scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

fix VertexAIPartnerModels

Browse source
This commit is contained in:
Ishaan Jaff 2025-09-03 13:01:47 -07:00
parent 8e9352fce7
commit c9f211f331
1 changed files with 22 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -1,5 +1,6 @@
# What is this?
## API Handler for calling Vertex AI Partner Models
from enum import Enum
from typing import Callable, Optional, Union
import httpx # type: ignore
@ -27,6 +28,16 @@ class VertexAIError(Exception):
self.message
) # Call the base class constructor with the parameters it needs
class PartnerModelPrefixes(Enum, str):
META_PREFIX = "meta/"
DEEPSEEK_PREFIX = "deepseek-ai"
MISTRAL_PREFIX = "mistral"
CODERESTAL_PREFIX = "codestral"
JAMBA_PREFIX = "jamba"
CLAUDE_PREFIX = "claude"
QWEN_PREFIX = "qwen"
GPT_OSS_PREFIX = "openai/gpt-oss-"
class VertexAIPartnerModels(VertexBase):
def __init__(self) -> None:
@ -42,14 +53,14 @@ class VertexAIPartnerModels(VertexBase):
bool: True if the model string is a Vertex AI Partner Model, False otherwise
"""
if (
model.startswith("meta/")
or model.startswith("deepseek-ai")
or model.startswith("mistral")
or model.startswith("codestral")
or model.startswith("jamba")
or model.startswith("claude")
or model.startswith("qwen")
or model.startswith("openai")
model.startswith(PartnerModelPrefixes.META_PREFIX)
or model.startswith(PartnerModelPrefixes.DEEPSEEK_PREFIX)
or model.startswith(PartnerModelPrefixes.MISTRAL_PREFIX)
or model.startswith(PartnerModelPrefixes.CODERESTAL_PREFIX)
or model.startswith(PartnerModelPrefixes.JAMBA_PREFIX)
or model.startswith(PartnerModelPrefixes.CLAUDE_PREFIX)
or model.startswith(PartnerModelPrefixes.QWEN_PREFIX)
or model.startswith(PartnerModelPrefixes.GPT_OSS_PREFIX)
):
return True
return False
@ -58,9 +69,9 @@ class VertexAIPartnerModels(VertexBase):
def should_use_openai_handler(model: str):
OPENAI_LIKE_VERTEX_PROVIDERS = [
"llama",
"deepseek-ai",
"qwen",
"openai",
PartnerModelPrefixes.DEEPSEEK_PREFIX,
PartnerModelPrefixes.QWEN_PREFIX,
PartnerModelPrefixes.GPT_OSS_PREFIX,
]
if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS):
return True