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