feat(providers): separate embedding providers, add custom-openai provider

This commit is contained in:
ItzCrazyKns 2024-05-04 10:51:06 +05:30
parent c710f4f88c
commit 9f45ecb98d
No known key found for this signature in database
GPG key ID: 8162927C7CCE3065
8 changed files with 360 additions and 72 deletions

View file

@ -8,7 +8,7 @@ import {
} from '../config';
import logger from '../utils/logger';
export const getAvailableProviders = async () => {
export const getAvailableChatModelProviders = async () => {
const openAIApiKey = getOpenaiApiKey();
const groqApiKey = getGroqApiKey();
const ollamaEndpoint = getOllamaApiEndpoint();
@ -33,10 +33,6 @@ export const getAvailableProviders = async () => {
modelName: 'gpt-4-turbo',
temperature: 0.7,
}),
embeddings: new OpenAIEmbeddings({
openAIApiKey,
modelName: 'text-embedding-3-large',
}),
};
} catch (err) {
logger.error(`Error loading OpenAI models: ${err}`);
@ -86,10 +82,6 @@ export const getAvailableProviders = async () => {
baseURL: 'https://api.groq.com/openai/v1',
},
),
embeddings: new OpenAIEmbeddings({
openAIApiKey: openAIApiKey,
modelName: 'text-embedding-3-large',
}),
};
} catch (err) {
logger.error(`Error loading Groq models: ${err}`);
@ -110,17 +102,56 @@ export const getAvailableProviders = async () => {
});
return acc;
}, {});
if (Object.keys(models['ollama']).length > 0) {
models['ollama']['embeddings'] = new OllamaEmbeddings({
baseUrl: ollamaEndpoint,
model: models['ollama'][Object.keys(models['ollama'])[0]].model,
});
}
} catch (err) {
logger.error(`Error loading Ollama models: ${err}`);
}
}
models['custom_openai'] = {};
return models;
};
export const getAvailableEmbeddingModelProviders = async () => {
const openAIApiKey = getOpenaiApiKey();
const ollamaEndpoint = getOllamaApiEndpoint();
const models = {};
if (openAIApiKey) {
try {
models['openai'] = {
'Text embedding 3 small': new OpenAIEmbeddings({
openAIApiKey,
modelName: 'text-embedding-3-small',
}),
'Text embedding 3 large': new OpenAIEmbeddings({
openAIApiKey,
modelName: 'text-embedding-3-large',
}),
};
} catch (err) {
logger.error(`Error loading OpenAI embeddings: ${err}`);
}
}
if (ollamaEndpoint) {
try {
const response = await fetch(`${ollamaEndpoint}/api/tags`);
const { models: ollamaModels } = (await response.json()) as any;
models['ollama'] = ollamaModels.reduce((acc, model) => {
acc[model.model] = new OllamaEmbeddings({
baseUrl: ollamaEndpoint,
model: model.model,
});
return acc;
}, {});
} catch (err) {
logger.error(`Error loading Ollama embeddings: ${err}`);
}
}
return models;
};