diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 19d161b8aa4..0309bfabcaa 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2854,20 +2854,27 @@ 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 + 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") - 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 isinstance(p, str) and _is_llm_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/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/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 7a5d5a1e259..d449f62af2c 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 collections.abc import Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final @@ -9,6 +8,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.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( GeminiSseChunkParser, @@ -207,11 +207,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 5010de46fcb..39da57081c2 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 @@ -19,6 +19,8 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import from litellm.llms.vertex_ai.common_utils import ( get_vertex_ai_lyria_generation_cost, 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, @@ -47,6 +49,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 @@ -684,11 +688,16 @@ 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 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: @@ -734,18 +743,36 @@ class VertexPassthroughLoggingHandler: return None @staticmethod - def _get_custom_llm_provider(url: str, logging_obj: LiteLLMLoggingObj) -> str: - hostname: Final = urlparse(url).hostname - if logging_obj.model_call_details.get("custom_llm_provider") == litellm.LlmProviders.GEMINI.value or ( - hostname is not None - and ( - hostname == "generativelanguage.googleapis.com" - or hostname.endswith(".generativelanguage.googleapis.com") - ) + 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 is_gemini_api_host(url: str) -> bool: + hostname: Final = urlparse(url).hostname + return hostname is not None and ( + hostname == GEMINI_API_HOST_SUFFIX or hostname.endswith(f".{GEMINI_API_HOST_SUFFIX}") + ) + + @staticmethod + def custom_llm_provider_from_url(url: str) -> str: + if VertexPassthroughLoggingHandler.is_gemini_api_host(url): return litellm.LlmProviders.GEMINI.value return litellm.LlmProviders.VERTEX_AI.value + @staticmethod + def _get_custom_llm_provider(url: str, logging_obj: LiteLLMLoggingObj) -> str: + if logging_obj.model_call_details.get("custom_llm_provider") == litellm.LlmProviders.GEMINI.value: + return litellm.LlmProviders.GEMINI.value + return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url) + @staticmethod def _is_multimodal_embedding_response(json_response: dict) -> bool: """ diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e1e8864678e..2e6e352ba83 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -72,6 +72,7 @@ from litellm.llms.base_llm.managed_resources.utils import ( ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.oss_decision import validate_oss_request +from litellm.llms.vertex_ai.common_utils import get_vertex_model_id_from_url from litellm.passthrough import BasePassthroughUtils from litellm.proxy._lazy_features import lazy_owned_routes from litellm.proxy._types import ( @@ -144,6 +145,7 @@ from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( is_tinyfish_agent_url, ) +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 ( @@ -846,12 +848,36 @@ 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 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 + + +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, 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``. @@ -868,10 +894,13 @@ 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 "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider + 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_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 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) @@ -1548,6 +1577,7 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, ), logging_obj=logging_obj, ) @@ -1586,6 +1616,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, @@ -1640,6 +1671,7 @@ async def pass_through_request( logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, + url=url, ), logging_obj=logging_obj, ) @@ -1678,6 +1710,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, @@ -1740,6 +1773,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( @@ -1960,19 +1994,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/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index bff89b6cbf1..33d920c9378 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -238,6 +238,60 @@ 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({"model": "gemini-3.8-flash", "custom_llm_provider": "vertex_ai"}) == "vertex_ai" + assert ( + extract( + { + "custom_llm_provider": "gemini", + "standard_logging_object": {"custom_llm_provider": "vertex_ai"}, + } + ) + == "vertex_ai" + ) + + +@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/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 db0a74d360a..d2819b9eda3 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 @@ -38,12 +38,16 @@ from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router +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, SafeRouteAdder, + _build_passthrough_failure_request_payload, _registered_pass_through_routes, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -429,53 +433,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 @@ -8442,6 +8411,279 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) +_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" +) +_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" +) +_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( + 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" + proxy_url = httpx.URL(target.replace("https://", "http://test-proxy.com/")) + mock_request.url = proxy_url + mock_request.scope = {"type": "http", "method": "POST", "path": proxy_url.path} + 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: + received = 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, + ) + 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 + + +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), + ) + + +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", + [ + (_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"), + (_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( + 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 +@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( + 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" + + +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("stream", [True, False]) +@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, stream: bool +): + failure = await _run_passthrough_failure( + 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" + assert failure["request_data"]["custom_llm_provider"] == expected_provider + + +@pytest.mark.parametrize( + "parsed_body, url, custom_llm_provider, expected_model, expected_provider", + [ + ({"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": []}, _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"), + ( + {}, + _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": []}, + _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": []}, _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), + ], +) +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, + 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 + 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 @pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"]) def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str): """A provider route that resolved a router deployment stashes its model_info on request.state. That