scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 117

Merge pull request #41307 from BerriAI/litellm_passthrough_deployment_model_id

Browse source 
fix(passthrough): attribute Vertex passthrough successes to the resolved router deployment
This commit is contained in:
Yassin Kortam 2026-09-15 15:44:38 -07:00 • committed by GitHub
commit 94a6dfaf52
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 171 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -77,6 +77,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
from litellm.secret_managers.main import get_secret_str, str_to_bool
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
@ -1322,7 +1323,7 @@ def _resolve_vertex_model_from_router(
endpoint: str,
vertex_project: str | None,
vertex_location: str | None,
) -> tuple[str, str, str | None, str | None]:
) -> tuple[str, str, str | None, str | None, Mapping[str, object] | None]:
"""
Resolve Vertex AI model configuration from router.
@ -1335,18 +1336,21 @@ def _resolve_vertex_model_from_router(
vertex_location: Current vertex location (may be from URL)
Returns:
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location)
with resolved values from router config
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info)
with resolved values from router config; deployment_model_info is the resolved
deployment's `model_info`, or None when no deployment matched
"""
if not llm_router:
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
try:
deployment: Final = llm_router.get_available_deployment_for_pass_through(model=model_id)
if not deployment:
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
litellm_params: Final = deployment.get("litellm_params", {})
model_info: Final = deployment.get("model_info")
deployment_model_info: Final = model_info if isinstance(model_info, Mapping) else None
# Always override with router config values (they take precedence over URL values)
config_vertex_project: Final = litellm_params.get("vertex_project")
@ -1387,10 +1391,11 @@ def _resolve_vertex_model_from_router(
encoded_endpoint = encoded_endpoint.replace(model_id, actual_model)
endpoint = endpoint.replace(model_id, actual_model)
return encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info
except Exception as e:
verbose_proxy_logger.debug("Error resolving vertex model from router for model %s: %s", model_id, e)
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
def _is_bedrock_agent_runtime_route(endpoint: str) -> bool:
@ -2134,6 +2139,7 @@ async def _base_vertex_proxy_route(
endpoint,
vertex_project,
vertex_location,
deployment_model_info,
) = _resolve_vertex_model_from_router(
model_id=model_id,
llm_router=llm_router,
@ -2142,6 +2148,8 @@ async def _base_vertex_proxy_route(
vertex_project=vertex_project,
vertex_location=vertex_location,
)
if deployment_model_info:
setattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, deployment_model_info)
vertex_credentials: Final = passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,

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

@ -96,6 +96,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
EndpointType,
@ -613,6 +614,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
_metadata.update(
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
)
_request_state: Final = getattr(request, "state", None)
deployment_model_info: Final = getattr(
_request_state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, None
)
if isinstance(deployment_model_info, Mapping):
_metadata["model_info"] = dict(deployment_model_info)
kwargs: Final = {
"litellm_params": {
@ -2002,6 +2009,8 @@ def create_pass_through_route(
delattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY):
delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY)
if hasattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY):
delattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY)
# The upstream withholds its response headers until its first token, so
# the whole time-to-first-token is spent inside _relay with nothing on

3
litellm/types/passthrough_endpoints/pass_through_endpoints.py
View file

@ -11,6 +11,9 @@ LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY: Final = "litellm_pass_through_custom
# exact byte/string body, such as AWS SigV4-signed requests.
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY: Final = "litellm_pass_through_raw_body"
# `model_info` of the router deployment a provider route (e.g. Vertex) resolved for this request.
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: Final = "litellm_pass_through_deployment_model_info"
# Attribute set on the FastAPI endpoint function of every user-defined pass-through
# route. Auth reads it off the dispatched endpoint (``request.scope["endpoint"]``) to
# decide whether a request body ``model`` names an upstream model rather than a

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

@ -34,6 +34,7 @@ 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_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
@ -5934,6 +5935,32 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
)
@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
deployment, not a model_info the client put in its own body, is what spend logs and metrics attribute
the call to (LIT-1761: passthrough successes carried model_id="")."""
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
mock_request.headers = Headers({})
mock_request.scope = {}
mock_request.state = SimpleNamespace(
**{LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: {"id": "vertex-gemini-38-flash-dep"}}
)
kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=mock_request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
passthrough_logging_payload=MagicMock(),
logging_obj=MagicMock(),
_parsed_body={client_metadata_key: {"model_info": {"id": "client-forged-id"}}},
litellm_call_id="lit-1761-call-id",
)
assert kwargs["litellm_params"]["metadata"]["model_info"] == {"id": "vertex-gemini-38-flash-dep"}
@pytest.mark.asyncio
async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_error_for_an_unknown_model(
monkeypatch: pytest.MonkeyPatch,

118
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py
View file

@ -1,13 +1,19 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import Request
from starlette.datastructures import Headers, State
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
VertexAIPassThroughHandler,
_base_vertex_proxy_route,
_resolve_vertex_model_from_router,
_upstream_headers_for_vertex_route,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
)
from litellm.types.router import DeploymentTypedDict
@ -758,3 +764,115 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
assert (
"gemini-3-pro" in target_url
), f"Actual Vertex AI model name should be in target URL. Got: {target_url}"
@pytest.mark.asyncio
async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment():
"""The router deployment that rewrote the upstream URL is the one the logging kwargs must name, so
the Prometheus model_id label (and SpendLogs.model_id) on a Vertex passthrough success reads the
deployment's id instead of "" (LIT-1761)."""
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
mock_request.headers = Headers({})
mock_request.scope = {}
mock_request.state = State()
mock_handler = MagicMock()
mock_handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
mock_router = MagicMock()
mock_router.get_available_deployment_for_pass_through.return_value = {
"model_name": "gemini-3.8-flash",
"litellm_params": {
"model": "vertex_ai/gemini-3.8-flash",
"vertex_project": "p",
"vertex_location": "global",
"use_in_pass_through": True,
},
"model_info": {"id": "vertex-gemini-38-flash-dep"},
}
async def relay_returning_logging_kwargs(
request: Request, fastapi_response: object, user_api_key_dict: UserAPIKeyAuth
) -> dict:
return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=request,
user_api_key_dict=user_api_key_dict,
passthrough_logging_payload=MagicMock(),
logging_obj=MagicMock(),
_parsed_body={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
litellm_call_id="lit-1761-call-id",
)
with (
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
"litellm.proxy.proxy_server.llm_router", mock_router
),
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
) as mock_pt_router,
patch( # test-quality-ok: the route offers no injection point for its header preparation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
new_callable=AsyncMock,
return_value=({}, False, "p", "global"),
),
patch( # test-quality-ok: the relay is captured here to read the logging kwargs, the route offers no seam
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
return_value=relay_returning_logging_kwargs,
),
patch( # test-quality-ok: the route calls auth directly rather than through Depends
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
return_value=UserAPIKeyAuth(api_key="hashed-key"),
),
):
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
logging_kwargs = await _base_vertex_proxy_route(
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
request=mock_request,
fastapi_response=MagicMock(),
get_vertex_pass_through_handler=mock_handler,
)
assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep"
def _router_without_deployment() -> MagicMock:
router = MagicMock()
router.get_available_deployment_for_pass_through.return_value = None
return router
def _router_raising_on_lookup() -> MagicMock:
router = MagicMock()
router.get_available_deployment_for_pass_through.side_effect = ValueError("no healthy deployment")
return router
@pytest.mark.parametrize(
"llm_router",
[None, _router_without_deployment(), _router_raising_on_lookup()],
ids=["no-router", "no-matching-deployment", "lookup-raises"],
)
def test_vertex_passthrough_without_a_resolved_deployment_keeps_the_url_and_reports_no_model_info(
llm_router: MagicMock | None,
):
"""A Vertex passthrough call that no router deployment serves must keep the URL-derived values and carry no
deployment model_info, so logging cannot attribute it to a deployment that never handled it."""
resolved = _resolve_vertex_model_from_router(
model_id="gemini-3.8-flash",
llm_router=llm_router,
encoded_endpoint="/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
vertex_project="url-project",
vertex_location="url-location",
)
assert resolved == (
"/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
"v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
"url-project",
"url-location",
None,
)