From f0f4f406c4db73c02b202f9dcdda9b0e533a1e41 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 30 Sep 2026 09:26:04 +0000 Subject: [PATCH 1/4] refactor(types): replace Any with proven types in 9 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 3 ++- litellm/litellm_core_utils/streaming_handler.py | 2 +- .../mcp_server/mcp_server_manager.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 16 ++++++++-------- .../guardrail_hooks/bedrock_guardrails.py | 2 +- .../guardrail_hooks/xecguard/xecguard.py | 16 +++++++++++++--- litellm/proxy/proxy_server.py | 2 +- .../responses/mcp/litellm_proxy_mcp_handler.py | 2 +- .../complexity_router/complexity_router.py | 2 +- 9 files changed, 29 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e292ab7b2ec2..8d3c2c9c4cf1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -111,6 +111,7 @@ SERVED_OUTPUT_TEXTS_KEY, overlay_served_output_texts, ) +from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.responses.utils import ResponseAPILoggingUtils @@ -173,7 +174,7 @@ Usage, ) from litellm.types.videos.main import VideoObject -from litellm.utils import _get_base_model_from_metadata, executor, print_verbose +from litellm.utils import _get_base_model_from_metadata, print_verbose from ..integrations.argilla import ArgillaLogger from ..integrations.arize.arize_phoenix import ArizePhoenixLogger diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index d2853a625c9b..dd5569a2504c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1106,7 +1106,7 @@ def _dispatch_provider_chunk( self, chunk: Any, model_response: ModelResponseStream, - completion_obj: dict[str, Any], + completion_obj: dict[str, object], ) -> _ProviderChunkResult: response_obj: dict[str, Any] = {} if ( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c0792c32de2f..e0cb50b1263b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1678,7 +1678,7 @@ def _create_elicitation_callback(): def _record_mcp_guardrail_evaluations( - synthetic_llm_data: dict[str, Any], # mutable-ok: `_sync_guardrail_info_to_logging_obj` takes a concrete dict + synthetic_llm_data: dict[str, object], # mutable-ok: `_sync_guardrail_info_to_logging_obj` takes a concrete dict litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> None: """Bridge guardrail decision records off an MCP synthetic request onto the request's logger. diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ed3ec7b4dde2..24ac37e01313 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -13,7 +13,7 @@ import secrets from collections.abc import Mapping from datetime import datetime, timezone -from typing import Any, Final, NamedTuple, Protocol, Union, cast +from typing import Final, NamedTuple, Protocol, Union, cast import fastapi import orjson @@ -207,7 +207,7 @@ def _get_model_from_request_context( request_data: dict, route: str, request: Request | None, - llm_router: Any | None = None, + llm_router: litellm.Router | None = None, team_id: str | None = None, ) -> str | list[str] | None: return get_model_from_request( @@ -512,8 +512,8 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str: def _routing_selector_matches_claim( - selector_value: Any | None, - claim_value: Any | None, + selector_value: str | list[str] | None, + claim_value: object, *, split_space_delimited: bool = False, ) -> bool: @@ -653,7 +653,7 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str # ``websocket.url``, which Starlette reconstructs from the (poisonable) # Host header. Carry the ASGI scope's path / root_path so the lookup # never reaches the fallback. - synthetic_scope: Final[dict[str, Any]] = { + synthetic_scope: Final[dict[str, object]] = { "type": "http", "headers": scope_headers, "path": ws_scope.get("path", ""), @@ -3099,7 +3099,7 @@ async def _reserve_budget_after_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, request_data: dict, route: str, - llm_router: Any | None, + llm_router: litellm.Router | None, team_object: LiteLLM_TeamTableCachedObj | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, @@ -3142,7 +3142,7 @@ def _should_skip_budget_checks( request_data: dict, route: str, request: Request | None, - llm_router: Any | None, + llm_router: litellm.Router | None, team_id: str | None = None, ) -> bool: model: Final = _get_model_from_request_context( @@ -3598,7 +3598,7 @@ async def _enforce_key_and_fallback_model_access( route: str, request: Request | None, llm_model_list: list | None, - llm_router: Any | None, + llm_router: litellm.Router | None, ) -> None: """ Key-level model allowlist and client fallbacks (same as standard auth). diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 228b31604a3b..bdcd9a51938e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1181,7 +1181,7 @@ async def _post_apply_guardrail_content_with_retry( or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES ): raise - await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt)) + await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2.0**attempt)) raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") async def _post_apply_guardrail_content( diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index ddf9cace8b91..d5a724f7458b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -28,6 +28,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional from fastapi.exceptions import HTTPException +from typing_extensions import ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -77,6 +78,12 @@ def _sanitize_scan_result_for_logging(scan_result: dict) -> dict: ] +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + supported_event_hooks: ReadOnly[list[GuardrailEventHooks] | None] + + class XecGuardMissingCredentials(Exception): pass @@ -90,7 +97,7 @@ def __init__( policy_names: list[str] | None = None, block_on_error: bool | None = None, grounding_strictness: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.api_key = api_key or os.environ.get("XECGUARD_API_KEY") if not self.api_key: @@ -122,9 +129,12 @@ def __init__( llm_provider=httpxSpecialProvider.GuardrailCallback, ) - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) + options: Final[_CustomGuardrailOptions] = { + "supported_event_hooks": list(self.get_supported_event_hooks()), + **kwargs, + } - super().__init__(**kwargs) + super().__init__(**options) @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a994293a4d89..4f2cc9e506e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3598,7 +3598,7 @@ async def _get_source_cache_base_spend( ) -> float: source_cache_keys: Final = [source_cache_key] if isinstance(source_cache_key, str) else source_cache_key for cache_key in source_cache_keys: - source = await user_api_key_cache.async_get_cache(key=cache_key) + source: object = await user_api_key_cache.async_get_cache(key=cache_key) if source is None: continue if isinstance(source, dict): diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 71f610791540..5bef5f319a15 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -158,7 +158,7 @@ def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: items: Final = tuple(tools or ()) gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] - other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + other_tools: Final[list[ToolParam]] = [tool for tool in items if not _names_gateway_explicitly(tool)] return gateway_tools, other_tools @staticmethod diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 76ee977bf281..ccd2fdf11dde 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2441,7 +2441,7 @@ async def _classify_with_llm( self, prompt: str, system_prompt: str | None = None, - request_kwargs: dict[str, Any] | None = None, + request_kwargs: Mapping[str, object] | None = None, messages: Sequence[Mapping[str, object]] | None = None, ) -> tuple[ComplexityTier | str, float | None]: """ From 18181bd0906c0a83ecc34455cd680c283d8472a9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 30 Sep 2026 10:29:11 +0000 Subject: [PATCH 2/4] test(integration): pin xecguard-adjacent guardrail retry, mcp mixed tools, jwt routing and complexity router paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_jwt_routing_override.py | 83 +++++++++++++++ .../mcp/test_responses_mcp_mixed_tools.py | 79 ++++++++++++++ .../test_bedrock_guardrail_throttle_retry.py | 100 ++++++++++++++++++ .../test_complexity_router_llm_classifier.py | 93 ++++++++++++++++ 4 files changed, 355 insertions(+) create mode 100644 tests/integration/authorization/test_jwt_routing_override.py create mode 100644 tests/integration/mcp/test_responses_mcp_mixed_tools.py create mode 100644 tests/integration/observability/test_bedrock_guardrail_throttle_retry.py create mode 100644 tests/integration/routing/test_complexity_router_llm_classifier.py diff --git a/tests/integration/authorization/test_jwt_routing_override.py b/tests/integration/authorization/test_jwt_routing_override.py new file mode 100644 index 000000000000..21b659a65d19 --- /dev/null +++ b/tests/integration/authorization/test_jwt_routing_override.py @@ -0,0 +1,83 @@ +import json +import time +import uuid +from pathlib import Path +from typing import Final + +import jwt +import yaml +from cryptography.hazmat.primitives.asymmetric import rsa + +from tests.integration._support.client import Gateway +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, wire_server + +KEY_ID: Final = "integration-routing-override-key" +ISSUER: Final = "https://idp.example.com" + + +def test_jwt_routing_override_sends_only_claim_matching_tokens_to_oauth2_introspection( + gateway: Gateway, tmp_path: Path +) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_jwk: Final = jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()) + subject: Final = f"oauth-machine-{uuid.uuid4().hex}" + + def jwks(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=json.dumps({"keys": [{**json.loads(public_jwk), "kid": KEY_ID}]}).encode()) + + def token_info(request: Request) -> Reply: + assert request.method == "GET", request + assert request.headers["authorization"].startswith("Bearer "), request.headers + return Reply(body=json.dumps({"sub": subject, "role": "proxy_admin"}).encode()) + + def token(**claims: object) -> str: + now: Final = int(time.time()) + return jwt.encode( + {"sub": subject, "iat": now, "exp": now + 300, **claims}, + private_key, + algorithm="RS256", + headers={"kid": KEY_ID}, + ) + + with wire_server(jwks) as keys, wire_server(token_info) as introspection, gateway.scenario() as scenario: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "routing_overrides": [ + { + "iss": ["https://other.example.com", "https://idp.*"], + "scope": "litellm.gateway", + "path": "oauth2", + } + ], + }, + } + path: Final = tmp_path / "jwt_routing_override.yaml" + path.write_text(yaml.safe_dump(config)) + overrides: Final = {"JWT_PUBLIC_KEY_URL": keys.url, "OAUTH_TOKEN_INFO_ENDPOINT": introspection.url} + with owned_proxy(gateway, tmp_path, overrides, config=path) as candidate: + model: Final = scenario.model() + + def chat(bearer: str) -> int: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "routing override"}]}, + key=bearer, + ).status_code + + assert chat(token(iss=ISSUER, scope="openid litellm.gateway")) == 200 + assert len(introspection.drain()) == 1, "space-delimited scope must match the override" + assert chat(token(iss=ISSUER, scope=["litellm.gateway", "openid"])) == 200 + assert len(introspection.drain()) == 1, "list scope must match the override" + + chat(token(iss=ISSUER, scope="openid")) + chat(token(iss=f"{ISSUER} evil.example.com", scope="litellm.gateway")) + chat(token(iss=ISSUER)) + assert introspection.drain() == (), "non-matching claims must stay on the JWT path" + assert keys.drain(), "non-matching tokens must be validated against the JWKS" diff --git a/tests/integration/mcp/test_responses_mcp_mixed_tools.py b/tests/integration/mcp/test_responses_mcp_mixed_tools.py new file mode 100644 index 000000000000..985a8b8a65ef --- /dev/null +++ b/tests/integration/mcp/test_responses_mcp_mixed_tools.py @@ -0,0 +1,79 @@ +import json +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.mcp import mcp_peer, register_mcp, tool_calls +from integration._support.wire import Reply, Request, wire_server + +_FUNCTION_TOOL: Final = { + "type": "function", + "name": "lookup_weather", + "description": "Look up the forecast for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, +} + + +def test_responses_with_gateway_mcp_and_caller_function_tool_hands_both_to_model_and_returns_the_function_call( + gateway: Gateway, +) -> None: + alias: Final = "mix" + uuid.uuid4().hex[:8] + upstream_tools: list[tuple[str, ...]] = [] + + def respond(request: Request) -> Reply: + assert request.target.endswith("/responses"), request.target + body: Final = json.loads(request.body) + upstream_tools.append(tuple(str(tool.get("name")) for tool in body.get("tools", ()))) + return Reply( + body=json.dumps( + { + "id": "resp_mixed", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "function_call", + "id": "fc_weather", + "call_id": "call_weather", + "name": "lookup_weather", + "arguments": json.dumps({"city": "Paris"}), + "status": "completed", + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with mcp_peer() as peer, wire_server(respond) as wire, gateway.scenario() as scenario: + server_id: Final = register_mcp(scenario, peer, alias) + model: Final = scenario.model(model="openai/responses/gpt-4o-mini", api_base=wire.url + "/v1") + key: Final = scenario.key(object_permission={"mcp_servers": [server_id]}) + peer.drain() + response: Final = gateway.client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {key}"}, + json={ + "model": model, + "input": "what is the weather in Paris", + "tools": [ + { + "type": "mcp", + "server_url": "litellm_proxy", + "server_label": "litellm", + "require_approval": "never", + }, + _FUNCTION_TOOL, + ], + }, + timeout=90, + ) + assert response.status_code == 200, response.text + assert upstream_tools, "model was never called" + assert all("lookup_weather" in names and f"{alias}-add" in names for names in upstream_tools), upstream_tools + calls: Final = [item for item in response.json()["output"] if item.get("type") == "function_call"] + assert [call["name"] for call in calls] == ["lookup_weather"], response.text + assert json.loads(calls[0]["arguments"]) == {"city": "Paris"} + assert tool_calls(peer.drain()) == (), "a caller-owned function call must never reach the MCP peer" diff --git a/tests/integration/observability/test_bedrock_guardrail_throttle_retry.py b/tests/integration/observability/test_bedrock_guardrail_throttle_retry.py new file mode 100644 index 000000000000..3381801b4900 --- /dev/null +++ b/tests/integration/observability/test_bedrock_guardrail_throttle_retry.py @@ -0,0 +1,100 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +_THROTTLED: Final = Reply( + status=429, + body=json.dumps({"message": "Rate exceeded"}).encode(), + headers={"x-amzn-ErrorType": "ThrottlingException"}, +) +_PASSED: Final = Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) +_PROVIDER: Final = Reply( + body=json.dumps( + { + "id": "chatcmpl-bedrock-retry", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "retried"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + } + ).encode() +) + + +def _config(tmp_path: Path, guardrail_id: str, endpoint: str) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + { + "guardrail_name": "bedrock" + uuid.uuid4().hex, + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "default_on": True, + "guardrailIdentifier": guardrail_id, + "guardrailVersion": "DRAFT", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIASYNTHETICTHROTTLE", + "aws_secret_access_key": "synthetic-secret", + "aws_bedrock_runtime_endpoint": endpoint, + }, + } + ] + path: Final = tmp_path / "bedrock-throttle.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def test_throttled_bedrock_guardrail_scan_is_retried_then_request_reaches_provider( + gateway: Gateway, tmp_path: Path +) -> None: + guardrail_id: Final = "throttle" + uuid.uuid4().hex[:8] + prompt: Final = f"retry marker {uuid.uuid4().hex}" + replies: Final = [_THROTTLED, _PASSED] + + def guardrail(request: Request) -> Reply: + assert request.target == f"/guardrail/{guardrail_id}/version/DRAFT/apply", request.target + assert json.loads(request.body)["source"] == "INPUT", request.body + return replies.pop(0) + + with wire_server(guardrail) as policy, wire_server(lambda _request: _PROVIDER) as upstream: + with ( + owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, guardrail_id, policy.url)) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key="synthetic") + response: Final = candidate.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": prompt}]} + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "retried" + scans: Final = policy.drain() + assert len(scans) == 2, scans + assert scans[0].body == scans[1].body, "the retry must resend the identical scan" + assert prompt in scans[1].body.decode() + assert len(upstream.drain()) == 1 + + +def test_bedrock_guardrail_throttle_that_never_clears_stops_after_the_retry_cap_without_calling_the_provider( + gateway: Gateway, tmp_path: Path +) -> None: + guardrail_id: Final = "throttle" + uuid.uuid4().hex[:8] + + with wire_server(lambda _request: _THROTTLED) as policy, wire_server(lambda _request: _PROVIDER) as upstream: + with ( + owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, guardrail_id, policy.url)) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key="synthetic") + response: Final = candidate.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "capped"}]} + ) + assert response.status_code == 429, response.text + assert len(policy.drain()) == 4 + assert upstream.drain() == () diff --git a/tests/integration/routing/test_complexity_router_llm_classifier.py b/tests/integration/routing/test_complexity_router_llm_classifier.py new file mode 100644 index 000000000000..fccc3a3a6d05 --- /dev/null +++ b/tests/integration/routing/test_complexity_router_llm_classifier.py @@ -0,0 +1,93 @@ +import json +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + + +def _completion(content: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex[:8], + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, + } + ).encode() + ) + + +def _tier_model(answer: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + assert request.target == "/chat/completions", request.target + return _completion(answer) + + return respond + + +def test_llm_classifier_verdict_routes_the_request_to_the_classified_tier_model( + gateway: Gateway, tmp_path: Path +) -> None: + prompt: Final = f"hi there {uuid.uuid4().hex}" + + def classifier(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + assert request.target == "/chat/completions", request.target + body: Final = json.loads(request.body) + assert "response_format" in body, body + assert [message["role"] for message in body["messages"]] == ["system", "user"], body + assert prompt in json.dumps(body["messages"][1]), body + return _completion(json.dumps({"tier": "COMPLEX"})) + + with ( + wire_server(classifier) as judge, + wire_server(_tier_model("simple answer")) as simple, + wire_server(_tier_model("complex answer")) as complex_tier, + ): + router: Final = "router-" + uuid.uuid4().hex[:8] + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + { + "model_name": name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": url, "api_key": "synthetic"}, + } + for name, url in (("judge", judge.url), ("simple", simple.url), ("complex", complex_tier.url)) + ] + [ + { + "model_name": router, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "judge", "timeout_ms": 20000}, + "tiers": {"SIMPLE": "simple", "MEDIUM": "simple", "COMPLEX": "complex", "REASONING": "complex"}, + }, + }, + } + ] + path: Final = tmp_path / "complexity_router.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {}, config=path) as candidate: + response: Final = candidate.request( + "POST", "/v1/chat/completions", {"model": router, "messages": [{"role": "user", "content": prompt}]} + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "complex answer", response.text + assert len([call for call in judge.drain() if call.method == "POST"]) == 1 + assert [call for call in simple.drain() if call.method == "POST"] == [] + forwarded: Final = [call for call in complex_tier.drain() if call.method == "POST"] + assert len(forwarded) == 1 + assert prompt in forwarded[0].body.decode() From 36a21709723040f7a5f4d063fd2b3bad6655e40d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 08:42:25 +0000 Subject: [PATCH 3/4] refactor(types): keep only live-provable Any removals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 16 +-- .../guardrail_hooks/bedrock_guardrails.py | 2 +- .../guardrail_hooks/xecguard/xecguard.py | 16 +-- .../test_jwt_routing_override.py | 83 --------------- .../test_bedrock_guardrail_throttle_retry.py | 100 ------------------ 5 files changed, 12 insertions(+), 205 deletions(-) delete mode 100644 tests/integration/authorization/test_jwt_routing_override.py delete mode 100644 tests/integration/observability/test_bedrock_guardrail_throttle_retry.py diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 24ac37e01313..ed3ec7b4dde2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -13,7 +13,7 @@ import secrets from collections.abc import Mapping from datetime import datetime, timezone -from typing import Final, NamedTuple, Protocol, Union, cast +from typing import Any, Final, NamedTuple, Protocol, Union, cast import fastapi import orjson @@ -207,7 +207,7 @@ def _get_model_from_request_context( request_data: dict, route: str, request: Request | None, - llm_router: litellm.Router | None = None, + llm_router: Any | None = None, team_id: str | None = None, ) -> str | list[str] | None: return get_model_from_request( @@ -512,8 +512,8 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str: def _routing_selector_matches_claim( - selector_value: str | list[str] | None, - claim_value: object, + selector_value: Any | None, + claim_value: Any | None, *, split_space_delimited: bool = False, ) -> bool: @@ -653,7 +653,7 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str # ``websocket.url``, which Starlette reconstructs from the (poisonable) # Host header. Carry the ASGI scope's path / root_path so the lookup # never reaches the fallback. - synthetic_scope: Final[dict[str, object]] = { + synthetic_scope: Final[dict[str, Any]] = { "type": "http", "headers": scope_headers, "path": ws_scope.get("path", ""), @@ -3099,7 +3099,7 @@ async def _reserve_budget_after_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, request_data: dict, route: str, - llm_router: litellm.Router | None, + llm_router: Any | None, team_object: LiteLLM_TeamTableCachedObj | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, @@ -3142,7 +3142,7 @@ def _should_skip_budget_checks( request_data: dict, route: str, request: Request | None, - llm_router: litellm.Router | None, + llm_router: Any | None, team_id: str | None = None, ) -> bool: model: Final = _get_model_from_request_context( @@ -3598,7 +3598,7 @@ async def _enforce_key_and_fallback_model_access( route: str, request: Request | None, llm_model_list: list | None, - llm_router: litellm.Router | None, + llm_router: Any | None, ) -> None: """ Key-level model allowlist and client fallbacks (same as standard auth). diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index bdcd9a51938e..228b31604a3b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1181,7 +1181,7 @@ async def _post_apply_guardrail_content_with_retry( or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES ): raise - await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2.0**attempt)) + await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt)) raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") async def _post_apply_guardrail_content( diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index d5a724f7458b..ddf9cace8b91 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -28,7 +28,6 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional from fastapi.exceptions import HTTPException -from typing_extensions import ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -78,12 +77,6 @@ def _sanitize_scan_result_for_logging(scan_result: dict) -> dict: ] -class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): - """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" - - supported_event_hooks: ReadOnly[list[GuardrailEventHooks] | None] - - class XecGuardMissingCredentials(Exception): pass @@ -97,7 +90,7 @@ def __init__( policy_names: list[str] | None = None, block_on_error: bool | None = None, grounding_strictness: str | None = None, - **kwargs: Unpack[_CustomGuardrailOptions], + **kwargs: Any, ) -> None: self.api_key = api_key or os.environ.get("XECGUARD_API_KEY") if not self.api_key: @@ -129,12 +122,9 @@ def __init__( llm_provider=httpxSpecialProvider.GuardrailCallback, ) - options: Final[_CustomGuardrailOptions] = { - "supported_event_hooks": list(self.get_supported_event_hooks()), - **kwargs, - } + kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - super().__init__(**options) + super().__init__(**kwargs) @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: diff --git a/tests/integration/authorization/test_jwt_routing_override.py b/tests/integration/authorization/test_jwt_routing_override.py deleted file mode 100644 index 21b659a65d19..000000000000 --- a/tests/integration/authorization/test_jwt_routing_override.py +++ /dev/null @@ -1,83 +0,0 @@ -import json -import time -import uuid -from pathlib import Path -from typing import Final - -import jwt -import yaml -from cryptography.hazmat.primitives.asymmetric import rsa - -from tests.integration._support.client import Gateway -from tests.integration._support.process import owned_proxy -from tests.integration._support.wire import Reply, Request, wire_server - -KEY_ID: Final = "integration-routing-override-key" -ISSUER: Final = "https://idp.example.com" - - -def test_jwt_routing_override_sends_only_claim_matching_tokens_to_oauth2_introspection( - gateway: Gateway, tmp_path: Path -) -> None: - private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) - public_jwk: Final = jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()) - subject: Final = f"oauth-machine-{uuid.uuid4().hex}" - - def jwks(request: Request) -> Reply: - assert request.method == "GET", request - return Reply(body=json.dumps({"keys": [{**json.loads(public_jwk), "kid": KEY_ID}]}).encode()) - - def token_info(request: Request) -> Reply: - assert request.method == "GET", request - assert request.headers["authorization"].startswith("Bearer "), request.headers - return Reply(body=json.dumps({"sub": subject, "role": "proxy_admin"}).encode()) - - def token(**claims: object) -> str: - now: Final = int(time.time()) - return jwt.encode( - {"sub": subject, "iat": now, "exp": now + 300, **claims}, - private_key, - algorithm="RS256", - headers={"kid": KEY_ID}, - ) - - with wire_server(jwks) as keys, wire_server(token_info) as introspection, gateway.scenario() as scenario: - config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - config["general_settings"] = { - **config["general_settings"], - "enable_jwt_auth": True, - "litellm_jwtauth": { - "user_id_jwt_field": "sub", - "routing_overrides": [ - { - "iss": ["https://other.example.com", "https://idp.*"], - "scope": "litellm.gateway", - "path": "oauth2", - } - ], - }, - } - path: Final = tmp_path / "jwt_routing_override.yaml" - path.write_text(yaml.safe_dump(config)) - overrides: Final = {"JWT_PUBLIC_KEY_URL": keys.url, "OAUTH_TOKEN_INFO_ENDPOINT": introspection.url} - with owned_proxy(gateway, tmp_path, overrides, config=path) as candidate: - model: Final = scenario.model() - - def chat(bearer: str) -> int: - return candidate.request( - "POST", - "/v1/chat/completions", - {"model": model, "messages": [{"role": "user", "content": "routing override"}]}, - key=bearer, - ).status_code - - assert chat(token(iss=ISSUER, scope="openid litellm.gateway")) == 200 - assert len(introspection.drain()) == 1, "space-delimited scope must match the override" - assert chat(token(iss=ISSUER, scope=["litellm.gateway", "openid"])) == 200 - assert len(introspection.drain()) == 1, "list scope must match the override" - - chat(token(iss=ISSUER, scope="openid")) - chat(token(iss=f"{ISSUER} evil.example.com", scope="litellm.gateway")) - chat(token(iss=ISSUER)) - assert introspection.drain() == (), "non-matching claims must stay on the JWT path" - assert keys.drain(), "non-matching tokens must be validated against the JWKS" diff --git a/tests/integration/observability/test_bedrock_guardrail_throttle_retry.py b/tests/integration/observability/test_bedrock_guardrail_throttle_retry.py deleted file mode 100644 index 3381801b4900..000000000000 --- a/tests/integration/observability/test_bedrock_guardrail_throttle_retry.py +++ /dev/null @@ -1,100 +0,0 @@ -import json -import uuid -from pathlib import Path -from typing import Final - -import yaml -from integration._support.client import Gateway -from integration._support.process import owned_proxy -from integration._support.wire import Reply, Request, wire_server - -_THROTTLED: Final = Reply( - status=429, - body=json.dumps({"message": "Rate exceeded"}).encode(), - headers={"x-amzn-ErrorType": "ThrottlingException"}, -) -_PASSED: Final = Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) -_PROVIDER: Final = Reply( - body=json.dumps( - { - "id": "chatcmpl-bedrock-retry", - "object": "chat.completion", - "created": 1700000000, - "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "retried"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, - } - ).encode() -) - - -def _config(tmp_path: Path, guardrail_id: str, endpoint: str) -> Path: - config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - config["guardrails"] = [ - { - "guardrail_name": "bedrock" + uuid.uuid4().hex, - "litellm_params": { - "guardrail": "bedrock", - "mode": "pre_call", - "default_on": True, - "guardrailIdentifier": guardrail_id, - "guardrailVersion": "DRAFT", - "aws_region_name": "us-east-1", - "aws_access_key_id": "AKIASYNTHETICTHROTTLE", - "aws_secret_access_key": "synthetic-secret", - "aws_bedrock_runtime_endpoint": endpoint, - }, - } - ] - path: Final = tmp_path / "bedrock-throttle.yaml" - path.write_text(yaml.safe_dump(config)) - return path - - -def test_throttled_bedrock_guardrail_scan_is_retried_then_request_reaches_provider( - gateway: Gateway, tmp_path: Path -) -> None: - guardrail_id: Final = "throttle" + uuid.uuid4().hex[:8] - prompt: Final = f"retry marker {uuid.uuid4().hex}" - replies: Final = [_THROTTLED, _PASSED] - - def guardrail(request: Request) -> Reply: - assert request.target == f"/guardrail/{guardrail_id}/version/DRAFT/apply", request.target - assert json.loads(request.body)["source"] == "INPUT", request.body - return replies.pop(0) - - with wire_server(guardrail) as policy, wire_server(lambda _request: _PROVIDER) as upstream: - with ( - owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, guardrail_id, policy.url)) as candidate, - candidate.scenario() as scenario, - ): - model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key="synthetic") - response: Final = candidate.request( - "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": prompt}]} - ) - assert response.status_code == 200, response.text - assert response.json()["choices"][0]["message"]["content"] == "retried" - scans: Final = policy.drain() - assert len(scans) == 2, scans - assert scans[0].body == scans[1].body, "the retry must resend the identical scan" - assert prompt in scans[1].body.decode() - assert len(upstream.drain()) == 1 - - -def test_bedrock_guardrail_throttle_that_never_clears_stops_after_the_retry_cap_without_calling_the_provider( - gateway: Gateway, tmp_path: Path -) -> None: - guardrail_id: Final = "throttle" + uuid.uuid4().hex[:8] - - with wire_server(lambda _request: _THROTTLED) as policy, wire_server(lambda _request: _PROVIDER) as upstream: - with ( - owned_proxy(gateway, tmp_path, {}, config=_config(tmp_path, guardrail_id, policy.url)) as candidate, - candidate.scenario() as scenario, - ): - model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url, api_key="synthetic") - response: Final = candidate.request( - "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "capped"}]} - ) - assert response.status_code == 429, response.text - assert len(policy.drain()) == 4 - assert upstream.drain() == () From b123b8a6d28ea7a04d73a33c3b707cc54c2a1370 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 09:18:56 +0000 Subject: [PATCH 4/4] refactor(types): type cache_hit as bool | None on the sync success path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- litellm/rust_bridge/callbacks_legacy_python.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8d3c2c9c4cf1..ab6a244eb1d1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3965,7 +3965,7 @@ def handle_sync_success_callbacks_for_async_calls( result: object, start_time: datetime.datetime, end_time: datetime.datetime, - cache_hit: object | None = None, + cache_hit: bool | None = None, ) -> None: """ Handles calling success callbacks for Async calls. diff --git a/litellm/rust_bridge/callbacks_legacy_python.py b/litellm/rust_bridge/callbacks_legacy_python.py index 25513666c43c..ce8cce2ebf9a 100644 --- a/litellm/rust_bridge/callbacks_legacy_python.py +++ b/litellm/rust_bridge/callbacks_legacy_python.py @@ -124,7 +124,7 @@ def post_call( ) -> object: ... def handle_sync_success_callbacks_for_async_calls( - self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: object = None + self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: bool | None = None ) -> None: ... def failure_handler(