blob: fb9615536bd980ae3d7e90064fec6ef125db642b (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
|
import OpenAI from "openai";
import type { EmbeddingProvider } from "./embedding-provider.js";
export type OpenAIEmbeddingModel =
| "text-embedding-3-small"
| "text-embedding-3-large";
export type OpenAIEmbeddingProviderConfiguration = {
apiKey: string;
model?: OpenAIEmbeddingModel;
};
const MODEL_DIMENSIONS: Record<OpenAIEmbeddingModel, number> = {
"text-embedding-3-small": 1536,
"text-embedding-3-large": 3072,
};
export class OpenAIEmbeddingProvider implements EmbeddingProvider {
private client: OpenAI;
private model: OpenAIEmbeddingModel;
readonly dimensions: number;
constructor(configuration: OpenAIEmbeddingProviderConfiguration) {
this.client = new OpenAI({ apiKey: configuration.apiKey });
this.model = configuration.model ?? "text-embedding-3-small";
this.dimensions = MODEL_DIMENSIONS[this.model];
}
async generate(text: string): Promise<number[]> {
const response = await this.client.embeddings.create({
model: this.model,
input: text,
});
return response.data[0]?.embedding ?? [];
}
async generateBatch(texts: string[]): Promise<number[][]> {
const response = await this.client.embeddings.create({
model: this.model,
input: texts,
});
return response.data.map((item) => item.embedding);
}
}
|