refactor(vertex_ai): tighten Interactions route match and type the usage boundary

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/cost_calculator.py

@@ -1460,7 +1460,6 @@ def completion_cost(
                     duration_seconds: float | None = None
                     video_resolution: str | None = None
                     provider_reported_cost: float | None = None
-                    video_count: int = 1
                     if completion_response is not None and usage_obj:
                         # Handle both dict and Pydantic Usage object
                         if isinstance(usage_obj, dict):
@@ -1475,8 +1474,7 @@ def completion_cost(
                             _vc = getattr(usage_obj, "video_count", None)
                         if _vr is not None:
                             video_resolution = str(_vr).strip().lower()
-                        if isinstance(_vc, int) and not isinstance(_vc, bool) and _vc > 1:
-                            video_count = _vc
+                        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)

litellm/llms/vertex_ai/videos/transformation.py

@@ -71,7 +71,6 @@ def _parse_veo_operation(raw_response: httpx.Response) -> _VeoOperation:
 
 
 def veo_video_count_from_parameters(parameters: Mapping[str, object]) -> int | None:
-    """Number of videos Veo generates for one request (``parameters.sampleCount``)."""
     sample_count: Final = parameters.get("sampleCount")
     if isinstance(sample_count, bool) or not isinstance(sample_count, int) or sample_count < 1:
         return None

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -6,6 +6,7 @@ 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
@@ -53,6 +54,9 @@ 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],
@@ -72,18 +76,22 @@ class VertexPassthroughLoggingHandler:
     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,
+        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[Mapping[str, object]] = httpx_response.json()
+        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):
@@ -95,6 +103,7 @@ class VertexPassthroughLoggingHandler:
                 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,
@@ -107,9 +116,10 @@ class VertexPassthroughLoggingHandler:
                 vertex_location=vertex_location,
             )
         )
-        logging_kwargs["custom_llm_provider"] = custom_llm_provider
-        logging_obj.custom_llm_provider = custom_llm_provider
-        return {"result": litellm_model_response, "kwargs": logging_kwargs}
+        return {
+            "result": litellm_model_response,
+            "kwargs": {**logging_kwargs, "custom_llm_provider": custom_llm_provider},
+        }
 
     @staticmethod
     def vertex_passthrough_handler(

litellm/proxy/pass_through_endpoints/success_handler.py

@@ -363,9 +363,7 @@ class PassThroughEndpointLogging:
             return True
         if any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES):
             return True
-        return "/locations/" in urlparse(url_route).path and VertexPassthroughLoggingHandler.is_interactions_route(
-            url_route
-        )
+        return VertexPassthroughLoggingHandler.is_vertex_interactions_route(url_route)
 
     def is_anthropic_route(self, url_route: str):
         for route in self.TRACKED_ANTHROPIC_ROUTES:

tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -510,6 +510,13 @@ def test_interactions_create_routes_are_tracked_for_vertex_and_gemini():
     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