From cd330f5f98a80e37ce0aad3f0eb49b5d66fcffea Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 17 Apr 2026 01:46:47 +0100 Subject: [PATCH] test: isolate gemini embedding request helpers --- .../src/host/embeddings-gemini-request.ts | 121 +++++++++++ .../src/host/embeddings-gemini.test.ts | 189 +----------------- .../src/host/embeddings-gemini.ts | 144 +++---------- .../host/embeddings-gemini-request.ts | 115 +++++++++++ .../host/embeddings-gemini.test.ts | 21 +- src/memory-host-sdk/host/embeddings-gemini.ts | 143 +++---------- 6 files changed, 300 insertions(+), 433 deletions(-) create mode 100644 packages/memory-host-sdk/src/host/embeddings-gemini-request.ts create mode 100644 src/memory-host-sdk/host/embeddings-gemini-request.ts diff --git a/packages/memory-host-sdk/src/host/embeddings-gemini-request.ts b/packages/memory-host-sdk/src/host/embeddings-gemini-request.ts new file mode 100644 index 00000000000..887376bbad3 --- /dev/null +++ b/packages/memory-host-sdk/src/host/embeddings-gemini-request.ts @@ -0,0 +1,121 @@ +import type { EmbeddingInput } from "./embedding-inputs.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 = + | "RETRIEVAL_QUERY" + | "RETRIEVAL_DOCUMENT" + | "SEMANTIC_SIMILARITY" + | "CLASSIFICATION" + | "CLUSTERING" + | "QUESTION_ANSWERING" + | "FACT_VERIFICATION"; + +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; +} diff --git a/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts b/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts index f1401dcf412..06c804c0bf1 100644 --- a/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts +++ b/packages/memory-host-sdk/src/host/embeddings-gemini.test.ts @@ -1,133 +1,13 @@ -import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; -import * as authModule from "../../../../src/agents/model-auth.js"; +import { describe, expect, it } from "vitest"; +import { + buildGeminiEmbeddingRequest, + DEFAULT_GEMINI_EMBEDDING_MODEL, + normalizeGeminiModel, + resolveGeminiOutputDimensionality, +} from "./embeddings-gemini-request.js"; -vi.mock("../../../../src/infra/net/fetch-guard.js", () => ({ - fetchWithSsrFGuard: async (params: { - url: string; - init?: RequestInit; - fetchImpl?: typeof fetch; - }) => { - const fetchImpl = params.fetchImpl ?? globalThis.fetch; - if (!fetchImpl) { - throw new Error("fetch is not available"); - } - const response = await fetchImpl(params.url, params.init); - return { - response, - finalUrl: params.url, - release: async () => {}, - }; - }, -})); - -const { resolveApiKeyForProviderMock } = vi.hoisted(() => ({ - resolveApiKeyForProviderMock: vi.fn(), -})); - -vi.mock("../../../../src/agents/model-auth.js", () => { - return { - resolveApiKeyForProvider: resolveApiKeyForProviderMock, - requireApiKey: (auth: { apiKey?: string; mode?: string }, provider: string) => { - if (auth.apiKey) { - return auth.apiKey; - } - throw new Error(`No API key resolved for provider "${provider}" (auth mode: ${auth.mode}).`); - }, - }; -}); - -const createGeminiFetchMock = (embeddingValues = [1, 2, 3]) => - vi.fn(async (_input?: unknown, _init?: unknown) => ({ - ok: true, - status: 200, - json: async () => ({ embedding: { values: embeddingValues } }), - })); - -const createGeminiBatchFetchMock = (count: number, embeddingValues = [1, 2, 3]) => - vi.fn(async (_input?: unknown, _init?: unknown) => ({ - ok: true, - status: 200, - json: async () => ({ - embeddings: Array.from({ length: count }, () => ({ values: embeddingValues })), - }), - })); - -function installFetchMock(fetchMock: typeof globalThis.fetch) { - vi.stubGlobal("fetch", fetchMock); -} - -function parseFetchBody(fetchMock: { mock: { calls: unknown[][] } }, callIndex = 0) { - const init = fetchMock.mock.calls[callIndex]?.[1] as RequestInit | undefined; - return JSON.parse((init?.body as string) ?? "{}") as Record; -} - -function magnitude(values: number[]) { - return Math.sqrt(values.reduce((sum, value) => sum + value * value, 0)); -} - -let buildGeminiEmbeddingRequest: typeof import("./embeddings-gemini.js").buildGeminiEmbeddingRequest; -let createGeminiEmbeddingProvider: typeof import("./embeddings-gemini.js").createGeminiEmbeddingProvider; -let DEFAULT_GEMINI_EMBEDDING_MODEL: typeof import("./embeddings-gemini.js").DEFAULT_GEMINI_EMBEDDING_MODEL; -let normalizeGeminiModel: typeof import("./embeddings-gemini.js").normalizeGeminiModel; -let resolveGeminiOutputDimensionality: typeof import("./embeddings-gemini.js").resolveGeminiOutputDimensionality; - -beforeAll(async () => { - vi.doUnmock("undici"); - ({ - buildGeminiEmbeddingRequest, - createGeminiEmbeddingProvider, - DEFAULT_GEMINI_EMBEDDING_MODEL, - normalizeGeminiModel, - resolveGeminiOutputDimensionality, - } = await import("./embeddings-gemini.js")); -}); - -beforeEach(() => { - vi.useRealTimers(); - vi.doUnmock("undici"); -}); - -afterEach(() => { - vi.doUnmock("undici"); - vi.resetAllMocks(); - vi.unstubAllGlobals(); -}); - -function mockResolvedProviderKey(apiKey = "test-key") { - vi.mocked(authModule.resolveApiKeyForProvider).mockResolvedValue({ - apiKey, - mode: "api-key", - source: "test", - }); -} - -type GeminiFetchMock = - | ReturnType - | ReturnType; - -async function createProviderWithFetch( - fetchMock: GeminiFetchMock, - options: Partial[0]> & { model: string }, -) { - installFetchMock(fetchMock as unknown as typeof globalThis.fetch); - mockResolvedProviderKey(); - const { provider } = await createGeminiEmbeddingProvider({ - config: {} as never, - provider: "gemini", - fallback: "none", - ...options, - }); - return provider; -} - -function expectNormalizedThreeFourVector(embedding: number[]) { - expect(embedding[0]).toBeCloseTo(0.6, 5); - expect(embedding[1]).toBeCloseTo(0.8, 5); - expect(magnitude(embedding)).toBeCloseTo(1, 5); -} - -describe("package Gemini embedding provider smoke", () => { - it("builds multimodal v2 requests and resolves dimensions", () => { +describe("package Gemini embedding request helpers", () => { + it("builds multimodal v2 requests and resolves model settings", () => { expect( buildGeminiEmbeddingRequest({ input: { @@ -158,57 +38,6 @@ describe("package Gemini embedding provider smoke", () => { expect(() => resolveGeminiOutputDimensionality("gemini-embedding-2-preview", 512)).toThrow( /Invalid outputDimensionality 512/, ); - }); - - it("handles legacy and v2 request/response behavior", async () => { - const legacyFetch = createGeminiBatchFetchMock(2); - const legacyProvider = await createProviderWithFetch(legacyFetch, { - model: "gemini-embedding-001", - }); - - await legacyProvider.embedQuery("test query"); - await legacyProvider.embedBatch(["text1", "text2"]); - - expect(parseFetchBody(legacyFetch, 0)).toMatchObject({ - taskType: "RETRIEVAL_QUERY", - content: { parts: [{ text: "test query" }] }, - }); - expect(parseFetchBody(legacyFetch, 0)).not.toHaveProperty("outputDimensionality"); - expect(parseFetchBody(legacyFetch, 1)).not.toHaveProperty("outputDimensionality"); - - const v2QueryFetch = createGeminiFetchMock([3, 4]); - const v2QueryProvider = await createProviderWithFetch(v2QueryFetch, { - model: "gemini-embedding-2-preview", - outputDimensionality: 768, - taskType: "SEMANTIC_SIMILARITY", - }); - await expect(v2QueryProvider.embedQuery(" ")).resolves.toEqual([]); - await expect(v2QueryProvider.embedBatch([])).resolves.toEqual([]); - expectNormalizedThreeFourVector(await v2QueryProvider.embedQuery("test query")); - - const v2BatchFetch = createGeminiBatchFetchMock(2, [3, 4]); - const v2BatchProvider = await createProviderWithFetch(v2BatchFetch, { - model: "gemini-embedding-2-preview", - outputDimensionality: 768, - taskType: "SEMANTIC_SIMILARITY", - }); - const batch = await v2BatchProvider.embedBatch(["text1", "text2"]); - expect(batch).toHaveLength(2); - for (const embedding of batch) { - expectNormalizedThreeFourVector(embedding); - } - - expect(parseFetchBody(v2QueryFetch)).toMatchObject({ - outputDimensionality: 768, - taskType: "SEMANTIC_SIMILARITY", - }); - expect(parseFetchBody(v2BatchFetch).requests).toEqual([ - expect.objectContaining({ outputDimensionality: 768 }), - expect.objectContaining({ outputDimensionality: 768 }), - ]); - }); - - it("normalizes known model prefixes and the default model", () => { expect(normalizeGeminiModel("models/gemini-embedding-2-preview")).toBe( "gemini-embedding-2-preview", ); diff --git a/packages/memory-host-sdk/src/host/embeddings-gemini.ts b/packages/memory-host-sdk/src/host/embeddings-gemini.ts index b4d4e0e9581..ed6bdaca899 100644 --- a/packages/memory-host-sdk/src/host/embeddings-gemini.ts +++ b/packages/memory-host-sdk/src/host/embeddings-gemini.ts @@ -12,10 +12,33 @@ import type { SsrFPolicy } from "../../../../src/infra/net/ssrf.js"; import type { EmbeddingInput } from "./embedding-inputs.js"; import { sanitizeAndNormalizeEmbedding } from "./embedding-vectors.js"; import { debugEmbeddingsLog } from "./embeddings-debug.js"; +import { + buildGeminiEmbeddingRequest, + buildGeminiTextEmbeddingRequest, + isGeminiEmbedding2Model, + normalizeGeminiModel, + resolveGeminiOutputDimensionality, +} from "./embeddings-gemini-request.js"; import type { EmbeddingProvider, EmbeddingProviderOptions } from "./embeddings.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; @@ -26,115 +49,9 @@ export type GeminiEmbeddingClient = { outputDimensionality?: number; }; -export const DEFAULT_GEMINI_EMBEDDING_MODEL = "gemini-embedding-001"; const GEMINI_MAX_INPUT_TOKENS: Record = { "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 = - | "RETRIEVAL_QUERY" - | "RETRIEVAL_DOCUMENT" - | "SEMANTIC_SIMILARITY" - | "CLASSIFICATION" - | "CLUSTERING" - | "QUESTION_ANSWERING" - | "FACT_VERIFICATION"; - -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, @@ -149,21 +66,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; diff --git a/src/memory-host-sdk/host/embeddings-gemini-request.ts b/src/memory-host-sdk/host/embeddings-gemini-request.ts new file mode 100644 index 00000000000..038843fe207 --- /dev/null +++ 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; +} diff --git a/src/memory-host-sdk/host/embeddings-gemini.test.ts b/src/memory-host-sdk/host/embeddings-gemini.test.ts index 13185ae2fbc..84cfbaae917 100644 --- 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[0]> & { model: string }, + options: Partial & { 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", ); diff --git a/src/memory-host-sdk/host/embeddings-gemini.ts b/src/memory-host-sdk/host/embeddings-gemini.ts index 91f678abd65..d4b989d5fcc 100644 --- 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; @@ -31,108 +50,9 @@ export type GeminiEmbeddingClient = { outputDimensionality?: number; }; -export const DEFAULT_GEMINI_EMBEDDING_MODEL = "gemini-embedding-001"; const GEMINI_MAX_INPUT_TOKENS: Record = { "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;