diff --git a/open-sse/handlers/embeddingProviders/gemini.js b/open-sse/handlers/embeddingProviders/gemini.js index 1a61954c..885ee306 100644 --- a/open-sse/handlers/embeddingProviders/gemini.js +++ b/open-sse/handlers/embeddingProviders/gemini.js @@ -13,12 +13,24 @@ export default { return `${BASE}/${path}:${op}?key=${encodeURIComponent(apiKey)}`; }, buildHeaders: () => ({ "Content-Type": "application/json" }), - buildBody: (model, { input }) => { + buildBody: (model, { input, dimensions }) => { const m = modelPath(model); + const outputDimensionality = Number(dimensions); + const hasOutputDimensionality = Number.isFinite(outputDimensionality) && outputDimensionality > 0; if (Array.isArray(input)) { - return { requests: input.map((text) => ({ model: m, content: { parts: [{ text: String(text) }] } })) }; + return { + requests: input.map((text) => ({ + model: m, + content: { parts: [{ text: String(text) }] }, + ...(hasOutputDimensionality ? { outputDimensionality } : {}), + })), + }; } - return { model: m, content: { parts: [{ text: String(input) }] } }; + return { + model: m, + content: { parts: [{ text: String(input) }] }, + ...(hasOutputDimensionality ? { outputDimensionality } : {}), + }; }, normalize: (responseBody, model) => { if (responseBody.object === "list" && Array.isArray(responseBody.data)) return responseBody; diff --git a/tests/unit/embeddingsCore.test.js b/tests/unit/embeddingsCore.test.js index cc5e0a1c..6450e4bb 100644 --- a/tests/unit/embeddingsCore.test.js +++ b/tests/unit/embeddingsCore.test.js @@ -141,6 +141,51 @@ describe("buildEmbeddingsBody", () => { const sent = JSON.parse(init.body); expect(sent.encoding_format).toBe("float"); }); + + it("gemini single input forwards dimensions as outputDimensionality", async () => { + vi.mocked(fetch).mockResolvedValueOnce(makeProviderResponse({ + embedding: { values: [0.1, 0.2, 0.3] }, + })); + + await handleEmbeddingsCore(makeOptions({ + body: { + model: "gemini/gemini-embedding-2-preview", + input: "test", + dimensions: 1536, + }, + modelInfo: { provider: "gemini", model: "gemini-embedding-2-preview" }, + credentials: { apiKey: "gemini-key" }, + })); + + const [, init] = vi.mocked(fetch).mock.calls[0]; + const sent = JSON.parse(init.body); + expect(sent.outputDimensionality).toBe(1536); + }); + + it("gemini batch input forwards dimensions on each request", async () => { + vi.mocked(fetch).mockResolvedValueOnce(makeProviderResponse({ + embeddings: [ + { values: [0.1, 0.2, 0.3] }, + { values: [0.4, 0.5, 0.6] }, + ], + })); + + await handleEmbeddingsCore(makeOptions({ + body: { + model: "gemini/gemini-embedding-2-preview", + input: ["hello", "world"], + dimensions: 1536, + }, + modelInfo: { provider: "gemini", model: "gemini-embedding-2-preview" }, + credentials: { apiKey: "gemini-key" }, + })); + + const [, init] = vi.mocked(fetch).mock.calls[0]; + const sent = JSON.parse(init.body); + expect(sent.requests).toHaveLength(2); + expect(sent.requests[0].outputDimensionality).toBe(1536); + expect(sent.requests[1].outputDimensionality).toBe(1536); + }); }); // ─── Test: buildEmbeddingsUrl ────────────────────────────────────────────────