diff --git a/litellm/__init__.py b/litellm/__init__.py index 959c9fb406d..9cc8965b628 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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": diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 13a53c55337..c52bc534058 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1196,6 +1196,9 @@ def _map_sagemaker_exception( ) +_VERTEX_403_STATUS_CODE_PATTERN: Final = re.compile(r"(? 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,...}}. diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index b68db22183a..073e3a3455e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -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 diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index aa3d7942174..2acf1c69594 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -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),) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 8b6946ee562..14ae59d4f85 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -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): diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index b5328b1ba4d..284a6ffea62 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -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 diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 94c5c5d016d..5bf19844287 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -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( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8545a5b568e..6ebb54c0035 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 449adaa73e9..3b2c2e09df5 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 02799b90566..96cf6ee11aa 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 961d32ab664..7198d74909c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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", diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 8dab693cd58..f5e8304c83c 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -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). diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index 242354ac417..44705c4d166 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 93b08831d84..199da7df30d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -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() diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d014ff6cf07..a014b9cf8a5 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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( diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 6f93bcd82dc..8bab7215b17 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -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) diff --git a/litellm/router.py b/litellm/router.py index 1b3857dfb71..8d77fd0dac5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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: diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 33db8cbc6f0..ff063d2e233 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index d34a3a13838..d38e54986b9 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -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=( diff --git a/litellm/utils.py b/litellm/utils.py index fc3acd3065d..b37827a3b13 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8545a5b568e..6ebb54c0035 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/tests/code_coverage_tests/ensure_async_clients_test.py b/tests/code_coverage_tests/ensure_async_clients_test.py index 5a534199cbc..93ae2e2837b 100644 --- a/tests/code_coverage_tests/ensure_async_clients_test.py +++ b/tests/code_coverage_tests/ensure_async_clients_test.py @@ -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", diff --git a/tests/e2e/perf/setup/locustfile.py b/tests/e2e/perf/setup/locustfile.py index 1609617bd73..ccc7cf698f8 100644 --- a/tests/e2e/perf/setup/locustfile.py +++ b/tests/e2e/perf/setup/locustfile.py @@ -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) diff --git a/tests/e2e/perf/setup/run.py b/tests/e2e/perf/setup/run.py index 9ba9b45104d..272c64754c4 100644 --- a/tests/e2e/perf/setup/run.py +++ b/tests/e2e/perf/setup/run.py @@ -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) diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py new file mode 100644 index 00000000000..37e5c7bc904 --- /dev/null +++ b/tests/e2e/perf/setup/upgrade.py @@ -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), + ) diff --git a/tests/e2e/perf/setup/upgrade_worker.py b/tests/e2e/perf/setup/upgrade_worker.py new file mode 100644 index 00000000000..e69e841814f --- /dev/null +++ b/tests/e2e/perf/setup/upgrade_worker.py @@ -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()) diff --git a/tests/e2e/perf/setup/wire.py b/tests/e2e/perf/setup/wire.py index 07dc5d170ae..ced79d6970c 100644 --- a/tests/e2e/perf/setup/wire.py +++ b/tests/e2e/perf/setup/wire.py @@ -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( diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py new file mode 100644 index 00000000000..f165b9b738f --- /dev/null +++ b/tests/e2e/perf/test_database_upgrade.py @@ -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) diff --git a/tests/e2e_harness/perf/test_run.py b/tests/e2e_harness/perf/test_run.py new file mode 100644 index 00000000000..23f62c9c560 --- /dev/null +++ b/tests/e2e_harness/perf/test_run.py @@ -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), + ) diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py new file mode 100644 index 00000000000..0904b7a7b5a --- /dev/null +++ b/tests/e2e_harness/perf/test_upgrade.py @@ -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) diff --git a/tests/integration/observability/test_generic_guardrail_num_retries.py b/tests/integration/observability/test_generic_guardrail_num_retries.py new file mode 100644 index 00000000000..df761d25968 --- /dev/null +++ b/tests/integration/observability/test_generic_guardrail_num_retries.py @@ -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() == () diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py new file mode 100644 index 00000000000..bb24130aa6c --- /dev/null +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -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) diff --git a/tests/integration/providers/_idle_upstream_support.py b/tests/integration/providers/_idle_upstream_support.py index ca808abc008..5fccd9804a6 100644 --- a/tests/integration/providers/_idle_upstream_support.py +++ b/tests/integration/providers/_idle_upstream_support.py @@ -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}), }, }, ), diff --git a/tests/integration/providers/test_google_genai_vertex_routes_wire.py b/tests/integration/providers/test_google_genai_vertex_routes_wire.py index d2c2f3bafb2..d7c2aa72c74 100644 --- a/tests/integration/providers/test_google_genai_vertex_routes_wire.py +++ b/tests/integration/providers/test_google_genai_vertex_routes_wire.py @@ -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]) diff --git a/tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py b/tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py new file mode 100644 index 00000000000..6b12e61aa4e --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py @@ -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 diff --git a/tests/integration/providers/test_vector_store_team_model_provider_wire.py b/tests/integration/providers/test_vector_store_team_model_provider_wire.py new file mode 100644 index 00000000000..5b986a64ff6 --- /dev/null +++ b/tests/integration/providers/test_vector_store_team_model_provider_wire.py @@ -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 diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 9eb99516e02..731e842090a 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -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.""" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 5fdc7669496..a43dacaf100 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -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" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 691da35b8c4..d47b2f73a48 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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 diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 21958e88008..7bdcb8afd37 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -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.""" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index f50da1bff64..8308a7ba386 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -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"]} diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 0d59c655835..65c4a8faeb7 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -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 diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 67f29220cb6..4e4e2f23950 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -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 diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 02c90e0d97b..35fbd28eae7 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -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 diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 50a9d98385d..bd3cf738af8 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -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: """ diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d8625ff9da3..b353dc86d5f 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -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 = { diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 87f9d58e247..a8e90fded83 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -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", [ diff --git a/tests/unit/types/test_guardrails.py b/tests/unit/types/test_guardrails.py index 5c652a10e16..97bab8295f4 100644 --- a/tests/unit/types/test_guardrails.py +++ b/tests/unit/types/test_guardrails.py @@ -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) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7fae3437f22..e5c1b9f2c26 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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