diff --git a/extensions/openai/embedding-provider.test.ts b/extensions/openai/embedding-provider.test.ts index 2d5f364b186..6afebd552c9 100644 --- a/extensions/openai/embedding-provider.test.ts +++ b/extensions/openai/embedding-provider.test.ts @@ -85,4 +85,22 @@ describe("OpenAI embedding provider", () => { }), ); }); + + it("sends outputDimensionality as OpenAI dimensions", async () => { + const { provider } = await createOpenAiEmbeddingProvider( + createOptions({ outputDimensionality: 512 }), + ); + + await provider.embedBatch(["doc"]); + + expect(mocks.fetchRemoteEmbeddingVectors).toHaveBeenCalledWith( + expect.objectContaining({ + body: { + model: "text-embedding-3-small", + input: ["doc"], + dimensions: 512, + }, + }), + ); + }); }); diff --git a/extensions/openai/embedding-provider.ts b/extensions/openai/embedding-provider.ts index 94a30383541..c5ebaab4946 100644 --- a/extensions/openai/embedding-provider.ts +++ b/extensions/openai/embedding-provider.ts @@ -16,6 +16,7 @@ export type OpenAiEmbeddingClient = { inputType?: string; queryInputType?: string; documentInputType?: string; + outputDimensionality?: number; }; const DEFAULT_OPENAI_BASE_URL = "https://api.openai.com/v1"; @@ -59,6 +60,9 @@ export async function createOpenAiEmbeddingProvider( body: { model: client.model, input, + ...(typeof client.outputDimensionality === "number" + ? { dimensions: client.outputDimensionality } + : {}), ...(inputType ? { input_type: inputType } : {}), }, errorPrefix: "openai embeddings failed", @@ -96,5 +100,6 @@ async function resolveOpenAiEmbeddingClient( inputType: options.inputType, queryInputType: options.queryInputType, documentInputType: options.documentInputType, + outputDimensionality: options.outputDimensionality, }; } diff --git a/extensions/openai/memory-embedding-adapter.test.ts b/extensions/openai/memory-embedding-adapter.test.ts index 2cdc26d4742..a4fb3f3c0b9 100644 --- a/extensions/openai/memory-embedding-adapter.test.ts +++ b/extensions/openai/memory-embedding-adapter.test.ts @@ -37,6 +37,7 @@ describe("OpenAI memory embedding adapter", () => { model: "text-embedding-3-small", inputType: "passage", documentInputType: "document", + outputDimensionality: 512, }, }); }); @@ -66,6 +67,7 @@ describe("OpenAI memory embedding adapter", () => { body: { model: "text-embedding-3-small", input: "doc one", + dimensions: 512, input_type: "document", }, }), diff --git a/extensions/openai/memory-embedding-adapter.ts b/extensions/openai/memory-embedding-adapter.ts index b592b6c005b..a9a0474122f 100644 --- a/extensions/openai/memory-embedding-adapter.ts +++ b/extensions/openai/memory-embedding-adapter.ts @@ -32,6 +32,7 @@ export const openAiMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapte provider: "openai", baseUrl: client.baseUrl, model: client.model, + outputDimensionality: client.outputDimensionality, documentInputType: client.documentInputType ?? client.inputType, headers: sanitizeEmbeddingCacheHeaders(client.headers, ["authorization"]), }, @@ -47,6 +48,9 @@ export const openAiMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapte body: { model: client.model, input: chunk.text, + ...(typeof client.outputDimensionality === "number" + ? { dimensions: client.outputDimensionality } + : {}), ...(inputType ? { input_type: inputType } : {}), }, })),