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