Fix openai embedding endpoint override (#344)

This commit is contained in:
Logan Yang 2024-03-08 19:22:20 -08:00 committed by GitHub
parent 6af3f4bdc5
commit 08eb5c8fcb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 79 additions and 45 deletions

View file

@ -15,69 +15,92 @@ export default class EmbeddingManager {
private static instance: EmbeddingManager;
private constructor(
private langChainParams: LangChainParams,
private encryptionService: EncryptionService,
private encryptionService: EncryptionService
) {}
static getInstance(
langChainParams: LangChainParams,
encryptionService: EncryptionService,
encryptionService: EncryptionService
): EmbeddingManager {
if (!EmbeddingManager.instance) {
EmbeddingManager.instance = new EmbeddingManager(langChainParams, encryptionService);
EmbeddingManager.instance = new EmbeddingManager(
langChainParams,
encryptionService
);
}
return EmbeddingManager.instance;
}
static getModelName(embeddingsInstance: Embeddings): string {
const emb = embeddingsInstance as any;
if ('model' in emb && emb.model) {
if ("model" in emb && emb.model) {
return emb.model as string;
} else if ('modelName' in emb && emb.modelName) {
} else if ("modelName" in emb && emb.modelName) {
return emb.modelName as string;
} else {
throw new Error(`Embeddings instance missing model or modelName properties: ${embeddingsInstance}`);
throw new Error(
`Embeddings instance missing model or modelName properties: ${embeddingsInstance}`
);
}
}
getEmbeddingsAPI(): Embeddings | undefined {
const decrypt = (key: string) => this.encryptionService.getDecryptedKey(key);
getOpenAIEmbeddingAPI(): OpenAIEmbeddings | undefined {
const decrypt = (key: string) =>
this.encryptionService.getDecryptedKey(key);
const {
openAIApiKey,
azureOpenAIApiKey,
azureOpenAIApiInstanceName,
azureOpenAIApiVersion,
azureOpenAIApiEmbeddingDeploymentName,
embeddingModel,
openAIEmbeddingProxyBaseUrl,
openAIEmbeddingProxyModelName,
} = this.langChainParams;
const OpenAIEmbeddingsAPI = openAIApiKey ? (
openAIEmbeddingProxyBaseUrl ?
new ProxyOpenAIEmbeddings({
modelName: openAIEmbeddingProxyModelName || this.langChainParams.embeddingModel,
openAIApiKey: decrypt(openAIApiKey),
maxRetries: 3,
maxConcurrency: 3,
timeout: 10000,
openAIEmbeddingProxyBaseUrl,
}) :
new OpenAIEmbeddings({
modelName: openAIEmbeddingProxyModelName || this.langChainParams.embeddingModel,
openAIApiKey: decrypt(openAIApiKey),
maxRetries: 3,
maxConcurrency: 3,
timeout: 10000,
})
) : null;
if (openAIEmbeddingProxyBaseUrl) {
return new ProxyOpenAIEmbeddings({
modelName: openAIEmbeddingProxyModelName || embeddingModel,
openAIApiKey: openAIApiKey? decrypt(openAIApiKey) : 'default-key',
maxRetries: 3,
maxConcurrency: 3,
timeout: 10000,
openAIEmbeddingProxyBaseUrl,
});
} else if (openAIApiKey) {
// No proxy URL; now check if the API key exists
return new OpenAIEmbeddings({
modelName:
openAIEmbeddingProxyModelName || embeddingModel,
openAIApiKey: decrypt(openAIApiKey),
maxRetries: 3,
maxConcurrency: 3,
timeout: 10000,
});
}
}
const embeddingProvder = EMBEDDING_MODEL_TO_PROVIDERS[this.langChainParams.embeddingModel];
getEmbeddingsAPI(): Embeddings | undefined {
const decrypt = (key: string) =>
this.encryptionService.getDecryptedKey(key);
const {
azureOpenAIApiKey,
azureOpenAIApiInstanceName,
azureOpenAIApiVersion,
azureOpenAIApiEmbeddingDeploymentName,
openAIEmbeddingProxyModelName,
} = this.langChainParams;
switch(embeddingProvder) {
const OpenAIEmbeddingsAPI = this.getOpenAIEmbeddingAPI();
const embeddingProvder =
EMBEDDING_MODEL_TO_PROVIDERS[this.langChainParams.embeddingModel];
switch (embeddingProvder) {
case ModelProviders.OPENAI:
if (OpenAIEmbeddingsAPI) {
return OpenAIEmbeddingsAPI;
}
console.error('OpenAI API key is not provided for the embedding model.');
console.error(
"OpenAI API key is not provided for the embedding model."
);
break;
case ModelProviders.COHEREAI:
return new CohereEmbeddings({
@ -96,23 +119,34 @@ export default class EmbeddingManager {
maxConcurrency: 3,
});
}
console.error('Azure OpenAI API key is not provided for the embedding model.');
console.error(
"Azure OpenAI API key is not provided for the embedding model."
);
break;
case ModelProviders.OLLAMA:
return new OllamaEmbeddings({
...(this.langChainParams.ollamaBaseUrl ? { baseUrl: this.langChainParams.ollamaBaseUrl } : {}),
...(this.langChainParams.ollamaBaseUrl
? { baseUrl: this.langChainParams.ollamaBaseUrl }
: {}),
// TODO: Add custom ollama embedding model setting once they have other models
model: NOMIC_EMBED_TEXT,
})
default:
console.error('No embedding provider set or no valid API key provided. Defaulting to OpenAI.');
return OpenAIEmbeddingsAPI || new OpenAIEmbeddings({
modelName: openAIEmbeddingProxyModelName || this.langChainParams.embeddingModel,
openAIApiKey: 'default-key',
maxRetries: 3,
maxConcurrency: 3,
timeout: 10000,
});
default:
console.error(
"No embedding provider set or no valid API key provided. Defaulting to OpenAI."
);
return (
OpenAIEmbeddingsAPI ||
new OpenAIEmbeddings({
modelName:
openAIEmbeddingProxyModelName ||
this.langChainParams.embeddingModel,
openAIApiKey: "default-key",
maxRetries: 3,
maxConcurrency: 3,
timeout: 10000,
})
);
}
}
}

View file

@ -26,7 +26,7 @@ export class ProxyOpenAIEmbeddings extends OpenAIEmbeddings {
// Reinitialize the client with the updated clientConfig
this["client"] = new OpenAI({
...this["clientConfig"],
baseURL: fields.openAIProxyBaseUrl,
baseURL: fields.openAIEmbeddingProxyBaseUrl,
});
}
}