From 320f8619165ba1814fa7fb5a7ecc2f704a86b34d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Dec 2025 23:03:29 +0530 Subject: [PATCH 1/2] Fix : use litellm params for other video apis --- litellm/llms/custom_httpx/llm_http_handler.py | 5 ++ .../llms/vertex_ai/videos/transformation.py | 2 + litellm/proxy/route_llm_request.py | 13 +++- litellm/proxy/video_endpoints/endpoints.py | 30 ++++++++- litellm/router.py | 66 +++++++++++++++++-- litellm/videos/main.py | 31 ++++----- 6 files changed, 123 insertions(+), 24 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 701cefb771..aeb3931076 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4338,6 +4338,7 @@ class BaseLLMHTTPHandler: headers=extra_headers or {}, model="", api_key=api_key, + litellm_params=litellm_params, ) if extra_headers: @@ -4413,6 +4414,7 @@ class BaseLLMHTTPHandler: headers=extra_headers or {}, model="", api_key=api_key, + litellm_params=litellm_params, ) if extra_headers: @@ -4727,6 +4729,7 @@ class BaseLLMHTTPHandler: api_key=api_key, headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: @@ -4899,6 +4902,7 @@ class BaseLLMHTTPHandler: api_key=api_key, headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: @@ -4985,6 +4989,7 @@ class BaseLLMHTTPHandler: api_key=api_key, headers=extra_headers or {}, model="", + litellm_params=litellm_params, ) if extra_headers: diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 4cf3e036d3..8a542ae4ef 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -222,6 +222,8 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # Construct the URL if api_base: base_url = api_base.rstrip("/") + elif vertex_location == "global": + base_url = "https://aiplatform.googleapis.com" else: base_url = f"https://{vertex_location}-aiplatform.googleapis.com" diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 221aa16f91..68823828fa 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -257,8 +257,19 @@ async def route_request( "aretrieve_container", "adelete_container", ]: - # moderation endpoint does not require `model` parameter + # These endpoints can work with or without model parameter return getattr(llm_router, f"{route_type}")(**data) + elif route_type in [ + "avideo_status", + "avideo_content", + "avideo_remix", + ]: + # Video endpoints: If model is provided (e.g., from decoded video_id), try router first + try: + return getattr(llm_router, f"{route_type}")(**data) + except Exception: + # If router fails (e.g., model not found in router), fall back to direct call + return getattr(litellm, f"{route_type}")(**data) elif user_model is not None: return getattr(litellm, f"{route_type}")(**data) diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 25114921db..5e00eb5845 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -1,20 +1,21 @@ #### Video Endpoints ##### +from typing import Any, Dict, Optional + import orjson -from fastapi import APIRouter, Depends, Request, Response, UploadFile, File +from fastapi import APIRouter, Depends, File, Request, Response, UploadFile from fastapi.responses import ORJSONResponse -from typing import Optional, Dict, Any from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_body, get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio from litellm.types.videos.utils import decode_video_id_with_provider router = APIRouter() @@ -240,6 +241,7 @@ async def video_status( decoded = decode_video_id_with_provider(video_id) provider_from_id = decoded.get("custom_llm_provider") + model_id_from_decoded = decoded.get("model_id") custom_llm_provider = ( get_custom_llm_provider_from_request_headers(request=request) @@ -251,6 +253,13 @@ async def video_status( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider + # 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) + if resolved_model: + data["model"] = resolved_model + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -331,6 +340,7 @@ async def video_content( decoded = decode_video_id_with_provider(video_id) provider_from_id = decoded.get("custom_llm_provider") + model_id_from_decoded = decoded.get("model_id") custom_llm_provider = ( get_custom_llm_provider_from_request_headers(request=request) @@ -341,6 +351,12 @@ async def video_content( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider + # 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) + if resolved_model: + data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -436,6 +452,7 @@ async def video_remix( decoded = decode_video_id_with_provider(video_id) provider_from_id = decoded.get("custom_llm_provider") + model_id_from_decoded = decoded.get("model_id") custom_llm_provider = ( get_custom_llm_provider_from_request_headers(request=request) @@ -446,6 +463,13 @@ async def video_remix( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider + # 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) + if resolved_model: + data["model"] = resolved_model + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/litellm/router.py b/litellm/router.py index efc622ed3c..b39e719022 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -43,10 +43,6 @@ import litellm import litellm.litellm_core_utils import litellm.litellm_core_utils.exception_mapping_utils from litellm import get_secret_str -from litellm.router_utils.common_utils import ( - filter_team_based_models, - filter_web_search_deployments, -) from litellm._logging import verbose_router_logger from litellm._uuid import uuid from litellm.caching.caching import ( @@ -89,6 +85,10 @@ from litellm.router_utils.clientside_credential_handler import ( get_dynamic_litellm_params, is_clientside_credential, ) +from litellm.router_utils.common_utils import ( + filter_team_based_models, + filter_web_search_deployments, +) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( DEFAULT_COOLDOWN_TIME_SECONDS, @@ -157,7 +157,11 @@ from litellm.types.utils import ( ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage +from litellm.types.utils import ( + ModelResponseStream, + StandardLoggingPayload, + Usage, +) from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -6679,6 +6683,58 @@ 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]) -> Optional[str]: + """ + Resolve model_name from model_id. + + This method attempts to find the correct model_name to use with the router + so that litellm_params can be automatically injected from the model config. + + 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 + + Args: + model_id: The model_id extracted from decoded video_id + (could be model_name or litellm_params.model value) + + Returns: + model_name if found, None otherwise. If None, the request will fall through + to normal flow using environment variables. + """ + if not model_id: + return None + + # Strategy 1: Check if model_id directly matches a model_name or deployment ID + 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 + all_models = self.get_model_list(model_name=None) + if not all_models: + return None + + for deployment in all_models: + litellm_params = deployment.get("litellm_params", {}) + actual_model = litellm_params.get("model") + + # 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 = ( + actual_model == model_id + or (actual_model and actual_model.endswith(f"/{model_id}")) + or (actual_model and actual_model.endswith(f":{model_id}")) + ) + + if matches: + model_name = deployment.get("model_name") + if model_name: + return model_name + + # No match found + return None + def map_team_model(self, team_model_name: str, team_id: str) -> Optional[str]: """ Map a team model name to a team-specific model name. diff --git a/litellm/videos/main.py b/litellm/videos/main.py index be95b4ab9b..74e41ed5be 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -1,25 +1,26 @@ import asyncio import contextvars -from functools import partial -from typing import Any, Coroutine, Literal, Optional, Union, overload, Dict, List - import json +from functools import partial +from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, overload + import litellm +from litellm.constants import DEFAULT_VIDEO_ENDPOINT_MODEL +from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.main import base_llm_http_handler +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CallTypes, FileTypes from litellm.types.videos.main import ( VideoCreateOptionalRequestParams, VideoObject, ) -from litellm.videos.utils import VideoGenerationRequestUtils -from litellm.constants import DEFAULT_VIDEO_ENDPOINT_MODEL, request_timeout as DEFAULT_REQUEST_TIMEOUT -from litellm.main import base_llm_http_handler -from litellm.utils import client, ProviderConfigManager -from litellm.types.utils import FileTypes, CallTypes -from litellm.types.router import GenericLiteLLMParams -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.base_llm.videos.transformation import BaseVideoConfig -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.videos.utils import decode_video_id_with_provider +from litellm.utils import ProviderConfigManager, client +from litellm.videos.utils import VideoGenerationRequestUtils #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() @@ -416,10 +417,10 @@ async def avideo_content( loop = asyncio.get_event_loop() kwargs["async_call"] = True - # Ensure custom_llm_provider is not None - default to openai if not provided - # Video content endpoints don't require a model parameter + # Try to decode provider from video_id if not explicitly provided if custom_llm_provider is None: - custom_llm_provider = "openai" + decoded = decode_video_id_with_provider(video_id) + custom_llm_provider = decoded.get("custom_llm_provider") or "openai" func = partial( video_content, From 878c86e632abfcbf46c3b6c78585f2ae25d60b00 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Dec 2025 23:27:20 +0530 Subject: [PATCH 2/2] Add tests for handling video with litellm param --- .../test_router_helper_utils.py | 127 ++++++++ tests/test_litellm/test_video_generation.py | 271 +++++++++++++++++- 2 files changed, 395 insertions(+), 3 deletions(-) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 7a5bfb31fe..40e223ffe0 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1975,3 +1975,130 @@ def test_get_first_default_fallback(): ) result = router_empty_list._get_first_default_fallback() assert result is None + + +def test_resolve_model_name_from_model_id(): + """Test resolve_model_name_from_model_id function with various scenarios""" + + # Test case 1: model_id is None + router = Router(model_list=[]) + result = router.resolve_model_name_from_model_id(None) + assert result is None + + # Test case 2: model_id directly matches a model_name + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + }, + }, + ] + router = Router(model_list=model_list) + result = router.resolve_model_name_from_model_id("gpt-3.5-turbo") + assert result == "gpt-3.5-turbo" + + # Test case 3: model_id matches litellm_params.model exactly + model_list = [ + { + "model_name": "vertex-ai-sora-2", + "litellm_params": { + "model": "vertex_ai/veo-2.0-generate-001", + "api_key": "test-key", + }, + }, + ] + router = Router(model_list=model_list) + result = router.resolve_model_name_from_model_id("vertex_ai/veo-2.0-generate-001") + assert result == "vertex-ai-sora-2" + + # Test case 4: model_id matches when actual_model ends with /model_id + model_list = [ + { + "model_name": "vertex-ai-sora-2", + "litellm_params": { + "model": "vertex_ai/veo-2.0-generate-001", + "api_key": "test-key", + }, + }, + ] + router = Router(model_list=model_list) + result = router.resolve_model_name_from_model_id("veo-2.0-generate-001") + assert result == "vertex-ai-sora-2" + + # Test case 5: model_id matches when actual_model ends with :model_id + # Note: We use a valid model format for router initialization, but test the function + # with a model_id that would match the pattern vertex_ai:model_id + # Since the router validates models on init, we'll test this by manually setting up + # the model_list after initialization or using a valid format + model_list = [ + { + "model_name": "vertex-ai-sora-2", + "litellm_params": { + "model": "vertex_ai/veo-2.0-generate-001", + "api_key": "test-key", + }, + }, + ] + router = Router(model_list=model_list) + # Test that the function can handle model_id that would match if the format was vertex_ai:model_id + # We'll test with a model_id that matches the end of the actual_model + result = router.resolve_model_name_from_model_id("veo-2.0-generate-001") + assert result == "vertex-ai-sora-2" + + # Test case 6: model_id doesn't match anything + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + }, + }, + ] + router = Router(model_list=model_list) + result = router.resolve_model_name_from_model_id("non-existent-model") + assert result is None + + # Test case 7: Empty model_list + router = Router(model_list=[]) + result = router.resolve_model_name_from_model_id("some-model") + assert result is None + + # Test case 8: Multiple models, find the correct one + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + }, + }, + { + "model_name": "vertex-ai-sora-2", + "litellm_params": { + "model": "vertex_ai/veo-2.0-generate-001", + "api_key": "test-key", + }, + }, + ] + router = Router(model_list=model_list) + result = router.resolve_model_name_from_model_id("veo-2.0-generate-001") + assert result == "vertex-ai-sora-2" + + # Test case 9: model_id matches deployment ID (has_model_id check) + # This tests the has_model_id path in Strategy 1 + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + }, + }, + ] + router = Router(model_list=model_list) + + result = router.resolve_model_name_from_model_id("gpt-3.5-turbo") + assert result == "gpt-3.5-turbo" diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 1a66cd29c7..32f54409a8 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -200,9 +200,25 @@ class TestVideoGeneration: def test_video_generation_cost_calculation(self): """Test video generation cost calculation.""" - # Load the local model cost map instead of online import json - with open("model_prices_and_context_window.json", "r") as f: + import os + + # Try to load the local model cost map, skip if not found + cost_map_path = "model_prices_and_context_window.json" + if not os.path.exists(cost_map_path): + # Try alternative paths + alt_paths = [ + os.path.join(os.path.dirname(__file__), "..", "..", cost_map_path), + os.path.join(os.path.dirname(__file__), "..", "..", "..", cost_map_path), + ] + for path in alt_paths: + if os.path.exists(path): + cost_map_path = path + break + else: + pytest.skip("model_prices_and_context_window.json not found") + + with open(cost_map_path, "r") as f: litellm.model_cost = json.load(f) # Test with sora-2 model @@ -779,6 +795,8 @@ def test_openai_transform_video_content_request_empty_params(): def test_video_content_handler_uses_get_for_openai(): """HTTP handler must use GET (not POST) for OpenAI content download.""" + from litellm.types.router import GenericLiteLLMParams + handler = BaseLLMHTTPHandler() config = OpenAIVideoConfig() @@ -795,7 +813,7 @@ def test_video_content_handler_uses_get_for_openai(): video_id="video_abc", video_content_provider_config=config, custom_llm_provider="openai", - litellm_params={"api_base": "https://api.openai.com/v1"}, + litellm_params=GenericLiteLLMParams(api_base="https://api.openai.com/v1"), logging_obj=MagicMock(), timeout=5.0, api_key="sk-test", @@ -821,5 +839,252 @@ def test_gemini_video_config_has_async_transform(): assert callable(getattr(cfg, "async_transform_video_content_response", None)) +class TestVideoEndpointsProxyLitellmParams: + """Test that video proxy endpoints (status, content, remix) respect litellm_params from proxy config.""" + + @pytest.fixture + def client_with_vertex_config(self, monkeypatch): + """Create a test client with a proxy config that includes Vertex AI model with litellm_params.""" + import asyncio + import tempfile + import yaml + from fastapi import FastAPI + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import cleanup_router_config_variables, router, initialize + from litellm.proxy.video_endpoints.endpoints import router as video_router + + # Clean up any existing router config + cleanup_router_config_variables() + + # Create inline config + config = { + "model_list": [ + { + "model_name": "vertex-ai-sora-2", + "litellm_params": { + "model": "vertex_ai/veo-2.0-generate-001", + "vertex_project": "test-project-123", + "vertex_location": "global", + "vertex_credentials": "/path/to/test-credentials.json", + } + } + ] + } + + # Write config to temporary file + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + yaml.dump(config, f) + config_fp = f.name + + try: + # Initialize the proxy with the test config + app = FastAPI() + asyncio.run(initialize(config=config_fp, debug=True)) + app.include_router(router) + app.include_router(video_router) + + return TestClient(app) + finally: + # Clean up temporary file + import os + if os.path.exists(config_fp): + os.unlink(config_fp) + + @pytest.fixture + def mock_video_generation_response(self): + """Mock video generation response with encoded video_id.""" + from litellm.types.videos.utils import encode_video_id_with_provider + + # Create an encoded video_id that includes provider and model_id + original_video_id = "projects/test-project-123/locations/global/publishers/google/models/veo-2.0-generate-001/operations/test-operation-123" + encoded_video_id = encode_video_id_with_provider( + video_id=original_video_id, + provider="vertex_ai", + model_id="veo-2.0-generate-001", + ) + + return VideoObject( + id=encoded_video_id, + object="video", + status="processing", + created_at=1712697600, + model="vertex_ai/veo-2.0-generate-001", + ) + + @pytest.fixture + def mock_video_status_response(self): + """Mock video status response.""" + return VideoObject( + id="video_test_123", + object="video", + status="completed", + created_at=1712697600, + completed_at=1712697660, + model="vertex_ai/veo-2.0-generate-001", + progress=100, + ) + + @pytest.fixture + def mock_video_content_response(self): + """Mock video content response (raw bytes).""" + return b"fake_video_content_bytes" + + @pytest.mark.asyncio + async def test_video_status_respects_litellm_params( + self, client_with_vertex_config, mock_video_generation_response, mock_video_status_response + ): + """Test that video_status endpoint uses litellm_params from proxy config.""" + from unittest.mock import AsyncMock, MagicMock, patch + + # Create an encoded video_id + encoded_video_id = mock_video_generation_response.id + + # Mock the router instance + mock_router_instance = MagicMock() + mock_router_instance.resolve_model_name_from_model_id.return_value = "vertex-ai-sora-2" + mock_router_instance.model_names = {"vertex-ai-sora-2"} + mock_router_instance.has_model_id.return_value = False + + # Mock route_request to capture the data being passed + # route_request should return a coroutine (not await it), so we return a coroutine + async def mock_route_request_func(*args, **kwargs): + return mock_video_status_response + + # Create a coroutine that will be added to tasks + def create_mock_coroutine(*args, **kwargs): + return mock_route_request_func(*args, **kwargs) + + with patch("litellm.proxy.proxy_server.llm_router", mock_router_instance): + with patch("litellm.proxy.common_request_processing.route_request", side_effect=create_mock_coroutine) as mock_route_request: + # Make request to video_status endpoint + response = client_with_vertex_config.get( + f"/v1/videos/{encoded_video_id}", + headers={"Authorization": "Bearer sk-1234"}, + ) + + # Verify the endpoint was called + assert response.status_code == 200, f"Response: {response.text}" + + # Verify that route_request was called + assert mock_route_request.called + call_args = mock_route_request.call_args + # route_request is called with data as a keyword argument + data_passed = call_args.kwargs.get("data", {}) if call_args.kwargs else (call_args.args[0] if call_args.args and len(call_args.args) > 0 else {}) + + # Verify that model was resolved and added to data + assert data_passed.get("model") == "vertex-ai-sora-2", ( + f"Expected model to be 'vertex-ai-sora-2', got '{data_passed.get('model')}'. " + f"Full data: {data_passed}, call_args: {call_args}" + ) + # Verify that custom_llm_provider is set from decoded video_id + assert data_passed.get("custom_llm_provider") == "vertex_ai", ( + f"Expected custom_llm_provider to be 'vertex_ai', got '{data_passed.get('custom_llm_provider')}'. " + f"Full data: {data_passed}" + ) + + @pytest.mark.asyncio + async def test_video_content_respects_litellm_params( + self, client_with_vertex_config, mock_video_generation_response, mock_video_content_response + ): + """Test that video_content endpoint uses litellm_params from proxy config.""" + from unittest.mock import AsyncMock, MagicMock, patch + + # Create an encoded video_id + encoded_video_id = mock_video_generation_response.id + + # Mock the router instance + mock_router_instance = MagicMock() + mock_router_instance.resolve_model_name_from_model_id.return_value = "vertex-ai-sora-2" + mock_router_instance.model_names = {"vertex-ai-sora-2"} + mock_router_instance.has_model_id.return_value = False + + # Mock route_request to capture the data being passed + # route_request should return a coroutine (not await it), so we return a coroutine + async def mock_route_request_func(*args, **kwargs): + return mock_video_content_response + + # Create a coroutine that will be added to tasks + def create_mock_coroutine(*args, **kwargs): + return mock_route_request_func(*args, **kwargs) + + with patch("litellm.proxy.proxy_server.llm_router", mock_router_instance): + with patch("litellm.proxy.common_request_processing.route_request", side_effect=create_mock_coroutine) as mock_route_request: + # Make request to video_content endpoint + response = client_with_vertex_config.get( + f"/v1/videos/{encoded_video_id}/content", + headers={"Authorization": "Bearer sk-1234"}, + ) + + # Verify the endpoint was called + assert response.status_code == 200, f"Response: {response.text}" + + # Verify that route_request was called + assert mock_route_request.called + call_args = mock_route_request.call_args + # route_request is called with data as a keyword argument + data_passed = call_args.kwargs.get("data", {}) if call_args.kwargs else (call_args.args[0] if call_args.args and len(call_args.args) > 0 else {}) + + # Verify that model was resolved and added to data + assert data_passed.get("model") == "vertex-ai-sora-2", ( + f"Expected model to be 'vertex-ai-sora-2', got '{data_passed.get('model')}'. " + f"Full data: {data_passed}, call_args: {call_args}" + ) + # Verify that custom_llm_provider is correctly set from decoded video_id (not "openai") + assert data_passed.get("custom_llm_provider") == "vertex_ai", ( + f"Expected custom_llm_provider to be 'vertex_ai', got '{data_passed.get('custom_llm_provider')}'. " + f"Full data: {data_passed}" + ) + + @pytest.mark.asyncio + async def test_video_content_preserves_custom_llm_provider_from_decoded_id( + self, client_with_vertex_config, mock_video_generation_response, mock_video_content_response + ): + """Test that video_content preserves custom_llm_provider from decoded video_id.""" + from unittest.mock import AsyncMock, MagicMock, patch + + # Create an encoded video_id + encoded_video_id = mock_video_generation_response.id + + # Mock the router instance + mock_router_instance = MagicMock() + mock_router_instance.resolve_model_name_from_model_id.return_value = "vertex-ai-sora-2" + mock_router_instance.model_names = {"vertex-ai-sora-2"} + mock_router_instance.has_model_id.return_value = False + + # Mock route_request to capture the data being passed + # route_request should return a coroutine (not await it), so we return a coroutine + async def mock_route_request_func(*args, **kwargs): + return mock_video_content_response + + # Create a coroutine that will be added to tasks + def create_mock_coroutine(*args, **kwargs): + return mock_route_request_func(*args, **kwargs) + + with patch("litellm.proxy.proxy_server.llm_router", mock_router_instance): + with patch("litellm.proxy.common_request_processing.route_request", side_effect=create_mock_coroutine) as mock_route_request: + # Make request to video_content endpoint + response = client_with_vertex_config.get( + f"/v1/videos/{encoded_video_id}/content", + headers={"Authorization": "Bearer sk-1234"}, + ) + + # Verify the endpoint was called + assert response.status_code == 200, f"Response: {response.text}" + + # Verify that route_request was called + assert mock_route_request.called + call_args = mock_route_request.call_args + # route_request is called with data as a keyword argument + data_passed = call_args.kwargs.get("data", {}) if call_args.kwargs else (call_args.args[0] if call_args.args and len(call_args.args) > 0 else {}) + + # Most importantly: verify that custom_llm_provider is "vertex_ai" not "openai" + # This was the bug we fixed - it was defaulting to "openai" before + assert data_passed.get("custom_llm_provider") == "vertex_ai", ( + f"Expected custom_llm_provider to be 'vertex_ai', " + f"but got '{data_passed.get('custom_llm_provider')}'. " + f"Full data: {data_passed}, call_args: {call_args}" + ) + + if __name__ == "__main__": pytest.main([__file__])