scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

Merge remote-tracking branch 'origin/main' into litellm_vertex_service_tier_enum
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details

Browse source
This commit is contained in:
Devin AI 2026-10-17 19:37:06 +00:00
commit 7bae1bce39
49 changed files with 4406 additions and 239 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/__init__.py
View file

@ -742,7 +742,7 @@ def is_openai_finetune_model(key: str) -> bool:
def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None:
for key, value in model_cost_map.items():
for key, value in tuple(model_cost_map.items()):
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
open_ai_chat_completion_models.add(key)
elif value.get("litellm_provider") == "text-completion-openai":

13
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -1196,6 +1196,9 @@ def _map_sagemaker_exception(
)
_VERTEX_403_STATUS_CODE_PATTERN: Final = re.compile(r"(?<![0-9A-Za-z])403(?![0-9A-Za-z])")
def _map_vertex_exception(
*,
model: str,
@ -1206,6 +1209,8 @@ def _map_vertex_exception(
exception_provider: str,
extra_information: str,
) -> None:
status_code: Final[object] = getattr(original_exception, "status_code", None)
body_code: Final[int | None] = _get_body_error_code(error_str)
if "Vertex AI API has not been used in project" in error_str or "Unable to find your project" in error_str:
raise BadRequestError(
message=f"litellm.BadRequestError: {custom_llm_provider}Exception - {error_str}",
@ -1252,7 +1257,11 @@ def _map_vertex_exception(
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
)
elif "403" in error_str:
elif (
(not isinstance(status_code, int) or status_code in (403, 500))
and body_code in (None, 403)
and _VERTEX_403_STATUS_CODE_PATTERN.search(error_str)
):
raise BadRequestError(
message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}",
model=model,
@ -1306,7 +1315,7 @@ def _map_vertex_exception(
elif (
isinstance(getattr(original_exception, "status_code", None), int)
and 500 <= original_exception.status_code < 600
and _get_body_error_code(error_str) == 429
and body_code == 429
):
# upstream gateway wraps a 429 inside a 5xx envelope
# e.g. HTTP 500/503 with {"error":{"code":429,...}}.

36
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -557,22 +557,39 @@ def get_model_cost_map_loaded_at() -> "datetime | None":
return _cost_map_source_info.loaded_at
def _expand_model_aliases(model_cost: dict) -> dict:
class AliasedCostEntry(dict[str, object]):
__slots__ = ("names",)
def __init__(self, entry: Mapping[str, object], names: tuple[str, ...]) -> None:
super().__init__(entry)
self.names: tuple[str, ...] = names
def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]:
return MappingProxyType(
{
owner: (owner, *(alias for alias, alias_owner in alias_owners.items() if alias_owner == owner))
for owner in dict.fromkeys(alias_owners.values())
}
)
def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]:
"""
Expand ``aliases`` lists in model cost entries into top-level entries.
Each alias gets a reference to the **same** dict object as the canonical
entry (zero memory overhead). The ``aliases`` key is removed from the
entry so downstream code never sees it.
Each alias gets a reference to the **same** ``AliasedCostEntry`` as the
canonical entry, which carries every name it is served under. The
``aliases`` key is removed from the entry so downstream code never sees it.
If an alias collides with an existing canonical entry the alias is
skipped and a warning is logged.
"""
aliases_to_add: Final[dict[str, dict]] = {}
alias_owners: Final[dict[str, str]] = {}
keys_with_aliases: Final[list[str]] = []
for model_name, model_info in model_cost.items():
aliases: list | None = model_info.get("aliases")
aliases: object = model_info.get("aliases")
if aliases is None:
continue
keys_with_aliases.append(model_name)
@ -594,7 +611,7 @@ def _expand_model_aliases(model_cost: dict) -> dict:
model_name,
)
continue
if alias in aliases_to_add:
if alias in alias_owners:
verbose_logger.warning(
"LiteLLM model alias conflict: alias '%s' (from '%s') "
"was already claimed by another entry — skipping.",
@ -602,13 +619,14 @@ def _expand_model_aliases(model_cost: dict) -> dict:
model_name,
)
continue
aliases_to_add[alias] = model_info # same dict reference
alias_owners[alias] = model_name
# Remove the ``aliases`` key from entries so it doesn't pollute model info
for key in keys_with_aliases:
model_cost[key].pop("aliases", None)
model_cost.update(aliases_to_add)
for owner, names in _alias_groups(alias_owners).items():
model_cost.update(dict.fromkeys(names, AliasedCostEntry(model_cost[owner], names)))
return model_cost

2
litellm/llms/fal_ai/cost_calculator.py
View file

@ -87,7 +87,7 @@ def _keyed_rows(model: str, quality: str) -> tuple[tuple[int, int, float], ...]:
suffix: Final = f"/{model}"
return tuple(
(width, height, float(raw_cost))
for key in litellm.model_cost
for key in (*litellm.model_cost,)
if isinstance(key, str) and key.startswith(prefix) and key.endswith(suffix)
for size in (key[len(prefix) : -len(suffix)],)
for dimensions in (_parse_keyed_dimensions(size),)

2
litellm/llms/fireworks_ai/chat/transformation.py
View file

@ -544,7 +544,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
return cached[2]
index: Final[list[tuple[str, dict]]] = []
for key, model_info in model_cost.items():
for key, model_info in (*model_cost.items(),):
if not key.startswith("fireworks_ai/"):
continue
if not isinstance(model_info, dict):

2
litellm/llms/ollama/common_utils.py
View file

@ -197,7 +197,7 @@ class OllamaModelInfo(BaseLLMModelInfo):
"ollama/" + stripped_model,
"ollama_chat/" + stripped_model,
}
model_cost_keys: Final = {key.lower() for key in model_cost}
model_cost_keys: Final = {key.lower() for key in (*model_cost,)}
return any(name.lower() in model_cost_keys for name in potential_model_names)
@staticmethod

4
litellm/llms/openai/chat/gpt_5_transformation.py
View file

@ -28,7 +28,9 @@ def _catalogue_declares_default_effort() -> bool:
the one litellm gave before it existed. Scanning costs ~80us on the largest published map and
only on the fallback path, which is noise beside the request it precedes.
"""
return any(isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in litellm.model_cost.values())
return any(
isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in (*litellm.model_cost.values(),)
)
def _normalize_reasoning_effort_for_chat_completion(

39
litellm/model_prices_and_context_window_backup.json
View file

@ -12877,6 +12877,7 @@
"supports_prompt_caching": true
},
"azure_ai/ministral-3b": {
"deprecation_date": "2026-12-09",
"input_cost_per_token": 4e-08,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12940,6 +12941,7 @@
"supports_vision": true
},
"azure_ai/mistral-medium-2505": {
"deprecation_date": "2026-12-09",
"input_cost_per_token": 4e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 131072,
@ -12977,6 +12979,7 @@
"supports_tool_choice": true
},
"azure_ai/mistral-small-2503": {
"deprecation_date": "2026-12-09",
"input_cost_per_token": 1e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -29699,6 +29702,7 @@
},
"vertex_ai/gemini-3.8-flash": {
"prompt_cache_min_tokens": 4096,
"deprecation_date": "2027-09-02",
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_batches": 3.75e-08,
"cache_read_input_token_cost_flex": 3.75e-08,
@ -29806,6 +29810,7 @@
"supports_native_streaming": true
},
"vertex_ai/gemini-3.8-live": {
"deprecation_date": "2027-09-25",
"input_cost_per_audio_token": 3e-06,
"input_cost_per_image_token": 1e-06,
"input_cost_per_token": 7.5e-07,
@ -31960,6 +31965,7 @@
},
"gemini-3.8-flash": {
"prompt_cache_min_tokens": 4096,
"deprecation_date": "2027-09-02",
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_batches": 3.75e-08,
"cache_read_input_token_cost_flex": 3.75e-08,
@ -45425,6 +45431,38 @@
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/getting-started/pricing"
},
"perplexity/pplx-decider-v1.2": {
"input_cost_per_token": 2e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/pplx-decider-latest": {
"input_cost_per_token": 2e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/pplx-decider-flash-v1.2": {
"input_cost_per_token": 1e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/pplx-decider-flash-latest": {
"input_cost_per_token": 1e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/sonar": {
"input_cost_per_token": 1e-06,
"litellm_provider": "perplexity",
@ -75357,6 +75395,7 @@
"web_search_billing_unit": "per_query"
},
"vertex_ai/gemini-4-argon": {
"deprecation_date": "2027-10-13",
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_batches": 5e-08,
"cache_read_input_token_cost_flex": 5e-08,

75
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -862,12 +862,17 @@ def _raise_unless_oauth2_discovery_server(
def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool:
"""True when a DCR-bridge server relays client registration to the upstream authorization
server instead of short-circuiting to an admin-configured OAuth client. In the relay arm the
server instead of short-circuiting to the server's own OAuth client. In the relay arm the
upstream holds each client's own registration, so the authorize and token relays pass the
client's ``client_id`` and ``redirect_uri`` through verbatim and the authorization code
returns directly to the client's redirect URI without transiting the gateway. Gateway-side
redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit
arm, where the upstream only knows the gateway's own callback."""
arm, where the upstream only knows the gateway's own callback.
Authorize and token decide per request from the server's current client: once the gateway
holds its shared client (registered at ``/register`` under ``PROXY_BASE_URL``, or entered by an
admin) every caller runs the gateway-callback arm, while a server that still has none keeps
relaying, so clients registered before the shared client existed keep their own binding."""
return bool(
mcp_server.is_dcr_bridge
and mcp_server.effective_registration_url
@ -880,6 +885,27 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool:
)
def _register_relays_bridge_client(mcp_server: MCPServer) -> bool:
"""``/register`` relays the caller's own registration only when the gateway cannot keep a shared
client: with ``PROXY_BASE_URL`` set it registers one client bound to its own ``/callback`` and
keeps it as the server's client, so an upstream that restricts redirect domains (Atlassian)
never sees a laptop loopback redirect. Without a database the shared client would live only in
this process, so a database-free gateway keeps relaying as well."""
return _dcr_bridge_relays_client_registration(mcp_server) and (
configured_proxy_base_url() is None or not _shared_client_store_available()
)
def _shared_client_store_available() -> bool:
from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import
try:
get_prisma_client_or_throw("Database not connected. Cannot keep the shared MCP OAuth client.")
except HTTPException:
return False
return True
def _require_s256_pkce(
code_challenge: str | None,
code_challenge_method: str | None,
@ -1028,6 +1054,17 @@ async def _managed_authorization_scope(
return selected_scope, scope_user_id
async def _hydrate_bridge_shared_client(mcp_server: MCPServer, request: Request) -> None:
"""Whichever worker served ``/register`` persisted the bridge's shared gateway client; a worker
whose catalog snapshot predates that write still holds no ``client_id``, so the gateway-callback
arms load the persisted client before presenting one upstream. Callers that carry their own
client (a gateway-minted ephemeral client sealed into the state or the code) skip this, since the
upstream bound their code to that client."""
if not mcp_server.is_dcr_bridge or mcp_server.client_id:
return
await _reuse_persisted_dcr_client_if_available(mcp_server, f"{get_request_base_url(request)}/callback")
async def authorize_with_server(
request: Request,
mcp_server: MCPServer,
@ -1069,6 +1106,8 @@ async def authorize_with_server(
# calling this for its enforcement side effect, then falls through to the gateway
# /callback flow below, which reads the original code_challenge names.
bridge_challenge, bridge_method = _require_s256_pkce(code_challenge, code_challenge_method)
if ephemeral_dcr_client is None:
await _hydrate_bridge_shared_client(resolved_server, request)
# A gateway-minted ephemeral client is registered against {base}/callback, so its
# flow must run the short-circuit arm; the relay arm is only for clients that
# registered themselves through the front door and hold their own redirect binding.
@ -1233,6 +1272,7 @@ async def exchange_token_with_server(
refresh_token: str | None = None,
scope: str | None = None,
client_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None,
sealed_client: bool = False,
):
raise_if_not_oauth2(mcp_server)
if grant_type not in ("authorization_code", "refresh_token"):
@ -1282,11 +1322,14 @@ async def exchange_token_with_server(
# recovered from a sealed code) must authenticate the way its own registration was granted,
# not the way the server row is configured; callers that carry no method keep the row's method
# as before.
resolved_client_id: Final = resolved_server.client_id if resolved_server.client_id else client_id
resolved_client_secret: Final = resolved_server.client_secret if resolved_server.client_id else client_secret
if not sealed_client:
await _hydrate_bridge_shared_client(resolved_server, request)
server_client_wins: Final = bool(resolved_server.client_id) and not sealed_client
resolved_client_id: Final = resolved_server.client_id if server_client_wins else client_id
resolved_client_secret: Final = resolved_server.client_secret if server_client_wins else client_secret
resolved_auth_method: Final = (
resolved_server.token_endpoint_auth_method
if resolved_server.client_id
if server_client_wins
else (client_token_endpoint_auth_method or resolved_server.token_endpoint_auth_method)
)
try:
@ -1830,10 +1873,13 @@ async def _persist_dcr_client_registration(
empty and otherwise adopts whatever a concurrent registration persisted first.
The client-forwarded token modes (``true_passthrough`` / ``oauth_delegate``) are skipped
unconditionally: the caller holds the upstream token and the gateway must hold no OAuth
client identity for these servers. Persisting here would stamp ``oauth2_flow`` and a
``client_id`` onto a server whose mode promises the gateway stores nothing, making a
fresh pass-through server read as gateway-authorized.
unless the server is a DCR bridge: the caller holds the upstream token and the gateway holds
no OAuth client identity for these servers, so persisting would stamp ``oauth2_flow`` and a
``client_id`` onto a server whose mode promises the gateway stores nothing. A DCR bridge is
the exception by design: its gateway-hosted sign-in registers one shared client bound to the
gateway ``/callback`` (the same app an admin may enter by hand), and that client must survive
for every later caller's authorize and token exchange. ``oauth2_flow`` is still never stamped
on a client-forwarded row.
``redirect_uris`` records what the client is bound to so a later origin change can be
detected as a positive mismatch and trigger re-registration instead of stranding the
@ -1842,7 +1888,7 @@ async def _persist_dcr_client_registration(
``update_mcp_server`` merges credential blobs: a re-registered public client must not
inherit the previous client's secret or auth method.
"""
if mcp_server.is_client_forwarded_token:
if mcp_server.is_client_forwarded_token and not mcp_server.is_dcr_bridge:
return "skipped"
try:
@ -1910,20 +1956,23 @@ async def _persist_dcr_client_registration(
):
await _refresh_persisted_dcr_server(stored)
return "reused"
stamped_flow: Final[Literal["authorization_code"] | None] = (
None if mcp_server.is_client_forwarded_token else "authorization_code"
)
updated_row: Final = await update_mcp_server(
prisma_client=prisma_client,
data=(
UpdateMCPServerRequest(
server_id=mcp_server.server_id,
credentials=credentials,
oauth2_flow="authorization_code",
oauth2_flow=stamped_flow,
token_url=mcp_server.token_url,
)
if mcp_server.token_url
else UpdateMCPServerRequest(
server_id=mcp_server.server_id,
credentials=credentials,
oauth2_flow="authorization_code",
oauth2_flow=stamped_flow,
)
),
touched_by="mcp_oauth_dcr",
@ -2211,7 +2260,7 @@ async def register_client_with_server(
)
return dummy_return
bridge_relay: Final = _dcr_bridge_relays_client_registration(resolved_server)
bridge_relay: Final = _register_relays_bridge_client(resolved_server)
if bridge_relay and not client_redirect_uris:
raise HTTPException(
status_code=400,

15
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -4531,7 +4531,7 @@ class MCPServerManager:
)
else:
page: Final = (
await client.list_tools_page(params)
await self._fetch_tools_page_with_timeout(client, server, params)
if params is not None
else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name))
)
@ -5621,6 +5621,19 @@ class MCPServerManager:
)
raise_classified_list_failure(e, server_name)
async def _fetch_tools_page_with_timeout(
self, client: MCPClient, server: MCPServer, params: PaginatedRequestParams
) -> ListToolsResult:
deadline: Final = (
MCP_TOOL_LISTING_TIMEOUT if server.timeout is None else max(server.timeout, MCP_TOOL_LISTING_TIMEOUT)
)
try:
with anyio.fail_after(deadline):
return await client.list_tools_page(params)
except TimeoutError as e:
verbose_logger.warning("Timeout while listing tools from %s", server.name)
raise MCPServerListError(ServerListFault(tag="timeout"), server.name) from e
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
def assign_unique_short_prefix(

28
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -14412,6 +14412,20 @@
"description": "Optional field if guardrail requires a 'model' parameter",
"title": "Model"
},
"num_retries": {
"anyOf": [
{
"maximum": 5.0,
"minimum": 0.0,
"type": "integer"
},
{
"type": "null"
}
],
"description": "How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it.",
"title": "Num Retries"
},
"on_sensitive_data": {
"anyOf": [
{
@ -17001,6 +17015,20 @@
"description": "If True, logs violations without blocking. Defaults to False if not provided",
"title": "Monitor Mode"
},
"num_retries": {
"anyOf": [
{
"maximum": 5.0,
"minimum": 0.0,
"type": "integer"
},
{
"type": "null"
}
],
"description": "How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it.",
"title": "Num Retries"
},
"on_disallowed_action": {
"default": "block",
"description": "Choose whether disallowed tools block the request or get rewritten out of the payload",

25
litellm/proxy/auth/route_checks.py
View file

@ -72,14 +72,22 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spe
_LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp"))
_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS))
# Method-scoped because the generic credential CRUD handlers share these paths:
# the route check cannot admit PATCH/DELETE on a name like "user_connections"
# or "*/user_connection*" without also opening the admin CRUD handlers.
_INTERNAL_USER_CREDENTIAL_CONNECTION_ROUTES: Final = (
# Writes that only touch the caller's own credential row (BYOK connections and
# per-user MCP credentials). Method-scoped because the generic credential CRUD
# handlers share these paths: the route check cannot admit PATCH/DELETE on a
# name like "user_connections" without also opening the admin CRUD handlers.
_SELF_SCOPED_CREDENTIAL_ROUTES: Final = (
("GET", "/credentials/user_connections"),
("POST", "/credentials/{credential_name:path}/user_connection/start"),
("POST", "/credentials/{credential_name:path}/user_connection/poll"),
("DELETE", "/credentials/{credential_name:path}/user_connection"),
("POST", "/v1/mcp/server/oauth/{server_id}/register"),
("POST", "/v1/mcp/server/{server_id}/oauth-user-credential"),
("DELETE", "/v1/mcp/server/{server_id}/oauth-user-credential"),
("POST", "/v1/mcp/server/{server_id}/user-credential"),
("DELETE", "/v1/mcp/server/{server_id}/user-credential"),
("POST", "/v1/mcp/server/{server_id}/user-env-vars"),
("DELETE", "/v1/mcp/server/{server_id}/user-env-vars"),
)
@ -349,7 +357,7 @@ class RouteChecks:
_user_role == LitellmUserRoles.INTERNAL_USER.value
and (
RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value)
or RouteChecks.is_internal_user_credential_connection_route(route=route, request=request)
or RouteChecks.is_self_scoped_credential_route(route=route, request=request)
)
or user_is_org_admin(request_data=request_data, user_object=user_obj)
and RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value)
@ -702,13 +710,13 @@ class RouteChecks:
_get_request_method = get_request_method
@staticmethod
def is_internal_user_credential_connection_route(route: str, request: Request | None) -> bool:
def is_self_scoped_credential_route(route: str, request: Request | None) -> bool:
method: Final = RouteChecks.get_request_method(request)
if method is None:
return False
return any(
method == expected_method and RouteChecks.check_route_access(route=route, allowed_routes=(pattern,))
for expected_method, pattern in _INTERNAL_USER_CREDENTIAL_CONNECTION_ROUTES
for expected_method, pattern in _SELF_SCOPED_CREDENTIAL_ROUTES
)
@staticmethod
@ -1033,6 +1041,9 @@ class RouteChecks:
if route == "/session/logout":
return
if RouteChecks.is_self_scoped_credential_route(route=route, request=request):
return
# Hard-block known write routes regardless of HTTP method (defensive
# — these are POSTs in practice, but pinning them here protects
# against future GET-shaped writes).

1
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py
View file

@ -44,6 +44,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
litellm_params, optional_params, "streaming_buffer_until_moderated"
),
timeout=litellm_params.timeout,
num_retries=litellm_params.num_retries,
)
litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback)

82
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
View file

@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import (
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
httpxSpecialProvider,
)
@ -63,6 +64,21 @@ _HEADER_VALUE_ALLOWLIST: Final = frozenset(
# Placeholder for headers that exist but are not on the allowlist (we don't expose their value).
_HEADER_PRESENT_PLACEHOLDER: Final = "[present]"
_AFTER_CONNECT_FAILURES: Final = (
httpx.ReadTimeout,
httpx.WriteTimeout,
httpx.ReadError,
httpx.WriteError,
httpx.RemoteProtocolError,
)
def _after_connect_failure(error: Timeout | httpx.TransportError) -> httpx.TransportError | None:
"""AsyncHTTPHandler re-raises connect and read timeouts as the same litellm Timeout, so the phase comes from
the httpx error it was raised while handling."""
origin: Final = error if isinstance(error, httpx.TransportError) else error.__cause__ or error.__context__
return origin if isinstance(origin, _AFTER_CONNECT_FAILURES) else None
def _header_value_allowed(
header_name: str,
@ -81,7 +97,7 @@ def _header_value_allowed(
def _sanitize_inbound_headers(
headers: object,
headers: Mapping[object, object] | None,
extra_allowlist: set[str] | None = None,
) -> dict[str, str] | None:
"""
@ -254,9 +270,18 @@ class GenericGuardrailAPI(CustomGuardrail):
streaming_sampling_rate: int | None = None,
streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
streaming_buffer_until_moderated: bool | None = None,
num_retries: int | None = None,
async_handler: AsyncHTTPHandler | None = None,
fresh_handler_factory: Callable[[], AsyncHTTPHandler] | None = None,
**kwargs,
):
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self.async_handler = async_handler or get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
self.fresh_handler_factory: Callable[[], AsyncHTTPHandler] = fresh_handler_factory or (
lambda: AsyncHTTPHandler(timeout=self.async_handler.timeout)
)
self.num_retries: int = num_retries or 0
self.headers = headers or {}
self.extra_headers = extra_headers or []
@ -384,12 +409,9 @@ class GenericGuardrailAPI(CustomGuardrail):
return_inputs.update(inputs)
return return_inputs
def _build_request_headers(self) -> dict:
def _build_request_headers(self) -> Mapping[str, str]:
"""Build HTTP headers for the guardrail API request."""
headers: Final = {"Content-Type": "application/json"}
if self.headers:
headers.update(self.headers)
return headers
return {"Content-Type": "application/json", **self.headers}
def _request_body(
self,
@ -463,6 +485,35 @@ class GenericGuardrailAPI(CustomGuardrail):
verbose_proxy_logger.error("Generic Guardrail API: failed to make request: %s", str(error))
raise Exception(f"Generic Guardrail API failed: {error}")
async def _post_check(
self,
payload: Mapping[str, object],
headers: Mapping[str, str],
retries_left: int,
fresh: bool = False,
) -> httpx.Response:
handler: Final = self.fresh_handler_factory() if fresh else self.async_handler
try:
return await handler.post(
url=self.api_base, json=dict(payload), headers=dict(headers), timeout=self.timeout
)
except (Timeout, httpx.TransportError) as error:
failure: Final = _after_connect_failure(error)
if retries_left < 1 or failure is None:
raise
verbose_proxy_logger.warning(
"Generic Guardrail API check timed out or lost its connection, retrying on a fresh connection. "
"guardrail_name=%s api_base=%s error_type=%s retries_left=%d",
getattr(self, "guardrail_name", None),
self.api_base,
type(failure).__name__,
retries_left - 1,
)
finally:
if fresh:
await handler.close()
return await self._post_check(payload, headers, retries_left - 1, fresh=True)
@log_guardrail_information
async def apply_guardrail(
self,
@ -508,12 +559,8 @@ class GenericGuardrailAPI(CustomGuardrail):
request_body: Final = request_data.get("body") or {}
# Merge additional provider specific params from config and dynamic params
additional_params: Final = {**self.additional_provider_specific_params}
# Get dynamic params from request if available
dynamic_params: Final = self.get_guardrail_dynamic_request_body_params(request_body)
if dynamic_params:
additional_params.update(dynamic_params)
dynamic_params: Final[Mapping[str, object]] = self.get_guardrail_dynamic_request_body_params(request_body) or {}
additional_params: Final = {**self.additional_provider_specific_params, **dynamic_params}
# Extract user API key metadata
user_metadata: Final = self._extract_user_api_key_metadata(request_data)
@ -546,11 +593,10 @@ class GenericGuardrailAPI(CustomGuardrail):
# Make the API request
# Use mode="json" to ensure all iterables are converted to lists
response: Final = await self.async_handler.post(
url=self.api_base,
json=dict(self._request_body(guardrail_request, structured_messages)),
headers=headers,
timeout=self.timeout,
response: Final = await self._post_check(
self._request_body(guardrail_request, structured_messages),
headers,
self.num_retries,
)
response.raise_for_status()

5
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -2517,7 +2517,9 @@ if MCP_AVAILABLE:
resolved_redirect_uri: Final = f"{get_request_base_url(request)}/callback" if sealed_code else redirect_uri
caller_client_id: Final = sealed_code.client_id if sealed_code else client_id
caller_client_secret: Final = sealed_code.client_secret if sealed_code else client_secret
resolved_client_id: Final = mcp_server.client_id or caller_client_id or ""
resolved_client_id: Final = (
caller_client_id if sealed_code else (mcp_server.client_id or caller_client_id or "")
)
if not resolved_client_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
@ -2541,6 +2543,7 @@ if MCP_AVAILABLE:
refresh_token=refresh_token,
scope=scope,
client_token_endpoint_auth_method=sealed_code.token_endpoint_auth_method if sealed_code else None,
sealed_client=sealed_code is not None,
)
@router.post(

18
litellm/proxy/route_llm_request.py
View file

@ -813,7 +813,23 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr
route=ROUTE_ENDPOINT_MAPPING.get(route_type, route_type),
param="model",
)
team_model_name: Final = llm_router.map_team_model(data["model"], team_id) if team_id is not None else None
custom_llm_provider: Final[object] = (
JSON_OBJECT_ADAPTER.validate_python(data).get("custom_llm_provider")
if route_type
in (
"avector_store_search",
"avector_store_retrieve",
"avector_store_update",
"avector_store_delete",
)
else None
)
team_model_provider: Final[str | None] = custom_llm_provider if isinstance(custom_llm_provider, str) else None
team_model_name: Final = (
llm_router.map_team_model(data["model"], team_id, custom_llm_provider=team_model_provider)
if team_id is not None
else None
)
if team_model_name is not None:
data["model"] = team_model_name
return getattr(llm_router, f"{route_type}")(**data)

58
litellm/router.py
View file

@ -12565,26 +12565,56 @@ class Router:
members: Final = routing_group.models if routing_group is not None else ()
return served_model if served_model == target or served_model in members else model_name
def map_team_model(self, team_model_name: str | None, team_id: str) -> str | None:
def map_team_model(
self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None
) -> str | None:
"""
Check if team_model_name resolves to team-specific deployments.
Resolve explicit team models or select a deployment for model-less calls.
Returns the public model name (unchanged) so the router can find all
sibling deployments via team_id filtering, instead of collapsing to a
single internal model_name.
When team_model_name is None (e.g. vector store / file endpoints that
don't include a model in their request), returns the first matching
team deployment's team_public_model_name so the router can inject BYOK
credentials from the team-scoped deployment.
Returns:
- str: the team_model_name if team deployments exist for this team
- None: if no team-specific model is found
Model-less calls with no provider use the first team deployment. With a provider, they
use the first team deployment of that provider, pinned by ID when its public name also
covers another provider.
"""
models: Final = self.get_model_list(model_name=team_model_name, team_id=team_id)
if not models:
return None
if team_model_name is None and custom_llm_provider is not None:
def team_deployment_public_name(deployment: DeploymentTypedDict) -> str | None:
model_info: Final = _configured_model_info(deployment)
if model_info.get("team_id") != team_id:
return None
team_public_model_name: Final = model_info.get("team_public_model_name")
return (
team_public_model_name
if isinstance(team_public_model_name, str) and team_public_model_name
else deployment["model_name"]
)
matching_deployment: Final = next(
(
model
for model in models
if team_deployment_public_name(model) is not None
and provider_for_generic_call(model["litellm_params"]) == custom_llm_provider
),
None,
)
if matching_deployment is None:
return None
public_name: Final = team_deployment_public_name(matching_deployment)
if public_name is None:
return None
same_public_name_deployments: Final = tuple(
candidate for candidate in models if team_deployment_public_name(candidate) == public_name
)
if all(
provider_for_generic_call(candidate["litellm_params"]) == custom_llm_provider
for candidate in same_public_name_deployments
):
return public_name
deployment_id: Final = _configured_model_info(matching_deployment).get("id")
return deployment_id if isinstance(deployment_id, str) else None
for model in models:
if model.get("model_info", {}).get("team_id") == team_id:
if team_model_name is None:

22
litellm/types/guardrails.py
View file

@ -2,7 +2,7 @@ from collections.abc import Mapping
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Final, Literal, cast
from typing import Annotated, Final, Literal, TypeAlias, cast
from pydantic import ConfigDict, Field, field_validator, model_validator
from pydantic.fields import FieldInfo
@ -1012,6 +1012,22 @@ def _default_if_blank(field: FieldInfo | None, value: object) -> object:
return field.get_default(call_default_factory=True)
GuardrailNumRetries: TypeAlias = Annotated[
int | None,
Field(
ge=0,
le=5,
description=(
"How many more times to send a guardrail check whose connection timed out waiting for the answer "
"or was dropped after the request went out. Each retry opens a fresh connection and gets the full "
"timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 "
"sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', "
"other guardrails ignore it."
),
),
]
class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails
api_key: str | None = Field(default=None, description="API key for the guardrail service")
api_base: str | None = Field(default=None, description="Base URL for the guardrail service API")
@ -1220,6 +1236,8 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
description="Python-like code containing the apply_guardrail function for custom guardrail logic",
)
num_retries: GuardrailNumRetries = None
timeout: float | None = Field(
default=None,
description=(
@ -1437,7 +1455,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
super().__init__(**kwargs) # pyright: ignore[reportArgumentType] # pydantic validates each field's type at runtime
def __contains__(self, key) -> bool:
def __contains__(self, key: str) -> bool:
# Define custom behavior for the 'in' operator
return hasattr(self, key)

3
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py
View file

@ -4,6 +4,7 @@ from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows h
from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, model_serializer
from typing_extensions import TypedDict
from litellm.types.guardrails import GuardrailNumRetries
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import (
AllMessageValues,
@ -70,6 +71,8 @@ class GenericGuardrailAPIOptionalParams(LiteLLMBaseModel):
),
)
num_retries: GuardrailNumRetries = None
forward_messages_as_sent: bool | None = Field(
default=False,
description=(

396
litellm/utils.py
View file

@ -44,7 +44,7 @@ from httpx import Proxy
from httpx._utils import get_environment_proxies
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # OpenAI parser module is private
from openai.types.chat.completion_create_params import ResponseFormat
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, with_config
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config
import litellm
import litellm.litellm_core_utils
@ -89,6 +89,7 @@ from litellm.litellm_core_utils.fallback_generalizations import (
match_capability_generalizations,
match_fill_missing_generalizations,
)
from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.litellm_core_utils.tokenizer import strip_special_tokens
@ -3303,46 +3304,176 @@ _BEDROCK_REGION_PREFIXES: Final = (
"ap-northeast-1.",
)
_CACHE_PRICING_FIELDS: Final = (
"cache_creation_input_token_cost",
"cache_creation_input_token_cost_above_1hr",
"cache_creation_input_token_cost_above_200k_tokens",
"cache_read_input_token_cost",
"cache_read_input_token_cost_above_200k_tokens",
_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token")
def _builtin_cost_map() -> dict[str, Mapping[str, object]]:
return cast(dict[str, Mapping[str, object]], litellm.model_cost)
def _replace_cost_map_entry(cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object]) -> None:
current: Final = cost_map.get(key)
merged: Final = {**(current or {}), **entry}
if not isinstance(current, AliasedCostEntry):
cost_map.update({key: merged})
return
sharing_names: Final = tuple(name for name in current.names if cost_map.get(name) is current)
cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, current.names)))
def _without_synthesized_costs(
model_info: Mapping[str, object], key: str, value: Mapping[str, object]
) -> Mapping[str, object]:
"""``_get_model_info_helper`` synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` = 0 when the raw
entry has neither (#30198). Persisting those zeros would turn a name-priced entry into a free one and make
``_is_cost_explicitly_configured`` disable budget enforcement on the next re-registration."""
builtin_cost_map: Final = _builtin_cost_map()
raw_fields: Final = frozenset(builtin_cost_map.get(str(model_info["key"]), builtin_cost_map.get(key, {})))
return MappingProxyType(
{
name: field_value
for name, field_value in model_info.items()
if name not in _TOKEN_COST_FIELDS or name in raw_fields or name in value
}
)
_DIRECTION_RATE_FIELDS: Final = MappingProxyType(
{
"input_cost_per_token": ("input_cost_per_token",),
"output_cost_per_token": ("output_cost_per_token", "output_cost_per_reasoning_token"),
}
)
_RATE_TABLE: Final = TypeAdapter(Mapping[str, object])
_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, object] | None:
def _is_overridden_rate(name: str, overridden_fields: tuple[str, ...]) -> bool:
is_variant: Final = name.startswith(tuple(f"{rate_field}_" for rate_field in overridden_fields))
return name in overridden_fields or (is_variant and "_cache_" not in name)
def _without_overridden_rates(rates: Mapping[str, object], overridden_fields: tuple[str, ...]) -> Mapping[str, object]:
return MappingProxyType(
{name: rate for name, rate in rates.items() if not _is_overridden_rate(name, overridden_fields)}
)
def _rate_table_with_flat_prices(
name: str, table: object, overridden_fields: tuple[str, ...], flat_prices: Mapping[str, object]
) -> object:
try:
if name == "tiered_pricing":
return [
{**_without_overridden_rates(tier, overridden_fields), **flat_prices}
for tier in _RATE_TIERS.validate_python(table)
]
if name == "off_peak_pricing":
return dict(_without_overridden_rates(_RATE_TABLE.validate_python(table), overridden_fields))
except ValidationError:
return table
return table
def _without_overridden_price_variants(
entry: Mapping[str, object], value: Mapping[str, object]
) -> Mapping[str, object]:
"""A caller's flat ``input_cost_per_token`` / ``output_cost_per_token`` is that direction's price above every
token threshold, for every service tier and batch, in every ``tiered_pricing`` tier and off-peak window, so each
inherited rate that would stand in for it there, the reasoning-token rate included for output, is dropped, and
each tier carries the flat price instead."""
flat_prices: Final = MappingProxyType(
{cost_field: value[cost_field] for cost_field in _TOKEN_COST_FIELDS if value.get(cost_field) is not None}
)
overridden_fields: Final = tuple(
itertools.chain.from_iterable(_DIRECTION_RATE_FIELDS[cost_field] for cost_field in flat_prices)
)
if not overridden_fields:
return entry
return MappingProxyType(
{
name: _rate_table_with_flat_prices(name, field_value, overridden_fields, flat_prices)
for name, field_value in entry.items()
if not _is_overridden_rate(name, overridden_fields)
}
)
@dataclass(frozen=True, slots=True)
class _BuiltinCostMatch:
key: str
entry: Mapping[str, object]
names_the_row: bool
@dataclass(frozen=True, slots=True)
class _LookupMiss:
tripped: bool
@dataclass(frozen=True, slots=True)
class _FallbackRegistration:
entry: Mapping[str, object]
matched_builtin: bool
def _fallback_registration(key: str, provider: str, value: Mapping[str, object]) -> _FallbackRegistration:
"""The key keeps its own entry, seeded from the capability rules an unregistered lookup resolves it to, with the
matched built-in entry on top. The matched row's provider stays behind, since the key's own lookup can reject it
(an ``azure_ai/`` key never resolves to an ``openai`` entry), and the key keeps the provider its caller gave."""
match: Final = _resolve_builtin_model_cost_entry(key=key, provider=provider)
generalized: Final[Mapping[str, object]] = match_capability_generalizations(key) or {}
if match is None:
return _FallbackRegistration(entry=generalized, matched_builtin=False)
inherited: Final = match.entry if match.names_the_row else _without_overridden_price_variants(match.entry, value)
own_fields: Final = {name: field_value for name, field_value in inherited.items() if name != "litellm_provider"}
return _FallbackRegistration(entry={**generalized, **own_fields}, matched_builtin=True)
def _resolve_builtin_model_cost_entry(key: str, provider: str) -> _BuiltinCostMatch | None:
"""Best-effort lookup of a built-in ``model_cost`` entry for a custom key
whose shape ``get_model_info`` cannot resolve (repeated provider prefixes
like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region
aliases).
aliases), or whose lookup tripped.
Returns a copy of the matching entry so the caller can inherit its defaults
(most importantly cache pricing) without mutating the shared built-in.
Returns ``None`` when no safe match exists.
The key itself is tried first, then the key behind its provider prefixes
and bedrock route, so a catalog row whose lookup tripped inherits from
itself rather than from a region-stripped sibling priced at the base rate.
``names_the_row`` is set when the match is the row the key names, which
``get_model_info`` resolves it to, so the caller keeps that row's rates as
an untripped registration does. Returns ``None`` when no candidate exists
under a provider ``get_model_info`` would accept for ``provider``, or, when
the value names none, for the provider prefix the key itself carries,
which is what ``get_model_info`` derives when it reads the key back.
"""
candidates: Final[list[str]] = []
from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix
segments: Final = key.split("/")
idx = 0
while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet:
idx += 1
candidates.append("/".join(segments[idx:]))
base: Final = candidates[-1] if candidates else key
for region_prefix in _BEDROCK_REGION_PREFIXES:
if base.startswith(region_prefix):
candidates.append(base[len(region_prefix) :])
if provider:
stripped: Final = _strip_model_name(model=base, custom_llm_provider=provider)
if stripped != base:
candidates.append(stripped)
match_provider: Final = provider or (segments[0] if len(segments) > 1 and segments[0] in LlmProvidersSet else "")
provider_prefix_count: Final = next(
(idx for idx, segment in enumerate(segments[:-1]) if segment not in LlmProvidersSet), len(segments) - 1
)
prefix_stripped: Final = tuple("/".join(segments[idx:]) for idx in range(1, provider_prefix_count + 1))
prefix_free: Final = prefix_stripped[-1] if prefix_stripped else key
base: Final = (
strip_bedrock_routing_prefix(prefix_free) if match_provider in ("bedrock", "bedrock_converse") else prefix_free
)
region_stripped: Final = tuple(
base[len(region_prefix) :] for region_prefix in _BEDROCK_REGION_PREFIXES if base.startswith(region_prefix)
)
stripped: Final = _strip_model_name(model=base, custom_llm_provider=match_provider) if match_provider else base
row_names: Final = (key, *prefix_stripped, base, stripped)
candidates: Final = (key, *prefix_stripped, base, *region_stripped, stripped)
builtin_cost_map: Final = _builtin_cost_map()
for candidate in candidates:
entry = litellm.model_cost.get(candidate)
if entry is not None and entry.get("litellm_provider") is not None:
return dict(entry)
entry = builtin_cost_map.get(candidate)
if (
entry is not None
and entry.get("litellm_provider") is not None
and _check_provider_match(dict(entry), match_provider)
):
return _BuiltinCostMatch(key=candidate, entry=dict(entry), names_the_row=candidate in row_names)
return None
@ -3360,19 +3491,115 @@ def is_generalized_model_info(model_info: ModelInfo) -> bool:
return key not in litellm.model_cost and match_capability_generalizations(key) is not None
def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | None:
"""Resolve ``model`` to its built-in cost-map entry for registration merging.
_REGISTRATION_ATTEMPTS: Final = 3
Returns ``None`` when the lookup raises or when it resolved via a
fallback-generalization capability rule. A rule-derived entry carries no
pricing, so treating it as a hit would skip the built-in cache-pricing
inheritance for prefix-mangled keys.
"""
def _lookup_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss:
try:
info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except ModelNotMappedError:
return _LookupMiss(tripped=False)
except Exception:
return None
return None if is_generalized_model_info(info) else info
return _LookupMiss(tripped=True)
return _LookupMiss(tripped=False) if is_generalized_model_info(info) else info
def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss:
"""Resolve ``model`` to its built-in cost-map entry for registration merging.
Misses when the lookup cannot map the model or resolved it via a
fallback-generalization capability rule, since a rule-derived entry carries no
pricing. A lookup that raised anything else, such as a concurrent registration
changing the map mid-walk, is retried, so a transient failure registers exactly
where the untripped lookup would.
"""
lookups: Final = (_lookup_for_registration(model, custom_llm_provider) for _ in range(_REGISTRATION_ATTEMPTS))
return next(
(lookup for lookup in lookups if not (isinstance(lookup, _LookupMiss) and lookup.tripped)),
_LookupMiss(tripped=True),
)
@dataclass(frozen=True, slots=True)
class _RegistrationWrite:
key: str
entry: Mapping[str, object]
def _registration_base(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
) -> _RegistrationWrite:
provider: Final = str(value.get("litellm_provider") or "")
if provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO or any(
key.startswith(f"{p}/") for p in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
):
return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {})))
builtin_model_info: Final = _get_builtin_model_info_for_registration(
model=key, custom_llm_provider=custom_llm_provider
)
if not isinstance(builtin_model_info, _LookupMiss):
return _RegistrationWrite(
key=builtin_model_info["key"],
entry=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)),
)
fallback: Final = _fallback_registration(key=key, provider=provider, value=value)
if not fallback.matched_builtin and (
value.get("cache_creation_input_token_cost") is None
and value.get("cache_read_input_token_cost") is None
and value.get("tiered_pricing") is None
and (value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None)
):
verbose_logger.warning(
"register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info",
warning_display_name or key,
)
return _RegistrationWrite(key=key, entry=dict(fallback.entry))
def _registration_write(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
) -> _RegistrationWrite:
base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name)
# ``get_model_info`` returns ``litellm_provider: None`` when the
# provider is unknown (e.g. custom deployments registered via
# ``Router.add_deployment``). Persisting that None into
# ``litellm.model_cost`` causes ``_check_provider_match`` to drop
# custom pricing on subsequent cost lookups.
existing_model: Final = {
name: field_value
for name, field_value in base.entry.items()
if name != "litellm_provider" or field_value is not None
}
return _RegistrationWrite(key=base.key, entry=update_dictionary(existing_model, value))
def _registration_attempt(
key: str,
value: Mapping[str, object],
custom_llm_provider: str | None,
warning_display_name: str | None,
final_attempt: bool,
) -> _RegistrationWrite | None:
source: Final = _builtin_cost_map()
write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name)
_replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry)
return write if final_attempt or litellm.model_cost is source else None
def _register_on_current_map(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
) -> _RegistrationWrite:
"""A price data reload that swaps the map while the entry is being built or written would otherwise get the old
map's prices written over its own, or drop the entry with the old map, so each attempt writes into the map it was
built from and the entry is rebuilt when a reload replaced that map before the write finished. The last attempt
writes into the current map, so the entry is never lost."""
attempts: Final = (
_registration_attempt(
key, value, custom_llm_provider, warning_display_name, final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1
)
for attempt in range(_REGISTRATION_ATTEMPTS)
)
return next(write for write in attempts if write is not None)
_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload
@ -3492,79 +3719,15 @@ def register_model(
for _registered_key, _registered_value in _registrations.items():
_runtime_registered_model_cost[_registered_key] = dict(_registered_value)
_skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
for key, value in loaded_model_cost.items():
## get model info ##
provider = value.get("litellm_provider", "")
_key_str = str(key)
if provider in _skip_get_model_info_providers or any(
_key_str.startswith(f"{p}/") for p in _skip_get_model_info_providers
):
existing_model = litellm.model_cost.get(key, {})
model_cost_key = key
else:
builtin_model_info = _get_builtin_model_info_for_registration(
model=_key_str, custom_llm_provider=custom_llm_provider
)
if builtin_model_info is not None:
existing_model = cast(dict, builtin_model_info)
model_cost_key = existing_model["key"]
else:
# An exact entry ends the lookup ladder before the capability rules are
# consulted, so seed from them: otherwise registering an unmapped model
# shadows the very defaults it would have resolved to unregistered.
existing_model = dict(match_capability_generalizations(_key_str) or {})
model_cost_key = key
builtin_entry = _resolve_builtin_model_cost_entry(key=_key_str, provider=provider)
if builtin_entry is not None:
for field in _CACHE_PRICING_FIELDS:
if value.get(field) is None and builtin_entry.get(field) is not None:
existing_model[field] = builtin_entry[field]
elif (
value.get("cache_creation_input_token_cost") is None
and value.get("cache_read_input_token_cost") is None
and value.get("tiered_pricing") is None
and (
value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None
)
):
verbose_logger.warning(
"register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info",
warning_display_name or key,
)
# ``get_model_info`` returns ``litellm_provider: None`` when the
# provider is unknown (e.g. custom deployments registered via
# ``Router.add_deployment``). Persisting that None into
# ``litellm.model_cost`` causes ``_check_provider_match`` to drop
# custom pricing on subsequent cost lookups.
if existing_model.get("litellm_provider") is None:
existing_model.pop("litellm_provider", None)
# Same pattern for cost fields (#30198): ``_get_model_info_helper``
# synthesizes ``input_cost_per_token`` / ``output_cost_per_token``
# = 0 when they are absent from the raw entry. Writing those zeros
# back flips a sparse entry from "no cost keys" (priced via name)
# to "cost keys = 0" (free), which makes
# ``_is_cost_explicitly_configured`` return True and silently
# disables budget enforcement on the next re-registration.
_raw_entry = litellm.model_cost.get(model_cost_key)
if _raw_entry is None:
_raw_entry = litellm.model_cost.get(key)
if _raw_entry is None:
_raw_entry = {}
for _cost_field in ("input_cost_per_token", "output_cost_per_token"):
if _cost_field not in _raw_entry and _cost_field not in value:
existing_model.pop(_cost_field, None)
## override / add new keys to the existing model cost dictionary
updated_dictionary = update_dictionary(existing_model, value)
litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary)
write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name)
# Invalidate case-insensitive lookup map since model_cost was modified
_invalidate_model_cost_lowercase_map()
verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", model_cost_key, model_cost_key)
verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", write.key, write.key)
if add_to_provider_model_sets and value.get("litellm_provider") == "openrouter":
litellm.openrouter_models.add(_key_str.split("/", 1)[-1])
litellm.openrouter_models.add(str(key).split("/", 1)[-1])
if add_to_provider_model_sets:
litellm.add_known_models(loaded_model_cost)
return model_cost
@ -5732,7 +5895,7 @@ def _rebuild_model_cost_lowercase_map() -> dict[str, str]:
The rebuilt map (guaranteed to be not None).
"""
global _model_cost_lowercase_map
_model_cost_lowercase_map = {k.lower(): k for k in litellm.model_cost}
_model_cost_lowercase_map = {k.lower(): k for k in tuple(litellm.model_cost)}
return _model_cost_lowercase_map
@ -5747,9 +5910,7 @@ def _handle_stale_map_entry_rebuild(
Returns:
The matched key if found after rebuild, None otherwise.
"""
global _model_cost_lowercase_map
_model_cost_lowercase_map = _rebuild_model_cost_lowercase_map()
matched_key: Final = _model_cost_lowercase_map.get(potential_key_lower)
matched_key: Final = _rebuild_model_cost_lowercase_map().get(potential_key_lower)
if matched_key is not None and matched_key in litellm.model_cost:
return matched_key
return None
@ -5767,7 +5928,7 @@ def _handle_new_key_with_scan(
The matched key if found, None otherwise.
"""
global _model_cost_lowercase_map
for key in litellm.model_cost:
for key in tuple(_builtin_cost_map()):
if key.lower() == potential_key_lower:
_model_cost_lowercase_map = _rebuild_model_cost_lowercase_map()
return key
@ -5789,18 +5950,19 @@ def get_model_cost_key(potential_key: str) -> str | None:
called and confirmed not to cause performance issues, add it to the allowed_helpers
list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py
"""
global _model_cost_lowercase_map
# Exact match (O(1))
if potential_key in litellm.model_cost:
return potential_key
# Case-insensitive lookup via map (O(1))
if _model_cost_lowercase_map is None:
_model_cost_lowercase_map = _rebuild_model_cost_lowercase_map()
# Case-insensitive lookup via map (O(1)); bound locally because another thread's
# registration can invalidate the module-level map at any point
cached_lowercase_map: Final = _model_cost_lowercase_map
lowercase_map: Final = (
cached_lowercase_map if cached_lowercase_map is not None else _rebuild_model_cost_lowercase_map()
)
potential_key_lower: Final = potential_key.lower()
matched_key = _model_cost_lowercase_map.get(potential_key_lower)
matched_key = lowercase_map.get(potential_key_lower)
# Verify key exists (O(1) - handles model_cost.pop() case)
if matched_key is not None and matched_key in litellm.model_cost:

39
model_prices_and_context_window.json
View file

@ -12877,6 +12877,7 @@
"supports_prompt_caching": true
},
"azure_ai/ministral-3b": {
"deprecation_date": "2026-12-09",
"input_cost_per_token": 4e-08,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12940,6 +12941,7 @@
"supports_vision": true
},
"azure_ai/mistral-medium-2505": {
"deprecation_date": "2026-12-09",
"input_cost_per_token": 4e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 131072,
@ -12977,6 +12979,7 @@
"supports_tool_choice": true
},
"azure_ai/mistral-small-2503": {
"deprecation_date": "2026-12-09",
"input_cost_per_token": 1e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -29699,6 +29702,7 @@
},
"vertex_ai/gemini-3.8-flash": {
"prompt_cache_min_tokens": 4096,
"deprecation_date": "2027-09-02",
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_batches": 3.75e-08,
"cache_read_input_token_cost_flex": 3.75e-08,
@ -29806,6 +29810,7 @@
"supports_native_streaming": true
},
"vertex_ai/gemini-3.8-live": {
"deprecation_date": "2027-09-25",
"input_cost_per_audio_token": 3e-06,
"input_cost_per_image_token": 1e-06,
"input_cost_per_token": 7.5e-07,
@ -31960,6 +31965,7 @@
},
"gemini-3.8-flash": {
"prompt_cache_min_tokens": 4096,
"deprecation_date": "2027-09-02",
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_batches": 3.75e-08,
"cache_read_input_token_cost_flex": 3.75e-08,
@ -45425,6 +45431,38 @@
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/getting-started/pricing"
},
"perplexity/pplx-decider-v1.2": {
"input_cost_per_token": 2e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/pplx-decider-latest": {
"input_cost_per_token": 2e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/pplx-decider-flash-v1.2": {
"input_cost_per_token": 1e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/pplx-decider-flash-latest": {
"input_cost_per_token": 1e-08,
"litellm_provider": "perplexity",
"max_input_tokens": 262144,
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://docs.perplexity.ai/docs/decisions/quickstart"
},
"perplexity/sonar": {
"input_cost_per_token": 1e-06,
"litellm_provider": "perplexity",
@ -75357,6 +75395,7 @@
"web_search_billing_unit": "per_query"
},
"vertex_ai/gemini-4-argon": {
"deprecation_date": "2027-10-13",
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_batches": 5e-08,
"cache_read_input_token_cost_flex": 5e-08,

3
tests/code_coverage_tests/ensure_async_clients_test.py
View file

@ -5,6 +5,9 @@ ALLOWED_FILES = [
# Lens data traffic owns one pool per app lifespan, isolated from model traffic and closed on shutdown.
"../../litellm/tracing/remote.py",
"./litellm/tracing/remote.py",
# Only a guardrail check retry opens a client, on purpose to get a new connection, and closes it right after.
"../../litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py",
"./litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py",
# local files
"../../litellm/__init__.py",
"../../litellm/llms/custom_httpx/http_handler.py",

25
tests/e2e/perf/setup/locustfile.py
View file

@ -19,9 +19,10 @@ from typing import Final, cast
import gevent
from locust import FastHttpUser, task
from locust.contrib.fasthttp import FastResponse, ResponseContextManager
from setup.wire import Endpoint, WireRecord, WireRequest, request_body
from pydantic import TypeAdapter
from setup.wire import Endpoint, PlanProgress, WireRecord, WireRequest, request_body, response_id
_PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest])
_JSON: Final = TypeAdapter[object](object)
_PLAN: Final = _PLAN_ADAPTER.validate_json(Path(os.environ["PERF_PLAN_FILE"]).read_bytes())
@ -33,7 +34,8 @@ _INTERVAL_S: Final = _USERS / _RPS
_JITTER: Final = 0.25
_RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan")
_RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8")
_written = 0
_PROGRESS: Final = PlanProgress(len(_PLAN))
_STOP_FILE: Final = os.environ.get("PERF_STOP_FILE")
BLOCK_REASON: Final = "perf-test-block"
_TOKEN_MAX_AGE_S: Final = 240.0
_tokens: dict[str, tuple[str, float]] = {}
@ -155,10 +157,9 @@ def _judge(item: WireRequest, status: int, body_text: str) -> tuple[bool, str |
def _record(record: WireRecord) -> None:
global _written
_ = _RECORDS.write(record.model_dump_json() + "\n")
_RECORDS.flush()
_written += 1
_PROGRESS.completed += 1
def _authorization(user: FastHttpUser, item: WireRequest) -> str:
@ -203,7 +204,7 @@ class PlanUser(FastHttpUser):
a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the
users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate."""
network_timeout = 120.0
network_timeout = float(os.environ.get("PERF_REQUEST_TIMEOUT", "120"))
connection_timeout = 30.0
def on_start(self) -> None:
@ -221,9 +222,10 @@ class PlanUser(FastHttpUser):
@task
def send(self) -> None:
item: Final = next(_NEXT, None)
take: Final = _PROGRESS.take(stop=_STOP_FILE is not None and Path(_STOP_FILE).exists())
item: Final = next(_NEXT, None) if take else None
if item is None:
if _written >= len(_PLAN):
if _PROGRESS.drained:
_quit(self)
return
started_at: Final = time.time()
@ -237,6 +239,7 @@ class PlanUser(FastHttpUser):
error: str | None = None
header_s = 0.0
response: FastResponse | None = None
identifier: str | None = None # rebind-ok: the ID arrives in a JSON body or the first SSE event
try:
headers: Final = {"Authorization": _authorization(self, item), "Content-Type": "application/json"}
context: ResponseContextManager
@ -265,6 +268,7 @@ class PlanUser(FastHttpUser):
buffer += line
events, buffer = _events(buffer)
for event in events:
identifier = identifier or response_id(event)
if first_token_s is None and token_in_event(item.endpoint, event):
first_token_s = now
text += delta_of_event(item.endpoint, event)
@ -272,7 +276,9 @@ class PlanUser(FastHttpUser):
body_text = context.text or ""
first_byte_s = time.perf_counter() - started
if status == 200:
text = text_of_json(item.endpoint, _as_dict(_JSON.validate_json(body_text)))
payload: Final = _as_dict(_JSON.validate_json(body_text))
identifier = response_id(payload)
text = text_of_json(item.endpoint, payload)
ok, error = _judge(item, status, body_text)
if ok:
context.success()
@ -302,7 +308,8 @@ class PlanUser(FastHttpUser):
mock_delay_s=float(delay_header) if delay_header else None,
call_id=_header(response, "x-litellm-call-id") if response is not None else None,
error=error,
response_id=identifier,
)
)
if _written >= len(_PLAN):
if _PROGRESS.drained:
_quit(self)

111
tests/e2e/perf/setup/run.py
View file

@ -9,9 +9,11 @@ import shutil
import subprocess
import sys
import tempfile
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Final
from typing import BinaryIO, Final
from e2e_config import PROXY_BASE_URL
from pydantic import TypeAdapter
@ -73,7 +75,59 @@ def _keep_records_artifact(records_path: Path) -> None:
_ = shutil.copyfile(records_path, Path(artifact_dir) / "perf-records.jsonl")
def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: CaseSetup) -> RunOutcome:
class RecordStream:
def __init__(self, source: BinaryIO) -> None:
self.source: Final = source
self.partial = b""
def read(self) -> tuple[WireRecord, ...]:
chunks: Final = (self.partial + self.source.read()).split(b"\n")
self.partial = chunks[-1]
return tuple(_RECORD_ADAPTER.validate_json(line) for line in chunks[:-1] if line)
def run_until(
command: list[str],
env: dict[str, str],
records_path: Path,
stop_path: Path,
stop_when: Callable[[tuple[WireRecord, ...]], bool],
timeout: int,
) -> subprocess.CompletedProcess[str]:
with records_path.open("rb") as records, records_path.with_suffix(".log").open("w+") as output:
with subprocess.Popen(command, env=env, stdout=output, stderr=subprocess.STDOUT, text=True) as process:
deadline: Final = time.monotonic() + timeout
reader: Final = RecordStream(records)
try:
while process.poll() is None:
if stop_when(reader.read()):
stop_path.touch()
if time.monotonic() >= deadline:
raise subprocess.TimeoutExpired(command, timeout)
try:
_ = process.wait(timeout=1)
except subprocess.TimeoutExpired:
pass
finally:
stop_path.touch()
try:
_ = process.wait(timeout=35)
except subprocess.TimeoutExpired:
process.kill()
_ = process.wait(timeout=10)
_keep_records_artifact(records_path)
_ = output.seek(0)
return subprocess.CompletedProcess(command, process.returncode, output.read(), "")
def run_plan(
case: PerformanceTestCase,
plan: tuple[PlannedRequest, ...],
setup: CaseSetup,
*,
stop_when: Callable[[tuple[WireRecord, ...]], bool] | None = None,
request_timeout: float = 120.0,
) -> RunOutcome:
"""Replay `plan` with Locust in a child process and collect its records. The run-time ceiling
is well past the planned duration so Locust's own clock never cuts a slow run short; the
locustfile quits itself once the last record is written."""
@ -82,6 +136,7 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup:
with tempfile.TemporaryDirectory(prefix="e2e-perf-") as directory:
plan_path: Final = Path(directory) / "plan.json"
records_path: Final = Path(directory) / "records.jsonl"
stop_path: Final = Path(directory) / "stop"
_ = plan_path.write_bytes(_PLAN_ADAPTER.dump_json(wire))
records_path.touch()
env: Final = {
@ -91,33 +146,33 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup:
"PERF_RECORDS_FILE": str(records_path),
"PERF_RPS": str(case.traffic.rps),
"PERF_USERS": str(case.traffic.concurrency),
"PERF_REQUEST_TIMEOUT": str(request_timeout),
**({"PERF_STOP_FILE": str(stop_path)} if stop_when is not None else {}),
}
completed: Final = subprocess.run(
[
sys.executable,
"-P",
"-m",
"locust",
"--headless",
"--only-summary",
"--locustfile",
str(_LOCUSTFILE),
"--host",
PROXY_BASE_URL,
"--users",
str(case.traffic.concurrency),
"--spawn-rate",
str(case.traffic.concurrency),
"--run-time",
f"{ceiling_s}s",
"--exit-code-on-error",
"0",
],
env=env,
capture_output=True,
text=True,
timeout=ceiling_s + 120,
check=False,
command: Final = [
sys.executable,
"-P",
"-m",
"locust",
"--headless",
"--only-summary",
"--locustfile",
str(_LOCUSTFILE),
"--host",
PROXY_BASE_URL,
"--users",
str(case.traffic.concurrency),
"--spawn-rate",
str(case.traffic.concurrency),
"--run-time",
f"{ceiling_s}s",
"--exit-code-on-error",
"0",
]
completed: Final = (
run_until(command, env, records_path, stop_path, stop_when, ceiling_s + 120)
if stop_when is not None
else subprocess.run(command, env=env, capture_output=True, text=True, timeout=ceiling_s + 120, check=False)
)
records: Final = tuple(
_RECORD_ADAPTER.validate_json(line)

209
tests/e2e/perf/setup/upgrade.py Normal file
View file

@ -0,0 +1,209 @@
from __future__ import annotations
import math
import time
from collections.abc import Callable
from dataclasses import replace
from typing import Final
import psycopg
from cases.small_bodies_5rps import CASE
from e2e_config import MASTER_KEY
from e2e_http import send_text_external
from prometheus_client.parser import text_string_to_metric_families
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from setup.performance_test_case import PerformanceTestCase
from setup.upgrade_worker import UpgradeState, fingerprints, missing_indexes, read_state, release_indexes
from setup.wire import WireRecord
MAX_TRAFFIC_SECONDS: Final = 6 * 3600
_ROWS: Final = TypeAdapter(tuple[tuple[str, str | None], ...])
_ACTIVITY: Final = TypeAdapter(tuple[list[int], int])
_COUNT: Final = TypeAdapter(tuple[int])
_SERVICES: Final = ("postgres", "batch_write_to_db")
def upgrade_case() -> PerformanceTestCase:
return replace(
CASE,
name="database_upgrade_5rps",
traffic=replace(CASE.traffic, total_requests=int(CASE.traffic.rps * MAX_TRAFFIC_SECONDS)),
limits=tuple(
limit
for limit in CASE.limits
if limit.metric
in {
"unexpected_status_count",
"gateway_added_seconds",
"time_to_first_token_seconds",
"memory_growth_megabytes",
}
),
)
class MetricsHeaders(BaseModel):
model_config = ConfigDict(frozen=True)
authorization: str = Field(default=f"Bearer {MASTER_KEY}", alias="Authorization")
def postgres_failures(text: str) -> float:
families: Final = tuple(text_string_to_metric_families(text))
names: Final = frozenset(family.name for family in families)
required: Final = frozenset(f"litellm_{service}_failed_requests" for service in _SERVICES)
samples: Final = tuple(sample for family in families for sample in family.samples)
has_database_activity: Final = any(
sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0
for sample in samples
)
assert required <= names or has_database_activity, f"Missing database service counters: {required - names}"
failures: Final = tuple(sample.value for sample in samples if sample.name in {f"{name}_total" for name in required})
assert all(math.isfinite(value) and value >= 0 for value in failures), "Invalid database failure counter"
return sum(failures)
def assert_database_metrics(urls: tuple[str, ...]) -> None:
assert urls, "E2E_DATABASE_METRICS_URLS must name every gateway/backend /metrics endpoint"
for url in urls:
assert_database_metrics_url(url)
def assert_database_metrics_url(url: str) -> None:
result: Final = send_text_external("GET", url, headers=MetricsHeaders(), timeout=10)
assert result.status_code == 200, f"Database metrics unavailable at {url}: {result.status_code}"
assert postgres_failures(result.body) == 0, f"Failed PostgreSQL operations reported at {url}"
def persisted(connection: psycopg.Connection[tuple[object, ...]], records: tuple[WireRecord, ...]) -> frozenset[str]:
ids: Final = tuple(record.call_id for record in records if record.call_id)
if not ids:
return frozenset()
rows: Final = _ROWS.validate_python(
connection.execute(
'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=ANY(%s)', (list(ids),)
).fetchall()
)
return matched_response_ids(records, rows)
def matched_response_ids(records: tuple[WireRecord, ...], rows: tuple[tuple[str, str | None], ...]) -> frozenset[str]:
call_ids: Final = frozenset(call_id for _, call_id in rows if call_id)
return frozenset(record.response_id for record in records if record.response_id and record.call_id in call_ids)
def activity(connection: psycopg.Connection[tuple[object, ...]]) -> tuple[list[int], int]:
return _ACTIVITY.validate_python(
connection.execute("""SELECT
ARRAY(SELECT pid FROM pg_stat_progress_create_index WHERE datname=current_database()),
(SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock')
""").fetchone()
)
def assert_finished(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None:
assert state.phase == "complete", f"Upgrade did not finish: {state.phase}"
assert not missing_indexes(connection, state.required_indexes), "Required target indexes are missing or invalid"
(missing,) = _COUNT.validate_python(
connection.execute(
"""
SELECT count(*) FROM unnest(%s::text[]) name WHERE NOT EXISTS (
SELECT 1 FROM "_prisma_migrations" WHERE migration_name=name
AND finished_at IS NOT NULL AND rolled_back_at IS NULL)""",
(list(state.required_migrations),),
).fetchone()
)
assert missing == 0, f"{missing} target migrations were not applied"
(unfinished,) = _COUNT.validate_python(
connection.execute("""
SELECT count(*) FROM "_prisma_migrations" WHERE finished_at IS NULL AND rolled_back_at IS NULL
""").fetchone()
)
assert unfinished == 0, f"{unfinished} unfinished migrations"
assert fingerprints(connection, tuple(state.baseline_rows)) == state.baseline_rows, (
"Baseline rows changed or disappeared"
)
class UpgradeMonitor:
def __init__(
self,
*,
state: Callable[[], UpgradeState],
check_metrics: Callable[[], None],
read_activity: Callable[[], tuple[list[int], int]],
read_persisted: Callable[[tuple[WireRecord, ...]], frozenset[str]],
start_indexes: Callable[[], None],
clock: Callable[[], float] = time.monotonic,
wall_clock: Callable[[], float] = time.time,
) -> None:
self.state: Final = state
self.check_metrics: Final = check_metrics
self.read_activity: Final = read_activity
self.read_persisted: Final = read_persisted
self.start_indexes: Final = start_indexes
self.clock: Final = clock
self.wall_clock: Final = wall_clock
self.started: Final = clock()
self.status_limit: Final = next(limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count")
self.pending: tuple[WireRecord, ...] = ()
self.processed = 0
self.successes = 0
self.persisted_count = 0
self.overlap_writes = 0
self.max_lock_waiters = 0
self.previous_builds: frozenset[int] = frozenset()
self.build_seen_at: float | None = None
self.released = False
self.completed_at: float | None = None
def poll(self, records: tuple[WireRecord, ...]) -> bool:
self.check_metrics()
current: Final = self.state()
assert current.phase != "failed", "Target migration/index worker failed"
assert self.clock() - self.started < MAX_TRAFFIC_SECONDS, "Upgrade exceeded its six-hour operational timeout"
successful: Final = tuple(record for record in records if record.status == 200)
self.processed += len(records)
self.successes += len(successful)
assert self.processed - self.successes <= self.status_limit, "Too many unexpected non-200 responses"
assert all(record.ok and record.text for record in successful), "Failed or empty inference response"
assert all(record.response_id and record.call_id for record in successful), (
"Response missing spend-log identifiers"
)
builds, lock_waiters = self.read_activity()
if not self.previous_builds.intersection(builds):
self.build_seen_at = self.wall_clock() if builds else None
self.max_lock_waiters = max(self.max_lock_waiters, lock_waiters)
self.pending += successful
found: Final = self.read_persisted(self.pending)
self.persisted_count += len(found)
still_building, _ = self.read_activity()
since: Final = self.build_seen_at
if since is not None and frozenset(builds).intersection(still_building):
self.overlap_writes += sum(
1 for record in self.pending if record.response_id in found and record.started_at >= since
)
self.previous_builds = frozenset(builds)
self.pending = tuple(record for record in self.pending if record.response_id not in found)
if self.successes and not self.released:
self.start_indexes()
self.released = True
if current.phase == "complete" and self.completed_at is None:
self.completed_at = self.clock()
return self.completed_at is not None and self.clock() - self.completed_at >= 60 and self.successes >= 300
def assert_complete(self, state: UpgradeState) -> None:
assert self.successes >= 300, f"Only {self.successes} requests completed"
assert not self.pending, f"{len(self.pending)} successful requests have no matching spend log"
assert self.persisted_count == self.successes, "Spend-log count does not match successful requests"
assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work"
def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor:
return UpgradeMonitor(
state=lambda: read_state(connection),
check_metrics=lambda: assert_database_metrics(urls),
read_activity=lambda: activity(connection),
read_persisted=lambda records: persisted(connection, records),
start_indexes=lambda: release_indexes(connection),
)

146
tests/e2e/perf/setup/upgrade_worker.py Normal file
View file

@ -0,0 +1,146 @@
from __future__ import annotations
import json
import os
import sys
import time
from pathlib import Path
from typing import Final, Literal
import psycopg
from pydantic import BaseModel, ConfigDict, TypeAdapter
class UpgradeState(BaseModel):
model_config = ConfigDict(frozen=True)
phase: Literal["schema", "schema_ready", "indexes", "complete", "failed"]
started_at: float
schema_seconds: float = 0
index_seconds: float = 0
required_indexes: tuple[str, ...]
required_migrations: tuple[str, ...]
pending_indexes: tuple[str, ...]
baseline_rows: dict[str, str]
_STATE: Final = TypeAdapter(tuple[UpgradeState])
_BOOL: Final = TypeAdapter(tuple[bool])
_STRINGS: Final = TypeAdapter(tuple[tuple[str], ...])
_FINGERPRINTS: Final = TypeAdapter(tuple[tuple[str, str], ...])
_FINGERPRINT: Final = """md5(ROW(request_id, api_key, spend, total_tokens, prompt_tokens,
completion_tokens, model, "startTime", "endTime")::text)"""
def connect() -> psycopg.Connection[tuple[object, ...]]:
return psycopg.connect(
os.environ["DATABASE_URL"],
autocommit=True,
connect_timeout=10,
options="-c statement_timeout=10000 -c timezone=UTC",
)
def fingerprints(connection: psycopg.Connection[tuple[object, ...]], ids: tuple[str, ...]) -> dict[str, str]:
rows: Final = connection.execute(
f'SELECT request_id, {_FINGERPRINT} FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', (list(ids),)
).fetchall()
return dict(_FINGERPRINTS.validate_python(rows))
def missing_indexes(connection: psycopg.Connection[tuple[object, ...]], names: tuple[str, ...]) -> tuple[str, ...]:
rows: Final = connection.execute(
"""SELECT name FROM unnest(%s::text[]) AS name WHERE NOT EXISTS (
SELECT 1 FROM pg_index WHERE indexrelid = to_regclass(format('public.%%I', name))
AND indisvalid AND indisready)""",
(list(names),),
).fetchall()
return tuple(row[0] for row in _STRINGS.validate_python(rows))
def read_state(connection: psycopg.Connection[tuple[object, ...]]) -> UpgradeState:
(state,) = _STATE.validate_python(
connection.execute('SELECT state FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone()
)
return state
def write_state(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None:
_ = connection.execute(
"""INSERT INTO "_LiteLLMPerfUpgrade" (id, state) VALUES (1, %s::jsonb)
ON CONFLICT (id) DO UPDATE SET state = EXCLUDED.state""",
(state.model_dump_json(),),
)
def release_indexes(connection: psycopg.Connection[tuple[object, ...]]) -> None:
_ = connection.execute('UPDATE "_LiteLLMPerfUpgrade" SET traffic_ready=true WHERE id=1')
def wait_for_traffic(connection: psycopg.Connection[tuple[object, ...]], deadline: float) -> bool:
while time.monotonic() < deadline:
(ready,) = _BOOL.validate_python(
connection.execute('SELECT traffic_ready FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone()
)
if ready:
return True
time.sleep(1)
return False
def main() -> int:
assert os.environ.get("E2E_DATABASE_UPGRADE") == "1", "Only run against a disposable restored baseline"
assert os.environ.get("DATABASE_PERF_CONFIRM_DISPOSABLE") == "1", "Disposable database confirmation required"
from litellm_proxy_extras import request_log_indexes
from litellm_proxy_extras.utils import ProxyExtrasDBManager
with connect() as connection:
_ = connection.execute("""CREATE TABLE "_LiteLLMPerfUpgrade" (
id integer PRIMARY KEY CHECK (id=1), state jsonb NOT NULL, traffic_ready boolean NOT NULL DEFAULT false)""")
sample: Final = connection.execute("""
(SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id LIMIT 32)
UNION (SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id DESC LIMIT 32)""").fetchall()
ids: Final = tuple(row[0] for row in _STRINGS.validate_python(sample))
assert ids, "Restored baseline has no spend logs"
names: Final = tuple(index.name for index in request_log_indexes.REQUEST_LOG_INDEXES)
migrations: Final = Path(request_log_indexes.__file__).parent / "migrations"
initial: Final = UpgradeState(
phase="schema",
started_at=time.time(),
required_indexes=names,
required_migrations=tuple(sorted(path.parent.name for path in migrations.glob("*/migration.sql"))),
pending_indexes=missing_indexes(connection, names),
baseline_rows=fingerprints(connection, ids),
)
assert initial.required_migrations and initial.required_indexes, "Target schema requirements are missing"
write_state(connection, initial)
try:
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True), (
"Schema migration failed"
)
ready: Final = initial.model_copy(
update={
"phase": "schema_ready",
"schema_seconds": time.time() - initial.started_at,
"pending_indexes": missing_indexes(connection, names),
}
)
write_state(connection, ready)
assert wait_for_traffic(connection, time.monotonic() + 1800), "Traffic did not start within 30 minutes"
write_state(connection, ready.model_copy(update={"phase": "indexes"}))
index_start: Final = time.monotonic()
assert ProxyExtrasDBManager.build_request_log_indexes(), "Request-log index build failed"
complete: Final = ready.model_copy(
update={"phase": "complete", "index_seconds": time.monotonic() - index_start}
)
write_state(connection, complete)
_ = sys.stdout.write(complete.model_dump_json() + "\n")
return 0
except Exception as exc:
write_state(connection, read_state(connection).model_copy(update={"phase": "failed"}))
_ = sys.stderr.write(json.dumps({"error": type(exc).__name__}) + "\n")
return 1
if __name__ == "__main__":
raise SystemExit(main())

30
tests/e2e/perf/setup/wire.py
View file

@ -9,12 +9,13 @@ from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, Field
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from setup.corpus.session import Endpoint as CorpusEndpoint
from setup.corpus.session import request_body as corpus_body
Endpoint = Literal["chat_completions", "responses", "messages"]
_PAYLOAD: Final = TypeAdapter(dict[str, object])
class TokenSource(BaseModel):
@ -70,6 +71,33 @@ class WireRecord(BaseModel):
mock_delay_s: float | None
call_id: str | None
error: str | None
response_id: str | None = None
def response_id(payload: Mapping[str, object]) -> str | None:
nested: Final = payload.get("response") or payload.get("message")
source: Final = _PAYLOAD.validate_python(nested) if isinstance(nested, Mapping) else payload
identifier: Final = source.get("id")
return identifier if isinstance(identifier, str) else None
class PlanProgress:
def __init__(self, total: int) -> None:
self.total: Final = total
self.issued = 0
self.completed = 0
self.stopped = False
def take(self, *, stop: bool) -> bool:
self.stopped = self.stopped or stop
if self.stopped or self.issued >= self.total:
return False
self.issued += 1
return True
@property
def drained(self) -> bool:
return (self.stopped or self.issued == self.total) and self.completed == self.issued
_CORPUS_ENDPOINT: Final[Mapping[Endpoint, CorpusEndpoint]] = MappingProxyType(

85
tests/e2e/perf/test_database_upgrade.py Normal file
View file

@ -0,0 +1,85 @@
from __future__ import annotations
import os
import sys
import time
from pathlib import Path
from typing import Final
import pytest
from idp import Keycloak
from lifecycle import ResourceManager
from memory_readings import read_rss_everywhere
from proxy_client import ProxyClient
from setup.datastore import read_db_counters, read_redis_counters
from setup.observe import Observation, evaluate, read_mock_requests, render, reset_mock, set_mock_latency
from setup.plan import build_plan
from setup.podcpu import gateway_cpu_sampler
from setup.provision import MOCK_PROVIDERS, build_setup, mock_url
from setup.run import run_plan
from setup.upgrade import assert_finished, monitor_for, upgrade_case
from setup.upgrade_worker import connect, read_state
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only")
def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None:
case: Final = upgrade_case()
urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS)
for url in urls:
_ = reset_mock(url)
_ = set_mock_latency(url, case.mock_latency)
setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources)
metrics: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip())
with connect() as connection:
assert read_state(connection).phase == "schema_ready", "Expected the target schema ready, index work waiting"
rss_before: Final = read_rss_everywhere(proxy)
db_before: Final = read_db_counters()
redis_before: Final = read_redis_counters()
monitor: Final = monitor_for(connection, metrics)
with gateway_cpu_sampler() as cpu:
outcome: Final = run_plan(case, build_plan(case), setup, stop_when=monitor.poll, request_timeout=30)
assert outcome.locust_exit == 0, outcome.locust_output
remaining: Final = outcome.records[monitor.processed :]
_ = monitor.poll(remaining)
deadline: Final = time.monotonic() + 120
while monitor.pending and time.monotonic() < deadline:
time.sleep(1)
_ = monitor.poll(())
state: Final = read_state(connection)
_ = sys.stdout.write(
f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} "
f"unmatched_requests={len(monitor.pending)}\n"
)
monitor.assert_complete(state)
assert_finished(connection, state)
rss_after: Final = read_rss_everywhere(proxy)
assert {r.worker for r in rss_before.readings} == {r.worker for r in rss_after.readings}, (
"Gateway workers restarted"
)
observed: Final = Observation(
records=outcome.records,
mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)),
guardrail_calls=(),
rss_before=rss_before,
rss_after=rss_after,
wall_seconds=outcome.wall_seconds,
spend_logged_share=1,
db_before=db_before,
db_after=read_db_counters(),
redis_before=redis_before,
redis_after=read_redis_counters(),
gateway_cpu_seconds=cpu.seconds,
)
rows, exceeded = evaluate(case, observed)
_ = sys.stdout.write(render(case, observed, rows) + "\n")
_ = sys.stdout.write(
f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} "
f"overlap_writes={monitor.overlap_writes} max_lock_waiters={monitor.max_lock_waiters} "
f"pending_indexes={state.pending_indexes}\n"
)
artifact: Final = os.environ.get("E2E_ARTIFACT_DIR")
if artifact:
_ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8")
assert not exceeded, "\n".join(exceeded)

66
tests/e2e_harness/perf/test_run.py Normal file
View file

@ -0,0 +1,66 @@
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import Final
import pytest
from setup.run import run_until
from setup.wire import WireRecord
from test_upgrade import record
class StopAfterRecord:
def __init__(self, *, partial: Path, resume: Path, fail: bool = False) -> None:
self.partial: Final = partial
self.resume: Final = resume
self.fail: Final = fail
self.received: tuple[WireRecord, ...] = ()
def __call__(self, records: tuple[WireRecord, ...]) -> bool:
self.received += records
if self.partial.exists():
self.resume.touch()
assert not (self.fail and records), "Database failed"
return bool(records)
@pytest.mark.parametrize("fail", [False, True])
def test_controlled_process_preserves_partial_lines_and_drains_on_stop_or_failure(tmp_path: Path, fail: bool) -> None:
records: Final = tmp_path / "records.jsonl"
stop: Final = tmp_path / "stop"
partial: Final = tmp_path / "partial"
resume: Final = tmp_path / "resume"
records.touch()
encoded: Final = record().model_dump_json() + "\n"
program: Final = f"""
import time
from pathlib import Path
with Path({str(records)!r}).open("w") as f:
text = {encoded!r}
f.write(text[:10])
f.flush()
Path({str(partial)!r}).touch()
while not Path({str(resume)!r}).exists():
time.sleep(0.01)
f.write(text[10:])
f.flush()
while not Path({str(stop)!r}).exists():
time.sleep(0.01)
f.write({record(1).model_dump_json()!r} + "\\n")
f.flush()
"""
callback: Final = StopAfterRecord(partial=partial, resume=resume, fail=fail)
command: Final = [sys.executable, "-c", program]
if fail:
with pytest.raises(AssertionError, match="Database failed"):
_ = run_until(command, dict(os.environ), records, stop, callback, timeout=10)
else:
result: Final = run_until(command, dict(os.environ), records, stop, callback, timeout=10)
assert result.returncode == 0, result.stdout
assert callback.received == (record(),)
assert tuple(WireRecord.model_validate_json(line) for line in records.read_text().splitlines()) == (
record(),
record(1),
)

325
tests/e2e_harness/perf/test_upgrade.py Normal file
View file

@ -0,0 +1,325 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Final
import pytest
from cases.small_bodies_5rps import CASE
from memory_readings import RssCapture, RssReading
from setup.observe import Observation, evaluate
from setup.upgrade import MAX_TRAFFIC_SECONDS, UpgradeMonitor, matched_response_ids, postgres_failures, upgrade_case
from setup.upgrade_worker import UpgradeState
from setup.wire import PlanProgress, WireRecord, response_id
def record(seq: int = 0, *, started_at: float = 100) -> WireRecord:
return WireRecord(
seq=seq,
marker=f"request-{seq}",
caller="key",
endpoint="chat_completions",
stream=False,
blocked=False,
started_at=started_at,
status=200,
header_s=1,
first_byte_s=1,
first_token_s=None,
total_s=1,
text="answer",
ok=True,
mock_delay_s=1,
call_id=f"call-{seq}",
error=None,
response_id=f"response-{seq}",
)
def test_spend_logs_match_call_ids_when_response_ids_are_transformed() -> None:
records: Final = (record(0), record(1), record(2))
rows: Final = (("provider-response-id", "call-0"), ("response-1", "unrelated-call"))
assert matched_response_ids(records, rows) == frozenset({"response-0"})
def state(phase: str = "schema_ready", *, pending: bool = True) -> UpgradeState:
return UpgradeState.model_validate(
{
"phase": phase,
"started_at": 0,
"required_indexes": ["index"],
"required_migrations": ["migration"],
"pending_indexes": ["index"] if pending else [],
"baseline_rows": {"seed-row": "hash"},
}
)
class Probe:
def __init__(self) -> None:
self.now = 0.0
self.current = state()
self.builds: list[int] = []
self.found: frozenset[str] = frozenset()
self.releases = 0
self.metric_checks = 0
self.errors = False
def check_metrics(self) -> None:
self.metric_checks += 1
assert not self.errors, "Failed PostgreSQL operation"
def persisted(self, records: tuple[WireRecord, ...]) -> frozenset[str]:
return frozenset(row.response_id for row in records if row.response_id in self.found and row.response_id)
def release(self) -> None:
self.releases += 1
def monitor(self) -> UpgradeMonitor:
return UpgradeMonitor(
state=lambda: self.current,
check_metrics=self.check_metrics,
read_activity=lambda: (self.builds, 2),
read_persisted=self.persisted,
start_indexes=self.release,
clock=lambda: self.now,
wall_clock=lambda: self.now,
)
def test_stop_drains_issued_requests_without_issuing_the_rest() -> None:
progress: Final = PlanProgress(100)
assert progress.take(stop=False)
assert progress.take(stop=False)
progress.completed = 1
assert not progress.take(stop=True)
assert not progress.drained
assert not progress.take(stop=False)
progress.completed = 2
assert progress.drained
assert progress.issued == 2
def test_finite_plan_still_waits_for_its_last_response() -> None:
progress: Final = PlanProgress(1)
assert progress.take(stop=False)
assert not progress.take(stop=False)
assert not progress.drained
progress.completed = 1
assert progress.drained
@pytest.mark.parametrize(
"payload,expected",
[
({"id": "chat-id"}, "chat-id"),
({"type": "response.created", "response": {"id": "response-id"}}, "response-id"),
({"type": "message_start", "message": {"id": "message-id"}}, "message-id"),
({"type": "response.output_item.added", "item": {"id": "item-id"}}, None),
({"type": "content_block_delta", "delta": {"text": "hello"}}, None),
],
)
def test_response_id_uses_response_not_output_item(payload: Mapping[str, object], expected: str | None) -> None:
assert response_id(payload) == expected
def metrics(*, failure: str = "0", batch: str = "0") -> str:
return f"""# TYPE litellm_postgres_failed_requests_total counter
litellm_postgres_failed_requests_total{{function_name="query",error_class="Timeout"}} {failure}
# TYPE litellm_batch_write_to_db_failed_requests_total counter
litellm_batch_write_to_db_failed_requests_total{{function_name="flush",error_class="DatabaseError"}} {batch}
"""
def test_failed_sql_and_async_batch_writes_are_both_counted() -> None:
assert postgres_failures(metrics(failure="2", batch="3")) == 5
assert postgres_failures(metrics()) == 0
def test_no_error_label_series_is_a_valid_zero() -> None:
assert (
postgres_failures("""# TYPE litellm_postgres_failed_requests_total counter
# TYPE litellm_batch_write_to_db_failed_requests_total counter
""")
== 0
)
@pytest.mark.parametrize("service", [None, "postgres", "batch_write_to_db"])
def test_multiprocess_metrics_omit_failure_families_until_a_failure(service: str | None) -> None:
total: Final = """# TYPE litellm_postgres_total_requests_total counter
litellm_postgres_total_requests_total{postgres="postgres"} 1
"""
failure: Final = (
f"# TYPE litellm_{service}_failed_requests_total counter\n"
f'litellm_{service}_failed_requests_total{{function_name="query",error_class="Timeout"}} 1\n'
if service is not None
else ""
)
assert postgres_failures(total + failure) == (1 if service is not None else 0)
@pytest.mark.parametrize("total", ["0", "-1", "NaN", "+Inf"])
def test_missing_failure_families_require_observed_database_activity(total: str) -> None:
with pytest.raises(AssertionError):
_ = postgres_failures(
"# TYPE litellm_postgres_total_requests_total counter\n"
f'litellm_postgres_total_requests_total{{postgres="postgres"}} {total}\n'
)
@pytest.mark.parametrize("body", ["", "# TYPE unrelated counter\n", metrics(failure="NaN"), metrics(batch="-1")])
def test_missing_or_invalid_counters_do_not_pass_as_zero(body: str) -> None:
with pytest.raises(AssertionError):
_ = postgres_failures(body)
def test_index_completion_stops_after_the_post_upgrade_window() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
assert not monitor.poll((record(),))
assert probe.releases == 1
probe.builds = [17]
probe.now = 101
assert not monitor.poll(())
probe.found = frozenset({"response-0", "response-1"})
probe.now = 103
assert not monitor.poll((record(1, started_at=102),))
assert monitor.overlap_writes == 1
assert monitor.persisted_count == 2
assert not monitor.pending
assert monitor.max_lock_waiters == 2
probe.current = state("complete")
assert not monitor.poll(())
probe.now = 162
assert not monitor.poll(())
probe.now = 163
assert not monitor.poll(())
probe.found = frozenset(f"response-{i}" for i in range(300))
assert monitor.poll(tuple(record(i) for i in range(2, 300)))
assert probe.releases == 1
assert probe.metric_checks == 7
def test_http_200_without_spend_row_cannot_complete() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
records: Final = tuple(record(i) for i in range(300))
_ = monitor.poll(records)
with pytest.raises(AssertionError, match="no matching spend log"):
monitor.assert_complete(state("complete", pending=False))
def test_no_pending_index_work_still_requires_300_persisted_requests() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
records: Final = tuple(record(i) for i in range(300))
probe.found = frozenset(f"response-{i}" for i in range(300))
_ = monitor.poll(records)
monitor.assert_complete(state("complete", pending=False))
assert monitor.successes == monitor.persisted_count == 300
with pytest.raises(AssertionError, match="during index work"):
monitor.assert_complete(state("complete", pending=True))
def test_postgres_failure_fails_even_when_inference_succeeded() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
probe.errors = True
with pytest.raises(AssertionError, match="PostgreSQL"):
_ = monitor.poll((record(),))
def test_failed_migration_and_operational_timeout_are_not_success() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
probe.current = state("failed")
with pytest.raises(AssertionError, match="worker failed"):
_ = monitor.poll(())
probe.current = state()
probe.now = MAX_TRAFFIC_SECONDS
with pytest.raises(AssertionError, match="operational timeout"):
_ = monitor.poll(())
def test_upgrade_case_uses_5rps_with_the_same_request_and_memory_limits() -> None:
case: Final = upgrade_case()
assert case.traffic.rps == 5
assert not {"cpu_seconds_per_request", "postgres_calls_per_request", "redis_calls_per_request"}.intersection(
limit.metric for limit in case.limits
)
assert case.limits == tuple(
limit
for limit in CASE.limits
if limit.metric
in {
"unexpected_status_count",
"gateway_added_seconds",
"time_to_first_token_seconds",
"memory_growth_megabytes",
}
)
def failed_record(seq: int) -> WireRecord:
return record(seq).model_copy(update={"status": 503, "ok": False, "text": "", "response_id": None, "call_id": None})
def test_three_non_200s_do_not_start_indexes_or_require_spend_rows() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
assert not monitor.poll(tuple(failed_record(i) for i in range(3)))
assert probe.releases == 0
assert monitor.processed == 3
assert monitor.successes == 0
assert not monitor.pending
probe.found = frozenset(f"response-{i}" for i in range(3, 303))
assert not monitor.poll(tuple(record(i) for i in range(3, 303)))
monitor.assert_complete(state("complete", pending=False))
assert probe.releases == 1
assert monitor.processed == 303
assert monitor.successes == monitor.persisted_count == 300
def test_fourth_non_200_fails_across_poll_batches() -> None:
monitor: Final = Probe().monitor()
assert not monitor.poll(tuple(failed_record(i) for i in range(3)))
with pytest.raises(AssertionError, match="Too many unexpected non-200"):
_ = monitor.poll((failed_record(3),))
@pytest.mark.parametrize(
"failed", [record().model_copy(update={"text": ""}), record().model_copy(update={"ok": False})]
)
def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None:
with pytest.raises(AssertionError, match="Failed or empty inference"):
_ = Probe().monitor().poll((failed,))
@pytest.mark.parametrize(
"failures,overhead,ttft,growth,expected",
[
(3, 0.29, 3.7, 100, ()),
(4, 0.29, 3.7, 100, ("unexpected_status_count",)),
(3, 0.30, 3.7, 100, ("gateway_added_seconds percentile 50",)),
(3, 0.29, 3.8, 100, ("time_to_first_token_seconds percentile 99",)),
(3, 0.29, 3.7, 101, ("memory_growth_megabytes",)),
],
)
def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples(
failures: int, overhead: float, ttft: float, growth: float, expected: tuple[str, ...]
) -> None:
nonstream: Final = tuple(
record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100)
)
stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft})
observed: Final = Observation(
records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))),
mock_requests=(),
guardrail_calls=(),
rss_before=RssCapture((RssReading("gateway", "host", 1, 1000),), ()),
rss_after=RssCapture((RssReading("gateway", "host", 1, 1000 + growth),), ()),
wall_seconds=60,
)
rows, exceeded = evaluate(upgrade_case(), observed)
assert tuple(row.label for row in rows if row.over or row.missing) == expected
assert len(exceeded) == len(expected)

179
tests/integration/observability/test_generic_guardrail_num_retries.py Normal file
View file

@ -0,0 +1,179 @@
"""generic_guardrail_api `num_retries` resends a check that timed out waiting for its answer, through a real proxy.
Each guardrail points at its own owned stub. A stub that holds a check never answers it before the 1s guardrail
timeout, and holding lasts until teardown so a resent check can only be answered by a later arrival.
"""
import itertools
import json
import threading
from collections.abc import Callable, Iterator, Mapping
from contextlib import ExitStack
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.openai_wire import chat_reply
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
MODEL: Final = "guarded-chat"
UPSTREAM_MODEL: Final = "openai/gpt-6.1-sol"
GUARDRAIL_TIMEOUT_SECONDS: Final = 1
HOLDS_FIRST_RETRIED: Final = "holds-first-check-retried-once"
HOLDS_FIRST_UNSET: Final = "holds-first-check-num-retries-unset"
HOLDS_ALL_RETRIED_TWICE: Final = "holds-every-check-retried-twice"
UNAVAILABLE_RETRIED_TWICE: Final = "unavailable-retried-twice"
ANSWER: Final = "upstream answer"
NO_ACTION: Final = Reply(body=json.dumps({"action": "NONE"}).encode())
_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
_JSON: Final = TypeAdapter(dict[str, JsonValue])
@dataclass(frozen=True, slots=True)
class Rig:
owned: OwnedProxy
upstream: Wire
guardrails: Mapping[str, Wire]
def _holds_first(release: threading.Event) -> Callable[[Request], Reply]:
arrivals: Final = itertools.count()
def respond(_request: Request) -> Reply:
if next(arrivals) == 0:
release.wait()
return NO_ACTION
return respond
def _holds_all(release: threading.Event) -> Callable[[Request], Reply]:
def respond(_request: Request) -> Reply:
release.wait()
return NO_ACTION
return respond
def _unavailable(_request: Request) -> Reply:
return Reply(status=503, body=b'{"error": "guardrail unavailable"}')
def _upstream(_request: Request) -> Reply:
return chat_reply("chatcmpl-num-retries", "gpt-6.1-sol", ANSWER, stream=False)
def _guardrail(name: str, wire: Wire, num_retries: int | None) -> dict[str, JsonValue]:
return {
"guardrail_name": name,
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": False,
"api_base": wire.url,
"api_key": "synthetic-guardrail-key",
"timeout": GUARDRAIL_TIMEOUT_SECONDS,
**({} if num_retries is None else {"num_retries": num_retries}),
},
}
@pytest.fixture(scope="module")
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("generic-guardrail-num-retries")
release: Final = threading.Event()
with ExitStack() as stack:
upstream: Final = stack.enter_context(wire_server(_upstream))
guardrails: Final = MappingProxyType(
{
HOLDS_FIRST_RETRIED: stack.enter_context(wire_server(_holds_first(release), keep_alive=True)),
HOLDS_FIRST_UNSET: stack.enter_context(wire_server(_holds_first(release), keep_alive=True)),
HOLDS_ALL_RETRIED_TWICE: stack.enter_context(wire_server(_holds_all(release), keep_alive=True)),
UNAVAILABLE_RETRIED_TWICE: stack.enter_context(wire_server(_unavailable, keep_alive=True)),
}
)
stack.callback(release.set)
retries: Final = MappingProxyType(
{HOLDS_FIRST_RETRIED: 1, HOLDS_FIRST_UNSET: None, HOLDS_ALL_RETRIED_TWICE: 2, UNAVAILABLE_RETRIED_TWICE: 2}
)
base: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text()))
path: Final = directory / "generic-guardrail-num-retries.yaml"
path.write_text(
yaml.safe_dump(
{
**base,
"model_list": [
{
"model_name": MODEL,
"litellm_params": {
"model": UPSTREAM_MODEL,
"api_base": f"{upstream.url}/v1",
"api_key": "synthetic-openai-key",
},
}
],
"guardrails": [_guardrail(name, wire, retries[name]) for name, wire in guardrails.items()],
}
)
)
owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, {}, config=path))
yield Rig(owned, upstream, guardrails)
def _chat(rig: Rig, guardrail: str) -> httpx.Response:
return rig.owned.gateway.request(
"POST",
"/v1/chat/completions",
{"model": MODEL, "messages": [{"role": "user", "content": "is this allowed?"}], "guardrails": [guardrail]},
)
def _checks(wire: Wire) -> tuple[bytes, ...]:
return tuple(request.body for request in wire.drain())
def test_check_held_past_the_timeout_is_resent_on_a_fresh_connection_and_the_request_is_served(rig: Rig) -> None:
response: Final = _chat(rig, HOLDS_FIRST_RETRIED)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == ANSWER
stub: Final = rig.guardrails[HOLDS_FIRST_RETRIED]
checks: Final = _checks(stub)
assert len(checks) == 2, checks
assert checks[0] == checks[1]
assert stub.connections() == 2
assert len(rig.upstream.drain()) == 1
def test_check_held_past_the_timeout_is_sent_once_when_num_retries_is_unset(rig: Rig) -> None:
response: Final = _chat(rig, HOLDS_FIRST_UNSET)
assert response.status_code != 200, response.text
assert "Generic Guardrail API failed" in response.text
assert len(_checks(rig.guardrails[HOLDS_FIRST_UNSET])) == 1
assert rig.upstream.drain() == ()
def test_num_retries_bounds_how_many_times_a_held_check_is_sent(rig: Rig) -> None:
response: Final = _chat(rig, HOLDS_ALL_RETRIED_TWICE)
assert response.status_code != 200, response.text
assert "Generic Guardrail API failed" in response.text
assert len(_checks(rig.guardrails[HOLDS_ALL_RETRIED_TWICE])) == 3
assert rig.upstream.drain() == ()
def test_error_status_from_the_guardrail_is_not_resent(rig: Rig) -> None:
response: Final = _chat(rig, UNAVAILABLE_RETRIED_TWICE)
assert response.status_code != 200, response.text
assert "Generic Guardrail API failed" in response.text
assert len(_checks(rig.guardrails[UNAVAILABLE_RETRIED_TWICE])) == 1
assert rig.upstream.drain() == ()

433
tests/integration/pricing/test_price_data_reload_prices.py Normal file
View file

@ -0,0 +1,433 @@
import itertools
import json
import re
import threading
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
_ROW: Final = "anthropic.claude-audit-reload-v1:0"
_GEO: Final = f"au.{_ROW}"
_FINETUNE: Final = "ft:gpt-audit-reload-2025-01-01"
_CLAUDE: Final = "claude-audit-reload-sonnet"
_MINI: Final = "gpt-audit-reload-mini"
_USAGE: Final = re.compile(r"bill (\d+) (\d+) (\d+)")
_CUSTOM: Final = MappingProxyType({"input_cost_per_token": 1.1e-06, "output_cost_per_token": 2.3e-06})
_COST_MAP_LIMITS: Final = MappingProxyType({"MODEL_COST_MAP_MIN_MODEL_COUNT": "1", "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0"})
_PUBLISHED_SCALE: Final = 2.0
_CATALOG: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType(
{
_ROW: {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token_above_200k_tokens": 6e-06,
"output_cost_per_token_above_200k_tokens": 2.25e-05,
},
_GEO: {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
},
_FINETUNE: {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 3.75e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 1.875e-06,
},
_CLAUDE: {
"litellm_provider": "anthropic",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 3e-07,
},
_MINI: {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 5e-07,
"output_cost_per_token": 2e-06,
},
}
)
@dataclass(frozen=True, slots=True)
class _Usage:
prompt: int
cached: int
completion: int
def text(self) -> str:
return f"bill {self.prompt} {self.cached} {self.completion}"
@dataclass(frozen=True, slots=True)
class _Rates:
input: float
output: float
cache_read: float
def cost(self, usage: _Usage) -> float:
return (
(usage.prompt - usage.cached) * self.input + usage.cached * self.cache_read + usage.completion * self.output
)
@dataclass(frozen=True, slots=True)
class _Deployment:
name: str
model: str
prices: Mapping[str, float]
usage: _Usage
rates: Callable[[float], _Rates | None]
@dataclass(frozen=True, slots=True)
class _Charge:
request_id: str
expected: float
def _catalog(scale: float) -> dict[str, dict[str, JsonValue]]:
return {
model: {name: value * scale if isinstance(value, float) else value for name, value in row.items()}
for model, row in _CATALOG.items()
}
def _rate(model: str, field: str, scale: float) -> float:
rate: Final = _CATALOG[model][field]
assert isinstance(rate, float), rate
return rate * scale
def _catalog_rates(model: str) -> Callable[[float], _Rates]:
return lambda scale: _Rates(
input=_rate(model, "input_cost_per_token", scale),
output=_rate(model, "output_cost_per_token", scale),
cache_read=_rate(model, "cache_read_input_token_cost", scale)
if "cache_read_input_token_cost" in _CATALOG[model]
else _rate(model, "input_cost_per_token", scale),
)
def _custom_rates(_scale: float) -> _Rates:
return _Rates(
input=_CUSTOM["input_cost_per_token"],
output=_CUSTOM["output_cost_per_token"],
cache_read=_CUSTOM["input_cost_per_token"],
)
def _unpriced(_scale: float) -> None:
return None
def _expected_cost(rates: _Rates | None, usage: _Usage) -> float:
return 0.0 if rates is None else rates.cost(usage)
def _expected_listing(rates: _Rates | None) -> tuple[object, object]:
if rates is None:
return None, None
return pytest.approx(rates.input), pytest.approx(rates.output)
def _connection(model: str, wire_url: str) -> dict[str, JsonValue]:
if model.startswith("bedrock/"):
return {"api_base": wire_url, "api_key": "synthetic-bedrock-bearer", "aws_region_name": "ap-northeast-1"}
if model.startswith("azure/"):
return {"api_base": wire_url, "api_key": "synthetic-azure-key", "api_version": "2025-01-01-preview"}
return {"api_base": f"{wire_url}/v1", "api_key": "synthetic-openai-key"}
def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str) -> Path:
config: Final = directory / "price_data_reload_config.yaml"
config.write_text(
json.dumps(
{
"model_list": [
{
"model_name": deployment.name,
"litellm_params": {
"model": deployment.model,
**_connection(deployment.model, wire_url),
**deployment.prices,
},
}
for deployment in deployments
],
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
},
"router_settings": {"disable_cooldowns": True},
}
)
)
return config
def _usage_in(request: Request) -> _Usage:
found: Final = _USAGE.search(request.body.decode())
assert found is not None, request
return _Usage(prompt=int(found.group(1)), cached=int(found.group(2)), completion=int(found.group(3)))
def _reply(request: Request) -> Reply:
usage: Final = _usage_in(request)
if "/converse" in request.target:
return Reply(
body=json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "priced"}]}},
"stopReason": "end_turn",
"usage": {
"inputTokens": usage.prompt,
"outputTokens": usage.completion,
"totalTokens": usage.prompt + usage.completion,
},
"metrics": {"latencyMs": 1},
}
).encode()
)
if "/invoke" in request.target:
return Reply(
body=json.dumps(
{
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": "claude-scripted",
"content": [{"type": "text", "text": "priced"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": usage.prompt, "output_tokens": usage.completion},
}
).encode()
)
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "scripted",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "priced"}, "finish_reason": "stop"}
],
"usage": {
"prompt_tokens": usage.prompt,
"completion_tokens": usage.completion,
"total_tokens": usage.prompt + usage.completion,
"prompt_tokens_details": {"cached_tokens": usage.cached},
},
}
).encode()
)
def _responder(published: threading.Event) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.target == "/model_prices.json":
return Reply(body=json.dumps(_catalog(_PUBLISHED_SCALE if published.is_set() else 1.0)).encode())
if request.method == "GET":
return Reply(body=json.dumps({"object": "list", "data": []}).encode())
return _reply(request)
return respond
def _listed_rates(proxy: Gateway, name: str) -> tuple[JsonValue, JsonValue]:
entries: Final = proxy.get("/model/info")["data"]
assert isinstance(entries, list), entries
info: Final = next(
object_value(object_value(entry)["model_info"])
for entry in entries
if object_value(entry)["model_name"] == name
)
return info.get("input_cost_per_token"), info.get("output_cost_per_token")
def _billed(client: httpx.Client, key: str, deployment: _Deployment) -> tuple[str, float | None]:
response: Final = client.post(
"/v1/chat/completions",
json={"model": deployment.name, "messages": [{"role": "user", "content": deployment.usage.text()}]},
headers={"Authorization": f"Bearer {key}"},
)
assert response.status_code == 200, response.text
reported: Final = response.headers.get("x-litellm-response-cost")
return string_value(object_value(response.json())["id"]), None if reported is None else float(reported)
def _charged_as_listed(proxy: Gateway, deployment: _Deployment, scale: float) -> _Charge:
rates: Final = deployment.rates(scale)
assert _listed_rates(proxy, deployment.name) == _expected_listing(rates), deployment.name
request_id, charged = _billed(proxy.client, proxy.key, deployment)
expected: Final = _expected_cost(rates, deployment.usage)
assert charged == (None if rates is None else pytest.approx(expected, rel=1e-6)), deployment.name
return _Charge(request_id=request_id, expected=expected)
def _spend(request_id: str) -> float:
rows: Final = eventually(
lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)),
lambda values: len(values) == 1,
seconds=70,
)
spend: Final = rows[0]["spend"]
assert isinstance(spend, (int, float)), rows
return float(spend)
def _proxy_environment(wire_url: str) -> dict[str, str]:
return {**_COST_MAP_LIMITS, "LITELLM_MODEL_COST_MAP_URL": f"{wire_url}/model_prices.json"}
def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments: tuple[_Deployment, ...]) -> None:
published: Final = threading.Event()
with (
wire_server(_responder(published)) as wire,
owned_proxy(
gateway,
tmp_path,
_proxy_environment(wire.url),
config=_config(tmp_path, deployments, wire.url),
remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",),
) as proxy,
):
before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in deployments)
published.set()
reload: Final = proxy.request("POST", "/reload/model_cost_map")
assert reload.status_code == 200, reload.text
after: Final = tuple(_charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in deployments)
for charge in (*before, *after):
assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6)
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
def test_price_data_reload_reprices_catalog_deployments_and_keeps_custom_prices(
gateway: Gateway, tmp_path: Path
) -> None:
_assert_priced_across_a_reload(
gateway,
tmp_path,
(
_Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)),
_Deployment("mini-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates),
),
)
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
@pytest.mark.parametrize(
"deployment",
(
_Deployment(
"region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW)
),
_Deployment(
"finetune-unpriced",
f"openai/{_FINETUNE}:org:suffix:y",
{},
_Usage(2000, 1000, 1000),
_catalog_rates(_FINETUNE),
),
_Deployment("finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates),
_Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced),
_Deployment(
"region-custom-above-threshold",
f"bedrock/ap-northeast-1.{_ROW}",
_CUSTOM,
_Usage(250_000, 0, 1000),
_custom_rates,
),
),
ids=lambda deployment: deployment.name,
)
def test_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload(
gateway: Gateway, tmp_path: Path, deployment: _Deployment
) -> None:
_assert_priced_across_a_reload(gateway, tmp_path, (deployment,))
def _storm_deployment(index: int) -> _Deployment:
prices: Final = MappingProxyType(
{"input_cost_per_token": (index + 1) * 1e-07, "output_cost_per_token": (index + 1) * 3e-07}
)
return _Deployment(
f"storm-{index}",
f"openai/gpt-audit-reload-storm-{index}",
prices,
_Usage(100, 0, 50),
lambda _scale: _Rates(
input=prices["input_cost_per_token"],
output=prices["output_cost_per_token"],
cache_read=prices["input_cost_per_token"],
),
)
def _storm_traffic(
base_url: str, key: str, deployments: tuple[_Deployment, ...], offset: int, finished: threading.Event
) -> tuple[tuple[float, float], ...]:
rotation: Final = (
deployments[(offset + turn) % len(deployments)] for turn, _ in enumerate(iter(finished.is_set, True))
)
with httpx.Client(base_url=base_url, timeout=60, trust_env=False) as client:
return tuple(
(_billed(client, key, deployment)[1], _expected_cost(deployment.rates(1.0), deployment.usage))
for deployment in rotation
)
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_priced(
gateway: Gateway, tmp_path: Path
) -> None:
storm: Final = tuple(_storm_deployment(index) for index in range(24))
geo: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO))
finished: Final = threading.Event()
with (
wire_server(_responder(threading.Event())) as wire,
owned_proxy(
gateway,
tmp_path,
_proxy_environment(wire.url),
config=_config(tmp_path, (geo, *storm), wire.url),
remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",),
) as proxy,
ThreadPoolExecutor(max_workers=8) as pool,
):
traffic: Final = tuple(
pool.submit(_storm_traffic, str(proxy.client.base_url), proxy.key, storm, offset, finished)
for offset in range(8)
)
reloads: Final = tuple(proxy.request("POST", "/reload/model_cost_map") for _ in range(40))
finished.set()
charges: Final = tuple(itertools.chain.from_iterable(calls.result() for calls in traffic))
assert tuple(reload.status_code for reload in reloads) == (200,) * 40, tuple(
reload.text for reload in reloads if reload.status_code != 200
)
assert len(charges) >= 8
assert tuple(charged for charged, _ in charges) == pytest.approx(tuple(expected for _, expected in charges))
settled: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (geo, *storm))
for charge in settled:
assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6)

11
tests/integration/providers/_idle_upstream_support.py
View file

@ -294,11 +294,11 @@ def chat_httpx(target: ProxyTarget, model: str, marker: str) -> Answer:
)
def chat_with_guardrail(target: ProxyTarget, model: str, marker: str) -> Answer:
def chat_with_guardrail(target: ProxyTarget, model: str, marker: str, guardrail: str = GUARDRAIL) -> Answer:
return _post(
target,
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": marker}], "guardrails": [GUARDRAIL]},
{"model": model, "messages": [{"role": "user", "content": marker}], "guardrails": [guardrail]},
_chat_body,
)
@ -517,18 +517,21 @@ def pass_through_deployment(_model: str, url: str, _steady: str) -> ConfigPart:
return ConfigPart(pass_through=({"path": PASS_THROUGH, "target": f"{url}/passthrough"},))
def guardrail_deployment(model: str, url: str, steady: str) -> ConfigPart:
def guardrail_deployment(
model: str, url: str, steady: str, guardrail: str = GUARDRAIL, num_retries: int | None = None
) -> ConfigPart:
return ConfigPart(
models=openai_deployment(model, steady, steady).models,
guardrails=(
{
"guardrail_name": GUARDRAIL,
"guardrail_name": guardrail,
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": False,
"api_base": url,
"api_key": PROVIDER_KEY,
**({} if num_retries is None else {"num_retries": num_retries}),
},
},
),

2
tests/integration/providers/test_google_genai_vertex_routes_wire.py
View file

@ -388,7 +388,7 @@ def test_adapter_route_upstream_error_reaches_the_caller_with_its_status(
gateway: Gateway, status: int, stream: bool
) -> None:
route: Final = _ROUTES["model_garden_publisher"]
message: Final = f"scripted vertex rejection {uuid.uuid4().hex}"
message: Final = f"scripted vertex rejection {uuid.uuid4().hex}f4039"
with wire_server(_vertex_peer(status=status, error=message)) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire, route, num_retries=0)
key: Final = scenario.key(models=[model])

134
tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py Normal file
View file

@ -0,0 +1,134 @@
import threading
from collections.abc import Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from functools import partial
from itertools import chain
from types import MappingProxyType
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.idle_gateway import Arrival
from integration._support.process import graceful_stop_seconds
from integration._support.wire import Reply, Request
from integration.providers._idle_upstream_support import (
REUSE_IDLE_SECONDS,
ProxyTarget,
Row,
Sent,
UpstreamProxy,
answer,
answered,
chat_with_guardrail,
expected_answers,
guardrail_deployment,
landed,
on_each_worker,
send,
spend_statuses,
stay_idle,
upstream_proxy,
)
CUT_AFTER_SECONDS: Final = 3.0
POOLED_PER_WORKER: Final = 2
RETRIED: Final = "guardrail-cut-num-retries-1"
UNSET: Final = "guardrail-cut-num-retries-unset"
pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 180 + REUSE_IDLE_SECONDS))
@dataclass(frozen=True, slots=True)
class PairedWarmup:
"""While warming, each check waits for a second one so every worker pools two connections to the guardrail."""
warming: threading.Event = field(default_factory=threading.Event)
pair: threading.Barrier = field(default_factory=lambda: threading.Barrier(POOLED_PER_WORKER))
def respond(self, request: Request) -> Reply:
if self.warming.is_set():
self.pair.wait(timeout=30)
return answer(request)
@dataclass(frozen=True, slots=True)
class StaleGuardrails:
proxy: UpstreamProxy
pooled: Mapping[str, frozenset[int]]
WARMUPS: Final = MappingProxyType({RETRIED: PairedWarmup(), UNSET: PairedWarmup()})
ROWS: Final = (
Row(
RETRIED,
partial(guardrail_deployment, guardrail=RETRIED, num_retries=1),
partial(chat_with_guardrail, guardrail=RETRIED),
respond=WARMUPS[RETRIED].respond,
),
Row(
UNSET,
partial(guardrail_deployment, guardrail=UNSET),
partial(chat_with_guardrail, guardrail=UNSET),
respond=WARMUPS[UNSET].respond,
),
)
def _pair(row: Row, target: ProxyTarget) -> tuple[Sent, ...]:
with ThreadPoolExecutor(max_workers=POOLED_PER_WORKER) as pool:
return tuple(pool.map(send, (row,) * POOLED_PER_WORKER, (target,) * POOLED_PER_WORKER))
def _warm(proxy: UpstreamProxy, row: Row) -> frozenset[int]:
warmup: Final = WARMUPS[row.name]
warmup.warming.set()
sents: Final = on_each_worker(proxy.owned, proxy.workers, lambda: _pair(row, proxy.target))
warmup.warming.clear()
assert answered(sents) == expected_answers(row, sents), sents
return frozenset(arrival.connection for arrival in proxy.gateways[row.name].drain())
@pytest.fixture(scope="module")
def stale(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StaleGuardrails]:
directory: Final = tmp_path_factory.mktemp("idle-upstream-guardrail-cut")
with upstream_proxy(session_gateway, directory, ROWS, {}, drop_after=CUT_AFTER_SECONDS) as proxy:
pooled: Final = MappingProxyType({row.name: _warm(proxy, row) for row in ROWS})
assert {name: len(connections) for name, connections in pooled.items()} == {
row.name: 2 * POOLED_PER_WORKER for row in ROWS
}, pooled
stay_idle(tuple(proxy.gateways.values()), REUSE_IDLE_SECONDS)
yield StaleGuardrails(proxy, pooled)
def _cut_once(sents: tuple[Sent, ...]) -> tuple[tuple[str, bool], ...]:
return tuple(sorted((sent.marker, True) for sent in sents))
def _cut_once_then_landed(sents: tuple[Sent, ...]) -> tuple[tuple[str, bool], ...]:
return tuple(sorted(chain.from_iterable(((sent.marker, False), (sent.marker, True)) for sent in sents)))
def _landed_on(arrivals: tuple[Arrival, ...]) -> frozenset[int]:
return frozenset(arrival.connection for arrival in arrivals if not arrival.dropped)
def test_retried_check_cut_on_a_stale_pooled_connection_lands_on_a_connection_that_was_never_pooled(
stale: StaleGuardrails,
) -> None:
sents, arrivals = stale.proxy.probe(RETRIED)
assert answered(sents) == expected_answers(stale.proxy.rows[RETRIED], sents), (sents, arrivals)
assert landed(arrivals) == _cut_once_then_landed(sents), arrivals
assert _landed_on(arrivals).isdisjoint(stale.pooled[RETRIED]), (arrivals, stale.pooled)
assert spend_statuses(sents) == ("success",) * len(sents), sents
def test_check_cut_on_a_stale_pooled_connection_is_not_resent_when_num_retries_is_unset(
stale: StaleGuardrails,
) -> None:
sents, arrivals = stale.proxy.probe(UNSET)
assert [sent.answer.status == 200 for sent in sents] == [False] * len(sents), sents
assert all("Generic Guardrail API failed" in sent.answer.text for sent in sents), sents
assert landed(arrivals) == _cut_once(sents), arrivals

192
tests/integration/providers/test_vector_store_team_model_provider_wire.py Normal file
View file

@ -0,0 +1,192 @@
import json
import uuid
from collections.abc import Callable
from typing import Final
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
QUERY: Final = "team provider vector search"
TEAM_ACCESS_KEY: Final = "AKIAINTEGRATIONTEAM01"
STORE_ACCESS_KEY: Final = "AKIASTOREKEY00000001"
_RETRIEVE_RESPONSE: Final = json.dumps(
{
"retrievalResults": [
{
"content": {"text": "matching Bedrock content"},
"score": 0.87,
"metadata": {
"x-amz-bedrock-kb-source-uri": "s3://synthetic-bucket/vector-store.txt",
"x-amz-bedrock-kb-chunk-id": "chunk-1",
},
}
]
}
).encode()
_OPENAI_RESULT: Final[dict[str, JsonValue]] = {
"file_id": "file-search-wire",
"filename": "search.txt",
"score": 0.91,
"attributes": {"tenant": "a"},
"content": [{"type": "text", "text": "matching OpenAI content"}],
}
_BEDROCK_RESULT: Final[dict[str, JsonValue]] = {
"score": 0.87,
"content": [{"text": "matching Bedrock content", "type": "text"}],
"file_id": "s3://synthetic-bucket/vector-store.txt",
"filename": "vector-store.txt",
"attributes": {
"x-amz-bedrock-kb-source-uri": "s3://synthetic-bucket/vector-store.txt",
"x-amz-bedrock-kb-chunk-id": "chunk-1",
},
}
def _bedrock_retrieve_peer(knowledge_base_id: str, access_key: str) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == f"/knowledgebases/{knowledge_base_id}/retrieve", request.target
assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={access_key}/")
assert JSON_OBJECT.validate_json(request.body) == {
"retrievalQuery": {"text": QUERY},
"retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 3}},
}
return Reply(body=_RETRIEVE_RESPONSE)
return respond
def _openai_search_peer(store_id: str) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == f"/v1/vector_stores/{store_id}/search", request.target
assert request.headers["authorization"] == "Bearer team-openai-key"
assert JSON_OBJECT.validate_json(request.body)["query"] == QUERY
return Reply(
body=json.dumps(
{
"object": "vector_store.search_results.page",
"search_query": QUERY,
"data": [_OPENAI_RESULT],
"has_more": False,
"next_page": None,
}
).encode()
)
return respond
def test_bedrock_vector_store_search_uses_bedrock_team_model_credentials(gateway: Gateway) -> None:
knowledge_base_id: Final = f"KB{uuid.uuid4().hex[:8].upper()}"
with wire_server(_bedrock_retrieve_peer(knowledge_base_id, TEAM_ACCESS_KEY)) as wire, gateway.scenario() as scenario:
team: Final = scenario.team()
scenario.model(
model="langgraph/repro-agent",
api_base="http://127.0.0.1:1",
model_info={"team_id": team},
)
scenario.model(
model="bedrock/us.anthropic.claude-sonnet-5-5",
api_base=wire.url,
aws_access_key_id=TEAM_ACCESS_KEY,
aws_secret_access_key="synthetic-team-secret-key",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
model_info={"team_id": team},
)
team_key: Final = scenario.key(team_id=team)
gateway.post(
"/vector_store/new",
{
"vector_store_id": knowledge_base_id,
"custom_llm_provider": "bedrock",
"litellm_params": {
"aws_region_name": "us-east-1",
"aws_bedrock_runtime_endpoint": wire.url,
},
},
)
scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": knowledge_base_id})
response: Final = gateway.request(
"POST",
f"/v1/vector_stores/{knowledge_base_id}/search",
{"query": QUERY, "max_num_results": 3},
key=team_key,
)
assert response.status_code == 200, response.text
assert JSON_OBJECT.validate_json(response.content)["data"] == [_BEDROCK_RESULT], response.text
assert len(wire.drain()) == 1
def test_bedrock_vector_store_search_uses_store_credentials_without_matching_team_model(
gateway: Gateway,
) -> None:
knowledge_base_id: Final = f"KB{uuid.uuid4().hex[:8].upper()}"
with wire_server(_bedrock_retrieve_peer(knowledge_base_id, STORE_ACCESS_KEY)) as wire, gateway.scenario() as scenario:
team: Final = scenario.team()
scenario.model(
model="langgraph/repro-agent",
api_base="http://127.0.0.1:1",
model_info={"team_id": team},
)
team_key: Final = scenario.key(team_id=team)
gateway.post(
"/vector_store/new",
{
"vector_store_id": knowledge_base_id,
"custom_llm_provider": "bedrock",
"litellm_params": {
"aws_access_key_id": STORE_ACCESS_KEY,
"aws_secret_access_key": "synthetic-store-secret-key",
"aws_region_name": "us-east-1",
"aws_bedrock_runtime_endpoint": wire.url,
},
},
)
scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": knowledge_base_id})
response: Final = gateway.request(
"POST",
f"/v1/vector_stores/{knowledge_base_id}/search",
{"query": QUERY, "max_num_results": 3},
key=team_key,
)
assert response.status_code == 200, response.text
assert JSON_OBJECT.validate_json(response.content)["data"] == [_BEDROCK_RESULT], response.text
assert len(wire.drain()) == 1
def test_openai_vector_store_search_without_store_key_uses_team_byok_key(gateway: Gateway) -> None:
store_id: Final = f"vs_team_{uuid.uuid4().hex}"
with wire_server(_openai_search_peer(store_id)) as wire, gateway.scenario() as scenario:
team: Final = scenario.team()
scenario.model(
model="openai/gpt-4o-mini",
api_key="team-openai-key",
api_base=f"{wire.url}/v1",
model_info={"team_id": team},
)
team_key: Final = scenario.key(team_id=team)
gateway.post(
"/vector_store/new",
{
"vector_store_id": store_id,
"custom_llm_provider": "openai",
"litellm_params": {},
},
)
scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id})
response: Final = gateway.request(
"POST",
f"/v1/vector_stores/{store_id}/search",
{"query": QUERY},
key=team_key,
)
assert response.status_code == 200, response.text
assert JSON_OBJECT.validate_json(response.content)["data"] == [_OPENAI_RESULT], response.text
assert len(wire.drain()) == 1

126
tests/unit/litellm_core_utils/test_exception_mapping_utils.py
View file

@ -20,6 +20,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import (
exception_type,
extract_and_raise_litellm_exception,
)
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai.common_utils import OpenAIError
@ -578,6 +579,131 @@ def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_lim
)
def test_vertex_status_code_429_ignores_403_token_in_message() -> None:
original_exception: Final = BaseLLMException(status_code=429, message="scripted vertex rejection f4039")
with pytest.raises(litellm.RateLimitError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=original_exception,
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 429
def test_vertex_status_code_403_maps_to_403() -> None:
original_exception: Final = BaseLLMException(status_code=403, message="permission denied")
with pytest.raises(litellm.PermissionDeniedError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=original_exception,
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 403
def test_vertex_status_code_403_with_text_token_maps_to_bad_request() -> None:
original_exception: Final = BaseLLMException(
status_code=403,
message='{"error":{"code":403,"message":"Permission denied"}}',
)
with pytest.raises(litellm.BadRequestError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=original_exception,
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 403
def test_vertex_text_403_without_status_code_still_maps_to_bad_request() -> None:
with pytest.raises(litellm.BadRequestError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=Exception("403 Forbidden: permission denied"),
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 403
def test_vertex_status_substring_without_status_code_does_not_map_to_403() -> None:
with pytest.raises(litellm.APIConnectionError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=Exception("scripted vertex rejection f4039a"),
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code != 403
def test_vertex_wrapped_500_with_403_body_maps_to_bad_request() -> None:
original_exception: Final = BaseLLMException(
status_code=500,
message='Error code: 403 - {"error": {"code": 403, "message": "Permission denied"}}',
)
with pytest.raises(litellm.BadRequestError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=original_exception,
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 403
def test_vertex_wrapped_500_with_embedded_403_text_does_not_map_to_403() -> None:
original_exception: Final = BaseLLMException(status_code=500, message="scripted vertex rejection f4039")
with pytest.raises(litellm.InternalServerError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=original_exception,
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 500
def test_vertex_wrapped_500_with_body_code_429_and_403_token_maps_to_rate_limit() -> None:
original_exception: Final = BaseLLMException(
status_code=500,
message='{"error":{"code":429,"message":"upstream unavailable (request id: 403-xyz)"}}',
)
with pytest.raises(litellm.RateLimitError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=original_exception,
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 429
def test_vertex_wrapped_500_with_body_code_403_maps_to_bad_request() -> None:
original_exception: Final = BaseLLMException(
status_code=500,
message='{"error":{"code":403,"message":"Permission denied"}}',
)
with pytest.raises(litellm.BadRequestError) as excinfo:
exception_type(
model="gemini/gemini-2.5-flash",
original_exception=original_exception,
custom_llm_provider="vertex_ai",
)
assert excinfo.value.status_code == 403
class TestGetBodyErrorCode:
"""Unit tests for _get_body_error_code helper."""

327
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -5,11 +5,12 @@ import json
import time
from base64 import urlsafe_b64encode
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException, Request
from fastapi import HTTPException, Request, Response
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._types import LiteLLM_MCPServerTable
from litellm.types.mcp import MCPAuth
@ -26,6 +27,14 @@ if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
class _BridgeRegistrationPayload(TypedDict, total=False):
client_name: ReadOnly[str]
grant_types: ReadOnly[list[str]]
response_types: ReadOnly[list[str]]
token_endpoint_auth_method: ReadOnly[str]
redirect_uris: ReadOnly[list[str]]
def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None):
credential = {"type": "oauth2", "access_token": access_token}
if refresh_token is not None:
@ -5849,7 +5858,12 @@ def test_oauth_authorization_server_404_for_non_bridge_client_forwarded_server()
assert exc.value.status_code == 404
async def _bridge_register_response(server, request_payload, replace_stale_client=False):
async def _bridge_register_response(
server: "MCPServer",
request_payload: _BridgeRegistrationPayload,
replace_stale_client: bool = False,
database_connected: bool = True,
) -> tuple[Response, MagicMock, AsyncMock]:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
register_client_with_server,
)
@ -5879,6 +5893,10 @@ async def _bridge_register_response(server, request_payload, replace_stale_clien
new_callable=AsyncMock,
return_value=False,
),
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
**({"return_value": MagicMock()} if database_connected else {"side_effect": HTTPException(status_code=500)}),
),
):
response = await register_client_with_server(
request=_bridge_mock_request(),
@ -6033,6 +6051,309 @@ async def test_register_bridge_relay_never_persists():
mock_persist.assert_not_called()
_LAPTOP_LOOPBACK_REDIRECT: Final = "http://127.0.0.1:33418/"
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type_value", ["oauth_delegate", "true_passthrough"])
async def test_register_bridge_with_proxy_base_url_registers_gateway_client_not_laptop_loopback(
monkeypatch: pytest.MonkeyPatch, auth_type_value: Literal["oauth_delegate", "true_passthrough"]
) -> None:
"""A laptop MCP client (VS Code) registers with its own loopback redirect. With PROXY_BASE_URL set
the gateway must not relay that registration upstream: it registers its one shared client bound
to the gateway /callback, keeps it, and hands the caller a client bound to the caller's loopback.
Relaying the loopback is the customer failure on upstreams that restrict redirect domains."""
import json
from litellm.types.mcp import MCPAuth
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
response, mock_async_client, mock_persist = await _bridge_register_response(
_bridge_server(auth_type=MCPAuth(auth_type_value)),
{"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]},
)
posted: Final = mock_async_client.post.call_args.kwargs["json"]
assert posted["redirect_uris"] == ["https://litellm.example.com/callback"], posted
assert posted["client_name"] == "LiteLLM Proxy", posted
mock_persist.assert_awaited_once()
payload: Final = json.loads(response.body.decode("utf-8"))
assert payload["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT], payload
@pytest.mark.asyncio
async def test_register_bridge_without_database_still_relays_client_registration(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A database-free gateway cannot keep the shared client past this process, so even with
PROXY_BASE_URL set it relays the caller's own registration instead of stranding the next authorize."""
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
_, mock_async_client, mock_persist = await _bridge_register_response(
_bridge_server(),
{"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]},
database_connected=False,
)
assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT]
mock_persist.assert_not_called()
@pytest.mark.asyncio
async def test_register_bridge_without_proxy_base_url_still_relays_client_registration(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Without PROXY_BASE_URL the gateway cannot keep a shared client, so the relay arm remains the
fallback and the caller's own redirect goes upstream as before."""
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
_, mock_async_client, mock_persist = await _bridge_register_response(
_bridge_server(),
{"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]},
)
assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT]
mock_persist.assert_not_called()
def _persisted_shared_gateway_client(
server: "MCPServer", client_id: str = "shared-gateway-client"
) -> LiteLLM_MCPServerTable:
request: Final = _bridge_mock_request()
credentials: Final = {
"client_id": client_id,
"client_secret": None,
"token_endpoint_auth_method": None,
"redirect_uris": [f"{request.base_url.rstrip('/')}/callback"],
"dcr_issuer": server.issuer,
"dcr_server_url": server.url,
}
return LiteLLM_MCPServerTable.model_validate({**server.model_dump(exclude_none=True), "credentials": credentials})
@pytest.mark.asyncio
@pytest.mark.parametrize(
("shared_client_persisted", "expected_redirect_uri", "expected_client_id"),
[
(True, "https://litellm.example.com/callback", "shared-gateway-client"),
(False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"),
],
)
async def test_authorize_clientless_bridge_loads_shared_client_before_choosing_the_arm(
monkeypatch: pytest.MonkeyPatch,
shared_client_persisted: bool,
expected_redirect_uri: str,
expected_client_id: str,
) -> None:
"""This worker's catalog snapshot has no client_id. When another worker already persisted the shared
client, authorize must load it and send it with the gateway /callback upstream, never the caller's
own client_id or loopback. When no shared client exists yet (a client registered before the gateway
kept one), the caller keeps its own upstream binding through the relay."""
from urllib.parse import parse_qs, urlparse
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding")
server: Final = _bridge_server(url="https://mcp.example.com/mcp")
persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(return_value=persisted)
mock_update_server: Final = AsyncMock()
with (
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=prisma_client,
),
patch(
"litellm.proxy._experimental.mcp_server.db.get_mcp_server",
new=mock_get_mcp_server,
),
patch.object(
global_mcp_server_manager,
"update_server",
new=mock_update_server,
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper",
return_value="mocked_encrypted_state",
),
):
response: Final = await authorize_with_server(
request=_bridge_mock_request(),
mcp_server=server,
client_id="dcr-client-123",
redirect_uri=_LAPTOP_LOOPBACK_REDIRECT,
state="s",
code_challenge="chal",
code_challenge_method="S256",
)
location: Final = response.headers["location"]
assert location.startswith("https://provider.com/oauth/authorize"), location
query: Final = parse_qs(urlparse(location).query)
assert query["redirect_uri"] == [expected_redirect_uri], query
assert query["client_id"] == [expected_client_id], query
@pytest.mark.asyncio
@pytest.mark.parametrize(
("shared_client_persisted", "expected_redirect_uri", "expected_client_id"),
[
(True, "https://litellm.example.com/callback", "shared-gateway-client"),
(False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"),
],
)
async def test_token_clientless_bridge_loads_shared_client_before_choosing_the_arm(
monkeypatch: pytest.MonkeyPatch,
shared_client_persisted: bool,
expected_redirect_uri: str,
expected_client_id: str,
) -> None:
"""The token leg mirrors authorize: with the shared client persisted it redeems the code with that
client and the gateway /callback, and with none yet it relays the caller's own client and redirect."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding")
server: Final = _bridge_server(url="https://mcp.example.com/mcp")
persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(return_value=persisted)
mock_update_server: Final = AsyncMock()
with (
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma_client),
patch(
"litellm.proxy._experimental.mcp_server.db.get_mcp_server",
new=mock_get_mcp_server,
),
patch.object(
global_mcp_server_manager,
"update_server",
new=mock_update_server,
),
):
data: Final = await _bridge_token_post_data(server, redirect_uri=_LAPTOP_LOOPBACK_REDIRECT)
assert data["redirect_uri"] == expected_redirect_uri, data
assert data["client_id"] == expected_client_id, data
@pytest.mark.asyncio
@pytest.mark.parametrize(
("dcr_bridge", "expected"),
[(True, "persisted"), (False, "skipped")],
)
async def test_persist_dcr_client_for_client_forwarded_server_only_when_it_is_a_bridge(
dcr_bridge: bool,
expected: Literal["persisted", "skipped"],
) -> None:
"""A DCR bridge keeps the shared gateway client on the server row (without stamping oauth2_flow on
a client-forwarded row), while a client-forwarded server without the bridge still stores nothing."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_persist_dcr_client_registration,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.types.mcp import MCPAuth
server: Final = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=dcr_bridge)
prisma: Final = MagicMock()
prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None)
mock_update: Final = AsyncMock(return_value=MagicMock())
with (
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma),
patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update),
patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()),
patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False),
):
result: Final = await _persist_dcr_client_registration(
server,
{"client_id": "shared-gateway-client", "token_endpoint_auth_method": "none"},
"https://litellm.example.com/callback",
replace_stale_client=False,
)
assert result == expected
if expected == "skipped":
mock_update.assert_not_awaited()
return
written: Final = mock_update.await_args.kwargs["data"]
assert written.oauth2_flow is None, written
assert written.credentials["client_id"] == "shared-gateway-client", written
assert server.client_id == "shared-gateway-client"
@pytest.mark.asyncio
@pytest.mark.parametrize("server_has_shared_client", [False, True])
async def test_token_sealed_ephemeral_client_is_not_replaced_by_persisted_gateway_client(
monkeypatch: pytest.MonkeyPatch,
server_has_shared_client: bool,
) -> None:
"""The token leg must use the sealed ephemeral client even if this worker's catalog has a shared client."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding")
server: Final = _bridge_server(
url="https://mcp.example.com/mcp",
client_id="shared-gateway-client" if server_has_shared_client else None,
client_secret="shared-secret" if server_has_shared_client else None,
token_endpoint_auth_method="client_secret_basic" if server_has_shared_client else None,
)
fake_http_response: Final = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client: Final = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(
return_value=_persisted_shared_gateway_client(server, client_id="persisted-gateway-client")
)
mock_update_server: Final = AsyncMock()
with (
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=prisma_client,
),
patch(
"litellm.proxy._experimental.mcp_server.db.get_mcp_server",
new=mock_get_mcp_server,
),
patch.object(
global_mcp_server_manager,
"update_server",
new=mock_update_server,
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
return_value=fake_http_client,
),
):
await exchange_token_with_server(
request=_bridge_mock_request(),
mcp_server=server,
grant_type="authorization_code",
code="upstream-code",
redirect_uri="https://litellm.example.com/callback",
client_id="ephemeral-client",
client_secret="ephemeral-secret" if server_has_shared_client else None,
code_verifier="verifier",
client_token_endpoint_auth_method="client_secret_post" if server_has_shared_client else None,
sealed_client=True,
)
data: Final = fake_http_client.post.call_args.kwargs["data"]
assert data["client_id"] == "ephemeral-client", data
assert data["redirect_uri"] == "https://litellm.example.com/callback", data
if server_has_shared_client:
headers: Final = fake_http_client.post.call_args.kwargs["headers"]
assert data["client_secret"] == "ephemeral-secret", data
assert "Authorization" not in headers, headers
_BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef"

68
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -20056,6 +20056,74 @@ async def test_disabled_stdio_tools_are_empty_initially_and_reject_continuation(
assert page.next_cursor is None
@pytest.mark.asyncio
async def test_aggregate_bounds_each_paginated_server_by_its_own_listing_deadline(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from mcp.types import JSONRPCRequest, PaginatedRequestParams
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import catalog, operations
from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module
from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR, get_server_prefix
monkeypatch.setenv("LITELLM_SALT_KEY", "paginated-listing-deadline-test")
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(manager_module, "MCP_TOOL_LISTING_TIMEOUT", 0.0)
extended: Final = MCPServer(
server_id="extended",
name="extended",
transport=MCPTransport.http,
url="https://extended.example/mcp",
timeout=60.0,
)
global_deadline: Final = MCPServer(
server_id="global", name="global", transport=MCPTransport.http, url="https://global.example/mcp"
)
servers: Final = (extended, global_deadline)
manager: Final = MCPServerManager()
manager.registry = {server.server_id: server for server in servers}
monkeypatch.setattr(operations, "global_mcp_server_manager", manager)
context: Final = operations.prepare_context(UserAPIKeyAuth(api_key="owned-caller", user_id="alice"))
async def respond(request: httpx2.Request) -> httpx2.Response:
if request.method == "DELETE":
return httpx2.Response(200)
payload: Final = _JSONRPC_ADAPTER.validate_json(request.content)
if not isinstance(payload, JSONRPCRequest):
return httpx2.Response(202)
result: Final = {
"initialize": {
"protocolVersion": "2025-03-26",
"serverInfo": {"name": "test", "version": "1"},
"capabilities": {"tools": {}},
},
"server/discover": {
"supportedVersions": ["2026-07-28"],
"capabilities": {"tools": {}},
"ttlMs": 60000,
"cacheScope": "private",
"resultType": "complete",
},
"tools/list": {
"tools": [{"name": "echo", "inputSchema": {"type": "object"}}],
"ttlMs": 60000,
"cacheScope": "private",
"resultType": "complete",
},
}[payload.method]
return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result})
with _mcp_upstream(respond):
listing: Final = await catalog.aggregate_gateway_tools(
context, PaginatedRequestParams(), servers, {}, record_listing=True, enforce_rate_limits=False
)
assert [tool.name for tool in listing.tools] == [f"{get_server_prefix(extended)}{MCP_TOOL_PREFIX_SEPARATOR}echo"]
assert listing.outcomes["extended"].tag == "ok"
assert listing.outcomes["global"].tag == "timeout"
@pytest.mark.asyncio
async def test_failed_aggregate_continuation_preserves_only_delivered_tool_metadata(monkeypatch):
from mcp.shared.exceptions import MCPError

45
tests/unit/proxy/auth/test_route_checks.py
View file

@ -4782,6 +4782,51 @@ def test_view_only_user_denied_credential_connection_routes(method: str, route:
_internal_user_check(route, method, role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY)
_SELF_SCOPED_CREDENTIAL_WRITES = (
("POST", "/credentials/x/user_connection/start"),
("POST", "/credentials/x/user_connection/poll"),
("DELETE", "/credentials/x/user_connection"),
("POST", "/v1/mcp/server/oauth/srv-1/register"),
("POST", "/v1/mcp/server/srv-1/oauth-user-credential"),
("DELETE", "/v1/mcp/server/srv-1/oauth-user-credential"),
("POST", "/v1/mcp/server/srv-1/user-credential"),
("DELETE", "/v1/mcp/server/srv-1/user-credential"),
("POST", "/v1/mcp/server/srv-1/user-env-vars"),
("DELETE", "/v1/mcp/server/srv-1/user-env-vars"),
)
@pytest.mark.parametrize(("method", "route"), _SELF_SCOPED_CREDENTIAL_WRITES)
def test_proxy_admin_viewer_allowed_self_scoped_credential_writes(method: str, route: str) -> None:
"""Regression for LIT-9320: these handlers only write the caller's own credential
row, so a read-only admin can use them like /user/password/change."""
assert _internal_user_check(route, method, role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) is None
@pytest.mark.parametrize(
("method", "route"),
[
("POST", "/v1/mcp/server"),
("PUT", "/v1/mcp/server"),
("DELETE", "/v1/mcp/server/srv-1"),
("POST", "/v1/mcp/server/import"),
("POST", "/v1/mcp/server/srv-1/pin"),
("PATCH", "/v1/mcp/server/srv-1/oauth-user-credential"),
("PUT", "/v1/mcp/server/srv-1/user-credential"),
("PATCH", "/v1/mcp/server/srv-1/user-env-vars"),
("POST", "/v1/mcp/server/srv-1/user-credential/rotate"),
("POST", "/v1/mcp/server/oauth/srv-1/token"),
("DELETE", "/v1/mcp/server/oauth/srv-1/register"),
("PATCH", "/credentials/x/user_connection"),
("DELETE", "/credentials/user_connections"),
("POST", "/key/generate"),
],
)
def test_proxy_admin_viewer_denied_mcp_admin_and_sibling_writes(method: str, route: str) -> None:
with pytest.raises(HTTPException, match="user not allowed to access this route"):
_internal_user_check(route, method, role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
def test_user_connection_route_registration_order_wins_over_generic_credential_routes() -> None:
"""The :path placeholders in the generic CRUD routes could swallow the
user_connection paths; the concrete routes must register (and match) first."""

153
tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py
View file

@ -8,6 +8,9 @@ specifically focusing on metadata extraction and passing.
import json
import logging
import os
from collections.abc import Callable
from functools import partial
from queue import SimpleQueue
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -26,7 +29,7 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
_HEADER_PRESENT_PLACEHOLDER,
)
from litellm.types.utils import Choices, Message
from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message
@pytest.fixture
@ -2470,3 +2473,151 @@ class TestFailOnError:
request_data={},
input_type="response",
)
Outcome = Exception | Callable[[], httpx.Response]
REWRITTEN: Final = partial(httpx.Response, 200, json={"action": "GUARDRAIL_INTERVENED", "texts": ["redacted"]})
def _scripted(*outcomes: Outcome) -> tuple[httpx.MockTransport, SimpleQueue[bytes]]:
"""Each check gets the next outcome; the last one repeats once the script runs out."""
sent: Final[SimpleQueue[bytes]] = SimpleQueue()
script: Final = iter(outcomes)
def answer(request: httpx.Request) -> httpx.Response:
sent.put(request.read())
outcome: Final = next(script, outcomes[-1])
if isinstance(outcome, Exception):
raise outcome
return outcome()
return httpx.MockTransport(answer), sent
def _guardrail_on(
transport: httpx.MockTransport, num_retries: int | None, fail_on_error: bool = True
) -> GenericGuardrailAPI:
return GenericGuardrailAPI(
api_base="https://api.test.guardrail.com",
guardrail_name="retrying-guardrail",
fail_on_error=fail_on_error,
num_retries=num_retries,
timeout=5,
async_handler=AsyncHTTPHandler(transport=transport),
fresh_handler_factory=lambda: AsyncHTTPHandler(transport=transport),
)
async def _check(guardrail: GenericGuardrailAPI) -> GenericGuardrailAPIInputs:
return await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request")
async def _attempts_until_it_raises(outcome: Outcome, num_retries: int | None) -> int:
transport, sent = _scripted(outcome)
with pytest.raises(Exception, match=r"Generic Guardrail API failed|policy violation"):
await _check(_guardrail_on(transport, num_retries))
return sent.qsize()
class _OpenCloseLoggingHandler(AsyncHTTPHandler):
def __init__(self, transport: httpx.MockTransport, events: SimpleQueue[str]) -> None:
super().__init__(transport=transport)
self._events: Final = events
events.put("open")
async def close(self) -> None:
self._events.put("close")
await super().close()
class TestNumRetries:
@pytest.mark.asyncio
@pytest.mark.parametrize(
("outcomes", "fresh_connections"),
[
((httpx.ReadTimeout("no answer before the read timeout"),), 3),
((httpx.ReadTimeout("no answer"), httpx.ReadTimeout("no answer"), REWRITTEN), 2),
],
ids=["every_retry_fails", "second_retry_succeeds"],
)
async def test_each_retry_connection_closes_before_the_next_attempt(
self, outcomes: tuple[Outcome, ...], fresh_connections: int
) -> None:
transport, _ = _scripted(*outcomes)
events: Final[SimpleQueue[str]] = SimpleQueue()
guardrail: Final = GenericGuardrailAPI(
api_base="https://api.test.guardrail.com",
guardrail_name="retrying-guardrail",
fail_on_error=False,
num_retries=3,
timeout=5,
async_handler=AsyncHTTPHandler(transport=transport),
fresh_handler_factory=lambda: _OpenCloseLoggingHandler(transport, events),
)
await _check(guardrail)
assert [events.get_nowait() for _ in range(events.qsize())] == ["open", "close"] * fresh_connections
@pytest.mark.asyncio
@pytest.mark.parametrize(
"failure",
[
httpx.ReadTimeout("no answer before the read timeout"),
httpx.WriteTimeout("request body stalled"),
httpx.ReadError("connection reset by peer"),
httpx.WriteError("broken pipe"),
httpx.RemoteProtocolError("Server disconnected without sending a response."),
],
ids=["read_timeout", "write_timeout", "read_error", "write_error", "remote_protocol_error"],
)
async def test_check_lost_after_connect_is_resent_and_the_retry_verdict_applies(self, failure: Exception) -> None:
transport, sent = _scripted(failure, REWRITTEN)
result: Final = await _check(_guardrail_on(transport, num_retries=1))
bodies: Final = [sent.get_nowait() for _ in range(sent.qsize())]
assert len(bodies) == 2
assert bodies[0] == bodies[1]
assert result == {"texts": ["redacted"]}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"outcome",
[
httpx.ConnectTimeout("no TCP handshake before the timeout"),
httpx.ConnectError("connection refused"),
httpx.PoolTimeout("no free pooled connection"),
partial(httpx.Response, 503, json={"error": "unavailable"}),
partial(httpx.Response, 500, json={"error": "boom"}),
partial(httpx.Response, 200, json={"action": "BLOCKED", "blocked_reason": "policy violation"}),
],
ids=["connect_timeout", "connect_error", "pool_timeout", "http_503", "http_500", "blocked"],
)
async def test_num_retries_does_not_resend_connect_failures_error_statuses_or_blocks(
self, outcome: Outcome
) -> None:
assert await _attempts_until_it_raises(outcome, num_retries=3) == await _attempts_until_it_raises(
outcome, num_retries=None
)
@pytest.mark.asyncio
@pytest.mark.parametrize("num_retries", [None, 0])
async def test_unset_num_retries_sends_a_timed_out_check_once(self, num_retries: int | None) -> None:
transport, sent = _scripted(httpx.ReadTimeout("no answer before the read timeout"), REWRITTEN)
with pytest.raises(Exception, match="Generic Guardrail API failed"):
await _check(_guardrail_on(transport, num_retries))
assert sent.qsize() == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("num_retries", [1, 3])
async def test_num_retries_bounds_the_attempts_before_fail_on_error_applies(self, num_retries: int) -> None:
transport, sent = _scripted(httpx.ReadTimeout("no answer before the read timeout"))
result: Final = await _check(_guardrail_on(transport, num_retries, fail_on_error=False))
assert sent.qsize() == num_retries + 1
assert result == {"texts": ["hi"]}

91
tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -3545,6 +3545,95 @@ class TestTemporaryMCPSessionEndpoints:
assert exchange_mock.await_args.kwargs["redirect_uri"] == "https://litellm.example.com/callback"
assert exchange_mock.await_args.kwargs["client_token_endpoint_auth_method"] == "client_secret_basic"
@pytest.mark.asyncio
async def test_mcp_token_prefers_sealed_client_when_server_has_shared_client(
self: "TestTemporaryMCPSessionEndpoints",
) -> None:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
seal_passthrough_authorization_code,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
mcp_token,
)
request: Final = MagicMock()
request.base_url = "https://litellm.example.com/"
request.headers = {}
server: Final = generate_mock_mcp_server_config_record(server_id="server-1").model_copy(
update={
"auth_type": MCPAuth.true_passthrough,
"client_id": "shared-gateway-client",
"client_secret": "shared-secret",
"token_endpoint_auth_method": "client_secret_basic",
"token_url": "https://provider.example/oauth/token",
}
)
admin_auth: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN)
fake_http_response: Final = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client: Final = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(return_value=None)
mock_update_server: Final = AsyncMock()
with (
patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation",
return_value=nullcontext(server),
),
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=prisma_client,
),
patch(
"litellm.proxy._experimental.mcp_server.db.get_mcp_server",
new=mock_get_mcp_server,
),
patch.object(
global_mcp_server_manager,
"update_server",
new=mock_update_server,
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
return_value=fake_http_client,
),
):
sealed: Final = seal_passthrough_authorization_code(
upstream_code="up-code",
client_id="ephemeral-client",
client_secret="ephemeral-secret",
mcp_server_id="server-1",
token_endpoint_auth_method="client_secret_post",
)
result: Final = await mcp_token(
request=request,
server_id="server-1",
user_api_key_dict=admin_auth,
grant_type="authorization_code",
code=sealed,
redirect_uri="https://example.com/callback",
client_id=None,
client_secret=None,
code_verifier="verifier",
refresh_token=None,
scope=None,
)
data: Final = fake_http_client.post.call_args.kwargs["data"]
headers: Final = fake_http_client.post.call_args.kwargs["headers"]
assert data["client_id"] == "ephemeral-client", data
assert data["client_secret"] == "ephemeral-secret", data
assert data["code"] == "up-code", data
assert data["redirect_uri"] == "https://litellm.example.com/callback", data
assert "Authorization" not in headers, headers
assert result.status_code == 200
assert result.body == b'{"access_token":"tok","token_type":"Bearer"}'
@pytest.mark.asyncio
async def test_mcp_token_refresh_grant_never_opens_sealed_code(self):
"""The minted client is unrecoverable outside the single authorization_code flow by
@ -3841,6 +3930,7 @@ class TestTemporaryMCPSessionEndpoints:
refresh_token=None,
scope=None,
client_token_endpoint_auth_method=None,
sealed_client=False,
)
@pytest.mark.asyncio
@ -3895,6 +3985,7 @@ class TestTemporaryMCPSessionEndpoints:
refresh_token="rt-123",
scope=None,
client_token_endpoint_auth_method=None,
sealed_client=False,
)
@pytest.mark.asyncio

92
tests/unit/proxy/test_route_llm_request.py
View file

@ -1,4 +1,7 @@
import httpx
import litellm
import pytest
import respx
from types import SimpleNamespace
@ -1890,3 +1893,92 @@ async def test_null_user_config_is_treated_as_absent() -> None:
)
assert response.choices[0].message.content == "from the main router"
def _registered_bedrock_kb_search(team_id: str) -> dict[str, object]:
"""The request a team key's /v1/vector_stores/{id}/search carries once the registry row is merged in."""
return {
"model": None,
"vector_store_id": "KB12345678",
"query": "litellm proxy",
"custom_llm_provider": "bedrock",
"aws_region_name": "us-east-1",
"aws_access_key_id": "AKIAREGISTRYKEY00000",
"aws_secret_access_key": "registry-secret",
"metadata": {"user_api_key_team_id": team_id},
}
def _team_a_router(*team_models: tuple[str, dict[str, object]]) -> litellm.Router:
return litellm.Router(
model_list=[
{
"model_name": f"internal-team-a-{public_name}",
"litellm_params": litellm_params,
"model_info": {
"id": f"team-a-{public_name}",
"team_id": "team-a",
"team_public_model_name": public_name,
},
}
for public_name, litellm_params in team_models
]
)
_LANGGRAPH_AGENT_PARAMS: Final[dict[str, object]] = {
"model": "langgraph/repro-agent",
"api_base": "http://127.0.0.1:2024",
}
_BEDROCK_CLAUDE_PARAMS: Final[dict[str, object]] = {
"model": "bedrock/anthropic.claude-sonnet-4-6",
"aws_region_name": "us-east-1",
}
@pytest.mark.parametrize(
"team_models",
[
pytest.param(
[("a-langgraph-agent", _LANGGRAPH_AGENT_PARAMS), ("a-bedrock-claude", _BEDROCK_CLAUDE_PARAMS)],
id="langgraph-first",
),
pytest.param(
[("a-bedrock-claude", _BEDROCK_CLAUDE_PARAMS), ("a-langgraph-agent", _LANGGRAPH_AGENT_PARAMS)],
id="bedrock-first",
),
pytest.param([("a-langgraph-agent", _LANGGRAPH_AGENT_PARAMS)], id="no-bedrock-team-model"),
],
)
@pytest.mark.asyncio
async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_registry_credentials(
monkeypatch: pytest.MonkeyPatch,
team_models: list[tuple[str, dict[str, object]]],
) -> None:
"""A team key's model-less search of a registered Bedrock store was routed through the team's first
deployment, so a langgraph team model turned it into "Vector store search is not supported for langgraph"."""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
with respx.mock(assert_all_called=True) as respx_mock:
route: Final = respx_mock.post(
"https://bedrock-agent-runtime.us-east-1.amazonaws.com/knowledgebases/KB12345678/retrieve"
).mock(return_value=httpx.Response(200, json={"retrievalResults": []}))
response: Final = await (
await route_request(
_registered_bedrock_kb_search("team-a"),
_team_a_router(*team_models),
None,
"avector_store_search",
)
)
request: Final = route.calls.last.request
posted_url: Final = str(request.url)
signed_auth: Final = request.headers["Authorization"]
assert response["object"] == "vector_store.search_results.page", response
assert route.call_count == 1
assert posted_url == "https://bedrock-agent-runtime.us-east-1.amazonaws.com/knowledgebases/KB12345678/retrieve", (
posted_url
)
assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth

608
tests/unit/test_register_model_custom_pricing.py
View file

@ -10,13 +10,16 @@ calculations for DB-sourced models with prompt caching pricing.
"""
import copy
import itertools
import os
from typing import Final
from collections.abc import Callable, Mapping
from typing import Final, NoReturn
import pytest
import litellm
from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map
from litellm.main import _build_custom_pricing_entry
from litellm.utils import _invalidate_model_cost_lowercase_map
@ -796,6 +799,609 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key():
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("route", ["", "converse/", "invoke/"])
def test_register_model_prices_a_geo_alias_from_its_builtin_when_the_lookup_fails(
monkeypatch: pytest.MonkeyPatch, route: str
) -> None:
from litellm.types.utils import ModelResponse, Usage
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
_invalidate_model_cost_lowercase_map()
canonical_key: Final = "au.anthropic.claude-sonnet-5"
alias: Final = f"bedrock/{route}{canonical_key}"
canonical: Final = litellm.model_cost[canonical_key]
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
litellm.register_model(
{alias: {"litellm_provider": "bedrock", "mode": "chat", "supports_vision": True}},
persist_across_reloads=False,
)
usage: Final = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100)
alias_cost: Final = litellm.completion_cost(
completion_response=ModelResponse(model=alias, usage=usage), model=alias, custom_llm_provider="bedrock"
)
canonical_cost: Final = litellm.completion_cost(
completion_response=ModelResponse(model=canonical_key, usage=usage),
model=canonical_key,
custom_llm_provider="bedrock",
)
try:
assert alias_cost == canonical_cost > 0
assert litellm.get_model_info(alias)["input_cost_per_token"] == canonical["input_cost_per_token"]
assert litellm.get_model_info(alias)["supports_vision"] is True
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None:
geo_key: Final = "au.anthropic.lit-5853-geo-model"
geo_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
"cache_read_input_token_cost": 3.3e-07,
}
sibling_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 3e-07,
}
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
geo_key: dict(geo_row),
"anthropic.lit-5853-geo-model": sibling_row,
},
)
_invalidate_model_cost_lowercase_map()
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
litellm.register_model({geo_key: {"mode": "chat"}}, persist_across_reloads=False)
try:
for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"):
assert litellm.model_cost[geo_key][field] == geo_row[field]
assert litellm.get_model_info(geo_key)[field] == geo_row[field]
assert litellm.model_cost[geo_key]["litellm_provider"] == geo_row["litellm_provider"]
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("route", ["", "converse/", "invoke/"])
def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once(
monkeypatch: pytest.MonkeyPatch, route: str
) -> None:
catalog_key: Final = "au.anthropic.lit-5853-geo-model"
alias: Final = f"bedrock/{route}{catalog_key}"
repriced_input_cost: Final = 9e-06
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
catalog_key: {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
},
},
)
_invalidate_model_cost_lowercase_map()
real_lookup: Final = litellm.utils.get_model_info
lookups: Final = itertools.count()
def lookup_tripped_once_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> object:
if next(lookups) == 0:
raise RuntimeError("dictionary changed size during iteration")
return real_lookup(model=model, custom_llm_provider=custom_llm_provider)
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_once_by_a_concurrent_registration)
litellm.register_model({alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False)
litellm.register_model({catalog_key: {"input_cost_per_token": repriced_input_cost}}, persist_across_reloads=False)
try:
assert litellm.get_model_info(alias)["key"] == catalog_key
assert litellm.get_model_info(alias)["input_cost_per_token"] == repriced_input_cost
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("key_prefix", ["", "bedrock/", "bedrock/converse/"])
def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup_fails(
monkeypatch: pytest.MonkeyPatch, key_prefix: str
) -> None:
from litellm.types.utils import ModelResponse, Usage
catalog_key: Final = "au.anthropic.lit-5853-tiered-geo-model"
registered_key: Final = f"{key_prefix}{catalog_key}"
catalog_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
"input_cost_per_token_above_200k_tokens": 6.6e-06,
"output_cost_per_token_above_200k_tokens": 2.475e-05,
}
builtin_cost_map: Final = litellm.get_model_cost_map(url="")
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
def long_prompt_cost_after_registering_with(lookup: Callable[..., object]) -> float:
monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)})
_invalidate_model_cost_lowercase_map()
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup)
litellm.register_model(
{
registered_key: {
"litellm_provider": "bedrock",
"input_cost_per_token": 1e-06,
"output_cost_per_token": 2e-06,
}
},
persist_across_reloads=False,
)
return litellm.completion_cost(
completion_response=ModelResponse(
model=registered_key,
usage=Usage(prompt_tokens=250_000, completion_tokens=1_000, total_tokens=251_000),
),
model=registered_key,
custom_llm_provider="bedrock",
)
try:
assert long_prompt_cost_after_registering_with(lookup_tripped_by_a_concurrent_registration) == pytest.approx(
long_prompt_cost_after_registering_with(litellm.utils.get_model_info)
)
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("lookup_trips", [False, True])
def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_the_openai_row(
monkeypatch: pytest.MonkeyPatch, lookup_trips: bool
) -> None:
openai_key: Final = "lit-5853-openai-model"
azure_ai_key: Final = f"azure_ai/{openai_key}"
openai_row: Final = {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 5e-06,
"output_cost_per_token": 1.5e-05,
}
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setenv("AZURE_AI_API_BASE", "https://lit-5853.openai.azure.com")
monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), openai_key: dict(openai_row)})
monkeypatch.setattr(
litellm, "open_ai_chat_completion_models", {*litellm.open_ai_chat_completion_models, openai_key}
)
_invalidate_model_cost_lowercase_map()
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
if lookup_trips:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
litellm.register_model(
{azure_ai_key: {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}},
persist_across_reloads=False,
)
try:
assert litellm.model_cost[openai_key] == openai_row
assert litellm.cost_per_token(model=azure_ai_key, prompt_tokens=1000, completion_tokens=1000) == pytest.approx(
(1e-03, 2e-03)
)
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monkeypatch: pytest.MonkeyPatch) -> None:
catalog_key: Final = "au.anthropic.lit-5853-reloaded-model"
deployment_key: Final = f"bedrock/{catalog_key}"
catalog_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
}
reloaded_input_cost: Final = 6.6e-06
builtin_cost_map: Final = litellm.get_model_cost_map(url="")
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)})
_invalidate_model_cost_lowercase_map()
reloaded_cost_map: Final = {
**builtin_cost_map,
catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost},
}
real_lookup: Final = litellm.utils.get_model_info
def lookup_racing_a_price_data_reload(model: str, custom_llm_provider: str | None = None) -> object:
info: Final = real_lookup(model=model, custom_llm_provider=custom_llm_provider)
if litellm.model_cost is not reloaded_cost_map:
adopt_model_cost_map(reloaded_cost_map)
return info
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_racing_a_price_data_reload)
litellm.register_model({deployment_key: {"litellm_provider": "bedrock"}}, persist_across_reloads=False)
try:
assert litellm.model_cost[catalog_key]["input_cost_per_token"] == reloaded_input_cost
assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_its_price_when_a_reload_adopts_a_map_as_it_writes(
monkeypatch: pytest.MonkeyPatch,
) -> None:
catalog_key: Final = "au.anthropic.lit-5853-reloaded-as-written-model"
deployment_key: Final = f"bedrock/{catalog_key}"
catalog_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
}
reloaded_input_cost: Final = 6.6e-06
registered_output_cost: Final = 2e-05
builtin_cost_map: Final = litellm.get_model_cost_map(url="")
reloaded_cost_map: Final = {
**builtin_cost_map,
catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost},
}
class CostMapAReloadReplacesAsItIsWritten(dict[str, dict[str, object]]):
def update(self, other: Mapping[str, dict[str, object]], /) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # registration writes pass one mapping
if litellm.model_cost is self:
adopt_model_cost_map(reloaded_cost_map)
super().update(other)
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
CostMapAReloadReplacesAsItIsWritten({**builtin_cost_map, catalog_key: dict(catalog_row)}),
)
_invalidate_model_cost_lowercase_map()
litellm.register_model(
{deployment_key: {"litellm_provider": "bedrock", "output_cost_per_token": registered_output_cost}},
persist_across_reloads=False,
)
try:
assert litellm.model_cost is reloaded_cost_map
assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost
assert litellm.get_model_info(deployment_key)["output_cost_per_token"] == registered_output_cost
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_leaves_the_entry_a_lookup_is_walking_intact(monkeypatch: pytest.MonkeyPatch) -> None:
catalog_key: Final = "lit-5853-live-entry-model"
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
catalog_key: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 2e-06},
},
)
_invalidate_model_cost_lowercase_map()
walk: Final = iter(litellm.model_cost[catalog_key].items())
first_field: Final = next(walk)
try:
litellm.register_model(
{catalog_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06, "supports_vision": True}},
persist_across_reloads=False,
)
assert dict((first_field, *walk)) == {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 2e-06,
}
assert litellm.get_model_info(catalog_key)["input_cost_per_token"] == 1e-06
assert litellm.get_model_info(catalog_key)["supports_vision"] is True
finally:
_invalidate_model_cost_lowercase_map()
_CANONICAL_KEY: Final = "lit-5853-canonical-model"
_ALIAS_KEY: Final = "lit-5853-alias-model"
def _install_aliased_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
_expand_model_aliases(
{
**litellm.get_model_cost_map(url=""),
_CANONICAL_KEY: {
"aliases": [_ALIAS_KEY],
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 2e-06,
"output_cost_per_token": 8e-06,
},
}
),
)
_invalidate_model_cost_lowercase_map()
def _token_prices(name: str) -> tuple[float | None, float | None]:
info: Final = litellm.get_model_info(name)
return info["input_cost_per_token"], info["output_cost_per_token"]
@pytest.mark.parametrize("registered_key", [_CANONICAL_KEY, _ALIAS_KEY])
def test_register_model_prices_every_alias_of_the_registered_entry(
monkeypatch: pytest.MonkeyPatch, registered_key: str
) -> None:
_install_aliased_cost_map(monkeypatch)
try:
litellm.register_model(
{registered_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}},
persist_across_reloads=False,
)
litellm.register_model(
{registered_key: {"litellm_provider": "openai", "output_cost_per_token": 4e-06}},
persist_across_reloads=False,
)
assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == {
_CANONICAL_KEY: (1e-06, 4e-06),
_ALIAS_KEY: (1e-06, 4e-06),
}
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_leaves_an_alias_repointed_at_its_own_entry(monkeypatch: pytest.MonkeyPatch) -> None:
_install_aliased_cost_map(monkeypatch)
litellm.model_cost[_ALIAS_KEY] = {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 5e-06,
"output_cost_per_token": 6e-06,
}
try:
litellm.register_model(
{_CANONICAL_KEY: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}},
persist_across_reloads=False,
)
assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == {
_CANONICAL_KEY: (1e-06, 8e-06),
_ALIAS_KEY: (5e-06, 6e-06),
}
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_inherits_builtin_token_pricing_for_unmapped_key(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
_invalidate_model_cost_lowercase_map()
builtin_key: Final = "us.anthropic.claude-sonnet-4-6"
registered_key: Final = f"bedrock/bedrock/bedrock/{builtin_key}"
builtin: Final = litellm.model_cost[builtin_key]
try:
litellm.register_model(
{registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False
)
info: Final = litellm.get_model_info(registered_key)
assert info["input_cost_per_token"] == builtin["input_cost_per_token"] > 0
assert info["output_cost_per_token"] == builtin["output_cost_per_token"] > 0
assert info["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"]
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_the_capabilities_an_unregistered_region_alias_resolves_to(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.litellm_core_utils.fallback_generalizations import (
get_fallback_generalization_rules,
set_fallback_generalizations,
)
sibling_key: Final = "anthropic.lit-5853-flagless-model"
region_alias: Final = f"bedrock/ap-northeast-1.{sibling_key}"
sibling_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
}
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), sibling_key: sibling_row})
_invalidate_model_cost_lowercase_map()
previous_rules: Final = list(get_fallback_generalization_rules())
set_fallback_generalizations(
[
{
"name": "lit-5853-flagless",
"pattern": "lit-5853-flagless",
"model_info": {"supports_system_messages": True, "supports_parallel_function_calling": True},
}
]
)
try:
unregistered: Final = litellm.get_model_info(region_alias)
litellm.register_model(
{region_alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False
)
registered: Final = litellm.get_model_info(region_alias)
assert registered["supports_system_messages"] is unregistered["supports_system_messages"] is True
assert (
registered["supports_parallel_function_calling"]
is unregistered["supports_parallel_function_calling"]
is True
)
assert registered["input_cost_per_token"] == sibling_row["input_cost_per_token"]
finally:
set_fallback_generalizations(previous_rules)
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize(
("builtin_prices", "custom_prices", "usage", "expected_cost"),
[
pytest.param(
{
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"input_cost_per_token_above_200k_tokens": 6e-06,
"output_cost_per_token_above_200k_tokens": 2.25e-05,
},
{"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06},
{"prompt_tokens": 250_000, "completion_tokens": 1_000},
250_000 * 1e-06 + 1_000 * 2e-06,
id="above-200k-tier",
),
pytest.param(
{
"tiered_pricing": [
{"input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07, "range": [0, 256_000]},
{"input_cost_per_token": 2.5e-07, "output_cost_per_token": 2e-06, "range": [256_000, 1_000_000]},
]
},
{"input_cost_per_token": 1e-06},
{"prompt_tokens": 300_000, "completion_tokens": 1_000},
300_000 * 1e-06 + 1_000 * 2e-06,
id="tiered-pricing-input-only",
),
pytest.param(
{"input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05, "output_cost_per_reasoning_token": 3e-05},
{"output_cost_per_token": 2e-06},
{"prompt_tokens": 100, "completion_tokens": 1_000, "reasoning_tokens": 600},
100 * 3e-06 + 1_000 * 2e-06,
id="reasoning-rate-output-only",
),
pytest.param(
{
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"off_peak_pricing": {
"input_cost_per_token": 1e-06,
"output_cost_per_token": 5e-06,
"hours_utc": "00:00-00:00",
},
},
{"output_cost_per_token": 2e-06},
{"prompt_tokens": 100, "completion_tokens": 1_000},
100 * 1e-06 + 1_000 * 2e-06,
id="off-peak-output-only",
),
],
)
def test_register_model_keeps_custom_flat_prices_over_inherited_rate_variants(
monkeypatch: pytest.MonkeyPatch,
builtin_prices: dict[str, object],
custom_prices: dict[str, float],
usage: dict[str, int],
expected_cost: float,
) -> None:
from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage
builtin_key: Final = "anthropic.lit-5853-tiered-model-v1:0"
alias: Final = f"bedrock/ap-northeast-1.{builtin_key}"
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
builtin_key: {"litellm_provider": "bedrock", "mode": "chat", **builtin_prices},
},
)
_invalidate_model_cost_lowercase_map()
try:
litellm.register_model(
{alias: {"litellm_provider": "bedrock", **custom_prices}},
persist_across_reloads=False,
)
response_usage: Final = Usage(
prompt_tokens=usage["prompt_tokens"],
completion_tokens=usage["completion_tokens"],
total_tokens=usage["prompt_tokens"] + usage["completion_tokens"],
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=usage.get("reasoning_tokens", 0)),
)
cost: Final = litellm.completion_cost(
completion_response=ModelResponse(model=alias, usage=response_usage),
model=alias,
custom_llm_provider="bedrock",
)
assert cost == pytest.approx(expected_cost)
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_never_inherits_pricing_across_providers(monkeypatch: pytest.MonkeyPatch) -> None:
openai_key: Final = "lit-5853-openai-only-model"
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
openai_key: {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 2.5e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-06,
},
},
)
_invalidate_model_cost_lowercase_map()
registered_key: Final = f"bedrock/{openai_key}"
try:
litellm.register_model(
{registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False
)
registered: Final = litellm.model_cost[registered_key]
assert "cache_read_input_token_cost" not in registered
assert "input_cost_per_token" not in registered
assert "output_cost_per_token" not in registered
finally:
_invalidate_model_cost_lowercase_map()
def test_update_dictionary_merges_nested_dicts_without_aliasing():
"""A nested dict must be merged copy-on-write: the pre-existing nested dict
object stays untouched, and the caller's incoming nested dict is never

99
tests/unit/test_router/test_router.py
View file

@ -23912,6 +23912,105 @@ def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name
result = router.map_team_model(team_model_name="team-model", team_id="team-1")
assert result == "team-model", f"Expected public name 'team-model', got {result}"
def _team_deployment(
public_name: str,
litellm_model: str,
team_id: str = "team-a",
deployment_id: str | None = None,
) -> DeploymentTypedDict:
return {
"model_name": f"internal-{team_id}-{public_name}",
"litellm_params": {"model": litellm_model},
"model_info": {
"id": deployment_id or f"{team_id}-{public_name}",
"team_id": team_id,
"team_public_model_name": public_name,
},
}
_LANGGRAPH_AGENT: Final[DeploymentTypedDict] = _team_deployment("a-langgraph-agent", "langgraph/repro-agent")
_BEDROCK_CLAUDE: Final[DeploymentTypedDict] = _team_deployment(
"a-bedrock-claude", "bedrock/anthropic.claude-sonnet-4-6"
)
_OPENAI_GPT: Final[DeploymentTypedDict] = _team_deployment("a-gpt", "openai/gpt-4o-mini")
@pytest.mark.parametrize(
"team_models, custom_llm_provider, expected",
[
pytest.param([_LANGGRAPH_AGENT, _BEDROCK_CLAUDE], "bedrock", "a-bedrock-claude", id="match-is-second"),
pytest.param([_BEDROCK_CLAUDE, _LANGGRAPH_AGENT], "bedrock", "a-bedrock-claude", id="match-is-first"),
pytest.param([_LANGGRAPH_AGENT, _OPENAI_GPT], "bedrock", None, id="no-match-returns-none"),
pytest.param([_LANGGRAPH_AGENT, _BEDROCK_CLAUDE], None, "a-langgraph-agent", id="no-provider-keeps-first"),
pytest.param(
[_team_deployment("b-bedrock-claude", "bedrock/anthropic.claude-sonnet-4-6", team_id="team-b")],
"bedrock",
None,
id="other-teams-match-ignored",
),
pytest.param(
[
_team_deployment("shared", "langgraph/repro-agent", deployment_id="shared-langgraph"),
_team_deployment(
"shared",
"bedrock/anthropic.claude-sonnet-4-6",
deployment_id="shared-bedrock",
),
],
"bedrock",
"shared-bedrock",
id="mixed-public-name-langgraph-first",
),
pytest.param(
[
_team_deployment(
"shared",
"bedrock/anthropic.claude-sonnet-4-6",
deployment_id="shared-bedrock",
),
_team_deployment("shared", "langgraph/repro-agent", deployment_id="shared-langgraph"),
],
"bedrock",
"shared-bedrock",
id="mixed-public-name-bedrock-first",
),
pytest.param(
[
_team_deployment("shared", "bedrock/anthropic.claude-sonnet-4-6", deployment_id="bedrock-one"),
_team_deployment("shared", "bedrock/anthropic.claude-sonnet-4-6", deployment_id="bedrock-two"),
],
"bedrock",
"shared",
id="same-provider-public-name-keeps-load-balancing",
),
],
)
def test_map_team_model_without_model_picks_deployment_serving_the_requested_provider(
team_models: list[DeploymentTypedDict], custom_llm_provider: str | None, expected: str | None
) -> None:
"""A model-less vector store call routed through the team's first deployment, so a langgraph
team model made a Bedrock vector store search fail with "not supported for langgraph"."""
router = Router(model_list=team_models)
result = router.map_team_model(None, "team-a", custom_llm_provider=custom_llm_provider)
assert result == expected, (
f"team models {[m['litellm_params']['model'] for m in team_models]}, provider {custom_llm_provider}"
)
def test_map_team_model_with_explicit_model_ignores_requested_provider() -> None:
router = Router(model_list=[_LANGGRAPH_AGENT, _BEDROCK_CLAUDE])
result = router.map_team_model("a-langgraph-agent", "team-a", custom_llm_provider="bedrock")
assert result == "a-langgraph-agent", (
"an explicit model is the caller's choice and must not be re-picked by provider"
)
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
class TestPreCallChecksOptimization:
"""

75
tests/unit/test_router_model_cost_isolation.py
View file

@ -9,9 +9,14 @@ should still use the built-in pricing.
import asyncio
import copy
import itertools
import logging
import os
import queue
import re
import sys
import threading
from types import MappingProxyType
from typing import Final
from unittest.mock import Mock, patch
@ -3241,6 +3246,76 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info(
assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price
def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
catalog: Final = litellm.get_model_cost_map(url="")
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
_invalidate_model_cost_lowercase_map()
backends: Final = tuple(
f"bedrock/au.anthropic.{model}"
for model in ("claude-sonnet-5", "claude-opus-5", "claude-haiku-4-5-20251001-v1:0", "claude-sonnet-4-6")
)
canonical_prices: Final = MappingProxyType(
{backend: catalog[backend.removeprefix("bedrock/")]["input_cost_per_token"] for backend in backends}
)
assert all(price > 0 for price in canonical_prices.values())
router: Final = Router(
model_list=[
{
"model_name": backend.rsplit(".", 1)[-1],
"litellm_params": {"model": backend, "aws_region_name": "ap-southeast-2"},
"model_info": {"id": f"lit-5853-dep-{index}"},
}
for index, backend in enumerate(backends)
]
)
stop: Final = threading.Event()
writer_outcome: Final[queue.SimpleQueue[int | Exception]] = queue.SimpleQueue()
def register_until_stopped() -> None:
try:
for i in itertools.count():
if stop.is_set():
writer_outcome.put(i)
return
litellm.register_model(
{
f"lit-5853-writer-{i}": {
"litellm_provider": "lit-5853-writer",
"mode": "chat",
"input_cost_per_token": 1e-07,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 1e-08,
}
},
persist_across_reloads=False,
)
except Exception as failure:
writer_outcome.put(failure)
writer: Final = threading.Thread(target=register_until_stopped, daemon=True)
saved_switch_interval: Final = sys.getswitchinterval()
sys.setswitchinterval(1e-4)
writer.start()
try:
for _ in range(10):
_simulate_price_data_reload_with_provider_sets(monkeypatch, dict(catalog))
finally:
stop.set()
writer.join()
sys.setswitchinterval(saved_switch_interval)
registrations_landed: Final = writer_outcome.get(timeout=5)
try:
assert isinstance(registrations_landed, int) and registrations_landed > 0, registrations_landed
assert tuple(backend for backend in backends if backend in litellm.model_cost) == ()
for backend in backends:
assert litellm.get_model_info(backend)["input_cost_per_token"] == canonical_prices[backend]
assert tuple(router.get_model_names()) == tuple(backend.rsplit(".", 1)[-1] for backend in backends)
finally:
_invalidate_model_cost_lowercase_map()
_COLLIDING_BUILTIN_KEY: Final = "baseten/zai-org/GLM-5.2"
_COLLIDING_SHARED_OPENAI_KEY: Final = "openai/zai-org/GLM-5.2"
_COLLIDING_MODEL_INFO: Final = {

119
tests/unit/test_utils.py
View file

@ -14,7 +14,7 @@ from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from datetime import datetime, timedelta, timezone
from pathlib import Path, PurePath
from typing import Final, cast
from typing import Final, SupportsIndex, cast
from unittest import mock
from unittest.mock import AsyncMock, MagicMock, Mock, patch
@ -4887,6 +4887,104 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey
_invalidate_model_cost_lowercase_map()
def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.utils import _invalidate_model_cost_lowercase_map, _rebuild_model_cost_lowercase_map
class _KeyRegisteredMidWalk(str):
def lower(self) -> str:
litellm.model_cost.setdefault("lit-5853-registered-during-the-rebuild", {"mode": "chat"})
return super().lower()
class _EntryRegisteredMidWalk(dict):
def get(self, field: str, default: object = None) -> object:
litellm.model_cost.setdefault("lit-5853-registered-during-add-known-models", {"mode": "chat"})
return super().get(field, default)
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
litellm.model_cost[_KeyRegisteredMidWalk("LIT-5853-WALKED-LAST")] = _EntryRegisteredMidWalk(mode="chat")
try:
assert _rebuild_model_cost_lowercase_map()["lit-5853-walked-last"] == "LIT-5853-WALKED-LAST"
litellm.add_known_models(model_cost_map=litellm.model_cost)
assert "lit-5853-registered-during-the-rebuild" in litellm.model_cost
assert "lit-5853-registered-during-add-known-models" in litellm.model_cost
finally:
_invalidate_model_cost_lowercase_map()
class _KeyWhoseLoweringLosesTheRace(str):
"""Lowering the key is the one Python-level step between the lookup map's rebuild and
its read, so another thread's registration invalidating the map lands right there."""
def lower(self) -> str:
from litellm.utils import _invalidate_model_cost_lowercase_map
_invalidate_model_cost_lowercase_map()
return super().lower()
def test_case_insensitive_model_cost_lookup_survives_an_invalidation_mid_lookup() -> None:
from litellm.utils import _get_model_cost_key, _invalidate_model_cost_lowercase_map
_invalidate_model_cost_lowercase_map()
try:
assert _get_model_cost_key(_KeyWhoseLoweringLosesTheRace("GPT-4O")) == "gpt-4o"
finally:
_invalidate_model_cost_lowercase_map()
def _register_a_model_mid_walk() -> None:
litellm.model_cost.setdefault("lit-5853-registered-mid-walk", {"mode": "chat"})
class _KeyReadDuringARegistration(str):
def startswith(
self,
prefix: str | tuple[str, ...],
start: SupportsIndex | None = None,
end: SupportsIndex | None = None,
) -> bool:
_register_a_model_mid_walk()
return super().startswith(prefix, start, end)
def lower(self) -> str:
_register_a_model_mid_walk()
return super().lower()
class _EntryReadDuringARegistration(dict[str, object]):
def __contains__(self, key: object) -> bool:
_register_a_model_mid_walk()
return super().__contains__(key)
@pytest.mark.parametrize("walk", ["fal_ai", "fireworks_ai", "ollama", "gpt-5", "case-insensitive scan"])
def test_request_path_cost_map_walks_survive_a_registration_mid_walk(monkeypatch: pytest.MonkeyPatch, walk: str) -> None:
from litellm.llms.fal_ai.cost_calculator import _keyed_rows
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
from litellm.llms.ollama.common_utils import OllamaModelInfo
from litellm.llms.openai.chat.gpt_5_transformation import _catalogue_declares_default_effort
from litellm.utils import _handle_new_key_with_scan, _invalidate_model_cost_lowercase_map
walks: Final[Mapping[str, Callable[[], object]]] = {
"fal_ai": lambda: _keyed_rows("lit-5853-model", "standard"),
"fireworks_ai": FireworksAIConfig._get_fireworks_index,
"ollama": lambda: OllamaModelInfo._is_static_ollama_model("ollama/lit-5853-model"),
"gpt-5": _catalogue_declares_default_effort,
"case-insensitive scan": lambda: _handle_new_key_with_scan("lit-5853-model"),
}
monkeypatch.setattr(
litellm,
"model_cost",
{_KeyReadDuringARegistration("lit-5853-catalog-model"): _EntryReadDuringARegistration(mode="chat")},
)
try:
walks[walk]()
assert "lit-5853-registered-mid-walk" in litellm.model_cost
finally:
_invalidate_model_cost_lowercase_map()
class _JsonCapture(logging.Handler):
def __init__(self):
super().__init__()
@ -7091,6 +7189,25 @@ def test_get_model_info_gemini(monkeypatch):
assert info.get("rpm") is not None, f"{model} does not have rpm"
def test_register_model_without_a_provider_never_inherits_another_providers_entry(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.utils import _invalidate_model_cost_lowercase_map
key: Final = "azure/claude-sonnet-4-6"
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
litellm.model_cost.pop(key, None)
try:
litellm.register_model(
{key: {"input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07}}, persist_across_reloads=False
)
assert litellm.model_cost[key].get("litellm_provider") is None
info: Final = litellm.get_model_info(key)
assert info["litellm_provider"] == "azure"
assert info["input_cost_per_token"] == 1e-07
assert info["output_cost_per_token"] == 4e-07
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize(
"model, litellm_params, expects_native_config",
[

14
tests/unit/types/test_guardrails.py
View file

@ -7,6 +7,7 @@ from typing import Any, List, Optional, Tuple
from unittest.mock import MagicMock, patch
import pytest
from pydantic import ValidationError
import litellm
import litellm.litellm_core_utils
@ -178,3 +179,16 @@ def test_litellm_params_treats_blank_values_as_unset(blank: str) -> None:
assert with_blank_values.model_dump() == omitted.model_dump()
assert BaseLitellmParams(fail_on_error=blank).model_dump() == BaseLitellmParams().model_dump()
@pytest.mark.parametrize(("configured", "loaded"), [("2", 2), (2, 2), ("0", 0), (5, 5), ("", None)])
def test_num_retries_loads_the_dashboard_string_and_blank_means_unset(configured: object, loaded: int | None) -> None:
assert (
LitellmParams(guardrail="generic_guardrail_api", mode="pre_call", num_retries=configured).num_retries == loaded
)
@pytest.mark.parametrize("configured", [-1, 6, "1.5", "twice"])
def test_num_retries_rejects_values_that_are_not_a_retry_count(configured: object) -> None:
with pytest.raises(ValidationError, match="num_retries"):
LitellmParams(guardrail="generic_guardrail_api", mode="pre_call", num_retries=configured)

10
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -27052,6 +27052,11 @@ export interface components {
* @description Optional field if guardrail requires a 'model' parameter
*/
model?: string | null;
/**
* Num Retries
* @description How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it.
*/
num_retries?: number | null;
/**
* On Sensitive Data
* @description Action to take when sensitive data is detected. 'block' raises an exception (default behavior). 'route' reroutes the request to the model specified in sensitive_data_route_to_model.
@ -36638,6 +36643,11 @@ export interface components {
* @description If True, logs violations without blocking. Defaults to False if not provided
*/
monitor_mode?: boolean | null;
/**
* Num Retries
* @description How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it.
*/
num_retries?: number | null;
/**
* On Disallowed Action
* @description Choose whether disallowed tools block the request or get rewritten out of the payload