Merge pull request #17296 from BerriAI/litellm_video_passthrough_cost_tracking

(Feat)Add passthrough cost tracking for veo

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,
 ]
 

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)
 

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()

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -737,7 +737,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", {})
 

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,
                 )
             )

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"