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