LibreChat/api/server/services/ModelService.js
Danny Avila c54a57019e
🕒 feat: Add 5-second timeout for Fetching Model Lists (#4423)
* refactor: add 5 second timeout for fetching AI provider model lists

* ci: fix test due to recent changes
2024-10-15 19:37:41 -04:00

270 lines
8 KiB
JavaScript

const axios = require('axios');
const { HttpsProxyAgent } = require('https-proxy-agent');
const { EModelEndpoint, defaultModels, CacheKeys } = require('librechat-data-provider');
const { inputSchema, logAxiosError, extractBaseURL, processModelData } = require('~/utils');
const { OllamaClient } = require('~/app/clients/OllamaClient');
const getLogStores = require('~/cache/getLogStores');
/**
* Splits a string by commas and trims each resulting value.
* @param {string} input - The input string to split.
* @returns {string[]} An array of trimmed values.
*/
const splitAndTrim = (input) => {
if (!input || typeof input !== 'string') {
return [];
}
return input
.split(',')
.map((item) => item.trim())
.filter(Boolean);
};
const { openAIApiKey, userProvidedOpenAI } = require('./Config/EndpointService').config;
/**
* Fetches OpenAI models from the specified base API path or Azure, based on the provided configuration.
*
* @param {Object} params - The parameters for fetching the models.
* @param {Object} params.user - The user ID to send to the API.
* @param {string} params.apiKey - The API key for authentication with the API.
* @param {string} params.baseURL - The base path URL for the API.
* @param {string} [params.name='OpenAI'] - The name of the API; defaults to 'OpenAI'.
* @param {boolean} [params.azure=false] - Whether to fetch models from Azure.
* @param {boolean} [params.userIdQuery=false] - Whether to send the user ID as a query parameter.
* @param {boolean} [params.createTokenConfig=true] - Whether to create a token configuration from the API response.
* @param {string} [params.tokenKey] - The cache key to save the token configuration. Uses `name` if omitted.
* @returns {Promise<string[]>} A promise that resolves to an array of model identifiers.
* @async
*/
const fetchModels = async ({
user,
apiKey,
baseURL,
name = 'OpenAI',
azure = false,
userIdQuery = false,
createTokenConfig = true,
tokenKey,
}) => {
let models = [];
if (!baseURL && !azure) {
return models;
}
if (!apiKey) {
return models;
}
if (name && name.toLowerCase().startsWith('ollama')) {
return await OllamaClient.fetchModels(baseURL);
}
try {
const options = {
headers: {
Authorization: `Bearer ${apiKey}`,
},
timeout: 5000,
};
if (process.env.PROXY) {
options.httpsAgent = new HttpsProxyAgent(process.env.PROXY);
}
if (process.env.OPENAI_ORGANIZATION && baseURL.includes('openai')) {
options.headers['OpenAI-Organization'] = process.env.OPENAI_ORGANIZATION;
}
const url = new URL(`${baseURL}${azure ? '' : '/models'}`);
if (user && userIdQuery) {
url.searchParams.append('user', user);
}
const res = await axios.get(url.toString(), options);
/** @type {z.infer<typeof inputSchema>} */
const input = res.data;
const validationResult = inputSchema.safeParse(input);
if (validationResult.success && createTokenConfig) {
const endpointTokenConfig = processModelData(input);
const cache = getLogStores(CacheKeys.TOKEN_CONFIG);
await cache.set(tokenKey ?? name, endpointTokenConfig);
}
models = input.data.map((item) => item.id);
} catch (error) {
const logMessage = `Failed to fetch models from ${azure ? 'Azure ' : ''}${name} API`;
logAxiosError({ message: logMessage, error });
}
return models;
};
/**
* Fetches models from the specified API path or Azure, based on the provided options.
* @async
* @function
* @param {object} opts - The options for fetching the models.
* @param {string} opts.user - The user ID to send to the API.
* @param {boolean} [opts.azure=false] - Whether to fetch models from Azure.
* @param {boolean} [opts.assistants=false] - Whether to fetch models from Azure.
* @param {boolean} [opts.plugins=false] - Whether to fetch models from the plugins.
* @param {string[]} [_models=[]] - The models to use as a fallback.
*/
const fetchOpenAIModels = async (opts, _models = []) => {
let models = _models.slice() ?? [];
let apiKey = openAIApiKey;
const openaiBaseURL = 'https://api.openai.com/v1';
let baseURL = openaiBaseURL;
let reverseProxyUrl = process.env.OPENAI_REVERSE_PROXY;
if (opts.assistants && process.env.ASSISTANTS_BASE_URL) {
reverseProxyUrl = process.env.ASSISTANTS_BASE_URL;
} else if (opts.azure) {
return models;
// const azure = getAzureCredentials();
// baseURL = (genAzureChatCompletion(azure))
// .split('/deployments')[0]
// .concat(`/models?api-version=${azure.azureOpenAIApiVersion}`);
// apiKey = azureOpenAIApiKey;
} else if (process.env.OPENROUTER_API_KEY) {
reverseProxyUrl = 'https://openrouter.ai/api/v1';
apiKey = process.env.OPENROUTER_API_KEY;
}
if (reverseProxyUrl) {
baseURL = extractBaseURL(reverseProxyUrl);
}
const modelsCache = getLogStores(CacheKeys.MODEL_QUERIES);
const cachedModels = await modelsCache.get(baseURL);
if (cachedModels) {
return cachedModels;
}
if (baseURL || opts.azure) {
models = await fetchModels({
apiKey,
baseURL,
azure: opts.azure,
user: opts.user,
name: baseURL,
});
}
if (models.length === 0) {
return _models;
}
if (baseURL === openaiBaseURL) {
const regex = /(text-davinci-003|gpt-|o1-)/;
models = models.filter((model) => regex.test(model));
const instructModels = models.filter((model) => model.includes('instruct'));
const otherModels = models.filter((model) => !model.includes('instruct'));
models = otherModels.concat(instructModels);
}
await modelsCache.set(baseURL, models);
return models;
};
/**
* Loads the default models for the application.
* @async
* @function
* @param {object} opts - The options for fetching the models.
* @param {string} opts.user - The user ID to send to the API.
* @param {boolean} [opts.azure=false] - Whether to fetch models from Azure.
* @param {boolean} [opts.plugins=false] - Whether to fetch models for the plugins endpoint.
* @param {boolean} [opts.assistants=false] - Whether to fetch models for the Assistants endpoint.
*/
const getOpenAIModels = async (opts) => {
let models = defaultModels[EModelEndpoint.openAI];
if (opts.assistants) {
models = defaultModels[EModelEndpoint.assistants];
} else if (opts.azure) {
models = defaultModels[EModelEndpoint.azureAssistants];
}
if (opts.plugins) {
models = models.filter(
(model) =>
!model.includes('text-davinci') &&
!model.includes('instruct') &&
!model.includes('0613') &&
!model.includes('0314') &&
!model.includes('0301'),
);
}
let key;
if (opts.assistants) {
key = 'ASSISTANTS_MODELS';
} else if (opts.azure) {
key = 'AZURE_OPENAI_MODELS';
} else if (opts.plugins) {
key = 'PLUGIN_MODELS';
} else {
key = 'OPENAI_MODELS';
}
if (process.env[key]) {
models = splitAndTrim(process.env[key]);
return models;
}
if (userProvidedOpenAI && !process.env.OPENROUTER_API_KEY) {
return models;
}
return await fetchOpenAIModels(opts, models);
};
const getChatGPTBrowserModels = () => {
let models = ['text-davinci-002-render-sha', 'gpt-4'];
if (process.env.CHATGPT_MODELS) {
models = splitAndTrim(process.env.CHATGPT_MODELS);
}
return models;
};
const getAnthropicModels = () => {
let models = defaultModels[EModelEndpoint.anthropic];
if (process.env.ANTHROPIC_MODELS) {
models = splitAndTrim(process.env.ANTHROPIC_MODELS);
}
return models;
};
const getGoogleModels = () => {
let models = defaultModels[EModelEndpoint.google];
if (process.env.GOOGLE_MODELS) {
models = splitAndTrim(process.env.GOOGLE_MODELS);
}
return models;
};
const getBedrockModels = () => {
let models = defaultModels[EModelEndpoint.bedrock];
if (process.env.BEDROCK_AWS_MODELS) {
models = splitAndTrim(process.env.BEDROCK_AWS_MODELS);
}
return models;
};
module.exports = {
fetchModels,
splitAndTrim,
getOpenAIModels,
getBedrockModels,
getChatGPTBrowserModels,
getAnthropicModels,
getGoogleModels,
};