Add openai embeddings (#915)

* Add OpenAI embedding compatibility

* Use OPENAI_API_KEY by default

* lint

* Add default OpenAI URL
replace `authorization` by `apiKey`

* Add a note in readme

---------

Co-authored-by: Nathan Sarrazin <[email protected]>

README.md

@@ -120,7 +120,7 @@ TEXT_EMBEDDING_MODELS = `[
 ```
 
 The required fields are `name`, `chunkCharLength` and `endpoints`.
-Supported text embedding backends are: [`transformers.js`](https://huggingface.co/docs/transformers.js) and [`TEI`](https://github.com/huggingface/text-embeddings-inference). `transformers.js` models run locally as part of `chat-ui`, whereas `TEI` models run in a different environment & accessed through an API endpoint.
+Supported text embedding backends are: [`transformers.js`](https://huggingface.co/docs/transformers.js), [`TEI`](https://github.com/huggingface/text-embeddings-inference) and [`OpenAI`](https://platform.openai.com/docs/guides/embeddings). `transformers.js` models run locally as part of `chat-ui`, whereas `TEI` models run in a different environment & accessed through an API endpoint. `openai` models are accessed through the [OpenAI API](https://platform.openai.com/docs/guides/embeddings).
 
 When more than one embedding models are supplied in `.env.local` file, the first will be used by default, and the others will only be used on LLM's which configured `embeddingModel` to the name of the model.
 

src/lib/server/embeddingEndpoints/embeddingEndpoints.ts

@@ -7,6 +7,10 @@ import {
 	embeddingEndpointTransformersJS,
 	embeddingEndpointTransformersJSParametersSchema,
 } from "./transformersjs/embeddingEndpoints";
+import {
+	embeddingEndpointOpenAI,
+	embeddingEndpointOpenAIParametersSchema,
+} from "./openai/embeddingEndpoints";
 
 // parameters passed when generating text
 interface EmbeddingEndpointParameters {
@@ -21,6 +25,7 @@ export type EmbeddingEndpoint = (params: EmbeddingEndpointParameters) => Promise
 export const embeddingEndpointSchema = z.discriminatedUnion("type", [
 	embeddingEndpointTeiParametersSchema,
 	embeddingEndpointTransformersJSParametersSchema,
+	embeddingEndpointOpenAIParametersSchema,
 ]);
 
 type EmbeddingEndpointTypeOptions = z.infer<typeof embeddingEndpointSchema>["type"];
@@ -36,6 +41,7 @@ export const embeddingEndpoints: {
 } = {
 	tei: embeddingEndpointTei,
 	transformersjs: embeddingEndpointTransformersJS,
+	openai: embeddingEndpointOpenAI,
 };
 
 export default embeddingEndpoints;

src/lib/server/embeddingEndpoints/openai/embeddingEndpoints.ts

@@ -0,0 +1,51 @@
+import { z } from "zod";
+import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
+import { chunk } from "$lib/utils/chunk";
+import { OPENAI_API_KEY } from "$env/static/private";
+
+export const embeddingEndpointOpenAIParametersSchema = z.object({
+	weight: z.number().int().positive().default(1),
+	model: z.any(),
+	type: z.literal("openai"),
+	url: z.string().url().default("https://api.openai.com/v1/embeddings"),
+	apiKey: z.string().default(OPENAI_API_KEY),
+});
+
+export async function embeddingEndpointOpenAI(
+	input: z.input<typeof embeddingEndpointOpenAIParametersSchema>
+): Promise<EmbeddingEndpoint> {
+	const { url, model, apiKey } = embeddingEndpointOpenAIParametersSchema.parse(input);
+
+	const maxBatchSize = model.maxBatchSize || 100;
+
+	return async ({ inputs }) => {
+		const requestURL = new URL(url);
+
+		const batchesInputs = chunk(inputs, maxBatchSize);
+
+		const batchesResults = await Promise.all(
+			batchesInputs.map(async (batchInputs) => {
+				const response = await fetch(requestURL, {
+					method: "POST",
+					headers: {
+						Accept: "application/json",
+						"Content-Type": "application/json",
+						...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
+					},
+					body: JSON.stringify({ input: batchInputs, model: model.name }),
+				});
+
+				const embeddings: Embedding[] = [];
+				const responseObject = await response.json();
+				for (const embeddingObject of responseObject.data) {
+					embeddings.push(embeddingObject.embedding);
+				}
+				return embeddings;
+			})
+		);
+
+		const flatAllEmbeddings = batchesResults.flat();
+
+		return flatAllEmbeddings;
+	};
+}

src/lib/server/embeddingModels.ts

@@ -22,6 +22,7 @@ const modelConfig = z.object({
 	modelUrl: z.string().url().optional(),
 	endpoints: z.array(embeddingEndpointSchema).nonempty(),
 	chunkCharLength: z.number().positive(),
+	maxBatchSize: z.number().positive().optional(),
 	preQuery: z.string().default(""),
 	prePassage: z.string().default(""),
 });
@@ -70,6 +71,8 @@ const addEndpoint = (m: Awaited<ReturnType<typeof processEmbeddingModel>>) => ({
 						return embeddingEndpoints.tei(args);
 					case "transformersjs":
 						return embeddingEndpoints.transformersjs(args);
+					case "openai":
+						return embeddingEndpoints.openai(args);
 				}
 			}