Merge pull request #41322 from BerriAI/litellm_gcp_video_usage

fix(vertex_ai): bill Gemini Omni Interactions usage and Veo sampleCount on passthrough

litellm/cost_calculator.py

@@ -1466,12 +1466,15 @@ def completion_cost(
                             duration_seconds = usage_obj.get("duration_seconds", None)
                             _vr = usage_obj.get("video_resolution", None)
                             provider_reported_cost = usage_obj.get("provider_reported_cost_usd", None)
+                            _vc = usage_obj.get("video_count", None)
                         else:
                             duration_seconds = getattr(usage_obj, "duration_seconds", None)
                             _vr = getattr(usage_obj, "video_resolution", None)
                             provider_reported_cost = getattr(usage_obj, "provider_reported_cost_usd", None)
+                            _vc = getattr(usage_obj, "video_count", None)
                         if _vr is not None:
                             video_resolution = str(_vr).strip().lower()
+                        video_count = _vc if isinstance(_vc, int) and not isinstance(_vc, bool) and _vc > 1 else 1
 
                         if _video_model_info is None and provider_reported_cost is not None:
                             return float(provider_reported_cost)
@@ -1482,12 +1485,15 @@ def completion_cost(
                                 video_generation_cost,
                             )
 
-                            return video_generation_cost(
-                                model=model,
-                                duration_seconds=duration_seconds,
-                                custom_llm_provider=custom_llm_provider,
-                                model_info=_video_model_info,
-                                video_resolution=video_resolution,
+                            return (
+                                video_generation_cost(
+                                    model=model,
+                                    duration_seconds=duration_seconds,
+                                    custom_llm_provider=custom_llm_provider,
+                                    model_info=_video_model_info,
+                                    video_resolution=video_resolution,
+                                )
+                                * video_count
                             )
                     # Fallback to default video cost calculation if no duration available
                     return default_video_cost_calculator(

litellm/llms/gemini/videos/transformation.py

@@ -9,6 +9,7 @@ import litellm
 from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
 from litellm.images.utils import ImageEditRequestUtils
 from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
+from litellm.llms.vertex_ai.videos.transformation import veo_video_count_from_parameters
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.gemini import (
     GeminiLongRunningOperationResponse,
@@ -354,6 +355,9 @@ class GeminiVideoConfig(BaseVideoConfig):
             video_resolution: Final = _usage_video_resolution_from_parameters(parameters)
             if video_resolution is not None:
                 usage_data["video_resolution"] = video_resolution
+            video_count: Final = veo_video_count_from_parameters(parameters)
+            if video_count is not None:
+                usage_data["video_count"] = video_count
 
         video_obj.usage = usage_data
         return video_obj

litellm/llms/vertex_ai/videos/transformation.py

@@ -70,10 +70,17 @@ def _parse_veo_operation(raw_response: httpx.Response) -> _VeoOperation:
     return operation
 
 
+def veo_video_count_from_parameters(parameters: Mapping[str, object]) -> int | None:
+    sample_count: Final = parameters.get("sampleCount")
+    if isinstance(sample_count, bool) or not isinstance(sample_count, int) or sample_count < 1:
+        return None
+    return sample_count
+
+
 def _build_vertex_video_usage_from_request_data(
     request_data: dict[str, Any] | None,
 ) -> dict[str, float | str]:
-    """Build usage metadata (duration, resolution) for video cost calculation."""
+    """Build usage metadata (duration, resolution, video count) for video cost calculation."""
     usage_data: Final[dict[str, float | str]] = {}
     if not request_data:
         return usage_data
@@ -88,6 +95,9 @@ def _build_vertex_video_usage_from_request_data(
     res: Final = parameters.get("resolution")
     if res is not None and str(res).strip() != "":
         usage_data["video_resolution"] = str(res).strip().lower()
+    video_count: Final = veo_video_count_from_parameters(parameters)
+    if video_count is not None:
+        usage_data["video_count"] = video_count
     return usage_data
 
 

litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py

@@ -12,6 +12,9 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     ModelResponseIterator as GeminiModelResponseIterator,
 )
 from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
+    VertexPassthroughLoggingHandler,
+)
 from litellm.types.utils import (
     ModelResponse,
     TextCompletionResponse,
@@ -40,6 +43,17 @@ class GeminiPassthroughLoggingHandler:
         request_body: dict,
         **kwargs,
     ) -> PassThroughEndpointLoggingTypedDict:
+        if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
+            return VertexPassthroughLoggingHandler.interactions_passthrough_handler(
+                httpx_response=httpx_response,
+                request_body=request_body,
+                logging_obj=logging_obj,
+                kwargs=kwargs,
+                start_time=start_time,
+                end_time=end_time,
+                custom_llm_provider="gemini",
+                vertex_location=None,
+            )
         if "predictLongRunning" in url_route:
             model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route)
 

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -1,15 +1,20 @@
 import asyncio
 import re
+from collections.abc import Mapping
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, cast
 from urllib.parse import urlparse
 
 import httpx
+from pydantic import TypeAdapter
 
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
+    InteractionsUsageObjectTransformation,
+)
 from litellm.llms.vertex_ai.common_utils import (
     get_vertex_ai_lyria_generation_cost,
     get_vertex_location_from_url,
@@ -49,8 +54,73 @@ else:
 
 EndpointType = Any
 
+_VERTEX_INTERACTIONS_PATH: Final = re.compile(r"/projects/[^/]+/locations/[^/]+/interactions/?$")
+_INTERACTIONS_RESPONSE_BODY: Final = TypeAdapter(dict[str, object])
+
+
+def _interactions_model(
+    response_body: Mapping[str, object],
+    request_body: Mapping[str, object] | None,
+) -> str | None:
+    response_model: Final = response_body.get("model")
+    if isinstance(response_model, str) and response_model:
+        return response_model
+    request_model: Final = (request_body or {}).get("model")
+    if isinstance(request_model, str) and request_model:
+        return request_model
+    return None
+
 
 class VertexPassthroughLoggingHandler:
+    @staticmethod
+    def is_interactions_route(url_route: str) -> bool:
+        return urlparse(url_route).path.rstrip("/").endswith("/interactions")
+
+    @staticmethod
+    def is_vertex_interactions_route(url_route: str) -> bool:
+        return _VERTEX_INTERACTIONS_PATH.search(urlparse(url_route).path) is not None
+
+    @staticmethod
+    def interactions_passthrough_handler(
+        httpx_response: httpx.Response,
+        request_body: Mapping[str, object] | None,
+        logging_obj: LiteLLMLoggingObj,
+        kwargs: dict[str, object],
+        start_time: datetime,
+        end_time: datetime,
+        custom_llm_provider: Literal["vertex_ai", "gemini"],
+        vertex_location: str | None,
+    ) -> PassThroughEndpointLoggingTypedDict:
+        response_body: Final = _INTERACTIONS_RESPONSE_BODY.validate_python(httpx_response.json())
+        usage_object: Final = response_body.get("usage")
+        model: Final = _interactions_model(response_body, request_body)
+        if model is None or not InteractionsUsageObjectTransformation.is_interactions_usage_object(usage_object):
+            return {"result": None, "kwargs": kwargs}
+
+        litellm_model_response: Final = ModelResponse(
+            model=model,
+            usage=InteractionsUsageObjectTransformation.transform_interactions_usage_object(
+                cast(Mapping[str, Any], usage_object)
+            ),
+        )
+        logging_obj.custom_llm_provider = custom_llm_provider
+        logging_kwargs: Final = (
+            VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content(
+                litellm_model_response=litellm_model_response,
+                model=model,
+                kwargs=kwargs,
+                start_time=start_time,
+                end_time=end_time,
+                logging_obj=logging_obj,
+                custom_llm_provider=custom_llm_provider,
+                vertex_location=vertex_location,
+            )
+        )
+        return {
+            "result": litellm_model_response,
+            "kwargs": {**logging_kwargs, "custom_llm_provider": custom_llm_provider},
+        }
+
     @staticmethod
     def vertex_passthrough_handler(
         httpx_response: httpx.Response,
@@ -66,6 +136,17 @@ class VertexPassthroughLoggingHandler:
         vertex_location: Final = get_vertex_location_from_url(url_route)
         if vertex_location is not None:
             logging_obj.optional_params["vertex_location"] = vertex_location
+        if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
+            return VertexPassthroughLoggingHandler.interactions_passthrough_handler(
+                httpx_response=httpx_response,
+                request_body=request_body,
+                logging_obj=logging_obj,
+                kwargs=kwargs,
+                start_time=start_time,
+                end_time=end_time,
+                custom_llm_provider="vertex_ai",
+                vertex_location=vertex_location,
+            )
         if "predictLongRunning" in url_route:
             model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
 

litellm/proxy/pass_through_endpoints/success_handler.py

@@ -361,7 +361,9 @@ class PassThroughEndpointLogging:
     def is_vertex_route(self, url_route: str) -> bool:
         if any(f":{method}" in url_route for method in self.TRACKED_VERTEX_METHOD_ROUTES):
             return True
-        return any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES)
+        if any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES):
+            return True
+        return VertexPassthroughLoggingHandler.is_vertex_interactions_route(url_route)
 
     def is_anthropic_route(self, url_route: str):
         for route in self.TRACKED_ANTHROPIC_ROUTES:
@@ -434,8 +436,12 @@ class PassThroughEndpointLogging:
 
     def is_gemini_route(self, url_route: str, custom_llm_provider: str | None = None):
         """Check if the URL route is a Gemini API route."""
+        if custom_llm_provider != "gemini":
+            return False
+        if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
+            return True
         for route in self.TRACKED_GEMINI_ROUTES:
-            if route in url_route and custom_llm_provider == "gemini":
+            if route in url_route:
                 return True
         return False
 

tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py

@@ -430,6 +430,25 @@ class TestGeminiVideoConfig:
         assert result.usage["video_resolution"] == "1080p"
         assert result.usage["duration_seconds"] == 8.0
 
+    def test_transform_video_create_response_usage_includes_video_count(self):
+        """Regression for LIT-6896: sampleCount (number of generated videos) is copied into usage for billing."""
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.json.return_value = {"name": "operations/generate_1234567890"}
+        request_data = {
+            "instances": [{"prompt": "Test"}],
+            "parameters": {"durationSeconds": 8, "sampleCount": 3},
+        }
+        result = self.config.transform_video_create_response(
+            model="gemini/veo-3.1-fast-generate-preview",
+            raw_response=mock_response,
+            logging_obj=self.mock_logging_obj,
+            custom_llm_provider="gemini",
+            request_data=request_data,
+        )
+        assert result.usage is not None
+        assert result.usage["video_count"] == 3
+        assert result.usage["duration_seconds"] == 8.0
+
     def test_transform_video_create_response_cost_tracking_with_different_durations(
         self,
     ):

tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py

@@ -9,7 +9,95 @@ import litellm
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
     VertexPassthroughLoggingHandler,
 )
-from litellm.types.utils import PassthroughCallTypes
+from litellm.types.utils import ModelResponse, PassthroughCallTypes
+
+_OMNI_INTERACTIONS_USAGE: Final = {
+    "total_tokens": 4041,
+    "total_input_tokens": 12,
+    "input_tokens_by_modality": [{"modality": "text", "tokens": 12}],
+    "total_output_tokens": 4009,
+    "output_tokens_by_modality": [
+        {"modality": "text", "tokens": 9},
+        {"modality": "video", "tokens": 4000},
+    ],
+    "total_tool_use_tokens": 0,
+    "total_thought_tokens": 20,
+}
+
+
+def test_interactions_create_response_logs_modality_usage_and_cost() -> None:
+    """
+    Regression for LIT-6896: gemini-omni Interactions passthrough rows were logged
+    with zero tokens and zero spend. Input, text-output and video-output tokens
+    must land in usage, priced with the model's per-modality rates, and the
+    response id must stay the litellm_call_id so SpendLogs keep their request_id.
+    """
+    logging_obj = MagicMock()
+    logging_obj.model_call_details = {}
+    logging_obj.optional_params = {}
+    logging_obj.litellm_call_id = "call-6896"
+    response = httpx.Response(
+        status_code=200,
+        json={
+            "id": "interactions/abc",
+            "model": "gemini-omni-flash-preview",
+            "status": "completed",
+            "outputs": [{"type": "text", "text": "hi"}],
+            "usage": _OMNI_INTERACTIONS_USAGE,
+        },
+    )
+
+    result = VertexPassthroughLoggingHandler.vertex_passthrough_handler(
+        httpx_response=response,
+        logging_obj=logging_obj,
+        url_route="https://aiplatform.googleapis.com/v1beta1/projects/p/locations/global/interactions",
+        result=response.text,
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+        cache_hit=False,
+        request_body={"model": "gemini-omni-flash-preview", "input": [{"type": "text", "text": "say hi"}]},
+    )
+
+    model_response = result["result"]
+    assert isinstance(model_response, ModelResponse)
+    assert model_response.id == "call-6896"
+    usage = model_response.usage
+    assert usage.prompt_tokens == 12
+    assert usage.completion_tokens == 4009 + 20
+    assert usage.completion_tokens_details.text_tokens == 9
+    assert usage.completion_tokens_details.video_tokens == 4000
+
+    model_info = litellm.get_model_info(model="gemini-omni-flash-preview", custom_llm_provider="vertex_ai")
+    expected_cost = (
+        12 * model_info["input_cost_per_token"]
+        + (9 + 20) * model_info["output_cost_per_token"]
+        + 4000 * model_info["output_cost_per_video_token"]
+    )
+    assert result["kwargs"]["response_cost"] == pytest.approx(expected_cost)
+    assert result["kwargs"]["custom_llm_provider"] == "vertex_ai"
+    assert logging_obj.model_call_details["model"] == "gemini-omni-flash-preview"
+    assert logging_obj.model_call_details["custom_llm_provider"] == "vertex_ai"
+
+
+def test_interactions_response_without_usage_falls_back_to_generic_logging() -> None:
+    logging_obj = MagicMock()
+    logging_obj.model_call_details = {}
+    logging_obj.optional_params = {}
+    response = httpx.Response(status_code=200, json={"id": "interactions/abc", "status": "in_progress"})
+
+    result = VertexPassthroughLoggingHandler.vertex_passthrough_handler(
+        httpx_response=response,
+        logging_obj=logging_obj,
+        url_route="https://aiplatform.googleapis.com/v1beta1/projects/p/locations/global/interactions",
+        result=response.text,
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+        cache_hit=False,
+        request_body={"agent": "projects/p/locations/global/reasoningEngines/1"},
+    )
+
+    assert result["result"] is None
+    assert "response_cost" not in result["kwargs"]
 
 
 def test_lyria_predict_response_preserves_audio_response_and_logs_cost(

tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py

@@ -717,6 +717,33 @@ class TestVertexAIVideoConfig:
         assert video_obj.usage["duration_seconds"] == 8.0
         assert video_obj.usage["video_resolution"] == "1080p"
 
+    @pytest.mark.parametrize(
+        "sample_count,expected_video_count",
+        [(2, 2), (1, 1), (None, None), (0, None), ("2", None)],
+        ids=["two", "one", "unset", "zero", "string"],
+    )
+    def test_transform_video_create_response_usage_includes_video_count(self, sample_count, expected_video_count):
+        """Regression for LIT-6896: sampleCount is the number of generated videos and must reach usage for billing."""
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.json.return_value = {
+            "name": "projects/p/locations/us-central1/publishers/google/models/veo-3.1-fast-generate-001/operations/op-1"
+        }
+        parameters = {"durationSeconds": 4, "resolution": "720p"}
+        if sample_count is not None:
+            parameters["sampleCount"] = sample_count
+
+        video_obj = self.config.transform_video_create_response(
+            model="veo-3.1-fast-generate-001",
+            raw_response=mock_response,
+            logging_obj=self.mock_logging_obj,
+            custom_llm_provider="vertex_ai",
+            request_data={"instances": [{"prompt": "a red ball"}], "parameters": parameters},
+        )
+
+        assert video_obj.usage is not None
+        assert video_obj.usage["duration_seconds"] == 4.0
+        assert video_obj.usage.get("video_count") == expected_video_count
+
     def test_transform_video_remix_request_not_supported(self):
         """Test that video remix raises NotImplementedError."""
         with pytest.raises(NotImplementedError, match="Video remix is not supported"):

tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py

@@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import httpx
 import pytest
 
-
+import litellm
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler import (
     GeminiPassthroughLoggingHandler,
@@ -397,3 +397,52 @@ class TestGeminiPassthroughLoggingHandler:
         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"
+
+    def test_interactions_create_response_is_priced_as_gemini(self):
+        """Regression for LIT-6896: Gemini API Interactions passthrough must not log zero usage."""
+        usage = {
+            "total_tokens": 1030,
+            "total_input_tokens": 10,
+            "input_tokens_by_modality": [{"modality": "text", "tokens": 10}],
+            "total_output_tokens": 1020,
+            "output_tokens_by_modality": [
+                {"modality": "text", "tokens": 20},
+                {"modality": "video", "tokens": 1000},
+            ],
+            "total_tool_use_tokens": 0,
+            "total_thought_tokens": 0,
+        }
+        mock_httpx_response = MagicMock(spec=httpx.Response)
+        mock_httpx_response.json.return_value = {
+            "id": "interactions/abc",
+            "model": "gemini-omni-flash-preview",
+            "status": "completed",
+            "usage": usage,
+        }
+        mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
+        mock_logging_obj.model_call_details = {}
+        mock_logging_obj.litellm_call_id = "call-6896"
+
+        result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler(
+            httpx_response=mock_httpx_response,
+            response_body=mock_httpx_response.json.return_value,
+            logging_obj=mock_logging_obj,
+            url_route="https://generativelanguage.googleapis.com/v1beta/interactions",
+            result="",
+            start_time=self.start_time,
+            end_time=self.end_time,
+            cache_hit=False,
+            request_body={"model": "gemini-omni-flash-preview", "input": "make a clip"},
+        )
+
+        model_info = litellm.get_model_info(model="gemini-omni-flash-preview", custom_llm_provider="gemini")
+        expected_cost = (
+            10 * model_info["input_cost_per_token"]
+            + 20 * model_info["output_cost_per_token"]
+            + 1000 * model_info["output_cost_per_video_token"]
+        )
+        assert result["result"].id == "call-6896"
+        assert result["result"].usage.completion_tokens_details.video_tokens == 1000
+        assert result["kwargs"]["response_cost"] == pytest.approx(expected_cost)
+        assert result["kwargs"]["custom_llm_provider"] == "gemini"
+        assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini"

tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -497,6 +497,33 @@ def test_is_vertex_route_ignores_plain_predict_path_segment():
     )
 
 
+def test_interactions_create_routes_are_tracked_for_vertex_and_gemini():
+    """
+    Regression for LIT-6896: Interactions API (gemini-omni) passthrough responses
+    were never handed to the Vertex/Gemini logging handlers, so SpendLogs rows
+    landed with zero tokens and zero spend. Only the create URL is billable;
+    GET/DELETE on an interaction id and non-Google `/interactions` URLs stay generic.
+    """
+    handler = PassThroughEndpointLogging()
+    vertex_create = "https://aiplatform.googleapis.com/v1beta1/projects/p/locations/global/interactions"
+    gemini_create = "https://generativelanguage.googleapis.com/v1beta/interactions"
+
+    assert handler.is_vertex_route(vertex_create) is True
+    assert handler.is_vertex_route(f"{vertex_create}/abc123") is False
+    assert handler.is_vertex_route("https://upstream.example.com/api/interactions") is False
+    assert handler.is_vertex_route("https://upstream.example.com/locations/eu/interactions") is False
+    assert (
+        handler.is_vertex_route(
+            "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/interactions"
+        )
+        is True
+    )
+
+    assert handler.is_gemini_route(gemini_create, custom_llm_provider="gemini") is True
+    assert handler.is_gemini_route(f"{gemini_create}/abc123", custom_llm_provider="gemini") is False
+    assert handler.is_gemini_route(gemini_create, custom_llm_provider=None) is False
+
+
 @pytest.mark.asyncio
 async def test_custom_passthrough_predict_path_logs_via_generic_handler():
     """

tests/test_litellm/test_cost_calculator.py

@@ -4141,6 +4141,31 @@ def test_completion_cost_bills_interactions_video_output_at_video_rate():
     assert cost == pytest.approx(expected)
 
 
+@pytest.mark.parametrize("video_count", [2, 3])
+def test_completion_cost_multiplies_video_cost_by_generated_video_count(video_count: int) -> None:
+    """Regression for LIT-6896: a Veo request for N samples generates N videos and must be billed N times."""
+    from litellm.types.videos.main import VideoObject
+
+    def _video(usage: dict[str, object]) -> VideoObject:
+        return VideoObject(id="v", object="video", status="processing", model="veo-3.1-fast-generate-001", usage=usage)
+
+    single_cost = completion_cost(
+        completion_response=_video({"duration_seconds": 4.0, "video_resolution": "720p"}),
+        model="veo-3.1-fast-generate-001",
+        custom_llm_provider="vertex_ai",
+        call_type="create_video",
+    )
+    multi_cost = completion_cost(
+        completion_response=_video({"duration_seconds": 4.0, "video_resolution": "720p", "video_count": video_count}),
+        model="veo-3.1-fast-generate-001",
+        custom_llm_provider="vertex_ai",
+        call_type="create_video",
+    )
+
+    assert single_cost > 0
+    assert multi_cost == pytest.approx(single_cost * video_count)
+
+
 @pytest.mark.parametrize(
     "batch_rate,expected_prompt,expected_completion",
     [