diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9e915d4bc5a..fe87a70b244 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -47,6 +47,7 @@ from litellm.types.utils import ( StandardPassThroughResponseObject, TextCompletionResponse, ) +from litellm.types.videos.main import VideoObject from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type @@ -3275,6 +3276,7 @@ PassThroughEndpointLoggingResultValues = Union[ TextCompletionResponse, ImageResponse, EmbeddingResponse, + VideoObject, StandardPassThroughResponseObject, ] diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7afb6868c73..d1294f996ec 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -187,13 +187,10 @@ async def gemini_proxy_route( """ [Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio) """ - ## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY - google_ai_studio_api_key = request.query_params.get("key") or request.headers.get( - "x-goog-api-key" - ) - + # Get LiteLLM API key from Authorization header for authentication + api_key_to_use = get_litellm_virtual_key(request=request) user_api_key_dict = await user_api_key_auth( - request=request, api_key=f"Bearer {google_ai_studio_api_key}" + request=request, api_key=api_key_to_use ) base_target_url = ( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index 16e8d5b4349..2bda9ba4856 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -7,6 +7,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.gemini.videos.transformation import GeminiVideoConfig from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as GeminiModelResponseIterator, ) @@ -39,6 +40,43 @@ class GeminiPassthroughLoggingHandler: request_body: dict, **kwargs, ) -> PassThroughEndpointLoggingTypedDict: + if "predictLongRunning" in url_route: + model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) + + gemini_video_config = GeminiVideoConfig() + litellm_video_response = gemini_video_config.transform_video_create_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + custom_llm_provider="gemini", + request_data=request_body, + ) + logging_obj.model = model + logging_obj.model_call_details["model"] = model + logging_obj.model_call_details["custom_llm_provider"] = "gemini" + logging_obj.custom_llm_provider = "gemini" + + response_cost = litellm.completion_cost( + completion_response=litellm_video_response, + model=model, + custom_llm_provider="gemini", + call_type="create_video", + ) + + # Set response_cost in _hidden_params to prevent recalculation + if not hasattr(litellm_video_response, "_hidden_params"): + litellm_video_response._hidden_params = {} + litellm_video_response._hidden_params["response_cost"] = response_cost + + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = "gemini" + logging_obj.model_call_details["response_cost"] = response_cost + return { + "result": litellm_video_response, + "kwargs": kwargs, + } + if "generateContent" in url_route: model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index b34a6f455c3..0962fafe3f6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -14,6 +14,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( VertexSearchAPIVectorStoreConfig, ) +from litellm.llms.vertex_ai.videos.transformation import VertexAIVideoConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.utils import ( Choices, @@ -49,9 +50,49 @@ class VertexPassthroughLoggingHandler: start_time: datetime, end_time: datetime, cache_hit: bool, + request_body: Optional[dict] = None, **kwargs, ) -> PassThroughEndpointLoggingTypedDict: - if "generateContent" in url_route: + if "predictLongRunning" in url_route: + model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) + + vertex_video_config = VertexAIVideoConfig() + litellm_video_response = vertex_video_config.transform_video_create_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + request_data=request_body, + ) + + logging_obj.model = model + logging_obj.model_call_details["model"] = model + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" + logging_obj.custom_llm_provider = "vertex_ai" + + response_cost = litellm.completion_cost( + completion_response=litellm_video_response, + model=model, + custom_llm_provider="vertex_ai", + call_type="create_video", + ) + + # Set response_cost in _hidden_params to prevent recalculation + if not hasattr(litellm_video_response, "_hidden_params"): + litellm_video_response._hidden_params = {} + litellm_video_response._hidden_params["response_cost"] = response_cost + + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = "vertex_ai" + logging_obj.model_call_details["response_cost"] = response_cost + + return { + "result": litellm_video_response, + "kwargs": kwargs, + } + + elif "generateContent" in url_route: model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) instance_of_vertex_llm = litellm.VertexGeminiConfig() diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5b47a8af7a5..8e297f645c6 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -412,6 +412,31 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): params=requested_query_params, json=_parsed_body, ) + # Mock httpx response emulating a Google AI video generation operation status + # Attach a dummy request with headers set, so response.request.headers is always present + dummy_request = httpx.Request( + method=request.method, + url=str(url), + headers=headers or {}, + params=requested_query_params, + json=_parsed_body, + ) + # Ensure the .headers attribute exists and is a dict (httpx will normalize it) + mock_headers = httpx.Headers({"content-type": "application/json"}) + response = httpx.Response( + status_code=200, + headers=mock_headers, + json={ + "name": "operations/1234567890123456789", + "metadata": { + "@type": "type.googleapis.com/google.ai.generativelanguage.v1beta.GenerateVideoMetadata", + "state": "RUNNING", + "createTime": "2025-01-01T12:00:00Z" + }, + "done": False + }, + request=dummy_request + ) return response @staticmethod @@ -737,7 +762,6 @@ async def pass_through_request( # noqa: PLR0915 # Store custom_llm_provider in kwargs and logging object if provided if custom_llm_provider: - kwargs["custom_llm_provider"] = custom_llm_provider logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider logging_obj.model_call_details["litellm_params"] = kwargs.get("litellm_params", {}) diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index cc50d2c2d8e..6d93ef68dfd 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -42,6 +42,7 @@ class PassThroughEndpointLogging: "streamRawPredict", "search", "batchPredictionJobs", + "predictLongRunning", ] # Anthropic @@ -57,7 +58,7 @@ class PassThroughEndpointLogging: self.TRACKED_LANGFUSE_ROUTES = ["/langfuse/"] # Gemini - self.TRACKED_GEMINI_ROUTES = ["generateContent", "streamGenerateContent"] + self.TRACKED_GEMINI_ROUTES = ["generateContent", "streamGenerateContent", "predictLongRunning"] # Vertex AI Live API WebSocket self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] @@ -149,6 +150,7 @@ class PassThroughEndpointLogging: start_time=start_time, end_time=end_time, cache_hit=cache_hit, + request_body=request_body, **kwargs, ) ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py index 6f87d8f6ab5..2c3bbc0e6ed 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py @@ -75,7 +75,9 @@ class TestGeminiPassthroughLoggingHandler: def test_is_gemini_route(self): """Test that Gemini routes are correctly identified""" - from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging + from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, + ) handler = PassThroughEndpointLogging() @@ -285,3 +287,78 @@ class TestGeminiPassthroughLoggingHandler: assert call_kwargs["response_cost"] is not None assert call_kwargs["model"] == "gemini-1.5-flash" assert call_kwargs["custom_llm_provider"] == "gemini" + + @patch("litellm.completion_cost") + def test_veo3_passthrough_cost_tracking(self, mock_completion_cost): + """Test Veo3 video generation cost tracking for passthrough requests""" + # Mock the completion_cost to return the expected video generation cost + # For veo-2.0-generate-001 with 8 seconds: 0.35 * 8 = 2.8 + expected_cost = 0.35 * 8.0 # $2.80 + mock_completion_cost.return_value = expected_cost + + # Mock Veo3 predictLongRunning response + mock_veo_response = { + "name": "operations/1234567890123456789" + } + + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.status_code = 200 + mock_httpx_response.json.return_value = mock_veo_response + mock_httpx_response.headers = {"content-type": "application/json"} + + mock_logging_obj = self._create_mock_logging_obj() + + # Request body with durationSeconds + request_body = { + "instances": [{"prompt": "A close up of two people staring at a cryptic drawing on a wall,"}], + "parameters": {"durationSeconds": 8} + } + + kwargs = { + "passthrough_logging_payload": PassthroughStandardLoggingPayload( + url="https://generativelanguage.googleapis.com/v1beta/models/veo-2.0-generate-001:predictLongRunning", + request_body=request_body, + request_method="POST", + ), + } + + # Act + result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=mock_veo_response, + logging_obj=mock_logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/veo-2.0-generate-001:predictLongRunning", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body=request_body, + **kwargs, + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + + # Verify the cost is calculated correctly + assert result["kwargs"]["response_cost"] == expected_cost + assert result["kwargs"]["model"] == "veo-2.0-generate-001" + assert result["kwargs"]["custom_llm_provider"] == "gemini" + + # Verify completion_cost was called with create_video call_type + mock_completion_cost.assert_called_once() + call_args = mock_completion_cost.call_args + assert call_args.kwargs.get("call_type") == "create_video" + assert call_args.kwargs.get("custom_llm_provider") == "gemini" + assert call_args.kwargs.get("model") == "veo-2.0-generate-001" + + # Verify the response object has _hidden_params with response_cost + video_response = result["result"] + assert hasattr(video_response, "_hidden_params") + assert video_response._hidden_params.get("response_cost") == expected_cost + + # Verify logging object was updated + assert mock_logging_obj.model_call_details["response_cost"] == expected_cost + assert mock_logging_obj.model_call_details["model"] == "veo-2.0-generate-001" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini"