From 74a2c2eab1f86ada6bc4ccc55dfab64bd6d39219 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:34:31 +0000 Subject: [PATCH] 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 | 4 +--- .../llms/vertex_ai/videos/transformation.py | 1 - .../vertex_passthrough_logging_handler.py | 20 ++++++++++++++----- .../pass_through_endpoints/success_handler.py | 4 +--- .../test_pass_through_endpoints.py | 7 +++++++ 5 files changed, 24 insertions(+), 12 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f81eb1539fd..e6bb8da4c6a 100644 --- a/litellm/cost_calculator.py +++ b/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) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 6525c6fa235..dc9caa13224 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/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 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 0c4bc2a5181..cd226e80c6e 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 @@ -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( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index dd8bd17eee9..76a471302f4 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/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: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 988c4782f6c..cc1d408176c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/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