scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 382

fix(passthrough): keep model and provider on mid-stream and Vertex failure logs

Browse source 
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
This commit is contained in:
mateo-berri 2026-09-05 01:01:17 -07:00
parent 3561ce241e
commit 006c166e5e
3 changed files with 74 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
View file

@ -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:

25
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -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,

73
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

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