feat: add support for Mistral and Codestral models

This commit is contained in:
dst03106
2025-04-18 12:02:55 +09:00
parent 58f6943b12
commit 869a179506
3 changed files with 50 additions and 0 deletions

View File

@ -112,6 +112,20 @@ MAX_TOKENS = {
"deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": 128000,
"deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": 128000,
"deepinfra/deepseek-ai/DeepSeek-R1": 128000,
"mistral/mistral-small-latest": 8191,
"mistral/mistral-medium-latest": 8191,
"mistral/mistral-large-2407": 128000,
"mistral/mistral-large-latest": 128000,
"mistral/open-mistral-7b": 8191,
"mistral/open-mixtral-8x7b": 8191,
"mistral/open-mixtral-8x22b": 8191,
"mistral/codestral-latest": 8191,
"mistral/open-mistral-nemo": 128000,
"mistral/open-mistral-nemo-2407": 128000,
"mistral/open-codestral-mamba": 256000,
"mistral/codestral-mamba-latest": 256000,
"codestral/codestral-latest": 8191,
"codestral/codestral-2405": 8191,
}
USER_MESSAGE_ONLY_MODELS = [

View File

@ -97,6 +97,14 @@ class LiteLLMAIHandler(BaseAiHandler):
if get_settings().get("DEEPINFRA.KEY", None):
os.environ['DEEPINFRA_API_KEY'] = get_settings().get("DEEPINFRA.KEY")
# Support mistral models
if get_settings().get("MISTRAL.KEY", None):
os.environ["MISTRAL_API_KEY"] = get_settings().get("MISTRAL.KEY")
# Support codestral models
if get_settings().get("CODESTRAL.KEY", None):
os.environ["CODESTRAL_API_KEY"] = get_settings().get("CODESTRAL.KEY")
# Check for Azure AD configuration
if get_settings().get("AZURE_AD.CLIENT_ID", None):
self.azure = True