test: isolate gemini embedding request helpers

2026-05-06 16:50:43 +00:00 · 2026-04-17 01:46:47 +01:00
parent fd48dfa68f
commit cd330f5f98
6 changed files with 300 additions and 433 deletions
--- a/src/memory-host-sdk/host/embeddings-gemini-request.ts
+++ b/src/memory-host-sdk/host/embeddings-gemini-request.ts
@@ -0,0 +1,115 @@
+import type { EmbeddingInput } from "./embedding-inputs.js";
+import type { GeminiTaskType } from "./embeddings.types.js";
+
+export const DEFAULT_GEMINI_EMBEDDING_MODEL = "gemini-embedding-001";
+
+export const GEMINI_EMBEDDING_2_MODELS = new Set([
+  "gemini-embedding-2-preview",
+  // Add the GA model name here once released.
+]);
+
+const GEMINI_EMBEDDING_2_DEFAULT_DIMENSIONS = 3072;
+const GEMINI_EMBEDDING_2_VALID_DIMENSIONS = [768, 1536, 3072] as const;
+
+export type { GeminiTaskType } from "./embeddings.types.js";
+
+export type GeminiTextPart = { text: string };
+export type GeminiInlinePart = {
+  inlineData: { mimeType: string; data: string };
+};
+export type GeminiPart = GeminiTextPart | GeminiInlinePart;
+export type GeminiEmbeddingRequest = {
+  content: { parts: GeminiPart[] };
+  taskType: GeminiTaskType;
+  outputDimensionality?: number;
+  model?: string;
+};
+export type GeminiTextEmbeddingRequest = GeminiEmbeddingRequest;
+
+/** Builds the text-only Gemini embedding request shape used across direct and batch APIs. */
+export function buildGeminiTextEmbeddingRequest(params: {
+  text: string;
+  taskType: GeminiTaskType;
+  outputDimensionality?: number;
+  modelPath?: string;
+}): GeminiTextEmbeddingRequest {
+  return buildGeminiEmbeddingRequest({
+    input: { text: params.text },
+    taskType: params.taskType,
+    outputDimensionality: params.outputDimensionality,
+    modelPath: params.modelPath,
+  });
+}
+
+export function buildGeminiEmbeddingRequest(params: {
+  input: EmbeddingInput;
+  taskType: GeminiTaskType;
+  outputDimensionality?: number;
+  modelPath?: string;
+}): GeminiEmbeddingRequest {
+  const request: GeminiEmbeddingRequest = {
+    content: {
+      parts: params.input.parts?.map((part) =>
+        part.type === "text"
+          ? ({ text: part.text } satisfies GeminiTextPart)
+          : ({
+              inlineData: { mimeType: part.mimeType, data: part.data },
+            } satisfies GeminiInlinePart),
+      ) ?? [{ text: params.input.text }],
+    },
+    taskType: params.taskType,
+  };
+  if (params.modelPath) {
+    request.model = params.modelPath;
+  }
+  if (params.outputDimensionality != null) {
+    request.outputDimensionality = params.outputDimensionality;
+  }
+  return request;
+}
+
+/**
+ * Returns true if the given model name is a gemini-embedding-2 variant that
+ * supports `outputDimensionality` and extended task types.
+ */
+export function isGeminiEmbedding2Model(model: string): boolean {
+  return GEMINI_EMBEDDING_2_MODELS.has(model);
+}
+
+/**
+ * Validate and return the `outputDimensionality` for gemini-embedding-2 models.
+ * Returns `undefined` for older models (they don't support the param).
+ */
+export function resolveGeminiOutputDimensionality(
+  model: string,
+  requested?: number,
+): number | undefined {
+  if (!isGeminiEmbedding2Model(model)) {
+    return undefined;
+  }
+  if (requested == null) {
+    return GEMINI_EMBEDDING_2_DEFAULT_DIMENSIONS;
+  }
+  const valid: readonly number[] = GEMINI_EMBEDDING_2_VALID_DIMENSIONS;
+  if (!valid.includes(requested)) {
+    throw new Error(
+      `Invalid outputDimensionality ${requested} for ${model}. Valid values: ${valid.join(", ")}`,
+    );
+  }
+  return requested;
+}
+
+export function normalizeGeminiModel(model: string): string {
+  const trimmed = model.trim();
+  if (!trimmed) {
+    return DEFAULT_GEMINI_EMBEDDING_MODEL;
+  }
+  const withoutPrefix = trimmed.replace(/^models\//, "");
+  if (withoutPrefix.startsWith("gemini/")) {
+    return withoutPrefix.slice("gemini/".length);
+  }
+  if (withoutPrefix.startsWith("google/")) {
+    return withoutPrefix.slice("google/".length);
+  }
+  return withoutPrefix;
+}
--- a/src/memory-host-sdk/host/embeddings-gemini.test.ts
+++ b/src/memory-host-sdk/host/embeddings-gemini.test.ts
@@ -3,13 +3,12 @@ import * as authModule from "../../agents/model-auth.js";
 import {
  buildGeminiEmbeddingRequest,
  buildGeminiTextEmbeddingRequest,
-  createGeminiEmbeddingProvider,
  DEFAULT_GEMINI_EMBEDDING_MODEL,
  GEMINI_EMBEDDING_2_MODELS,
  isGeminiEmbedding2Model,
  normalizeGeminiModel,
  resolveGeminiOutputDimensionality,
-} from "./embeddings-gemini.js";
+} from "./embeddings-gemini-request.js";
 import {
  createGeminiBatchFetchMock,
  createJsonResponseFetchMock,
@@ -47,12 +46,17 @@ afterEach(() => {
  vi.unstubAllGlobals();
 });

+type GeminiProviderOptions = Parameters<
+  typeof import("./embeddings-gemini.js").createGeminiEmbeddingProvider
+>[0];
+
 async function createProviderWithFetch(
  fetchMock: JsonFetchMock,
-  options: Partial<Parameters<typeof createGeminiEmbeddingProvider>[0]> & { model: string },
+  options: Partial<GeminiProviderOptions> & { model: string },
 ) {
  installFetchMock(fetchMock as unknown as typeof globalThis.fetch);
  mockResolvedProviderKey(authModule.resolveApiKeyForProvider);
+  const { createGeminiEmbeddingProvider } = await import("./embeddings-gemini.js");
  const { provider } = await createGeminiEmbeddingProvider({
    config: {} as never,
    provider: "gemini",
@@ -63,7 +67,7 @@ async function createProviderWithFetch(
 }

 describe("Gemini embedding request helpers", () => {
-  it("builds text and multimodal requests", () => {
+  it("builds requests and resolves model settings", () => {
    expect(
      buildGeminiTextEmbeddingRequest({
        text: "hello",
@@ -101,16 +105,10 @@ describe("Gemini embedding request helpers", () => {
      taskType: "RETRIEVAL_DOCUMENT",
      outputDimensionality: 1536,
    });
-  });
-
-  it("detects v2 model names", () => {
    expect(GEMINI_EMBEDDING_2_MODELS.has("gemini-embedding-2-preview")).toBe(true);
    expect(isGeminiEmbedding2Model("gemini-embedding-2-preview")).toBe(true);
    expect(isGeminiEmbedding2Model("gemini-embedding-001")).toBe(false);
    expect(isGeminiEmbedding2Model("text-embedding-004")).toBe(false);
-  });
-
-  it("resolves v2 dimensions and rejects invalid values", () => {
    expect(resolveGeminiOutputDimensionality("gemini-embedding-001")).toBeUndefined();
    expect(resolveGeminiOutputDimensionality("text-embedding-004")).toBeUndefined();
    expect(resolveGeminiOutputDimensionality("gemini-embedding-2-preview")).toBe(3072);
@@ -123,9 +121,6 @@ describe("Gemini embedding request helpers", () => {
    expect(() => resolveGeminiOutputDimensionality("gemini-embedding-2-preview", 1024)).toThrow(
      /Valid values: 768, 1536, 3072/,
    );
-  });
-
-  it("normalizes known model prefixes and default model", () => {
    expect(normalizeGeminiModel("models/gemini-embedding-2-preview")).toBe(
      "gemini-embedding-2-preview",
    );
--- a/src/memory-host-sdk/host/embeddings-gemini.ts
+++ b/src/memory-host-sdk/host/embeddings-gemini.ts
@@ -13,14 +13,33 @@ import { normalizeOptionalString } from "../../shared/string-coerce.js";
 import type { EmbeddingInput } from "./embedding-inputs.js";
 import { sanitizeAndNormalizeEmbedding } from "./embedding-vectors.js";
 import { debugEmbeddingsLog } from "./embeddings-debug.js";
-import type {
-  EmbeddingProvider,
-  EmbeddingProviderOptions,
-  GeminiTaskType,
-} from "./embeddings.types.js";
+import {
+  buildGeminiEmbeddingRequest,
+  buildGeminiTextEmbeddingRequest,
+  isGeminiEmbedding2Model,
+  normalizeGeminiModel,
+  resolveGeminiOutputDimensionality,
+} from "./embeddings-gemini-request.js";
+import type { EmbeddingProvider, EmbeddingProviderOptions } from "./embeddings.types.js";
 import { buildRemoteBaseUrlPolicy, withRemoteHttpResponse } from "./remote-http.js";
 import { resolveMemorySecretInputString } from "./secret-input.js";

+export {
+  buildGeminiEmbeddingRequest,
+  buildGeminiTextEmbeddingRequest,
+  DEFAULT_GEMINI_EMBEDDING_MODEL,
+  GEMINI_EMBEDDING_2_MODELS,
+  isGeminiEmbedding2Model,
+  normalizeGeminiModel,
+  resolveGeminiOutputDimensionality,
+  type GeminiEmbeddingRequest,
+  type GeminiInlinePart,
+  type GeminiPart,
+  type GeminiTaskType,
+  type GeminiTextEmbeddingRequest,
+  type GeminiTextPart,
+} from "./embeddings-gemini-request.js";
+
 export type GeminiEmbeddingClient = {
  baseUrl: string;
  headers: Record<string, string>;
@@ -31,108 +50,9 @@ export type GeminiEmbeddingClient = {
  outputDimensionality?: number;
 };

-export const DEFAULT_GEMINI_EMBEDDING_MODEL = "gemini-embedding-001";
 const GEMINI_MAX_INPUT_TOKENS: Record<string, number> = {
  "text-embedding-004": 2048,
 };
-
-// --- gemini-embedding-2-preview support ---
-
-export const GEMINI_EMBEDDING_2_MODELS = new Set([
-  "gemini-embedding-2-preview",
-  // Add the GA model name here once released.
-]);
-
-const GEMINI_EMBEDDING_2_DEFAULT_DIMENSIONS = 3072;
-const GEMINI_EMBEDDING_2_VALID_DIMENSIONS = [768, 1536, 3072] as const;
-
-export type { GeminiTaskType } from "./embeddings.types.js";
-
-export type GeminiTextPart = { text: string };
-export type GeminiInlinePart = {
-  inlineData: { mimeType: string; data: string };
-};
-export type GeminiPart = GeminiTextPart | GeminiInlinePart;
-export type GeminiEmbeddingRequest = {
-  content: { parts: GeminiPart[] };
-  taskType: GeminiTaskType;
-  outputDimensionality?: number;
-  model?: string;
-};
-export type GeminiTextEmbeddingRequest = GeminiEmbeddingRequest;
-
-/** Builds the text-only Gemini embedding request shape used across direct and batch APIs. */
-export function buildGeminiTextEmbeddingRequest(params: {
-  text: string;
-  taskType: GeminiTaskType;
-  outputDimensionality?: number;
-  modelPath?: string;
-}): GeminiTextEmbeddingRequest {
-  return buildGeminiEmbeddingRequest({
-    input: { text: params.text },
-    taskType: params.taskType,
-    outputDimensionality: params.outputDimensionality,
-    modelPath: params.modelPath,
-  });
-}
-
-export function buildGeminiEmbeddingRequest(params: {
-  input: EmbeddingInput;
-  taskType: GeminiTaskType;
-  outputDimensionality?: number;
-  modelPath?: string;
-}): GeminiEmbeddingRequest {
-  const request: GeminiEmbeddingRequest = {
-    content: {
-      parts: params.input.parts?.map((part) =>
-        part.type === "text"
-          ? ({ text: part.text } satisfies GeminiTextPart)
-          : ({
-              inlineData: { mimeType: part.mimeType, data: part.data },
-            } satisfies GeminiInlinePart),
-      ) ?? [{ text: params.input.text }],
-    },
-    taskType: params.taskType,
-  };
-  if (params.modelPath) {
-    request.model = params.modelPath;
-  }
-  if (params.outputDimensionality != null) {
-    request.outputDimensionality = params.outputDimensionality;
-  }
-  return request;
-}
-
-/**
- * Returns true if the given model name is a gemini-embedding-2 variant that
- * supports `outputDimensionality` and extended task types.
- */
-export function isGeminiEmbedding2Model(model: string): boolean {
-  return GEMINI_EMBEDDING_2_MODELS.has(model);
-}
-
-/**
- * Validate and return the `outputDimensionality` for gemini-embedding-2 models.
- * Returns `undefined` for older models (they don't support the param).
- */
-export function resolveGeminiOutputDimensionality(
-  model: string,
-  requested?: number,
-): number | undefined {
-  if (!isGeminiEmbedding2Model(model)) {
-    return undefined;
-  }
-  if (requested == null) {
-    return GEMINI_EMBEDDING_2_DEFAULT_DIMENSIONS;
-  }
-  const valid: readonly number[] = GEMINI_EMBEDDING_2_VALID_DIMENSIONS;
-  if (!valid.includes(requested)) {
-    throw new Error(
-      `Invalid outputDimensionality ${requested} for ${model}. Valid values: ${valid.join(", ")}`,
-    );
-  }
-  return requested;
-}
 function resolveRemoteApiKey(remoteApiKey: unknown): string | undefined {
  const trimmed = resolveMemorySecretInputString({
    value: remoteApiKey,
@@ -147,21 +67,6 @@ function resolveRemoteApiKey(remoteApiKey: unknown): string | undefined {
  return trimmed;
 }

-export function normalizeGeminiModel(model: string): string {
-  const trimmed = model.trim();
-  if (!trimmed) {
-    return DEFAULT_GEMINI_EMBEDDING_MODEL;
-  }
-  const withoutPrefix = trimmed.replace(/^models\//, "");
-  if (withoutPrefix.startsWith("gemini/")) {
-    return withoutPrefix.slice("gemini/".length);
-  }
-  if (withoutPrefix.startsWith("google/")) {
-    return withoutPrefix.slice("google/".length);
-  }
-  return withoutPrefix;
-}
-
 async function fetchGeminiEmbeddingPayload(params: {
  client: GeminiEmbeddingClient;
  endpoint: string;