fix VertexAIPartnerModels

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