From 1ffd96b348c2c0fd523640d0993a3523fa760b00 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 1 Aug 2024 15:22:32 -0700 Subject: [PATCH] fix(google.py): fix cost tracking for vertex ai mistral models --- litellm/litellm_core_utils/llm_cost_calc/google.py | 4 +++- litellm/tests/test_amazing_vertex_completion.py | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/google.py b/litellm/litellm_core_utils/llm_cost_calc/google.py index 26eeb7b7a7..0642f101e9 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/google.py +++ b/litellm/litellm_core_utils/llm_cost_calc/google.py @@ -44,7 +44,9 @@ def cost_router( Returns - str, the specific google cost calc function it should route to. """ - if custom_llm_provider == "vertex_ai" and ("claude" in model or "llama" in model): + if custom_llm_provider == "vertex_ai" and ( + "claude" in model or "llama" in model or "mistral" in model + ): return "cost_per_token" elif custom_llm_provider == "gemini": return "cost_per_token" diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 96ea323282..438997b383 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -904,6 +904,7 @@ from litellm.tests.test_completion import response_format_tests "model", [ "vertex_ai/mistral-large@2407", + "vertex_ai/mistral-nemo@2407", "vertex_ai/meta/llama3-405b-instruct-maas", ], # ) # "vertex_ai",