diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 5e00eb58455..a3c4af9ae5d 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -256,7 +256,9 @@ async def video_status( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) + resolved_model = llm_router.resolve_model_name_from_model_id( + model_id_from_decoded, custom_llm_provider=provider_from_id + ) if resolved_model: data["model"] = resolved_model @@ -354,7 +356,9 @@ async def video_content( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) + resolved_model = llm_router.resolve_model_name_from_model_id( + model_id_from_decoded, custom_llm_provider=provider_from_id + ) if resolved_model: data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing @@ -466,7 +470,9 @@ async def video_remix( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) + resolved_model = llm_router.resolve_model_name_from_model_id( + model_id_from_decoded, custom_llm_provider=provider_from_id + ) if resolved_model: data["model"] = resolved_model diff --git a/litellm/router.py b/litellm/router.py index bd02e8e019c..f73d907c8c7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6972,7 +6972,7 @@ class Router: return candidate_id in self.model_id_to_deployment_index_map def resolve_model_name_from_model_id( - self, model_id: Optional[str] + self, model_id: Optional[str], custom_llm_provider: Optional[str] = None ) -> Optional[str]: """ Resolve model_name from model_id. @@ -6982,12 +6982,15 @@ class Router: Strategy: 1. First, check if model_id directly matches a model_name or deployment ID - 2. If not, search through router's model_list to find a match by litellm_params.model - 3. Return the model_name if found, None otherwise + 2. If custom_llm_provider is provided, check with provider prefix + 3. Search through router's model_list to find a match by litellm_params.model + 4. If custom_llm_provider is provided, try to find a wildcard pattern match + 5. Return the model_name if found, None otherwise Args: model_id: The model_id extracted from decoded video_id (could be model_name or litellm_params.model value) + custom_llm_provider: The provider name (e.g., "vertex_ai") for wildcard matching Returns: model_name if found, None otherwise. If None, the request will fall through @@ -7000,15 +7003,26 @@ class Router: if model_id in self.model_names or self.has_model_id(model_id): return model_id - # Strategy 2: Search through router's model_list to find by litellm_params.model + # Strategy 2: Check with provider prefix (e.g., "vertex_ai/veo-3.0-generate-preview") + if custom_llm_provider: + full_model_name = f"{custom_llm_provider}/{model_id}" + if full_model_name in self.model_names or self.has_model_id(full_model_name): + return full_model_name + + # Strategy 3: Search through router's model_list to find by litellm_params.model all_models = self.get_model_list(model_name=None) if not all_models: return None + # First pass: exact matches (non-wildcard) for deployment in all_models: litellm_params = deployment.get("litellm_params", {}) actual_model = litellm_params.get("model") + # Skip wildcard patterns in first pass + if actual_model and actual_model.endswith("/*"): + continue + # Match by exact match or by checking if actual_model ends with /model_id or :model_id # e.g., model_id="veo-2.0-generate-001" matches actual_model="vertex_ai/veo-2.0-generate-001" matches = ( @@ -7022,6 +7036,19 @@ class Router: if model_name: return model_name + # Strategy 4: Wildcard patterns using PatternMatchRouter + # For video status/content, we need to match model_id like "veo-3.0-generate-preview" + # to wildcard patterns like "vertex_ai/*" + if custom_llm_provider: + full_model_name = f"{custom_llm_provider}/{model_id}" + pattern_deployments = self.pattern_router.route(full_model_name) + if pattern_deployments: + # Return the first matching wildcard model_name + for pattern_deployment in pattern_deployments: + matched_model_name = pattern_deployment.get("model_name") + if matched_model_name: + return matched_model_name + # No match found return None diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 6279e96305f..7201b961588 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2054,3 +2054,190 @@ async def test_aguardrail(): assert result["result"] == "success" assert result["selected_guardrail"]["id"] == "guardrail-1" + + +def test_resolve_model_name_from_model_id_wildcard_pattern(): + """ + Test that resolve_model_name_from_model_id correctly resolves model names + for wildcard patterns using PatternMatchRouter. + + This is critical for video status/content endpoints where model_id extracted + from video_id (e.g., "veo-3.0-generate-preview") needs to match wildcard + patterns like "vertex_ai/*" to inject credentials from the model config. + """ + # Set up router with wildcard pattern + router = litellm.Router( + model_list=[ + { + "model_name": "vertex_ai/*", + "litellm_params": { + "model": "vertex_ai/*", + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + }, + { + "model_name": "specific-model", + "litellm_params": { + "model": "vertex_ai/gemini-pro", + "vertex_project": "specific-project", + "vertex_location": "us-east1", + }, + }, + ], + ) + + # Test Case 1: Wildcard pattern matching with custom_llm_provider + # This simulates video_id like "vertex_ai:veo-3.0-generate-preview:..." + result = router.resolve_model_name_from_model_id( + model_id="veo-3.0-generate-preview", + custom_llm_provider="vertex_ai", + ) + assert result == "vertex_ai/*", f"Expected 'vertex_ai/*', got '{result}'" + + # Test Case 2: Different model name should also match wildcard + result = router.resolve_model_name_from_model_id( + model_id="gemini-2.0-flash", + custom_llm_provider="vertex_ai", + ) + assert result == "vertex_ai/*", f"Expected 'vertex_ai/*', got '{result}'" + + # Test Case 3: Without custom_llm_provider, should not match wildcard + result = router.resolve_model_name_from_model_id( + model_id="veo-3.0-generate-preview", + custom_llm_provider=None, + ) + assert result is None, f"Expected None without provider, got '{result}'" + + # Test Case 4: Exact model_name match should take precedence + result = router.resolve_model_name_from_model_id( + model_id="specific-model", + custom_llm_provider="vertex_ai", + ) + assert result == "specific-model", f"Expected 'specific-model', got '{result}'" + + +def test_resolve_model_name_from_model_id_exact_match(): + """ + Test that resolve_model_name_from_model_id correctly resolves exact model names. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "my-gpt-model", + "litellm_params": { + "model": "azure/gpt-4", + "api_key": "test-key", + }, + }, + { + "model_name": "veo-model", + "litellm_params": { + "model": "vertex_ai/veo-2.0-generate-001", + "vertex_project": "test-project", + }, + }, + ], + ) + + # Test Case 1: Direct model_name match + result = router.resolve_model_name_from_model_id(model_id="my-gpt-model") + assert result == "my-gpt-model", f"Expected 'my-gpt-model', got '{result}'" + + # Test Case 2: Match by litellm_params.model suffix + result = router.resolve_model_name_from_model_id(model_id="veo-2.0-generate-001") + assert result == "veo-model", f"Expected 'veo-model', got '{result}'" + + # Test Case 3: Non-existent model should return None + result = router.resolve_model_name_from_model_id(model_id="non-existent-model") + assert result is None, f"Expected None, got '{result}'" + + +def test_resolve_model_name_from_model_id_provider_prefix(): + """ + Test that resolve_model_name_from_model_id handles provider prefix correctly. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "vertex_ai/gemini-pro", + "litellm_params": { + "model": "vertex_ai/gemini-pro", + "vertex_project": "test-project", + }, + }, + ], + ) + + # Test Case 1: Full model name with provider prefix as model_name + result = router.resolve_model_name_from_model_id( + model_id="vertex_ai/gemini-pro", + custom_llm_provider=None, + ) + assert result == "vertex_ai/gemini-pro", f"Expected 'vertex_ai/gemini-pro', got '{result}'" + + # Test Case 2: Model ID with provider prefix constructed from custom_llm_provider + result = router.resolve_model_name_from_model_id( + model_id="gemini-pro", + custom_llm_provider="vertex_ai", + ) + assert result == "vertex_ai/gemini-pro", f"Expected 'vertex_ai/gemini-pro', got '{result}'" + + +def test_resolve_model_name_from_model_id_multiple_wildcards(): + """ + Test that resolve_model_name_from_model_id works with multiple wildcard patterns. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "vertex_ai/*", + "litellm_params": { + "model": "vertex_ai/*", + "vertex_project": "vertex-project", + }, + }, + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "openai-key", + }, + }, + { + "model_name": "anthropic/*", + "litellm_params": { + "model": "anthropic/*", + "api_key": "anthropic-key", + }, + }, + ], + ) + + # Test Case 1: Match vertex_ai wildcard + result = router.resolve_model_name_from_model_id( + model_id="veo-3.0-generate-preview", + custom_llm_provider="vertex_ai", + ) + assert result == "vertex_ai/*", f"Expected 'vertex_ai/*', got '{result}'" + + # Test Case 2: Match openai wildcard + result = router.resolve_model_name_from_model_id( + model_id="gpt-4o", + custom_llm_provider="openai", + ) + assert result == "openai/*", f"Expected 'openai/*', got '{result}'" + + # Test Case 3: Match anthropic wildcard + result = router.resolve_model_name_from_model_id( + model_id="claude-3-opus", + custom_llm_provider="anthropic", + ) + assert result == "anthropic/*", f"Expected 'anthropic/*', got '{result}'" + + # Test Case 4: Non-matching provider should return None + result = router.resolve_model_name_from_model_id( + model_id="some-model", + custom_llm_provider="bedrock", + ) + assert result is None, f"Expected None for non-matching provider, got '{result}'"