diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index d64a5899519..e5dc05e4797 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -118,15 +118,20 @@ class HeliconeLogger: f"Helicone Logging - Enters logging function for model {model}" ) litellm_params = kwargs.get("litellm_params", {}) + custom_llm_provider = litellm_params.get("custom_llm_provider", "") kwargs.get("litellm_call_id", None) metadata = litellm_params.get("metadata", {}) or {} metadata = self.add_metadata_from_header(litellm_params, metadata) + + # Check if model is a vertex_ai model + is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/") + model = ( model if any( accepted_model in model for accepted_model in self.helicone_model_list - ) + ) or is_vertex_ai else "gpt-3.5-turbo" ) provider_request = {"model": model, "messages": messages} @@ -135,7 +140,7 @@ class HeliconeLogger: ): response_obj = response_obj.json() - if "claude" in model: + if "claude" in model and not is_vertex_ai: response_obj = self.claude_mapping( model=model, messages=messages, response_obj=response_obj ) @@ -149,12 +154,15 @@ class HeliconeLogger: # Code to be executed provider_url = self.provider_url url = f"{self.api_base}/oai/v1/log" - if "claude" in model: + if "claude" in model and not is_vertex_ai: url = f"{self.api_base}/anthropic/v1/log" provider_url = "https://api.anthropic.com/v1/messages" elif "gemini" in model: url = f"{self.api_base}/custom/v1/log" provider_url = "https://generativelanguage.googleapis.com/v1beta" + elif is_vertex_ai: + url = f"{self.api_base}/custom/v1/log" + provider_url = "https://aiplatform.googleapis.com/v1" headers = { "Authorization": f"Bearer {self.key}", "Content-Type": "application/json", diff --git a/tests/litellm/integrations/helicone/test_helicone_gemini.py b/tests/litellm/integrations/helicone/test_helicone_gemini.py index 37c50241a14..f42a7016131 100644 --- a/tests/litellm/integrations/helicone/test_helicone_gemini.py +++ b/tests/litellm/integrations/helicone/test_helicone_gemini.py @@ -33,3 +33,32 @@ def test_helicone_gemini_models_recognized(): for accepted_model in logger.helicone_model_list ) assert is_recognized, f"{model} should be recognized by helicone_model_list" + + +def test_helicone_vertex_ai_models_recognized(): + """ + Test that Vertex AI models (GLM, DeepSeek, etc.) are recognized via custom_llm_provider. + """ + # Test models that don't contain "gemini" but are vertex_ai + test_models = [ + "vertex_ai/zai-org/glm-4.7-maas", + "vertex_ai/deepseek-ai/deepseek-v3", + "vertex_ai/meta/llama-3.1-405b", + ] + for model in test_models: + is_vertex_ai = model.startswith("vertex_ai/") + assert is_vertex_ai, f"{model} should be recognized as vertex_ai model" + + +def test_helicone_vertex_ai_via_custom_llm_provider(): + """ + Test that vertex_ai models are recognized when custom_llm_provider is set. + """ + # Models without vertex_ai/ prefix but with custom_llm_provider="vertex_ai" + test_cases = [ + ("zai-org/glm-4.7-maas", "vertex_ai"), + ("deepseek-ai/deepseek-v3", "vertex_ai"), + ] + for model, custom_llm_provider in test_cases: + is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/") + assert is_vertex_ai, f"{model} with custom_llm_provider={custom_llm_provider} should be recognized as vertex_ai"