From fe1fb7634355debd36980d7cbc3966f1051538e7 Mon Sep 17 00:00:00 2001 From: Marc J Saint-jour <82672745+Junebugg1214@users.noreply.github.com> Date: Thu, 12 Mar 2026 22:03:16 -0400 Subject: [PATCH] fix: restore Cortex and Gemini CI regressions --- src/memory/embeddings-gemini.ts | 163 ++++++++++++++++++++++++++++---- 1 file changed, 145 insertions(+), 18 deletions(-) diff --git a/src/memory/embeddings-gemini.ts b/src/memory/embeddings-gemini.ts index 5786ead73fd..1fc7b21b6e2 100644 --- a/src/memory/embeddings-gemini.ts +++ b/src/memory/embeddings-gemini.ts @@ -5,11 +5,41 @@ import { import { requireApiKey, resolveApiKeyForProvider } from "../agents/model-auth.js"; import { parseGeminiAuth } from "../infra/gemini-auth.js"; import type { SsrFPolicy } from "../infra/net/ssrf.js"; +import { + hasNonTextEmbeddingParts, + isInlineDataEmbeddingInputPart, + type EmbeddingInput, +} from "./embedding-inputs.js"; +import { sanitizeAndNormalizeEmbedding } from "./embedding-vectors.js"; import { debugEmbeddingsLog } from "./embeddings-debug.js"; import type { EmbeddingProvider, EmbeddingProviderOptions } from "./embeddings.js"; import { buildRemoteBaseUrlPolicy, withRemoteHttpResponse } from "./remote-http.js"; import { resolveMemorySecretInputString } from "./secret-input.js"; +export const GEMINI_EMBEDDING_2_MODELS = new Set(["gemini-embedding-2-preview"]); +const GEMINI_EMBEDDING_2_DIMENSIONS = new Set([768, 1536, 3072]); + +export type GeminiTaskType = + | "RETRIEVAL_QUERY" + | "RETRIEVAL_DOCUMENT" + | "SEMANTIC_SIMILARITY" + | (string & {}); + +export type GeminiEmbeddingRequestPart = + | { text: string } + | { inlineData: { mimeType: string; data: string } }; + +export type GeminiEmbeddingContent = { + parts: GeminiEmbeddingRequestPart[]; +}; + +export type GeminiTextEmbeddingRequest = { + model?: string; + content: GeminiEmbeddingContent; + taskType?: GeminiTaskType; + outputDimensionality?: number; +}; + export type GeminiEmbeddingClient = { baseUrl: string; headers: Record; @@ -17,6 +47,7 @@ export type GeminiEmbeddingClient = { model: string; modelPath: string; apiKeys: string[]; + outputDimensionality?: number; }; const DEFAULT_GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"; @@ -24,6 +55,29 @@ export const DEFAULT_GEMINI_EMBEDDING_MODEL = "gemini-embedding-001"; const GEMINI_MAX_INPUT_TOKENS: Record = { "text-embedding-004": 2048, }; + +export function isGeminiEmbedding2Model(model: string): boolean { + return GEMINI_EMBEDDING_2_MODELS.has(normalizeGeminiModel(model)); +} + +export function resolveGeminiOutputDimensionality( + model: string, + outputDimensionality?: number, +): number | undefined { + if (!isGeminiEmbedding2Model(model)) { + return undefined; + } + if (typeof outputDimensionality !== "number") { + return 3072; + } + if (!GEMINI_EMBEDDING_2_DIMENSIONS.has(outputDimensionality)) { + throw new Error( + `Invalid outputDimensionality ${outputDimensionality} for ${normalizeGeminiModel(model)}. Valid values: 768, 1536, 3072.`, + ); + } + return outputDimensionality; +} + function resolveRemoteApiKey(remoteApiKey: unknown): string | undefined { const trimmed = resolveMemorySecretInputString({ value: remoteApiKey, @@ -66,6 +120,46 @@ function buildGeminiModelPath(model: string): string { return model.startsWith("models/") ? model : `models/${model}`; } +export function buildGeminiTextEmbeddingRequest(params: { + text: string; + taskType?: GeminiTaskType; + modelPath?: string; + outputDimensionality?: number; +}): GeminiTextEmbeddingRequest { + return { + ...(params.modelPath ? { model: params.modelPath } : {}), + content: { parts: [{ text: params.text }] }, + ...(params.taskType ? { taskType: params.taskType } : {}), + ...(typeof params.outputDimensionality === "number" + ? { outputDimensionality: params.outputDimensionality } + : {}), + }; +} + +export function buildGeminiEmbeddingRequest(params: { + input: EmbeddingInput; + taskType?: GeminiTaskType; + modelPath?: string; + outputDimensionality?: number; +}): GeminiTextEmbeddingRequest { + const parts = + params.input.parts?.length && hasNonTextEmbeddingParts(params.input) + ? params.input.parts.map((part) => + isInlineDataEmbeddingInputPart(part) + ? { inlineData: { mimeType: part.mimeType, data: part.data } } + : { text: part.text }, + ) + : [{ text: params.input.text }]; + return { + ...(params.modelPath ? { model: params.modelPath } : {}), + content: { parts }, + ...(params.taskType ? { taskType: params.taskType } : {}), + ...(typeof params.outputDimensionality === "number" + ? { outputDimensionality: params.outputDimensionality } + : {}), + }; +} + export async function createGeminiEmbeddingProvider( options: EmbeddingProviderOptions, ): Promise<{ provider: EmbeddingProvider; client: GeminiEmbeddingClient }> { @@ -110,29 +204,32 @@ export async function createGeminiEmbeddingProvider( provider: "google", apiKeys: client.apiKeys, execute: (apiKey) => - fetchWithGeminiAuth(apiKey, embedUrl, { - content: { parts: [{ text }] }, - taskType: "RETRIEVAL_QUERY", - ...(typeof options.outputDimensionality === "number" - ? { outputDimensionality: options.outputDimensionality } - : {}), - }), + fetchWithGeminiAuth( + apiKey, + embedUrl, + buildGeminiTextEmbeddingRequest({ + text, + taskType: options.taskType ?? "RETRIEVAL_QUERY", + modelPath: client.modelPath, + outputDimensionality: client.outputDimensionality, + }), + ), }); - return payload.embedding?.values ?? []; + return sanitizeAndNormalizeEmbedding(payload.embedding?.values ?? []); }; const embedBatch = async (texts: string[]): Promise => { if (texts.length === 0) { return []; } - const requests = texts.map((text) => ({ - model: client.modelPath, - content: { parts: [{ text }] }, - taskType: "RETRIEVAL_DOCUMENT", - ...(typeof options.outputDimensionality === "number" - ? { outputDimensionality: options.outputDimensionality } - : {}), - })); + const requests = texts.map((text) => + buildGeminiTextEmbeddingRequest({ + text, + taskType: "RETRIEVAL_DOCUMENT", + modelPath: client.modelPath, + outputDimensionality: client.outputDimensionality, + }), + ); const payload = await executeWithApiKeyRotation({ provider: "google", apiKeys: client.apiKeys, @@ -142,7 +239,31 @@ export async function createGeminiEmbeddingProvider( }), }); const embeddings = Array.isArray(payload.embeddings) ? payload.embeddings : []; - return texts.map((_, index) => embeddings[index]?.values ?? []); + return texts.map((_, index) => sanitizeAndNormalizeEmbedding(embeddings[index]?.values ?? [])); + }; + + const embedBatchInputs = async (inputs: EmbeddingInput[]): Promise => { + if (inputs.length === 0) { + return []; + } + const requests = inputs.map((input) => + buildGeminiEmbeddingRequest({ + input, + taskType: "RETRIEVAL_DOCUMENT", + modelPath: client.modelPath, + outputDimensionality: client.outputDimensionality, + }), + ); + const payload = await executeWithApiKeyRotation({ + provider: "google", + apiKeys: client.apiKeys, + execute: (apiKey) => + fetchWithGeminiAuth(apiKey, batchUrl, { + requests, + }), + }); + const embeddings = Array.isArray(payload.embeddings) ? payload.embeddings : []; + return inputs.map((_, index) => sanitizeAndNormalizeEmbedding(embeddings[index]?.values ?? [])); }; return { @@ -152,6 +273,7 @@ export async function createGeminiEmbeddingProvider( maxInputTokens: GEMINI_MAX_INPUT_TOKENS[client.model], embedQuery, embedBatch, + embedBatchInputs, }, client, }; @@ -189,13 +311,18 @@ export async function resolveGeminiEmbeddingClient( }); const model = normalizeGeminiModel(options.model); const modelPath = buildGeminiModelPath(model); + const outputDimensionality = resolveGeminiOutputDimensionality( + model, + options.outputDimensionality, + ); debugEmbeddingsLog("memory embeddings: gemini client", { rawBaseUrl, baseUrl, model, modelPath, + outputDimensionality, embedEndpoint: `${baseUrl}/${modelPath}:embedContent`, batchEndpoint: `${baseUrl}/${modelPath}:batchEmbedContents`, }); - return { baseUrl, headers, ssrfPolicy, model, modelPath, apiKeys }; + return { baseUrl, headers, ssrfPolicy, model, modelPath, apiKeys, outputDimensionality }; }