mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(passthrough): log the URL model and declared provider on Gemini and Vertex failures (#49075)
* fix(passthrough): log the URL model and declared provider on Gemini and Vertex failures Pass-through failures on Gemini and Vertex routes carry the model only in the URL, so the failure spend row and the Prometheus failure counter logged an empty model. The failure payload now falls back to the model parsed from the URL on those routes, and the Prometheus provider label honors the provider the request declares before inferring one from the model name * fix(passthrough): keep model and provider on mid-stream and Vertex failure logs 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 * 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. * fix(passthrough): name Google failures from the URL on untracked methods too A Gemini or Vertex pass-through failure on a method the success handler does not track (countTokens, for one) still logged model "" and no provider, although the model sits in the URL. The failure payload now falls back to the Google AI host for the provider and reads /models/<name> from the URL, so countTokens 400s name their model and provider like generateContent ones do. Non-Google hosts keep the tracked-route rule and still log nothing extra. * fix(passthrough): trust the declared Google provider and Vertex-shaped paths on custom bases The failure path only named the Google provider from the URL when the method was one the success handler tracks or the host was a Google API host, so a Gemini countTokens failure behind a custom GEMINI_API_BASE (or a Vertex one behind a custom base) still logged an empty model. A route that declares gemini or vertex_ai now wins outright, and a URL whose path carries /projects/<p>/locations/<l>/ counts as Vertex on any host, while a tracked method on a non-Google host no longer passes as Vertex. * fix(pass_through): stop the URL model parser at the query string so a configured key never lands in a spend row's model * test(pass_through): run the URL-model failure tests through the real ProxyLogging with a recording callback * fix(passthrough): scope declared provider to pass-through failures and match Google hosts only Prometheus trusts the resolved provider as before and reads the request-declared provider only on pass-through failures. Vertex attribution matches Google hosts, strips operation poll suffixes, prefixes vertex_ai/ only for rawPredict, reads the URL path only, covers Vertex Live WebSocket failures, and labels a router alias failure with the alias the client requested * fix(passthrough): attribute failures from the URL only on built-in Google routes and keep other routes' metric labels A user-defined pass-through route to a Google host keeps its empty model and provider, since its URL names an upstream model rather than a LiteLLM-managed one. Prometheus failure labels read the route's provider and model_group only for gemini and vertex_ai, so Azure, OpenRouter, Bedrock and other pass-through routes keep their earlier labels. The Gemini URL parser drops the query string, Google host matching needs a dot or dash boundary, and the Vertex Claude prefix follows alt=sse like the success handler. * fix(passthrough): log the URL model on guardrail blocks and the named model on Vertex Live failures A pre-call guardrail block on a pass-through route now builds its failure payload the same way upstream failures do, so the model and provider read from the URL reach spend logs and metrics. Vertex Live websocket failures (connect errors, upstream rejections, and the credential failure) now carry the model from the ?model= query parameter under the vertex_ai provider * fix(passthrough): mark guardrail-block and Vertex Live credential failures as pass-through calls Prometheus reads a failure's declared provider only on pass-through calls, so without call_type it fell back to guessing from the model name. That labeled a guardrail block on the Gemini route as vertex_ai and the Live credential failure as None * fix(passthrough): let the route name a failure's provider and give guardrail blocks their own spend row * fix(pass-through): keep URL model attribution off routes that name another provider A built-in route that names a non-Google provider, like /openai pointed at Gemini's OpenAI-compatible host, no longer takes its failure model from the Google URL. A guardrail block response now returns the x-litellm-call-id of its spend row. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tessellark
GitHub
parent
b195d15b93
commit
a084d13635
10 changed files with 1025 additions and 174 deletions
|
|
@ -82,6 +82,7 @@ from litellm.types.proxy.carried_budget_state import (
|
|||
UserBudgetSnapshot,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
StandardLoggingGuardrailInformation,
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingZeroCostDiagnostic,
|
||||
|
|
@ -279,6 +280,11 @@ def _logged_fields(value: object) -> Mapping[str, object]:
|
|||
return _NO_LOGGED_FIELDS
|
||||
|
||||
|
||||
_GOOGLE_PASS_THROUGH_PROVIDERS: Final = frozenset(
|
||||
(litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value)
|
||||
)
|
||||
|
||||
|
||||
def _is_llm_provider(name: str) -> bool:
|
||||
routable: Final = (
|
||||
name in litellm.provider_list or JSONProviderRegistry.exists(name) or is_registered_custom_provider(name)
|
||||
|
|
@ -2942,6 +2948,22 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _google_pass_through_provider(request_data: Mapping[str, object]) -> str | None:
|
||||
declared_provider: Final = request_data.get("custom_llm_provider")
|
||||
if request_data.get("call_type") != CallTypes.pass_through.value or not isinstance(declared_provider, str):
|
||||
return None
|
||||
return declared_provider if declared_provider in _GOOGLE_PASS_THROUGH_PROVIDERS else None
|
||||
|
||||
@staticmethod
|
||||
def _failed_request_requested_model(request_data: Mapping[str, object]) -> object:
|
||||
requested_model: Final = request_data.get("model", "")
|
||||
if PrometheusLogger._google_pass_through_provider(request_data) is None:
|
||||
return requested_model
|
||||
request_metadata: Final = _logged_fields(_logged_fields(request_data.get("litellm_params")).get("metadata"))
|
||||
model_group: Final = request_metadata.get("model_group")
|
||||
return model_group if isinstance(model_group, str) and model_group else requested_model
|
||||
|
||||
@staticmethod
|
||||
def _extract_api_provider_from_request_data(request_data: dict) -> str | None:
|
||||
"""
|
||||
|
|
@ -2950,22 +2972,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`` ->
|
||||
``openai``) since the proxy's failure ``request_data`` usually carries
|
||||
only the client-supplied model. Return ``None`` when it cannot be
|
||||
determined so the label emits empty rather than a guess.
|
||||
``standard_logging_object`` (e.g. a stream that broke mid-flight). On a
|
||||
Gemini or Vertex pass-through failure only, next take the provider the
|
||||
failure payload names (``gemini`` or ``vertex_ai``), so every other
|
||||
pass-through route keeps its earlier label. 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 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
|
||||
resolved_provider: Final = litellm_params.get("custom_llm_provider") or standard_logging_object.get(
|
||||
"custom_llm_provider"
|
||||
)
|
||||
if resolved_provider:
|
||||
return resolved_provider
|
||||
google_provider: Final = PrometheusLogger._google_pass_through_provider(request_data)
|
||||
if google_provider is not None:
|
||||
return google_provider
|
||||
model: Final = litellm_params.get("model") or request_data.get("model")
|
||||
if not model:
|
||||
if not isinstance(model, str) or not model:
|
||||
return None
|
||||
try:
|
||||
return litellm.get_llm_provider(model=model)[1] or None
|
||||
|
|
@ -3052,7 +3079,7 @@ class PrometheusLogger(CustomLogger):
|
|||
team_alias=user_api_key_dict.team_alias,
|
||||
org_id=user_api_key_dict.org_id,
|
||||
org_alias=user_api_key_dict.organization_alias,
|
||||
requested_model=_bounded_requested_model_label(request_data.get("model", "")),
|
||||
requested_model=_bounded_requested_model_label(self._failed_request_requested_model(request_data)),
|
||||
status_code=str(status_code),
|
||||
exception_status=str(status_code),
|
||||
exception_class=self._get_exception_class_name(original_exception),
|
||||
|
|
|
|||
|
|
@ -1257,7 +1257,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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm import get_llm_provider
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS,
|
||||
AWS_ERROR_TYPE_HEADER,
|
||||
|
|
@ -149,7 +150,7 @@ from litellm.types.passthrough_endpoints.tinyfish import (
|
|||
)
|
||||
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
|
||||
from litellm.types.router import LiteLLMParamsTypedDict
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.types.utils import CallTypes, LlmProviders
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
|
|
@ -4105,7 +4106,12 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data={},
|
||||
request_data={
|
||||
"model": model or "",
|
||||
"custom_llm_provider": litellm.LlmProviders.VERTEX_AI.value,
|
||||
"call_type": CallTypes.pass_through.value,
|
||||
"litellm_call_id": str(uuid.uuid4()),
|
||||
},
|
||||
)
|
||||
if websocket.client_state != WebSocketState.DISCONNECTED:
|
||||
await websocket.close(code=1011, reason=VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON)
|
||||
|
|
@ -4140,6 +4146,8 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
builtin_provider_route=True,
|
||||
model=model,
|
||||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project=resolved_project,
|
||||
vertex_location=resolved_location,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import re
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast
|
||||
from urllib.parse import urlparse
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
|
@ -19,6 +19,7 @@ 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,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator as VertexModelResponseIterator,
|
||||
|
|
@ -47,6 +48,19 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
vertex_search_api_config: Final = VertexSearchAPIVectorStoreConfig()
|
||||
GEMINI_API_HOST_SUFFIX: Final = "generativelanguage.googleapis.com"
|
||||
GOOGLE_AI_HOST_SUFFIXES: Final = (
|
||||
GEMINI_API_HOST_SUFFIX,
|
||||
"aiplatform.googleapis.com",
|
||||
"discoveryengine.googleapis.com",
|
||||
)
|
||||
_VERTEX_MULTI_REGION_HOST: Final = re.compile(r"aiplatform\.[a-z0-9-]+\.rep\.googleapis\.com")
|
||||
|
||||
|
||||
def _is_host_or_subdomain(hostname: str, suffix: str) -> bool:
|
||||
return hostname == suffix or hostname.endswith((f".{suffix}", f"-{suffix}"))
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
|
|
@ -684,11 +698,18 @@ 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:
|
||||
parsed_url: Final = urlparse(url_route)
|
||||
model: Final = get_vertex_model_id_from_url(parsed_url.path)
|
||||
if model is None:
|
||||
return "vertex_ai/search_api" if ":search" in parsed_url.path else None
|
||||
streamed: Final = "sse" in parse_qs(parsed_url.query).get("alt", ())
|
||||
if parsed_url.path.endswith((":rawPredict", ":streamRawPredict")) and not streamed:
|
||||
return f"vertex_ai/{model}"
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def extract_model_name_from_vertex_path(vertex_model_path: str) -> str:
|
||||
|
|
@ -734,18 +755,32 @@ class VertexPassthroughLoggingHandler:
|
|||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_custom_llm_provider(url: str, logging_obj: LiteLLMLoggingObj) -> str:
|
||||
def is_google_ai_url(url: str) -> bool:
|
||||
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")
|
||||
)
|
||||
):
|
||||
return hostname is not None and (
|
||||
any(_is_host_or_subdomain(hostname, suffix) for suffix in GOOGLE_AI_HOST_SUFFIXES)
|
||||
or _VERTEX_MULTI_REGION_HOST.fullmatch(hostname) 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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import asyncio
|
|||
import copy
|
||||
import json
|
||||
import posixpath
|
||||
import re
|
||||
import traceback
|
||||
from base64 import b64encode
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence
|
||||
|
|
@ -77,6 +78,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.eu_jev.common_utils import eu_jev_scoped_model, eu_jev_upstream_body
|
||||
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 (
|
||||
|
|
@ -147,11 +149,12 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
EndpointType,
|
||||
PassthroughStandardLoggingPayload,
|
||||
)
|
||||
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage
|
||||
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, 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 (
|
||||
|
|
@ -874,13 +877,45 @@ 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 is not None:
|
||||
return custom_llm_provider if custom_llm_provider in _GOOGLE_PASSTHROUGH_PROVIDERS else None
|
||||
if VertexPassthroughLoggingHandler.is_google_ai_url(url_route):
|
||||
return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route)
|
||||
return None
|
||||
|
||||
|
||||
_OPERATION_POLL_SUFFIX: Final = re.compile(r"/operations/[^/]+$")
|
||||
|
||||
|
||||
def _without_operation_poll_suffix(model: str | None) -> str | None:
|
||||
return _OPERATION_POLL_SUFFIX.sub("", model) if model else model
|
||||
|
||||
|
||||
def _passthrough_google_model(url_route: str, google_provider: str | None) -> str | None:
|
||||
match google_provider:
|
||||
case "gemini":
|
||||
return _without_operation_poll_suffix(get_vertex_model_id_from_url(urlparse(url_route).path))
|
||||
case "vertex_ai":
|
||||
return _without_operation_poll_suffix(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,
|
||||
upstream_model: str | None = None,
|
||||
builtin_provider_route: bool = False,
|
||||
) -> dict:
|
||||
"""Build the ``request_data`` dict passed to ``post_call_failure_hook``.
|
||||
|
||||
|
|
@ -893,16 +928,25 @@ def _build_passthrough_failure_request_payload(
|
|||
comes paired with a usage object, so both keys are written together.
|
||||
``upstream_model`` is the model name the upstream reported, which replaces
|
||||
the body's ``model`` on the failure row just as it does on a success row.
|
||||
``builtin_provider_route`` lets a Gemini or Vertex URL name the failure's
|
||||
model and provider. A user-defined route never gets that, since its URL
|
||||
names an upstream model rather than a LiteLLM-managed one.
|
||||
"""
|
||||
request_payload: Final[dict] = dict(cast(Mapping[str, object], parsed_body or {})) # cast-ok: json body
|
||||
if kwargs:
|
||||
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 builtin_provider_route else None
|
||||
)
|
||||
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
|
||||
request_payload.pop("custom_llm_provider", None)
|
||||
if 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)
|
||||
|
|
@ -1203,6 +1247,7 @@ async def pass_through_request(
|
|||
# Initialize variables
|
||||
#########################################################
|
||||
litellm_call_id: Final = str(uuid.uuid4())
|
||||
builtin_provider_route: Final = not request_dispatched_to_pass_through_endpoint(request)
|
||||
url: httpx.URL | None = None
|
||||
|
||||
# parsed request body
|
||||
|
|
@ -1612,6 +1657,8 @@ async def pass_through_request(
|
|||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
upstream_usage=upstream_usage,
|
||||
url=url,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
upstream_model=upstream_model,
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -1651,6 +1698,8 @@ async def pass_through_request(
|
|||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
url=url,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
),
|
||||
),
|
||||
managed_id_provider=_managed_id_provider,
|
||||
|
|
@ -1710,6 +1759,8 @@ async def pass_through_request(
|
|||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
upstream_usage=upstream_usage,
|
||||
url=url,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
upstream_model=upstream_model,
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -1749,6 +1800,8 @@ async def pass_through_request(
|
|||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
url=url,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
),
|
||||
),
|
||||
managed_id_provider=_managed_id_provider,
|
||||
|
|
@ -1811,6 +1864,8 @@ async def pass_through_request(
|
|||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
upstream_usage=upstream_usage,
|
||||
url=url,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
upstream_model=upstream_model,
|
||||
)
|
||||
failure_request_payload["response_body"] = response_body
|
||||
|
|
@ -1991,7 +2046,14 @@ async def pass_through_request(
|
|||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data=e.request_data,
|
||||
request_data=_build_passthrough_failure_request_payload(
|
||||
parsed_body=e.request_data,
|
||||
kwargs={"call_type": CallTypes.pass_through.value, "litellm_call_id": litellm_call_id},
|
||||
logging_obj=None,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
url=url,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
),
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -2010,6 +2072,7 @@ async def pass_through_request(
|
|||
content=json.dumps(error_body),
|
||||
status_code=200,
|
||||
media_type="application/json",
|
||||
headers={"x-litellm-call-id": litellm_call_id},
|
||||
)
|
||||
except Exception as e:
|
||||
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
|
|
@ -2033,19 +2096,14 @@ 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,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
)
|
||||
_carry_guardrail_logging_info(request_payload, post_call_guardrail_data)
|
||||
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
|
|
@ -2547,6 +2605,8 @@ async def websocket_passthrough_request(
|
|||
cost_per_request: float | None = None,
|
||||
accept_websocket: bool = True,
|
||||
setup_model_rewriter: Callable[[str], str] | None = None,
|
||||
builtin_provider_route: bool = False,
|
||||
model: str | None = None,
|
||||
):
|
||||
"""
|
||||
WebSocket passthrough request handler.
|
||||
|
|
@ -2560,6 +2620,8 @@ async def websocket_passthrough_request(
|
|||
endpoint: The endpoint path (for logging purposes)
|
||||
cost_per_request: Optional field - cost per request to the target endpoint
|
||||
setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
|
||||
builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider
|
||||
model: The model the caller named before connecting, logged on a failure before the setup frame arrives
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj
|
||||
|
|
@ -2912,19 +2974,17 @@ async def websocket_passthrough_request(
|
|||
except InvalidStatus as exc:
|
||||
verbose_proxy_logger.exception("WebSocket passthrough (%s): upstream rejected WebSocket connection", endpoint)
|
||||
|
||||
# Prepare request payload for logging
|
||||
request_payload = {}
|
||||
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
|
||||
|
||||
# Log the connection failure using the same pattern as HTTP
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=exc,
|
||||
request_data=request_payload,
|
||||
request_data=_build_passthrough_failure_request_payload(
|
||||
parsed_body=None if model is None else {"model": model},
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=None,
|
||||
url=httpx.URL(target),
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
),
|
||||
traceback_str=traceback.format_exc(
|
||||
limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
|
||||
),
|
||||
|
|
@ -2940,19 +3000,17 @@ async def websocket_passthrough_request(
|
|||
"WebSocket passthrough (%s): unexpected error while proxying WebSocket", endpoint
|
||||
)
|
||||
|
||||
# Prepare request payload for logging
|
||||
request_payload = {}
|
||||
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
|
||||
|
||||
# Log the unexpected error using the same pattern as HTTP
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data=request_payload,
|
||||
request_data=_build_passthrough_failure_request_payload(
|
||||
parsed_body=None if model is None else {"model": model},
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=None,
|
||||
url=httpx.URL(target),
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
),
|
||||
traceback_str=traceback.format_exc(
|
||||
limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -238,6 +238,53 @@ def test_extract_api_provider_from_request_data_failure_path():
|
|||
assert extract({"model": "some-unknown-model-xyz"}) is None
|
||||
|
||||
|
||||
_PASS_THROUGH_CALL_TYPE: Final = "pass_through_endpoint"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_data, expected_provider",
|
||||
[
|
||||
({"model": "gpt-4o-mini", "call_type": "acompletion", "custom_llm_provider": "anthropic"}, "openai"),
|
||||
({"model": "gpt-4o-mini", "litellm_params": {"custom_llm_provider": "humanloop"}}, "humanloop"),
|
||||
({"model": "gemini-3.8-flash", "call_type": _PASS_THROUGH_CALL_TYPE}, "vertex_ai"),
|
||||
(
|
||||
{"model": "gemini-3.8-flash", "call_type": _PASS_THROUGH_CALL_TYPE, "custom_llm_provider": "gemini"},
|
||||
"gemini",
|
||||
),
|
||||
(
|
||||
{
|
||||
"call_type": _PASS_THROUGH_CALL_TYPE,
|
||||
"custom_llm_provider": "gemini",
|
||||
"standard_logging_object": {"custom_llm_provider": "vertex_ai"},
|
||||
},
|
||||
"vertex_ai",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_extract_api_provider_reads_the_request_declared_provider_only_on_pass_through_failures(
|
||||
request_data: dict[str, object], expected_provider: str
|
||||
):
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
assert PrometheusLogger._extract_api_provider_from_request_data(request_data) == expected_provider
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"declared_provider",
|
||||
["azure", "bedrock", "openrouter", "client-supplied-value", "", {"nested": "object"}, 42],
|
||||
)
|
||||
def test_extract_api_provider_ignores_pass_through_declared_values_outside_gemini_and_vertex(
|
||||
declared_provider: object,
|
||||
):
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
extract = PrometheusLogger._extract_api_provider_from_request_data
|
||||
declared = {"call_type": _PASS_THROUGH_CALL_TYPE, "custom_llm_provider": declared_provider}
|
||||
|
||||
assert extract({"model": "gpt-4o-mini", **declared}) == "openai"
|
||||
assert extract({"model": "no-such-model", **declared}) is None
|
||||
|
||||
|
||||
def test_extract_api_provider_swallows_unknown_model_but_logs_unexpected_errors():
|
||||
"""
|
||||
get_llm_provider raises BadRequestError for a model that maps to no
|
||||
|
|
@ -849,6 +896,56 @@ async def test_failure_hook_prefers_request_data_provider_over_exception_provide
|
|||
) == ["azure"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call_type, custom_llm_provider, expected_requested_model",
|
||||
[
|
||||
(_PASS_THROUGH_CALL_TYPE, "vertex_ai", "gcp/google/gemini-2.5-flash"),
|
||||
("acompletion", "vertex_ai", "other"),
|
||||
(_PASS_THROUGH_CALL_TYPE, "bedrock", "other"),
|
||||
],
|
||||
)
|
||||
async def test_failure_hook_labels_a_pass_through_failure_with_the_router_model_the_client_requested(
|
||||
call_type: str, custom_llm_provider: str, expected_requested_model: str
|
||||
):
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gcp/google/gemini-2.5-flash",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-flash",
|
||||
"vertex_project": "p",
|
||||
"vertex_location": "us-central1",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
_clear_prometheus_registry()
|
||||
try:
|
||||
with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
|
||||
await PrometheusLogger().async_post_call_failure_hook(
|
||||
request_data={
|
||||
"model": "gemini-2.5-flash",
|
||||
"call_type": call_type,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"metadata": {},
|
||||
"litellm_params": {"metadata": {"model_group": "gcp/google/gemini-2.5-flash"}},
|
||||
},
|
||||
original_exception=Exception("upstream 400"),
|
||||
user_api_key_dict=UserAPIKeyAuth(token="tok"),
|
||||
)
|
||||
assert {
|
||||
s.labels.get("requested_model") for s in _collected_samples("litellm_proxy_failed_requests_metric_total")
|
||||
} == {expected_requested_model}
|
||||
finally:
|
||||
_clear_prometheus_registry()
|
||||
|
||||
|
||||
def test_model_group_in_deployment_metrics():
|
||||
"""
|
||||
Test that model_group label is present on the deployment-scoped metrics
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -25,12 +25,14 @@ from starlette.datastructures import FormData
|
|||
import litellm
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from tests._master_key import MASTER_KEY as SHARED_MASTER_KEY
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
from litellm.types.utils import CallTypes, CallTypesLiteral
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException
|
||||
from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.passthrough.main import raise_for_provider_status
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
|
|
@ -6490,6 +6492,55 @@ class TestVertexAILiveWebsocketPassthrough:
|
|||
assert "default_vertex_config" in close_kwargs["reason"]
|
||||
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")])
|
||||
async def test_credential_failure_logs_the_named_model_under_vertex_ai(
|
||||
self, monkeypatch: pytest.MonkeyPatch, model: str | None, expected_model: str
|
||||
):
|
||||
from litellm.proxy.pass_through_endpoints import (
|
||||
llm_passthrough_endpoints as passthrough_module,
|
||||
)
|
||||
|
||||
class _FailureRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.request_data: tuple[dict[str, object], ...] = ()
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
self.request_data = (*self.request_data, request_data)
|
||||
|
||||
recorder: Final = _FailureRecorder()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())
|
||||
)
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None)
|
||||
self._clear_vertex_env(monkeypatch)
|
||||
ensure_token: Final = AsyncMock(side_effect=Exception("Unable to find your credentials"))
|
||||
|
||||
with patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token):
|
||||
await passthrough_module.vertex_ai_live_websocket_passthrough(
|
||||
websocket=self._websocket(),
|
||||
model=model,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
assert [
|
||||
(
|
||||
d["model"],
|
||||
PrometheusLogger._extract_api_provider_from_request_data(d),
|
||||
get_spend_logs_id(CallTypes.pass_through.value, {}, d) is not None,
|
||||
)
|
||||
for d in recorder.request_data
|
||||
] == [(expected_model, "vertex_ai", True)]
|
||||
|
||||
|
||||
class TestAnthropicProxyRoute:
|
||||
"""The /anthropic passthrough route: custom auth headers must not clobber the
|
||||
|
|
|
|||
|
|
@ -30,20 +30,26 @@ from starlette.datastructures import UploadFile as StarletteUploadFile
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
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,
|
||||
|
|
@ -61,8 +67,11 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
|
|||
PassThroughEndpointLogging,
|
||||
)
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types import utils as types_utils
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CallTypes
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
|
||||
|
|
@ -140,9 +149,7 @@ def test_is_multipart():
|
|||
),
|
||||
],
|
||||
)
|
||||
def test_get_endpoint_type_classifies_anthropic_messages_routes(
|
||||
url: str, expected_endpoint_type: EndpointType
|
||||
) -> None:
|
||||
def test_get_endpoint_type_classifies_anthropic_messages_routes(url: str, expected_endpoint_type: EndpointType) -> None:
|
||||
assert HttpPassThroughEndpointHelpers.get_endpoint_type(url) == expected_endpoint_type
|
||||
|
||||
|
||||
|
|
@ -430,53 +437,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: Final = 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 "Request failed" in (failure.traceback_str or "")
|
||||
assert failure.request_data["contents"] == "not a list"
|
||||
assert failure.request_data["model"] == ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3656,9 +3628,7 @@ async def test_pass_through_request_preserves_target_query_without_client_query(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_pass_through_without_configured_headers_preserves_no_authorization(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
async def test_config_pass_through_without_configured_headers_preserves_no_authorization(tmp_path, monkeypatch):
|
||||
proxy = await _boot_db_backed_proxy(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
|
|
@ -3725,9 +3695,7 @@ async def test_pass_through_rpm_limit_is_sequential_and_scoped_per_key(tmp_path,
|
|||
proxy = await _boot_db_backed_proxy(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
config_pass_through_endpoints=[
|
||||
{"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True}
|
||||
],
|
||||
config_pass_through_endpoints=[{"path": "/rpm-rerank", "target": "http://upstream.test/rerank", "auth": True}],
|
||||
db_pass_through_endpoints=[],
|
||||
master_key=MASTER_KEY,
|
||||
)
|
||||
|
|
@ -4902,8 +4870,10 @@ from fastapi import HTTPException as _FastAPIHTTPException
|
|||
from litellm.exceptions import (
|
||||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
ModifyResponseException,
|
||||
)
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
||||
_PT_MODULE = "litellm.proxy.pass_through_endpoints.pass_through_endpoints"
|
||||
|
||||
|
||||
|
|
@ -8182,8 +8152,10 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata
|
|||
}
|
||||
)
|
||||
body: Final = {
|
||||
"model": "upstream-only-model", metadata_slot: {
|
||||
"model_group": "managed-model", "customer_label": "retained",
|
||||
"model": "upstream-only-model",
|
||||
metadata_slot: {
|
||||
"model_group": "managed-model",
|
||||
"customer_label": "retained",
|
||||
"user_api_key_team_model_max_budget": budget,
|
||||
},
|
||||
}
|
||||
|
|
@ -8191,48 +8163,84 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata
|
|||
assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model")
|
||||
start: Final = datetime.now()
|
||||
logging_obj: Final = LiteLLMLoggingObj(
|
||||
model="upstream-only-model", messages=[], stream=False, call_type="pass_through_endpoint",
|
||||
start_time=start, litellm_call_id="custom-budget", function_id="custom-budget", kwargs={},
|
||||
model="upstream-only-model",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="pass_through_endpoint",
|
||||
start_time=start,
|
||||
litellm_call_id="custom-budget",
|
||||
function_id="custom-budget",
|
||||
kwargs={},
|
||||
dynamic_async_success_callbacks=[limiter],
|
||||
)
|
||||
payload: Final = {
|
||||
"url": "https://upstream.test/echo", "request_body": body, "request_method": "POST", "cost_per_request": 0.25,
|
||||
"url": "https://upstream.test/echo",
|
||||
"request_body": body,
|
||||
"request_method": "POST",
|
||||
"cost_per_request": 0.25,
|
||||
}
|
||||
kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
|
||||
request=request, user_api_key_dict=auth, passthrough_logging_payload=payload, logging_obj=logging_obj,
|
||||
_parsed_body=body, litellm_call_id="custom-budget",
|
||||
request=request,
|
||||
user_api_key_dict=auth,
|
||||
passthrough_logging_payload=payload,
|
||||
logging_obj=logging_obj,
|
||||
_parsed_body=body,
|
||||
litellm_call_id="custom-budget",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="upstream-only-model", user="unknown", optional_params={},
|
||||
litellm_params=kwargs["litellm_params"], call_type="pass_through_endpoint",
|
||||
model="upstream-only-model",
|
||||
user="unknown",
|
||||
optional_params={},
|
||||
litellm_params=kwargs["litellm_params"],
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
response: Final = httpx.Response(
|
||||
200, request=httpx.Request("POST", "https://upstream.test/echo"), json={"ok": True},
|
||||
200,
|
||||
request=httpx.Request("POST", "https://upstream.test/echo"),
|
||||
json={"ok": True},
|
||||
)
|
||||
await PassThroughEndpointLogging().pass_through_async_success_handler(
|
||||
httpx_response=response, response_body={"ok": True}, request_body=body, logging_obj=logging_obj,
|
||||
url_route="https://upstream.test/echo", result=response.text, start_time=start, end_time=datetime.now(),
|
||||
cache_hit=False, **kwargs,
|
||||
httpx_response=response,
|
||||
response_body={"ok": True},
|
||||
request_body=body,
|
||||
logging_obj=logging_obj,
|
||||
url_route="https://upstream.test/echo",
|
||||
result=response.text,
|
||||
start_time=start,
|
||||
end_time=datetime.now(),
|
||||
cache_hit=False,
|
||||
**kwargs,
|
||||
)
|
||||
assert logging_obj.model_call_details["response_cost"] == 0.25
|
||||
assert await limiter.is_team_within_model_budget("shared-team", budget, None, "managed-model")
|
||||
metadata: Final = kwargs["litellm_params"]["metadata"]
|
||||
assert (metadata["model_group"], metadata["customer_label"]) == ("managed-model", "retained")
|
||||
assert metadata.keys().isdisjoint({
|
||||
"user_api_key_model_max_budget", "user_api_key_team_model_max_budget",
|
||||
"user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget",
|
||||
})
|
||||
assert metadata.keys().isdisjoint(
|
||||
{
|
||||
"user_api_key_model_max_budget",
|
||||
"user_api_key_team_model_max_budget",
|
||||
"user_api_key_user_model_max_budget",
|
||||
"user_api_key_end_user_model_max_budget",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_slot", ["metadata", "litellm_metadata"])
|
||||
def test_builtin_passthrough_pins_model_group_to_the_resolved_model(metadata_slot: str) -> None:
|
||||
request: Final = Request({
|
||||
"type": "http", "method": "POST", "path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent",
|
||||
"headers": [], "query_string": b"",
|
||||
})
|
||||
request: Final = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/gemini/v1beta/models/gemini-2.5-flash:generateContent",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
}
|
||||
)
|
||||
kwargs: Final = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
|
||||
request=request, user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"),
|
||||
passthrough_logging_payload=MagicMock(), logging_obj=MagicMock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(token="hash", user_id="u-1"),
|
||||
passthrough_logging_payload=MagicMock(),
|
||||
logging_obj=MagicMock(),
|
||||
_parsed_body={"contents": [], metadata_slot: {"model_group": "unbounded-client-choice"}},
|
||||
)
|
||||
assert kwargs["litellm_params"]["metadata"]["model_group"] == "gemini-2.5-flash"
|
||||
|
|
@ -8591,6 +8599,581 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
|
|||
)
|
||||
|
||||
|
||||
_GEMINI_FLASH_TARGET: Final = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash:generateContent"
|
||||
_VERTEX_FLASH_TARGET: Final = (
|
||||
"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: Final = (
|
||||
"https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5"
|
||||
"/publishers/anthropic/models/claude-opus-5:rawPredict"
|
||||
)
|
||||
_VERTEX_SEARCH_TARGET: Final = (
|
||||
"https://discoveryengine.googleapis.com/v1/projects/p/locations/global"
|
||||
"/dataStores/default/servingConfigs/default:search"
|
||||
)
|
||||
_RELAYED_GEMINI_COUNT_TOKENS_TARGET: Final = "http://127.0.0.1:30732/v1beta/models/gemini-3.8-flash:countTokens"
|
||||
_TRANSLATE_TARGET: Final = "https://translate.googleapis.com/v3/projects/p/locations/global:translateText"
|
||||
_NON_GOOGLE_VERTEX_SHAPED_TARGET: Final = "https://llm.example.com/v1/projects/a/locations/b/models/my-model:predict"
|
||||
_VERTEX_LIVE_TARGET: Final = (
|
||||
"wss://us-central1-aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
)
|
||||
_VEO_OPERATION_SUFFIX: Final = "/models/veo-3.0-generate-001/operations/op-8f2c1"
|
||||
_RELAYED_VERTEX_COUNT_TOKENS_TARGET: Final = (
|
||||
"http://127.0.0.1:30732/v1/projects/p/locations/us-east5/publishers/google/models/gemini-3.8-flash:countTokens"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RecordedFailure:
|
||||
request_data: dict[str, object]
|
||||
original_exception: Exception
|
||||
user_api_key_dict: UserAPIKeyAuth
|
||||
traceback_str: str | None
|
||||
|
||||
|
||||
class _FailureHookRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.failures: tuple[_RecordedFailure, ...] = ()
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
self.failures = (
|
||||
*self.failures,
|
||||
_RecordedFailure(request_data, original_exception, user_api_key_dict, traceback_str),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PassthroughRun:
|
||||
failure: _RecordedFailure
|
||||
response: Response | None
|
||||
|
||||
|
||||
async def _drained_pass_through_response(
|
||||
request: Request,
|
||||
target: str,
|
||||
custom_llm_provider: str | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
stream: bool,
|
||||
) -> Response | None:
|
||||
try:
|
||||
received: Final = await pass_through_request(
|
||||
request=request,
|
||||
target=target,
|
||||
custom_headers={},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
)
|
||||
if isinstance(received, StreamingResponse):
|
||||
async for _ in received.body_iterator:
|
||||
pass
|
||||
except (ProxyException, httpx.HTTPError):
|
||||
return None
|
||||
return received if isinstance(received, Response) else None
|
||||
|
||||
|
||||
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,
|
||||
dispatched_endpoint: Callable[..., object] | None = None,
|
||||
guardrail: CustomGuardrail | None = None,
|
||||
body: bytes = b'{"contents": "not a list"}',
|
||||
) -> _RecordedFailure:
|
||||
run: Final = await _run_passthrough(
|
||||
target, custom_llm_provider, upstream, stream, user_api_key_dict, dispatched_endpoint, guardrail, body
|
||||
)
|
||||
return run.failure
|
||||
|
||||
|
||||
async def _run_passthrough(
|
||||
target: str,
|
||||
custom_llm_provider: str | None,
|
||||
upstream: httpx.Response | Exception,
|
||||
stream: bool = False,
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
dispatched_endpoint: Callable[..., object] | None = None,
|
||||
guardrail: CustomGuardrail | None = None,
|
||||
body: bytes = b'{"contents": "not a list"}',
|
||||
) -> _PassthroughRun:
|
||||
recorder: Final = _FailureHookRecorder()
|
||||
with (
|
||||
pytest.MonkeyPatch.context() as litellm_patch,
|
||||
patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request; there is no injection seam
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())
|
||||
),
|
||||
respx.mock(assert_all_called=False) as upstream_mock,
|
||||
):
|
||||
litellm_patch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm_patch.setattr(litellm, "callbacks", [recorder] if guardrail is None else [guardrail, recorder])
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
upstream_route: Final = upstream_mock.route(method="POST")
|
||||
if isinstance(upstream, Exception):
|
||||
upstream_route.mock(side_effect=upstream)
|
||||
else:
|
||||
upstream_route.mock(return_value=upstream)
|
||||
mock_request: Final = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
proxy_url: Final = 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,
|
||||
**({"endpoint": dispatched_endpoint} if dispatched_endpoint is not None else {}),
|
||||
}
|
||||
mock_request.body = AsyncMock(return_value=body)
|
||||
mock_request.headers = Headers({"content-type": "application/json"})
|
||||
mock_request.query_params = QueryParams({})
|
||||
response: Final = await _drained_pass_through_response(
|
||||
mock_request,
|
||||
target,
|
||||
custom_llm_provider,
|
||||
user_api_key_dict or UserAPIKeyAuth(api_key="sk-test"),
|
||||
stream,
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
assert len(recorder.failures) == 1
|
||||
return _PassthroughRun(recorder.failures[0], response)
|
||||
|
||||
|
||||
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"),
|
||||
],
|
||||
)
|
||||
async def test_passthrough_upstream_error_logs_model_from_url(
|
||||
target: str, custom_llm_provider: str | None, stream: bool, expected_provider: str
|
||||
):
|
||||
failure: Final = 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",
|
||||
[
|
||||
_RELAYED_VERTEX_COUNT_TOKENS_TARGET,
|
||||
_TRANSLATE_TARGET,
|
||||
_NON_GOOGLE_VERTEX_SHAPED_TARGET,
|
||||
],
|
||||
)
|
||||
async def test_passthrough_upstream_error_to_a_host_outside_gemini_and_vertex_gets_no_url_attribution(target: str):
|
||||
failure: Final = await _run_passthrough_failure(target, None, _upstream_400(target))
|
||||
assert failure.request_data["model"] == ""
|
||||
assert failure.request_data.get("custom_llm_provider") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("target", [_GEMINI_FLASH_TARGET, _VERTEX_FLASH_TARGET])
|
||||
async def test_passthrough_upstream_error_on_a_user_defined_route_to_google_gets_no_url_attribution(target: str):
|
||||
failure: Final = await _run_passthrough_failure(
|
||||
target, None, _upstream_400(target), dispatched_endpoint=_marked_pass_through_endpoint()
|
||||
)
|
||||
assert failure.request_data["model"] == ""
|
||||
assert failure.request_data.get("custom_llm_provider") is None
|
||||
|
||||
|
||||
@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: Final = 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: Final = 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: Final = 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
|
||||
|
||||
|
||||
class _ModifyResponseOnPreCallGuardrail(CustomGuardrail):
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict[str, object],
|
||||
call_type: str,
|
||||
) -> None:
|
||||
self.raise_passthrough_exception(violation_message="blocked by policy", request_data=data, detection_info=None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"target, custom_llm_provider, body, expected_model, expected_provider",
|
||||
[
|
||||
(_GEMINI_FLASH_TARGET, "gemini", b'{"contents": "not a list"}', "gemini-3.8-flash", "gemini"),
|
||||
(_VERTEX_FLASH_TARGET, None, b'{"contents": "not a list"}', "gemini-3.8-flash", "vertex_ai"),
|
||||
(_VERTEX_FLASH_TARGET, None, b'{"custom_llm_provider": "gemini"}', "gemini-3.8-flash", "vertex_ai"),
|
||||
],
|
||||
)
|
||||
async def test_passthrough_guardrail_block_logs_model_from_url(
|
||||
target: str, custom_llm_provider: str | None, body: bytes, expected_model: str, expected_provider: str
|
||||
):
|
||||
guardrail: Final = _ModifyResponseOnPreCallGuardrail(
|
||||
guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True
|
||||
)
|
||||
failure: Final = await _run_passthrough_failure(
|
||||
target, custom_llm_provider, _upstream_400(target), guardrail=guardrail, body=body
|
||||
)
|
||||
assert isinstance(failure.original_exception, ModifyResponseException)
|
||||
assert failure.request_data["model"] == expected_model
|
||||
assert failure.request_data["custom_llm_provider"] == expected_provider
|
||||
assert PrometheusLogger._extract_api_provider_from_request_data(failure.request_data) == expected_provider
|
||||
assert get_spend_logs_id(CallTypes.pass_through.value, {}, failure.request_data) is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("body", [b'{"contents": "not a list"}', b'{"litellm_call_id": "client-chosen-id"}'])
|
||||
async def test_passthrough_guardrail_block_response_names_the_call_id_of_its_spend_row(body: bytes):
|
||||
guardrail: Final = _ModifyResponseOnPreCallGuardrail(
|
||||
guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True
|
||||
)
|
||||
run: Final = await _run_passthrough(
|
||||
_GEMINI_FLASH_TARGET, "gemini", _upstream_400(_GEMINI_FLASH_TARGET), guardrail=guardrail, body=body
|
||||
)
|
||||
assert isinstance(run.failure.original_exception, ModifyResponseException)
|
||||
assert run.response is not None
|
||||
assert run.response.headers["x-litellm-call-id"] == get_spend_logs_id(
|
||||
CallTypes.pass_through.value, {}, run.failure.request_data
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("body", [b'{"contents": "not a list"}', b'{"custom_llm_provider": "vertex_ai"}'])
|
||||
async def test_passthrough_guardrail_block_on_a_user_defined_route_gets_no_url_attribution(body: bytes):
|
||||
guardrail: Final = _ModifyResponseOnPreCallGuardrail(
|
||||
guardrail_name="block-policy", event_hook=GuardrailEventHooks.pre_call, default_on=True
|
||||
)
|
||||
failure: Final = await _run_passthrough_failure(
|
||||
_GEMINI_FLASH_TARGET,
|
||||
None,
|
||||
_upstream_400(_GEMINI_FLASH_TARGET),
|
||||
dispatched_endpoint=_marked_pass_through_endpoint(),
|
||||
guardrail=guardrail,
|
||||
body=body,
|
||||
)
|
||||
assert isinstance(failure.original_exception, ModifyResponseException)
|
||||
assert failure.request_data["model"] == ""
|
||||
assert "custom_llm_provider" not in failure.request_data
|
||||
|
||||
|
||||
def _websocket_connect_failure(kind: str) -> Exception:
|
||||
from websockets.datastructures import Headers as WebSocketHeaders
|
||||
from websockets.exceptions import InvalidStatus
|
||||
from websockets.http11 import Response as WebSocketHandshakeResponse
|
||||
|
||||
if kind == "rejected_handshake":
|
||||
return InvalidStatus(WebSocketHandshakeResponse(403, "Forbidden", WebSocketHeaders()))
|
||||
return OSError("connection refused")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure_kind", ["rejected_handshake", "unreachable_host"])
|
||||
@pytest.mark.parametrize(
|
||||
"target, builtin_provider_route, expected_provider",
|
||||
[
|
||||
(_VERTEX_LIVE_TARGET, True, "vertex_ai"),
|
||||
(_VERTEX_LIVE_TARGET, False, None),
|
||||
("wss://api.openai.com/v1/realtime?model=gpt-realtime", True, None),
|
||||
],
|
||||
)
|
||||
async def test_websocket_passthrough_connect_failure_logs_the_provider_of_the_upstream_url(
|
||||
failure_kind: str, target: str, builtin_provider_route: bool, expected_provider: str | None
|
||||
):
|
||||
failure: Final = await _run_websocket_connect_failure(failure_kind, target, builtin_provider_route)
|
||||
assert failure.request_data.get("custom_llm_provider") == expected_provider
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure_kind", ["rejected_handshake", "unreachable_host"])
|
||||
@pytest.mark.parametrize("model, expected_model", [("gemini-3.8-live", "gemini-3.8-live"), (None, "")])
|
||||
async def test_websocket_passthrough_connect_failure_logs_the_model_the_caller_named(
|
||||
failure_kind: str, model: str | None, expected_model: str
|
||||
):
|
||||
failure: Final = await _run_websocket_connect_failure(failure_kind, _VERTEX_LIVE_TARGET, True, model=model)
|
||||
assert failure.request_data["model"] == expected_model
|
||||
|
||||
|
||||
async def _run_websocket_connect_failure(
|
||||
failure_kind: str, target: str, builtin_provider_route: bool, model: str | None = None
|
||||
) -> _RecordedFailure:
|
||||
recorder: Final = _FailureHookRecorder()
|
||||
with (
|
||||
pytest.MonkeyPatch.context() as litellm_patch,
|
||||
patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside websocket_passthrough_request; there is no injection seam
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
|
||||
side_effect=_websocket_connect_failure(failure_kind),
|
||||
),
|
||||
):
|
||||
litellm_patch.setattr(litellm, "callbacks", [recorder])
|
||||
await websocket_passthrough_request(
|
||||
websocket=_client_websocket(_pending_receive),
|
||||
target=target,
|
||||
custom_headers={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
builtin_provider_route=builtin_provider_route,
|
||||
model=model,
|
||||
)
|
||||
assert len(recorder.failures) == 1
|
||||
return recorder.failures[0]
|
||||
|
||||
|
||||
@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", "", "anthropic"),
|
||||
({}, "https://generativelanguage.googleapis.com/v1beta/openai/models/gemini-3.8-flash", "openai", "", "openai"),
|
||||
({"custom_llm_provider": "from-body"}, _VERTEX_FLASH_TARGET, None, "gemini-3.8-flash", "vertex_ai"),
|
||||
({"custom_llm_provider": "gemini"}, _VERTEX_FLASH_TARGET, "anthropic", "", "anthropic"),
|
||||
({"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?alt=sse"),
|
||||
None,
|
||||
"claude-opus-5",
|
||||
"vertex_ai",
|
||||
),
|
||||
({}, "https://generativelanguage.googleapis.com/v1beta/models?pageToken=/models/x", "gemini", "", "gemini"),
|
||||
(
|
||||
{"contents": []},
|
||||
"https://evilaiplatform.googleapis.com/v1/projects/p/locations/l"
|
||||
"/publishers/google/models/gemini-3.8-flash:generateContent",
|
||||
None,
|
||||
"",
|
||||
None,
|
||||
),
|
||||
({}, _VERTEX_SEARCH_TARGET, None, "vertex_ai/search_api", "vertex_ai"),
|
||||
({}, f"{_VERTEX_FLASH_TARGET}?tag=rawPredict", None, "gemini-3.8-flash", "vertex_ai"),
|
||||
(
|
||||
{},
|
||||
"https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs"
|
||||
"?filter=:search",
|
||||
None,
|
||||
"",
|
||||
"vertex_ai",
|
||||
),
|
||||
(
|
||||
{},
|
||||
f"https://generativelanguage.googleapis.com/v1beta{_VEO_OPERATION_SUFFIX}",
|
||||
"gemini",
|
||||
"veo-3.0-generate-001",
|
||||
"gemini",
|
||||
),
|
||||
(
|
||||
{},
|
||||
"https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1"
|
||||
f"/publishers/google{_VEO_OPERATION_SUFFIX}",
|
||||
None,
|
||||
"veo-3.0-generate-001",
|
||||
"vertex_ai",
|
||||
),
|
||||
(
|
||||
{"contents": []},
|
||||
"https://aiplatform.us.rep.googleapis.com/v1/projects/p/locations/us"
|
||||
"/publishers/google/models/gemini-3.8-flash:generateContent",
|
||||
None,
|
||||
"gemini-3.8-flash",
|
||||
"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, "", None),
|
||||
({"contents": []}, _TRANSLATE_TARGET, None, "", None),
|
||||
({"contents": []}, _NON_GOOGLE_VERTEX_SHAPED_TARGET, None, "", None),
|
||||
({"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: Final = _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,
|
||||
builtin_provider_route=True,
|
||||
)
|
||||
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_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":streamRawPredict?alt=sse"),
|
||||
"claude-opus-5",
|
||||
),
|
||||
(_VERTEX_CLAUDE_RAW_PREDICT_TARGET.replace(":rawPredict", ":rawPredict?alt=sse"), "claude-opus-5"),
|
||||
(f"{_VERTEX_FLASH_TARGET}?tag=rawPredict", "gemini-3.8-flash"),
|
||||
(_VERTEX_SEARCH_TARGET, "vertex_ai/search_api"),
|
||||
("https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs", None),
|
||||
(
|
||||
"https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5/batchPredictionJobs"
|
||||
"?filter=:search",
|
||||
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
|
||||
|
|
@ -8598,7 +9181,9 @@ def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_bo
|
|||
the call to (LIT-1761: passthrough successes carried model_id="")."""
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.url = httpx.URL("http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent")
|
||||
mock_request.url = httpx.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(
|
||||
|
|
@ -9668,12 +10253,14 @@ async def _drain_logging_worker():
|
|||
await GLOBAL_LOGGING_WORKER.stop()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def _vcr_outcome_gate(request, vcr):
|
||||
install_live_call_probe(request, vcr)
|
||||
yield
|
||||
record_vcr_outcome(request, vcr)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate")
|
||||
def test_update_pass_through_route_updates_registry():
|
||||
"""
|
||||
|
|
@ -9763,9 +10350,7 @@ def mock_request():
|
|||
return self._dict.values()
|
||||
|
||||
class MockRequest:
|
||||
def __init__(
|
||||
self, headers=None, method="POST", request_body: Optional[dict] = None
|
||||
):
|
||||
def __init__(self, headers=None, method="POST", request_body: Optional[dict] = None):
|
||||
self.headers = headers or {}
|
||||
self.query_params = QueryParams()
|
||||
self.method = method
|
||||
|
|
@ -9873,9 +10458,9 @@ def test_pass_through_routes_support_all_methods():
|
|||
path = route.path
|
||||
methods = set(route.methods)
|
||||
allowed = PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES.get(path, expected_methods)
|
||||
assert (
|
||||
methods == allowed
|
||||
), f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}"
|
||||
assert methods == allowed, (
|
||||
f"Route {path} does not support all methods. Supported: {methods}, Expected: {allowed}"
|
||||
)
|
||||
|
||||
check_router_methods(llm_router)
|
||||
|
||||
|
|
@ -9905,19 +10490,11 @@ def test_is_bedrock_agent_runtime_route():
|
|||
|
||||
# Test agent runtime endpoints (should return True)
|
||||
assert _is_bedrock_agent_runtime_route("/knowledgebases/kb-123/retrieve") is True
|
||||
assert (
|
||||
_is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve")
|
||||
is True
|
||||
)
|
||||
assert _is_bedrock_agent_runtime_route("/agents/knowledgebases/kb-123/retrieve") is True
|
||||
|
||||
# Test regular bedrock runtime endpoints (should return False)
|
||||
assert (
|
||||
_is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False
|
||||
)
|
||||
assert (
|
||||
_is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse")
|
||||
is False
|
||||
)
|
||||
assert _is_bedrock_agent_runtime_route("/guardrail/test-id/version/1/apply") is False
|
||||
assert _is_bedrock_agent_runtime_route("/model/cohere.command-r-v1:0/converse") is False
|
||||
assert _is_bedrock_agent_runtime_route("/some/random/endpoint") is False
|
||||
|
||||
|
||||
|
|
@ -9984,9 +10561,7 @@ def test_custom_pricing_used_in_cost_calculation():
|
|||
|
||||
# Verify custom cost is different from standard cost (unless prices happen to match)
|
||||
# This confirms custom pricing is actually being applied
|
||||
assert (
|
||||
custom_cost != standard_cost
|
||||
), "Custom pricing should produce different cost than standard pricing"
|
||||
assert custom_cost != standard_cost, "Custom pricing should produce different cost than standard pricing"
|
||||
|
||||
# Test 3: Custom pricing with cache_read_input_token_cost and input_cost_per_token_batches
|
||||
# This specifically tests the parameters that were causing the original issue
|
||||
|
|
@ -10562,6 +11137,7 @@ def test_init_kwargs_client_metadata_cannot_mark_a_pass_through_call_as_a_silent
|
|||
assert "is_silent_experiment" not in json.loads(row["metadata"])
|
||||
assert result["litellm_params"]["metadata"]["client_key"] == "kept"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_direct_passthrough_keeps_mapped_baggage_tags_outside_upstream_body():
|
||||
from litellm.proxy.common_utils.http_parsing_utils import read_request_body
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue