diff --git a/src/types/types.ts b/src/types/types.ts index 7406b88c0..5660935b7 100644 --- a/src/types/types.ts +++ b/src/types/types.ts @@ -366,6 +366,7 @@ export type OpenAiEmbedder = { documentTemplate?: string; dimensions?: number; distribution?: Distribution; + url?: string; }; export type HuggingFaceEmbedder = { @@ -388,12 +389,10 @@ export type RestEmbedder = { apiKey?: string; dimensions?: number; documentTemplate?: string; - inputField?: string[] | null; - inputType?: 'text' | 'textArray'; - query?: Record | null; - pathToEmbeddings?: string[] | null; - embeddingObject?: string[] | null; distribution?: Distribution; + request: Record; + response: Record; + headers?: Record; }; export type OllamaEmbedder = { @@ -403,6 +402,7 @@ export type OllamaEmbedder = { model?: string; documentTemplate?: string; distribution?: Distribution; + dimensions?: number; }; export type Embedder = diff --git a/tests/embedders.test.ts b/tests/embedders.test.ts index 73d58745c..68d927d9c 100644 --- a/tests/embedders.test.ts +++ b/tests/embedders.test.ts @@ -118,6 +118,7 @@ describe.each([{ permission: 'Master' }, { permission: 'Admin' }])( mean: 0.7, sigma: 0.3, }, + url: 'https://api.openai.com/v1/embeddings', }, }; const task: EnqueuedTask = await client @@ -169,17 +170,25 @@ describe.each([{ permission: 'Master' }, { permission: 'Admin' }])( dimensions: 1536, documentTemplate: "A movie titled '{{doc.title}}' whose description starts with {{doc.overview|truncatewords: 20}}", - inputField: ['input'], - inputType: 'textArray', - query: { - model: 'text-embedding-ada-002', - }, - pathToEmbeddings: ['data'], - embeddingObject: ['embedding'], distribution: { mean: 0.7, sigma: 0.3, }, + request: { + model: 'text-embedding-3-small', + input: ['{{text}}', '{{..}}'], + }, + response: { + data: [ + { + embedding: '{{embedding}}', + }, + '{{..}}', + ], + }, + headers: { + 'Custom-Header': 'CustomValue', + }, }, }; const task: EnqueuedTask = await client @@ -197,7 +206,7 @@ describe.each([{ permission: 'Master' }, { permission: 'Admin' }])( }); }); - test.skip(`${permission} key: Update embedders with 'ollama' source`, async () => { + test(`${permission} key: Update embedders with 'ollama' source`, async () => { const client = await getClient(permission); const newEmbedder: Embedders = { default: { @@ -210,6 +219,7 @@ describe.each([{ permission: 'Master' }, { permission: 'Admin' }])( mean: 0.7, sigma: 0.3, }, + dimensions: 512, }, }; const task: EnqueuedTask = await client