scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

feat(helicone): add Vertex AI support for non-Gemini models

Browse source 
Extends HeliconeLogger to properly log Vertex AI partner models (GLM, DeepSeek, etc.)
that don't contain "gemini" in their name. Uses custom_llm_provider to detect vertex_ai.
This commit is contained in:
Chesars 2026-01-19 11:35:34 -03:00
parent cc39c71ac6
commit f76719750b
2 changed files with 40 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/integrations/helicone.py
View file

@ -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",

29
tests/litellm/integrations/helicone/test_helicone_gemini.py
View file

@ -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"