diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index ee30c1749a..19bd75596f 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -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