phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(vertex_ai_partner.py): add /chat/completion codestral support

Browse source 
Closes https://github.com/BerriAI/litellm/issues/4984
This commit is contained in:
Krrish Dholakia 2024-08-01 16:12:05 -07:00
parent 1ffd96b348
commit 246b3227a9
5 changed files with 30 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/vertex_ai_partner.py
View file

@ -192,7 +192,7 @@ class VertexAIPartnerModels(BaseLLM):
if "llama" in model: if "llama" in model:
partner = "llama" partner = "llama"
elif "mistral" in model: elif "mistral" in model or "codestral" in model:
partner = "mistralai" partner = "mistralai"
optional_params["custom_endpoint"] = True optional_params["custom_endpoint"] = True

6
litellm/main.py
View file

@ -2074,7 +2074,11 @@ def completion(
timeout=timeout, timeout=timeout,
client=client, client=client,
) )
elif model.startswith("meta/") or model.startswith("mistral"): elif (
model.startswith("meta/")
or model.startswith("mistral")
or model.startswith("codestral")
):
model_response = vertex_partner_models_chat_completion.completion( model_response = vertex_partner_models_chat_completion.completion(
model=model, model=model,
messages=messages, messages=messages,

10
litellm/model_prices_and_context_window_backup.json
View file

@ -2092,6 +2092,16 @@
"mode": "chat", "mode": "chat",
"supports_function_calling": true "supports_function_calling": true
}, },
"vertex_ai/codestral@latest": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 0.000001,
"output_cost_per_token": 0.000003,
"litellm_provider": "vertex_ai-mistral_models",
"mode": "chat",
"supports_function_calling": true
},
"vertex_ai/codestral@2405": { "vertex_ai/codestral@2405": {
"max_tokens": 128000, "max_tokens": 128000,
"max_input_tokens": 128000, "max_input_tokens": 128000,

7
litellm/tests/test_amazing_vertex_completion.py
View file

@ -903,9 +903,10 @@ from litellm.tests.test_completion import response_format_tests
@pytest.mark.parametrize( @pytest.mark.parametrize(
"model", "model",
[ [
"vertex_ai/mistral-large@2407", # "vertex_ai/mistral-large@2407",
"vertex_ai/mistral-nemo@2407", # "vertex_ai/mistral-nemo@2407",
"vertex_ai/meta/llama3-405b-instruct-maas", "vertex_ai/codestral@2405",
# "vertex_ai/meta/llama3-405b-instruct-maas",
], # ], #
) # "vertex_ai", ) # "vertex_ai",
@pytest.mark.parametrize( @pytest.mark.parametrize(

10
model_prices_and_context_window.json
View file

@ -2092,6 +2092,16 @@
"mode": "chat", "mode": "chat",
"supports_function_calling": true "supports_function_calling": true
}, },
"vertex_ai/codestral@latest": {
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"input_cost_per_token": 0.000001,
"output_cost_per_token": 0.000003,
"litellm_provider": "vertex_ai-mistral_models",
"mode": "chat",
"supports_function_calling": true
},
"vertex_ai/codestral@2405": { "vertex_ai/codestral@2405": {
"max_tokens": 128000, "max_tokens": 128000,
"max_input_tokens": 128000, "max_input_tokens": 128000,