From 3561ce241ebb8d763ae4b116cad5a970d075e332 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:45:26 -0700 Subject: [PATCH 1/6] fix(passthrough): log the URL model and declared provider on Gemini and Vertex failures Pass-through failures on Gemini and Vertex routes carry the model only in the URL, so the failure spend row and the Prometheus failure counter logged an empty model. The failure payload now falls back to the model parsed from the URL on those routes, and the Prometheus provider label honors the provider the request declares before inferring one from the model name --- litellm/integrations/prometheus.py | 7 +- .../pass_through_endpoints.py | 39 +++-- .../integrations/test_prometheus_labels.py | 18 ++ .../test_pass_through_endpoints.py | 164 +++++++++++++----- 4 files changed, 165 insertions(+), 63 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6766d246894..a31b0f082e3 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2493,8 +2493,9 @@ class PrometheusLogger(CustomLogger): A request can fail before a deployment is resolved, so the provider is not always known. Prefer the resolved ``custom_llm_provider`` on ``litellm_params``, then any provider recovered onto a partial - ``standard_logging_object`` (e.g. a stream that broke mid-flight), and - finally infer it from the requested model name (e.g. ``gpt-4o-mini`` -> + ``standard_logging_object`` (e.g. a stream that broke mid-flight), then + the provider the request itself declares (e.g. a Gemini pass-through + route), and finally infer it from the requested model name (e.g. ``gpt-4o-mini`` -> ``openai``) since the proxy's failure ``request_data`` usually carries only the client-supplied model. Return ``None`` when it cannot be determined so the label emits empty rather than a guess. @@ -2504,7 +2505,7 @@ class PrometheusLogger(CustomLogger): if provider: return provider standard_logging_object: Final = request_data.get("standard_logging_object") or {} - provider = standard_logging_object.get("custom_llm_provider") + provider = standard_logging_object.get("custom_llm_provider") or request_data.get("custom_llm_provider") if provider: return provider model: Final = litellm_params.get("model") or request_data.get("model") diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3cb9acc6110..023d8eeb961 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -58,6 +58,7 @@ from litellm.llms.base_llm.managed_resources.utils import ( resolve_passthrough_managed_id_provider, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url from litellm.passthrough import BasePassthroughUtils from litellm.proxy._types import ( ConfigFieldInfo, @@ -730,12 +731,23 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non metadata.setdefault("standard_logging_guardrail_information", list(entries)) +def _passthrough_model_from_url(url: httpx.URL | None, custom_llm_provider: str | None) -> str | None: + if url is None: + return None + url_route: Final = str(url) + tracked_google_route: Final = pass_through_endpoint_logging.is_gemini_route( + url_route, custom_llm_provider + ) or pass_through_endpoint_logging.is_vertex_route(url_route) + return get_vertex_model_id_from_url(url_route) if tracked_google_route else None + + def _build_passthrough_failure_request_payload( parsed_body: dict | None, kwargs: dict | None, logging_obj: LiteLLMLoggingObj | None, custom_llm_provider: str | None, upstream_usage: UpstreamReportedUsage | None = None, + url: httpx.URL | None = None, ) -> dict: """Build the ``request_data`` dict passed to ``post_call_failure_hook``. @@ -752,8 +764,8 @@ def _build_passthrough_failure_request_payload( request_payload.update(kwargs) if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj - if "model" not in request_payload and parsed_body and isinstance(parsed_body, dict): - request_payload["model"] = parsed_body.get("model", "") + if "model" not in request_payload: + request_payload["model"] = _passthrough_model_from_url(url, custom_llm_provider) or "" if "custom_llm_provider" not in request_payload and custom_llm_provider: request_payload["custom_llm_provider"] = custom_llm_provider if upstream_usage is not None: @@ -1302,6 +1314,7 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, ), ) @@ -1393,6 +1406,7 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, ), ) @@ -1492,6 +1506,7 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, ) failure_request_payload["response_body"] = response_body await _log_passthrough_upstream_failure( @@ -1697,19 +1712,13 @@ async def pass_through_request( # Monitoring: Trigger post_call_failure_hook # for pass through endpoint failure ######################################################### - request_payload: Final[dict] = _parsed_body or {} - # add user_api_key_dict, litellm_call_id, passthrough_logging_payloa for logging - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - if "model" not in request_payload and _parsed_body and isinstance(_parsed_body, dict): - request_payload["model"] = _parsed_body.get("model", "") - if "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider - + request_payload: Final = _build_passthrough_failure_request_payload( + parsed_body=_parsed_body, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + url=url, + ) _carry_guardrail_logging_info(request_payload, post_call_guardrail_data) await proxy_logging_obj.post_call_failure_hook( diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 859cdd30c11..93191e7d888 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -235,6 +235,24 @@ def test_extract_api_provider_from_request_data_failure_path(): assert extract({"model": "some-unknown-model-xyz"}) is None +def test_extract_api_provider_prefers_declared_provider_over_model_name_inference(): + from litellm.integrations.prometheus import PrometheusLogger + + extract = PrometheusLogger._extract_api_provider_from_request_data + + assert extract({"model": "gemini-3.8-flash"}) == "vertex_ai" + assert extract({"model": "gemini-3.8-flash", "custom_llm_provider": "gemini"}) == "gemini" + assert ( + extract( + { + "custom_llm_provider": "gemini", + "standard_logging_object": {"custom_llm_provider": "vertex_ai"}, + } + ) + == "vertex_ai" + ) + + def test_extract_api_provider_swallows_unknown_model_but_logs_unexpected_errors(): """ get_llm_provider raises BadRequestError for a model that maps to no 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 fb4ed3db4d2..d43e66fb956 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 @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from fastapi import Request, UploadFile from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile @@ -21,6 +22,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + _build_passthrough_failure_request_payload, _registered_pass_through_routes, create_pass_through_route, initialize_pass_through_endpoints, @@ -344,53 +346,18 @@ async def test_pass_through_request_failure_handler(): Critical Test: When a users pass through endpoint request fails, we must log the failure code, exception in litellm spend logs. """ - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client: - with patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processing: - # Setup mock for post_call_failure_hook and pre_call_hook - mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.pre_call_hook = AsyncMock() + request_failure = httpx.HTTPError("Request failed") + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") - # Setup mock for httpx client - mock_client = MagicMock() - mock_client.client = MagicMock() - mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed")) - mock_get_client.return_value = mock_client + failure = await _run_passthrough_failure( + "http://test.com", None, request_failure, user_api_key_dict=user_api_key_dict + ) - # Mock headers for custom headers - mock_processing.get_custom_headers.return_value = {} - - # Create mock request - mock_request = MagicMock(spec=Request) - mock_request.method = "POST" - mock_request.body = AsyncMock(return_value=b'{"test": "data"}') - mock_request.headers = Headers({}) - - # Create a simple empty QueryParams - mock_request.query_params = QueryParams({}) - - # Create mock user API key dict - mock_user_api_key_dict = MagicMock() - - # Call the function with a target that will trigger an HTTPError - with pytest.raises(ProxyException): - await pass_through_request( - request=mock_request, - target="http://test.com", - custom_headers={}, - user_api_key_dict=mock_user_api_key_dict, - ) - - # Assert post_call_failure_hook was called - mock_proxy_logging.post_call_failure_hook.assert_called_once() - - # Verify the arguments to post_call_failure_hook - call_args = mock_proxy_logging.post_call_failure_hook.call_args[1] - assert call_args["user_api_key_dict"] == mock_user_api_key_dict - assert isinstance(call_args["original_exception"], TypeError) # Now expecting TypeError - assert "traceback_str" in call_args + assert failure["user_api_key_dict"] is user_api_key_dict + assert failure["original_exception"] is request_failure + assert "traceback_str" in failure + assert failure["request_data"]["contents"] == "not a list" + assert failure["request_data"]["model"] == "" @pytest.mark.asyncio @@ -5837,3 +5804,110 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) == "per-call-random-trace-id" ) + + +_GEMINI_FLASH_TARGET = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash:generateContent" +_VERTEX_FLASH_TARGET = ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/google/models/gemini-3.8-flash:generateContent" +) + + +async def _run_passthrough_failure( + target: str, + custom_llm_provider: str | None, + upstream: httpx.Response | Exception, + stream: bool = False, + user_api_key_dict: UserAPIKeyAuth | None = None, +) -> dict: + with ( + pytest.MonkeyPatch.context() as transport_patch, + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + respx.mock(assert_all_called=False) as upstream_mock, + ): + transport_patch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream_route = upstream_mock.route(method="POST") + if isinstance(upstream, Exception): + upstream_route.mock(side_effect=upstream) + else: + upstream_route.mock(return_value=upstream) + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = target.replace("https://", "http://test-proxy.com/") + mock_request.body = AsyncMock(return_value=b'{"contents": "not a list"}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + try: + await pass_through_request( + request=mock_request, + target=target, + custom_headers={}, + user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(api_key="sk-test"), + custom_llm_provider=custom_llm_provider, + stream=stream, + ) + except ProxyException: + pass + litellm.in_memory_llm_clients_cache.flush_cache() + mock_proxy_logging.post_call_failure_hook.assert_called_once() + return mock_proxy_logging.post_call_failure_hook.call_args.kwargs + + +def _upstream_400(target: str) -> httpx.Response: + return httpx.Response( + status_code=400, + headers={"content-type": "application/json"}, + content=json.dumps({"error": {"code": 400, "status": "INVALID_ARGUMENT"}}).encode(), + request=httpx.Request("POST", target), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, stream", + [ + (_GEMINI_FLASH_TARGET, "gemini", False), + (_GEMINI_FLASH_TARGET, "gemini", True), + (_VERTEX_FLASH_TARGET, None, False), + ], +) +async def test_passthrough_upstream_error_logs_model_from_url( + target: str, custom_llm_provider: str | None, stream: bool +): + failure = await _run_passthrough_failure(target, custom_llm_provider, _upstream_400(target), stream=stream) + assert failure["request_data"]["model"] == "gemini-3.8-flash" + + +@pytest.mark.asyncio +async def test_passthrough_internal_failure_logs_model_from_url(): + failure = await _run_passthrough_failure( + _GEMINI_FLASH_TARGET, "gemini", httpx.ReadTimeout("upstream timed out") + ) + assert failure["request_data"]["model"] == "gemini-3.8-flash" + + +@pytest.mark.parametrize( + "parsed_body, url, custom_llm_provider, expected_model", + [ + ({"model": "from-body"}, _GEMINI_FLASH_TARGET, "gemini", "from-body"), + ({"contents": []}, _GEMINI_FLASH_TARGET, "gemini", "gemini-3.8-flash"), + ({"contents": []}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash"), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, ""), + (None, None, None, ""), + ], +) +def test_build_passthrough_failure_request_payload_model_precedence( + parsed_body: dict | None, url: str | None, custom_llm_provider: str | None, expected_model: str +): + request_payload = _build_passthrough_failure_request_payload( + parsed_body=parsed_body, + kwargs=None, + logging_obj=None, + custom_llm_provider=custom_llm_provider, + url=httpx.URL(url) if url else None, + ) + assert request_payload["model"] == expected_model From 006c166e5e57ab8eb5cfbaf3ea5e289173cbd4a5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:01:17 -0700 Subject: [PATCH 2/6] fix(passthrough): keep model and provider on mid-stream and Vertex failure logs Build the failure payload from the URL at both streamed relay sites so a 200 stream that dies mid-body still logs the model, derive vertex_ai for Vertex routes that never declare a provider, and have the Vertex logging handler reuse get_vertex_model_id_from_url instead of its own regex --- .../vertex_passthrough_logging_handler.py | 9 +-- .../pass_through_endpoints.py | 25 ++++--- .../test_pass_through_endpoints.py | 73 +++++++++++++++---- 3 files changed, 74 insertions(+), 33 deletions(-) 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 49ec18013b5..a2cb61845dd 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 @@ -1,5 +1,4 @@ import asyncio -import re from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse @@ -10,7 +9,7 @@ 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.llms.vertex_ai.common_utils import get_vertex_location_from_url +from litellm.llms.vertex_ai.common_utils import get_vertex_location_from_url, get_vertex_model_id_from_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as VertexModelResponseIterator, ) @@ -507,11 +506,7 @@ class VertexPassthroughLoggingHandler: @staticmethod def extract_model_from_url(url: str) -> str: - pattern: Final = r"/models/([^:]+)" - match: Final = re.search(pattern, url) - if match: - return match.group(1) - return "unknown" + return get_vertex_model_id_from_url(url) or "unknown" @staticmethod def extract_model_name_from_vertex_path(vertex_model_path: str) -> str: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 023d8eeb961..d2d9811824a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -731,14 +731,12 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non metadata.setdefault("standard_logging_guardrail_information", list(entries)) -def _passthrough_model_from_url(url: httpx.URL | None, custom_llm_provider: str | None) -> str | None: - if url is None: - return None - url_route: Final = str(url) - tracked_google_route: Final = pass_through_endpoint_logging.is_gemini_route( - url_route, custom_llm_provider - ) or pass_through_endpoint_logging.is_vertex_route(url_route) - return get_vertex_model_id_from_url(url_route) if tracked_google_route else None +def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None) -> str | None: + if pass_through_endpoint_logging.is_gemini_route(url_route, custom_llm_provider): + return "gemini" + if pass_through_endpoint_logging.is_vertex_route(url_route): + return "vertex_ai" + return None def _build_passthrough_failure_request_payload( @@ -764,10 +762,13 @@ def _build_passthrough_failure_request_payload( request_payload.update(kwargs) if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj + url_route: Final = str(url) if url is not None else "" + google_provider: Final = _passthrough_google_provider(url_route, custom_llm_provider) if "model" not in request_payload: - request_payload["model"] = _passthrough_model_from_url(url, custom_llm_provider) or "" - if "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider + request_payload["model"] = (get_vertex_model_id_from_url(url_route) if google_provider else None) or "" + resolved_provider: Final = custom_llm_provider or google_provider + if "custom_llm_provider" not in request_payload and resolved_provider: + request_payload["custom_llm_provider"] = resolved_provider if upstream_usage is not None: request_payload["response_cost"] = upstream_usage.response_cost or 0.0 request_payload["combined_usage_object"] = Usage(total_tokens=upstream_usage.total_tokens or 0) @@ -1352,6 +1353,7 @@ async def pass_through_request( kwargs=kwargs, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + url=url, ), ), managed_id_provider=_managed_id_provider, @@ -1444,6 +1446,7 @@ async def pass_through_request( kwargs=kwargs, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + url=url, ), ), managed_id_provider=_managed_id_provider, 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 d43e66fb956..b9d67a7dfdd 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 @@ -13,6 +13,7 @@ import httpx import pytest import respx from fastapi import Request, UploadFile +from fastapi.responses import StreamingResponse from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile @@ -5842,7 +5843,7 @@ async def _run_passthrough_failure( mock_request.headers = Headers({"content-type": "application/json"}) mock_request.query_params = QueryParams({}) try: - await pass_through_request( + received = await pass_through_request( request=mock_request, target=target, custom_headers={}, @@ -5850,8 +5851,12 @@ async def _run_passthrough_failure( custom_llm_provider=custom_llm_provider, stream=stream, ) - except ProxyException: + if isinstance(received, StreamingResponse): + async for _ in received.body_iterator: + pass + except (ProxyException, httpx.HTTPError): pass + await asyncio.sleep(0) litellm.in_memory_llm_clients_cache.flush_cache() mock_proxy_logging.post_call_failure_hook.assert_called_once() return mock_proxy_logging.post_call_failure_hook.call_args.kwargs @@ -5868,18 +5873,20 @@ def _upstream_400(target: str) -> httpx.Response: @pytest.mark.asyncio @pytest.mark.parametrize( - "target, custom_llm_provider, stream", + "target, custom_llm_provider, stream, expected_provider", [ - (_GEMINI_FLASH_TARGET, "gemini", False), - (_GEMINI_FLASH_TARGET, "gemini", True), - (_VERTEX_FLASH_TARGET, None, False), + (_GEMINI_FLASH_TARGET, "gemini", False, "gemini"), + (_GEMINI_FLASH_TARGET, "gemini", True, "gemini"), + (_VERTEX_FLASH_TARGET, None, False, "vertex_ai"), + (_VERTEX_FLASH_TARGET, None, True, "vertex_ai"), ], ) async def test_passthrough_upstream_error_logs_model_from_url( - target: str, custom_llm_provider: str | None, stream: bool + target: str, custom_llm_provider: str | None, stream: bool, expected_provider: str ): failure = await _run_passthrough_failure(target, custom_llm_provider, _upstream_400(target), stream=stream) assert failure["request_data"]["model"] == "gemini-3.8-flash" + assert failure["request_data"]["custom_llm_provider"] == expected_provider @pytest.mark.asyncio @@ -5888,20 +5895,55 @@ async def test_passthrough_internal_failure_logs_model_from_url(): _GEMINI_FLASH_TARGET, "gemini", httpx.ReadTimeout("upstream timed out") ) assert failure["request_data"]["model"] == "gemini-3.8-flash" + assert failure["request_data"]["custom_llm_provider"] == "gemini" + + +def _upstream_dropping_stream(target: str) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + stream=_UpstreamDroppingMidStream(), + request=httpx.Request("POST", target), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target, custom_llm_provider, expected_provider", + [ + (_GEMINI_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), "gemini", "gemini"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), None, "vertex_ai"), + ], +) +async def test_passthrough_mid_stream_drop_logs_model_from_url( + target: str, custom_llm_provider: str | None, expected_provider: str +): + failure = await _run_passthrough_failure( + target, custom_llm_provider, _upstream_dropping_stream(target), stream=True + ) + assert isinstance(failure["original_exception"], httpx.ReadError) + assert failure["request_data"]["model"] == "gemini-3.8-flash" + assert failure["request_data"]["custom_llm_provider"] == expected_provider @pytest.mark.parametrize( - "parsed_body, url, custom_llm_provider, expected_model", + "parsed_body, url, custom_llm_provider, expected_model, expected_provider", [ - ({"model": "from-body"}, _GEMINI_FLASH_TARGET, "gemini", "from-body"), - ({"contents": []}, _GEMINI_FLASH_TARGET, "gemini", "gemini-3.8-flash"), - ({"contents": []}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash"), - ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, ""), - (None, None, None, ""), + ({"model": "from-body"}, _GEMINI_FLASH_TARGET, "gemini", "from-body", "gemini"), + ({"contents": []}, _GEMINI_FLASH_TARGET, "gemini", "gemini-3.8-flash", "gemini"), + ({"contents": []}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "vertex_ai"), + ({"contents": []}, _VERTEX_FLASH_TARGET, "anthropic", "gemini-3.8-flash", "anthropic"), + ({"custom_llm_provider": "from-body"}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "from-body"), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None), + (None, None, None, "", None), ], ) -def test_build_passthrough_failure_request_payload_model_precedence( - parsed_body: dict | None, url: str | None, custom_llm_provider: str | None, expected_model: str +def test_build_passthrough_failure_request_payload_model_and_provider_precedence( + parsed_body: dict | None, + url: str | None, + custom_llm_provider: str | None, + expected_model: str, + expected_provider: str | None, ): request_payload = _build_passthrough_failure_request_payload( parsed_body=parsed_body, @@ -5911,3 +5953,4 @@ def test_build_passthrough_failure_request_payload_model_precedence( url=httpx.URL(url) if url else None, ) assert request_payload["model"] == expected_model + assert request_payload.get("custom_llm_provider") == expected_provider From 60ae16db8cc74a9622cb6d28c844d3b3aa3e6b1e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:03:46 -0700 Subject: [PATCH 3/6] fix(passthrough): bound the Prometheus provider label and name Vertex failures like the success handler Declared providers only reach the api_provider label when they name a registered provider, otherwise the label falls back to inference. Vertex failure rows now use the success handler's model naming (vertex_ai/ prefix on rawPredict and streamRawPredict, vertex_ai/search_api on :search) and its host-based provider detection, and the detected-stream failure sites are covered by tests. --- litellm/integrations/prometheus.py | 25 +++-- .../gemini_passthrough_logging_handler.py | 8 +- .../vertex_passthrough_logging_handler.py | 17 +++- .../pass_through_endpoints.py | 15 ++- .../integrations/test_prometheus_labels.py | 36 +++++++ .../test_pass_through_endpoints.py | 99 ++++++++++++++++--- 6 files changed, 166 insertions(+), 34 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index a31b0f082e3..6ef3fd3d280 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -35,6 +35,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, ) +from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.proxy._types import ( LiteLLM_DeletedVerificationToken, LiteLLM_TeamTable, @@ -2485,6 +2486,12 @@ class PrometheusLogger(CustomLogger): return False + @staticmethod + def _is_known_provider(provider: object) -> bool: + return isinstance(provider, str) and ( + provider in litellm.provider_list or JSONProviderRegistry.exists(provider) + ) + @staticmethod def _extract_api_provider_from_request_data(request_data: dict) -> str | None: """ @@ -2497,17 +2504,21 @@ class PrometheusLogger(CustomLogger): the provider the request itself declares (e.g. a Gemini pass-through route), and finally infer it from the requested model name (e.g. ``gpt-4o-mini`` -> ``openai``) since the proxy's failure ``request_data`` usually carries - only the client-supplied model. Return ``None`` when it cannot be + only the client-supplied model. Declared values outside the known + provider set are ignored, since the request body can carry any string + and this label must stay bounded. Return ``None`` when it cannot be determined so the label emits empty rather than a guess. """ litellm_params: Final = request_data.get("litellm_params") or {} - provider = litellm_params.get("custom_llm_provider") - if provider: - return provider standard_logging_object: Final = request_data.get("standard_logging_object") or {} - provider = standard_logging_object.get("custom_llm_provider") or request_data.get("custom_llm_provider") - if provider: - return provider + declared_providers: Final = ( + litellm_params.get("custom_llm_provider"), + standard_logging_object.get("custom_llm_provider"), + request_data.get("custom_llm_provider"), + ) + known_provider: Final = next((p for p in declared_providers if PrometheusLogger._is_known_provider(p)), None) + if known_provider: + return known_provider model: Final = litellm_params.get("model") or request_data.get("model") if not model: return None 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 64d8b2929b6..c4b102101fb 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 @@ -1,4 +1,3 @@ -import re from datetime import datetime from typing import TYPE_CHECKING, Any, Final @@ -8,6 +7,7 @@ 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.common_utils import get_vertex_model_id_from_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as GeminiModelResponseIterator, ) @@ -200,11 +200,7 @@ class GeminiPassthroughLoggingHandler: @staticmethod def extract_model_from_url(url: str) -> str: - pattern: Final = r"/models/([^:]+)" - match: Final = re.search(pattern, url) - if match: - return match.group(1) - return "unknown" + return get_vertex_model_id_from_url(url) or "unknown" @staticmethod def _create_gemini_response_logging_payload_for_generate_content( 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 a2cb61845dd..b667a7916fb 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 @@ -126,7 +126,7 @@ class VertexPassthroughLoggingHandler: start_time=start_time, end_time=end_time, logging_obj=logging_obj, - custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route), + custom_llm_provider=VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route), vertex_location=vertex_location, ) @@ -377,7 +377,7 @@ class VertexPassthroughLoggingHandler: response_json=response_json, ) - custom_llm_provider: Final = VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route) + custom_llm_provider: Final = VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route) litellm_embedding_response.model = model logging_obj.model = model @@ -449,7 +449,7 @@ class VertexPassthroughLoggingHandler: start_time=start_time, end_time=end_time, logging_obj=litellm_logging_obj, - custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route), + custom_llm_provider=VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route), vertex_location=vertex_location, ) @@ -508,6 +508,15 @@ class VertexPassthroughLoggingHandler: def extract_model_from_url(url: str) -> str: return get_vertex_model_id_from_url(url) or "unknown" + @staticmethod + def model_from_url_route(url_route: str) -> str | None: + model: Final = get_vertex_model_id_from_url(url_route) + if model is None: + return "vertex_ai/search_api" if ":search" in url_route else None + if "rawPredict" in url_route or "streamRawPredict" in url_route: + return f"vertex_ai/{model}" + return model + @staticmethod def extract_model_name_from_vertex_path(vertex_model_path: str) -> str: """ @@ -552,7 +561,7 @@ class VertexPassthroughLoggingHandler: return None @staticmethod - def _get_custom_llm_provider_from_url(url: str) -> str: + def custom_llm_provider_from_url(url: str) -> str: parsed_url: Final = urlparse(url) if parsed_url.hostname and parsed_url.hostname.endswith("generativelanguage.googleapis.com"): return litellm.LlmProviders.GEMINI.value diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d2d9811824a..e59907bfa4c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -99,6 +99,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage +from .llm_provider_handlers.vertex_passthrough_logging_handler import VertexPassthroughLoggingHandler from .streaming_handler import PassThroughStreamingHandler from .success_handler import PassThroughEndpointLogging from .upstream_usage_headers import ( @@ -735,10 +736,20 @@ def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None if pass_through_endpoint_logging.is_gemini_route(url_route, custom_llm_provider): return "gemini" if pass_through_endpoint_logging.is_vertex_route(url_route): - return "vertex_ai" + return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route) return None +def _passthrough_google_model(url_route: str, google_provider: str | None) -> str | None: + match google_provider: + case "gemini": + return get_vertex_model_id_from_url(url_route) + case "vertex_ai": + return VertexPassthroughLoggingHandler.model_from_url_route(url_route) + case _: + return None + + def _build_passthrough_failure_request_payload( parsed_body: dict | None, kwargs: dict | None, @@ -765,7 +776,7 @@ def _build_passthrough_failure_request_payload( url_route: Final = str(url) if url is not None else "" google_provider: Final = _passthrough_google_provider(url_route, custom_llm_provider) if "model" not in request_payload: - request_payload["model"] = (get_vertex_model_id_from_url(url_route) if google_provider else None) or "" + request_payload["model"] = _passthrough_google_model(url_route, google_provider) or "" resolved_provider: Final = custom_llm_provider or google_provider if "custom_llm_provider" not in request_payload and resolved_provider: request_payload["custom_llm_provider"] = resolved_provider diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 93191e7d888..3506614768d 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -242,6 +242,7 @@ def test_extract_api_provider_prefers_declared_provider_over_model_name_inferenc assert extract({"model": "gemini-3.8-flash"}) == "vertex_ai" assert extract({"model": "gemini-3.8-flash", "custom_llm_provider": "gemini"}) == "gemini" + assert extract({"model": "gemini-3.8-flash", "custom_llm_provider": "vertex_ai"}) == "vertex_ai" assert ( extract( { @@ -253,6 +254,41 @@ def test_extract_api_provider_prefers_declared_provider_over_model_name_inferenc ) +@pytest.mark.parametrize( + "declared_provider", + ["client-supplied-value", "", {"nested": "object"}, 42], +) +def test_extract_api_provider_ignores_declared_values_outside_the_known_provider_set(declared_provider: object): + from litellm.integrations.prometheus import PrometheusLogger + + extract = PrometheusLogger._extract_api_provider_from_request_data + + assert extract({"model": "gpt-4o-mini", "custom_llm_provider": declared_provider}) == "openai" + assert extract({"model": "gpt-4o-mini", "litellm_params": {"custom_llm_provider": declared_provider}}) == "openai" + assert extract({"model": "no-such-model", "custom_llm_provider": declared_provider}) is None + assert ( + extract( + { + "litellm_params": {"custom_llm_provider": declared_provider}, + "standard_logging_object": {"custom_llm_provider": "azure"}, + } + ) + == "azure" + ) + + +def test_extract_api_provider_keeps_registered_custom_providers(monkeypatch: pytest.MonkeyPatch): + import litellm + from litellm.integrations.prometheus import PrometheusLogger + + extract = PrometheusLogger._extract_api_provider_from_request_data + assert extract({"custom_llm_provider": "my-custom-llm"}) is None + + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "my-custom-llm"]) + assert extract({"custom_llm_provider": "my-custom-llm"}) == "my-custom-llm" + assert extract({"litellm_params": {"custom_llm_provider": "my-custom-llm"}}) == "my-custom-llm" + + def test_extract_api_provider_swallows_unknown_model_but_logs_unexpected_errors(): """ get_llm_provider raises BadRequestError for a model that maps to no 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 b9d67a7dfdd..79c2e7fb3e9 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 @@ -17,32 +17,33 @@ from fastapi.responses import StreamingResponse from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile - +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( + VertexPassthroughLoggingHandler, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, + LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, - LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, _build_passthrough_failure_request_payload, _registered_pass_through_routes, create_pass_through_route, initialize_pass_through_endpoints, pass_through_request, - resolve_pass_through_request_timeout, resolve_llm_passthrough_timeout, + resolve_pass_through_request_timeout, websocket_passthrough_request, ) -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy._types import ProxyException, UserAPIKeyAuth -from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, -) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) - -import litellm +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, +) MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start"}\n\n' @@ -5812,6 +5813,14 @@ _VERTEX_FLASH_TARGET = ( "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" "/publishers/google/models/gemini-3.8-flash:generateContent" ) +_VERTEX_CLAUDE_RAW_PREDICT_TARGET = ( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/anthropic/models/claude-opus-5:rawPredict" +) +_VERTEX_SEARCH_TARGET = ( + "https://discoveryengine.googleapis.com/v1/projects/p/locations/global" + "/dataStores/default/servingConfigs/default:search" +) async def _run_passthrough_failure( @@ -5871,6 +5880,15 @@ def _upstream_400(target: str) -> httpx.Response: ) +def _upstream_400_as_event_stream(target: str) -> httpx.Response: + return httpx.Response( + status_code=400, + headers={"content-type": "text/event-stream"}, + content=b'data: {"error": {"code": 400, "status": "INVALID_ARGUMENT"}}\n\n', + request=httpx.Request("POST", target), + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "target, custom_llm_provider, stream, expected_provider", @@ -5890,11 +5908,27 @@ async def test_passthrough_upstream_error_logs_model_from_url( @pytest.mark.asyncio -async def test_passthrough_internal_failure_logs_model_from_url(): +@pytest.mark.parametrize( + "target, custom_llm_provider, expected_provider", + [ + (_GEMINI_FLASH_TARGET, "gemini", "gemini"), + (_VERTEX_FLASH_TARGET, None, "vertex_ai"), + ], +) +async def test_passthrough_detected_stream_upstream_error_logs_model_from_url( + target: str, custom_llm_provider: str | None, expected_provider: str +): failure = await _run_passthrough_failure( - _GEMINI_FLASH_TARGET, "gemini", httpx.ReadTimeout("upstream timed out") + target, custom_llm_provider, _upstream_400_as_event_stream(target), stream=False ) assert failure["request_data"]["model"] == "gemini-3.8-flash" + assert failure["request_data"]["custom_llm_provider"] == expected_provider + + +@pytest.mark.asyncio +async def test_passthrough_internal_failure_logs_model_from_url(): + failure = await _run_passthrough_failure(_GEMINI_FLASH_TARGET, "gemini", httpx.ReadTimeout("upstream timed out")) + assert failure["request_data"]["model"] == "gemini-3.8-flash" assert failure["request_data"]["custom_llm_provider"] == "gemini" @@ -5908,6 +5942,7 @@ def _upstream_dropping_stream(target: str) -> httpx.Response: @pytest.mark.asyncio +@pytest.mark.parametrize("stream", [True, False]) @pytest.mark.parametrize( "target, custom_llm_provider, expected_provider", [ @@ -5916,10 +5951,10 @@ def _upstream_dropping_stream(target: str) -> httpx.Response: ], ) async def test_passthrough_mid_stream_drop_logs_model_from_url( - target: str, custom_llm_provider: str | None, expected_provider: str + target: str, custom_llm_provider: str | None, expected_provider: str, stream: bool ): failure = await _run_passthrough_failure( - target, custom_llm_provider, _upstream_dropping_stream(target), stream=True + target, custom_llm_provider, _upstream_dropping_stream(target), stream=stream ) assert isinstance(failure["original_exception"], httpx.ReadError) assert failure["request_data"]["model"] == "gemini-3.8-flash" @@ -5934,6 +5969,23 @@ async def test_passthrough_mid_stream_drop_logs_model_from_url( ({"contents": []}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "vertex_ai"), ({"contents": []}, _VERTEX_FLASH_TARGET, "anthropic", "gemini-3.8-flash", "anthropic"), ({"custom_llm_provider": "from-body"}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "from-body"), + ({"contents": []}, _GEMINI_FLASH_TARGET, None, "gemini-3.8-flash", "gemini"), + ({}, _VERTEX_CLAUDE_RAW_PREDICT_TARGET, None, "vertex_ai/claude-opus-5", "vertex_ai"), + ( + {}, + _VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict"), + None, + "vertex_ai/claude-opus-5", + "vertex_ai", + ), + ({}, _VERTEX_SEARCH_TARGET, None, "vertex_ai/search_api", "vertex_ai"), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs", + None, + "", + "vertex_ai", + ), ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None), (None, None, None, "", None), ], @@ -5954,3 +6006,20 @@ def test_build_passthrough_failure_request_payload_model_and_provider_precedence ) assert request_payload["model"] == expected_model assert request_payload.get("custom_llm_provider") == expected_provider + + +@pytest.mark.parametrize( + "url_route, expected_model", + [ + (_VERTEX_FLASH_TARGET, "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":streamGenerateContent?alt=sse"), "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":predict"), "gemini-3.8-flash"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":embedContent"), "gemini-3.8-flash"), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET, "vertex_ai/claude-opus-5"), + (_VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict"), "vertex_ai/claude-opus-5"), + (_VERTEX_SEARCH_TARGET, "vertex_ai/search_api"), + ("https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs", None), + ], +) +def test_vertex_failure_model_name_matches_success_handler_naming(url_route: str, expected_model: str | None): + assert VertexPassthroughLoggingHandler.model_from_url_route(url_route) == expected_model From 571927705a7866dc358d01b2bdef68e65f253b5a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 03:19:27 -0700 Subject: [PATCH 4/6] fix(passthrough): name Google failures from the URL on untracked methods too A Gemini or Vertex pass-through failure on a method the success handler does not track (countTokens, for one) still logged model "" and no provider, although the model sits in the URL. The failure payload now falls back to the Google AI host for the provider and reads /models/ from the URL, so countTokens 400s name their model and provider like generateContent ones do. Non-Google hosts keep the tracked-route rule and still log nothing extra. --- .../vertex_passthrough_logging_handler.py | 9 +++++- .../pass_through_endpoints.py | 4 ++- .../test_pass_through_endpoints.py | 32 +++++++++++++++++++ 3 files changed, 43 insertions(+), 2 deletions(-) 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 b667a7916fb..960e4d5882b 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 @@ -35,6 +35,8 @@ from litellm.types.utils import ( ) vertex_search_api_config: Final = VertexSearchAPIVectorStoreConfig() +GEMINI_API_HOST_SUFFIX: Final = "generativelanguage.googleapis.com" +VERTEX_API_HOST_SUFFIX: Final = "aiplatform.googleapis.com" if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch @@ -560,10 +562,15 @@ class VertexPassthroughLoggingHandler: return "openapi" return None + @staticmethod + def is_google_ai_host(url: str) -> bool: + hostname: Final = urlparse(url).hostname + return hostname is not None and hostname.endswith((GEMINI_API_HOST_SUFFIX, VERTEX_API_HOST_SUFFIX)) + @staticmethod def custom_llm_provider_from_url(url: str) -> str: parsed_url: Final = urlparse(url) - if parsed_url.hostname and parsed_url.hostname.endswith("generativelanguage.googleapis.com"): + if parsed_url.hostname and parsed_url.hostname.endswith(GEMINI_API_HOST_SUFFIX): return litellm.LlmProviders.GEMINI.value return litellm.LlmProviders.VERTEX_AI.value diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e59907bfa4c..eb81a02871e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -735,7 +735,9 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None) -> str | None: if pass_through_endpoint_logging.is_gemini_route(url_route, custom_llm_provider): return "gemini" - if pass_through_endpoint_logging.is_vertex_route(url_route): + if pass_through_endpoint_logging.is_vertex_route(url_route) or VertexPassthroughLoggingHandler.is_google_ai_host( + url_route + ): return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route) return None 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 79c2e7fb3e9..f79b755f84e 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 @@ -5897,6 +5897,8 @@ def _upstream_400_as_event_stream(target: str) -> httpx.Response: (_GEMINI_FLASH_TARGET, "gemini", True, "gemini"), (_VERTEX_FLASH_TARGET, None, False, "vertex_ai"), (_VERTEX_FLASH_TARGET, None, True, "vertex_ai"), + (_GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), "gemini", False, "gemini"), + (_VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), None, False, "vertex_ai"), ], ) async def test_passthrough_upstream_error_logs_model_from_url( @@ -5986,6 +5988,36 @@ async def test_passthrough_mid_stream_drop_logs_model_from_url( "", "vertex_ai", ), + ( + {"contents": []}, + _GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), + "gemini", + "gemini-3.8-flash", + "gemini", + ), + ( + {"contents": []}, + _VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {"contents": []}, + "https://aiplatform.googleapis.com/v1/projects/p/locations/global" + "/publishers/google/models/gemini-3.8-flash:countTokens", + None, + "gemini-3.8-flash", + "vertex_ai", + ), + ( + {}, + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/cachedContents", + None, + "", + "vertex_ai", + ), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:countTokens", None, "", None), ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None), (None, None, None, "", None), ], From f80801985abeb3cdefc08b2b6434d971e3e29382 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 03:58:30 -0700 Subject: [PATCH 5/6] fix(passthrough): trust the declared Google provider and Vertex-shaped paths on custom bases The failure path only named the Google provider from the URL when the method was one the success handler tracks or the host was a Google API host, so a Gemini countTokens failure behind a custom GEMINI_API_BASE (or a Vertex one behind a custom base) still logged an empty model. A route that declares gemini or vertex_ai now wins outright, and a URL whose path carries /projects/

/locations// counts as Vertex on any host, while a tracked method on a non-Google host no longer passes as Vertex. --- .../vertex_passthrough_logging_handler.py | 19 +++++++++++++++---- .../pass_through_endpoints.py | 13 ++++++++----- .../test_pass_through_endpoints.py | 9 +++++++++ 3 files changed, 32 insertions(+), 9 deletions(-) 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 960e4d5882b..7f13364dc23 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 @@ -9,7 +9,11 @@ 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.llms.vertex_ai.common_utils import get_vertex_location_from_url, get_vertex_model_id_from_url +from litellm.llms.vertex_ai.common_utils import ( + get_vertex_location_from_url, + get_vertex_model_id_from_url, + get_vertex_project_id_from_url, +) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as VertexModelResponseIterator, ) @@ -563,9 +567,16 @@ class VertexPassthroughLoggingHandler: return None @staticmethod - def is_google_ai_host(url: str) -> bool: - hostname: Final = urlparse(url).hostname - return hostname is not None and hostname.endswith((GEMINI_API_HOST_SUFFIX, VERTEX_API_HOST_SUFFIX)) + def is_google_ai_url(url: str) -> bool: + parsed_url: Final = urlparse(url) + if parsed_url.hostname is not None and parsed_url.hostname.endswith( + (GEMINI_API_HOST_SUFFIX, VERTEX_API_HOST_SUFFIX) + ): + return True + return ( + get_vertex_project_id_from_url(parsed_url.path) is not None + and get_vertex_location_from_url(parsed_url.path) is not None + ) @staticmethod def custom_llm_provider_from_url(url: str) -> str: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index eb81a02871e..d9f294ea268 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -732,12 +732,15 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non metadata.setdefault("standard_logging_guardrail_information", list(entries)) +_GOOGLE_PASSTHROUGH_PROVIDERS: Final = frozenset( + (litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value) +) + + def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None) -> str | None: - if pass_through_endpoint_logging.is_gemini_route(url_route, custom_llm_provider): - return "gemini" - if pass_through_endpoint_logging.is_vertex_route(url_route) or VertexPassthroughLoggingHandler.is_google_ai_host( - url_route - ): + if custom_llm_provider in _GOOGLE_PASSTHROUGH_PROVIDERS: + return custom_llm_provider + if VertexPassthroughLoggingHandler.is_google_ai_url(url_route): return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route) return None 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 f79b755f84e..4998ffcc2a3 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 @@ -5821,6 +5821,10 @@ _VERTEX_SEARCH_TARGET = ( "https://discoveryengine.googleapis.com/v1/projects/p/locations/global" "/dataStores/default/servingConfigs/default:search" ) +_RELAYED_GEMINI_COUNT_TOKENS_TARGET = "http://127.0.0.1:30732/v1beta/models/gemini-3.8-flash:countTokens" +_RELAYED_VERTEX_COUNT_TOKENS_TARGET = ( + "http://127.0.0.1:30732/v1/projects/p/locations/us-east5/publishers/google/models/gemini-3.8-flash:countTokens" +) async def _run_passthrough_failure( @@ -5899,6 +5903,8 @@ def _upstream_400_as_event_stream(target: str) -> httpx.Response: (_VERTEX_FLASH_TARGET, None, True, "vertex_ai"), (_GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), "gemini", False, "gemini"), (_VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), None, False, "vertex_ai"), + (_RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", False, "gemini"), + (_RELAYED_VERTEX_COUNT_TOKENS_TARGET, None, False, "vertex_ai"), ], ) async def test_passthrough_upstream_error_logs_model_from_url( @@ -6017,7 +6023,10 @@ async def test_passthrough_mid_stream_drop_logs_model_from_url( "", "vertex_ai", ), + ({"contents": []}, _RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", "gemini-3.8-flash", "gemini"), + ({"contents": []}, _RELAYED_VERTEX_COUNT_TOKENS_TARGET, None, "gemini-3.8-flash", "vertex_ai"), ({"contents": []}, "https://api.example.com/v1/models/gpt-x:countTokens", None, "", None), + ({"contents": []}, "https://api.example.com/v1/models/gpt-x:predict", None, "", None), ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None), (None, None, None, "", None), ], From 0b9fd507d25c4e135fa7cf403e83317bd818fef8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:20:31 -0700 Subject: [PATCH 6/6] fix(pass_through): stop the URL model parser at the query string so a configured key never lands in a spend row's model --- litellm/llms/vertex_ai/common_utils.py | 2 +- .../llms/vertex_ai/test_vertex_ai_common_utils.py | 3 +++ .../test_pass_through_endpoints.py | 15 +++++++++++++++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 27f2aedb300..921bd21023a 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1226,7 +1226,7 @@ def get_vertex_model_id_from_url(url: str) -> str | None: `https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:streamGenerateContent` """ - match: Final = re.search(r"/models/([^:]+)", url) + match: Final = re.search(r"/models/([^:?#]+)", url) return match.group(1) if match else None diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 152ab2ffc13..3b222be6e5e 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1293,6 +1293,9 @@ def test_get_vertex_model_id_from_url(): model_id = get_vertex_model_id_from_url(url) assert model_id == "gemini-pro" + url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro?key=provider-secret#frag" + assert get_vertex_model_id_from_url(url) == "gemini-pro" + # Test with invalid URL url = "https://invalid-url.com" model_id = get_vertex_model_id_from_url(url) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index a16b268fc23..9ac5717ec19 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8583,6 +8583,21 @@ async def test_passthrough_mid_stream_drop_logs_model_from_url( ({"contents": []}, _VERTEX_FLASH_TARGET, "anthropic", "gemini-3.8-flash", "anthropic"), ({"custom_llm_provider": "from-body"}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "from-body"), ({"contents": []}, _GEMINI_FLASH_TARGET, None, "gemini-3.8-flash", "gemini"), + ( + {}, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash?key=provider-secret", + "gemini", + "gemini-3.8-flash", + "gemini", + ), + ( + {}, + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/google/models/gemini-3.8-flash?key=provider-secret", + None, + "gemini-3.8-flash", + "vertex_ai", + ), ({}, _VERTEX_CLAUDE_RAW_PREDICT_TARGET, None, "vertex_ai/claude-opus-5", "vertex_ai"), ( {},