fix: restore Cortex and Gemini CI regressions

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<string, string>;
@@ -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<string, number> = {
   "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<number[][]> => {
     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<number[][]> => {
+    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 };
 }