@@ -1,5 +1,5 @@
11import {
2-createRemoteEmbeddingProvider,
2+fetchRemoteEmbeddingVectors,
33resolveRemoteEmbeddingClient,
44type MemoryEmbeddingProvider,
55type MemoryEmbeddingProviderCreateOptions,
@@ -13,6 +13,9 @@ export type OpenAiEmbeddingClient = {
1313ssrfPolicy?: SsrFPolicy;
1414fetchImpl?: typeof fetch;
1515model: string;
16+inputType?: string;
17+queryInputType?: string;
18+documentInputType?: string;
1619};
17201821const DEFAULT_OPENAI_BASE_URL = "https://api.openai.com/v1";
@@ -35,25 +38,63 @@ export async function createOpenAiEmbeddingProvider(
3538options: MemoryEmbeddingProviderCreateOptions,
3639): Promise<{ provider: MemoryEmbeddingProvider; client: OpenAiEmbeddingClient }> {
3740const client = await resolveOpenAiEmbeddingClient(options);
41+const url = `${client.baseUrl.replace(/\/$/, "")}/embeddings`;
42+43+const resolveInputType = (kind: "query" | "document"): string | undefined => {
44+const explicit = kind === "query" ? client.queryInputType : client.documentInputType;
45+const value = explicit ?? client.inputType;
46+return typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
47+};
48+49+const embed = async (input: string[], kind: "query" | "document"): Promise<number[][]> => {
50+if (input.length === 0) {
51+return [];
52+}
53+const inputType = resolveInputType(kind);
54+return await fetchRemoteEmbeddingVectors({
55+ url,
56+headers: client.headers,
57+ssrfPolicy: client.ssrfPolicy,
58+fetchImpl: client.fetchImpl,
59+body: {
60+model: client.model,
61+ input,
62+ ...(inputType ? { input_type: inputType } : {}),
63+},
64+errorPrefix: "openai embeddings failed",
65+});
66+};
38673968return {
40-provider: createRemoteEmbeddingProvider({
69+provider: {
4170id: "openai",
42- client,
43-errorPrefix: "openai embeddings failed",
44-maxInputTokens: OPENAI_MAX_INPUT_TOKENS[client.model],
45-}),
71+model: client.model,
72+ ...(typeof OPENAI_MAX_INPUT_TOKENS[client.model] === "number"
73+ ? { maxInputTokens: OPENAI_MAX_INPUT_TOKENS[client.model] }
74+ : {}),
75+embedQuery: async (text) => {
76+const [vec] = await embed([text], "query");
77+return vec ?? [];
78+},
79+embedBatch: async (texts) => await embed(texts, "document"),
80+},
4681 client,
4782};
4883}
49845085export async function resolveOpenAiEmbeddingClient(
5186options: MemoryEmbeddingProviderCreateOptions,
5287): Promise<OpenAiEmbeddingClient> {
53-return await resolveRemoteEmbeddingClient({
88+const client = await resolveRemoteEmbeddingClient({
5489provider: "openai",
5590 options,
5691defaultBaseUrl: DEFAULT_OPENAI_BASE_URL,
5792normalizeModel: normalizeOpenAiModel,
5893});
94+return {
95+ ...client,
96+inputType: options.inputType,
97+queryInputType: options.queryInputType,
98+documentInputType: options.documentInputType,
99+};
59100}
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。