From f94bd6d903e1163803323b721c8677ffa8365057 Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Wed, 2 Sep 2026 09:11:36 +0000 Subject: [PATCH 01/38] refactor(typing): replace Any with proven types in 65 backend files Typing-only pass over backend modules that carried the most reportAny and reportExplicitAny errors. Every new annotation is backed by a construction site, a call site, or an isinstance narrowing that already existed; untyped JSON boundaries were left alone rather than declared without validation. Tree-wide basedpyright errors drop 138,481 to 138,007. reportAny drops 8,854 to 8,645 and reportExplicitAny drops 3,119 to 2,814. --- .../bedrock_agentcore/transformation.py | 12 ++++-- litellm/caching/redis_semantic_cache.py | 14 +++---- .../compression/scoring/embedding_scorer.py | 3 +- litellm/experimental_mcp_client/client.py | 24 +++++++++--- litellm/files/main.py | 8 ++-- litellm/integrations/newrelic/newrelic.py | 24 ++++++------ litellm/integrations/opentelemetry.py | 10 ++--- litellm/integrations/prometheus.py | 6 +-- .../websearch_interception/tools.py | 11 +++--- .../websearch_interception/transformation.py | 8 ++-- litellm/interactions/agents/http_handler.py | 16 ++++---- litellm/interactions/agents/main.py | 18 ++++----- .../transformation.py | 14 +++---- litellm/interactions/main.py | 8 ++-- litellm/litellm_core_utils/core_helpers.py | 10 ++--- .../llm_response_utils/response_metadata.py | 2 +- .../prompt_templates/factory.py | 6 +-- .../adapters/streaming_iterator.py | 8 ++-- .../messages/transformation.py | 6 +-- .../azure_ai/vector_stores/transformation.py | 7 ++-- .../guardrail_translation/base_translation.py | 10 ++--- litellm/llms/bedrock/common_utils.py | 19 +++++---- ...n_nova_canvas_image_edit_transformation.py | 16 ++++---- .../bedrock/vector_stores/transformation.py | 4 +- .../image_edit/transformation.py | 4 +- litellm/llms/gemini/agents/transformation.py | 27 ++++++------- .../milvus/vector_stores/transformation.py | 7 ++-- .../minimax/text_to_speech/transformation.py | 9 +++-- .../responses/count_tokens/transformation.py | 16 ++++---- .../guardrail_translation/handler.py | 8 ++-- litellm/llms/openai/videos/transformation.py | 4 +- .../openrouter/image_edit/transformation.py | 4 +- .../guardrail_translation/handler.py | 4 +- .../perplexity/embedding/transformation.py | 6 +-- .../ragflow/vector_stores/transformation.py | 5 ++- litellm/llms/vertex_ai/fine_tuning/handler.py | 4 +- .../mcp_server/discoverable_endpoints.py | 12 +++--- .../mcp_server/elicitation_handler.py | 32 +++++++++------ .../proxy/_experimental/mcp_server/server.py | 2 +- litellm/proxy/a2a/version_convert.py | 13 ++++--- litellm/proxy/client/cli/commands/models.py | 4 +- .../proxy/common_utils/cache_coordinator.py | 26 ++++++------- .../proxy/common_utils/http_parsing_utils.py | 18 +++++---- .../container_endpoints/handler_factory.py | 6 +-- litellm/proxy/db/prisma_client.py | 4 +- litellm/proxy/guardrails/_content_utils.py | 16 ++++---- .../guardrail_hooks/qualifire/qualifire.py | 16 ++++---- .../guardrail_hooks/singulr/singulr.py | 4 +- .../unified_guardrail/unified_guardrail.py | 2 +- .../team_callback_endpoints.py | 8 ++-- .../proxy/openai_evals_endpoints/endpoints.py | 22 +++++------ litellm/proxy/policy_engine/init_policies.py | 5 ++- .../management_endpoints.py | 4 +- litellm/rag/ingestion/gemini_ingestion.py | 8 ++-- litellm/realtime_api/main.py | 9 +++-- litellm/repositories/config_repository.py | 17 +++++--- .../router_strategy/adaptive_router/hooks.py | 24 +++++++----- .../quality_router/quality_router.py | 11 +++--- .../router_utils/fallback_event_handlers.py | 14 +++---- .../io_token_rate_limit_check.py | 12 +++--- litellm/router_utils/search_api_router.py | 17 ++++++-- .../secret_managers/aws_secret_manager_v2.py | 4 +- litellm/skills/main.py | 26 ++++++------- litellm/types/vector_stores.py | 39 ++++++++++--------- litellm/vector_store_files/main.py | 14 +++---- 65 files changed, 411 insertions(+), 340 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 32252711997..4c5abf596cb 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -7,7 +7,7 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT). import json from collections.abc import AsyncIterator, Mapping -from typing import Any, Final +from typing import Any, Final, Protocol from litellm._logging import verbose_logger from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig @@ -35,6 +35,12 @@ _RESERVED_PREFIX_HEADERS: Final[tuple[str, ...]] = ( ) +class _SSELineSource(Protocol): + """Minimal streaming-response surface used to read SSE lines.""" + + def aiter_lines(self) -> AsyncIterator[str]: ... + + def _filter_reserved_headers( agent_extra_headers: Mapping[str, str] | None, ) -> dict[str, str] | None: @@ -77,7 +83,7 @@ class BedrockAgentCoreA2ATransformation: @staticmethod def get_url_and_signed_request( request_id: str, - params: dict[str, Any], + params: Mapping[str, object], litellm_params: dict[str, Any], method: str = "message/send", stream: bool = False, @@ -170,7 +176,7 @@ class BedrockAgentCoreA2ATransformation: return url, signed_headers, signed_body @staticmethod - async def parse_sse_events(response: Any) -> AsyncIterator[dict[str, Any]]: + async def parse_sse_events(response: _SSELineSource) -> AsyncIterator[dict[str, Any]]: """ Parse SSE events from an httpx streaming response. diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index f5264e28124..9a70bfc1418 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -116,7 +116,7 @@ class RedisSemanticCache(BaseCache): password = password or os.environ["REDIS_PASSWORD"] except KeyError as e: # Raise a more informative exception if any of the required keys are missing - missing_var: Final = e.args[0] + missing_var: Final[object] = e.args[0] raise ValueError( f"Missing required Redis configuration: {missing_var}. Provide {missing_var} or redis_url." ) from e @@ -273,7 +273,7 @@ class RedisSemanticCache(BaseCache): return prompt or None @classmethod - def _collect_responses_input_text(cls, value: Any, prompt_parts: list[str]) -> None: + def _collect_responses_input_text(cls, value: object, prompt_parts: list[str]) -> None: value = cls._coerce_response_input_value(value) if value is None: return @@ -334,7 +334,7 @@ class RedisSemanticCache(BaseCache): resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router), ) - def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]: + def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> list[float]: """ Routes through the proxy Router when the embedding model is a Router deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies, @@ -425,7 +425,7 @@ class RedisSemanticCache(BaseCache): prompt_embedding: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata")) - store_kwargs: Final[dict[str, Any]] = { + store_kwargs: Final[dict[str, object]] = { "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -504,7 +504,7 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Error retrieving from Redis semantic cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]: + async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> list[float]: """ Asynchronously generate an embedding for the given prompt. @@ -571,7 +571,7 @@ class RedisSemanticCache(BaseCache): # Generate embedding for the value (response) to cache prompt_embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) - store_kwargs: Final[dict[str, Any]] = { + store_kwargs: Final[dict[str, object]] = { "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -665,7 +665,7 @@ class RedisSemanticCache(BaseCache): aindex: Final = await self.llmcache._get_async_index() return await aindex.info() - async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: object) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None: """ Asynchronously store multiple values in the semantic cache. diff --git a/litellm/compression/scoring/embedding_scorer.py b/litellm/compression/scoring/embedding_scorer.py index aab1371e097..7e645ba3f9c 100644 --- a/litellm/compression/scoring/embedding_scorer.py +++ b/litellm/compression/scoring/embedding_scorer.py @@ -5,6 +5,7 @@ Computes cosine similarity between the query embedding and each message embeddin """ import math +from collections.abc import Mapping from typing import Any, Final from litellm.caching.dual_cache import DualCache @@ -49,7 +50,7 @@ def embedding_score_messages( messages: list[dict], model: str, cache: DualCache | None = None, - embedding_model_params: dict[str, Any] | None = None, + embedding_model_params: Mapping[str, object] | None = None, ) -> list[float]: """ Score each message's semantic similarity to the query using embeddings. diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ea81e323da4..34af6fcffba 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -5,18 +5,28 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 import os -from collections.abc import Awaitable, Callable, Generator +from collections.abc import Awaitable, Callable, Generator, Sequence +from contextlib import AbstractAsyncContextManager from datetime import timedelta from functools import partial from importlib import metadata -from typing import Any, Final, TypeVar +from typing import Any, Final, Protocol, TypeAlias, TypeVar import httpx from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client -streamable_http_client: Any | None = None +_TransportContext: TypeAlias = AbstractAsyncContextManager[Sequence[Any]] + + +class _StreamableHttpClientFactory(Protocol): + """The ``streamable_http_client`` entry point this module calls on the installed MCP SDK.""" + + def __call__(self, *, url: str, http_client: httpx.AsyncClient | None) -> _TransportContext: ... + + +streamable_http_client: _StreamableHttpClientFactory | None = None try: import mcp.client.streamable_http as streamable_http_module @@ -217,10 +227,12 @@ class MCPSigV4Auth(httpx.Auth): aws_region_name: str, ): """Call STS AssumeRole and return temporary credentials.""" + import time + import boto3 from botocore.credentials import Credentials - session_name: Final = aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" + session_name: Final = aws_session_name or f"litellm-mcp-{int(time.time())}" sts_kwargs: Final[dict] = {"region_name": aws_region_name} if aws_access_key_id and aws_secret_access_key: sts_kwargs["aws_access_key_id"] = aws_access_key_id @@ -316,7 +328,7 @@ class MCPClient: def _create_transport_context( self, - ) -> tuple[Any, httpx.AsyncClient | None]: + ) -> tuple[_TransportContext, httpx.AsyncClient | None]: """ Create the appropriate transport context based on transport type. Returns: @@ -409,7 +421,7 @@ class MCPClient: async def _execute_session_operation( self, - transport_ctx: Any, + transport_ctx: _TransportContext, operation: Callable[[ClientSession], Awaitable[TSessionResult]], ) -> TSessionResult: """ diff --git a/litellm/files/main.py b/litellm/files/main.py index 294c62f3d80..e769a0a0508 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -431,7 +431,7 @@ async def afile_delete( extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, **kwargs, -) -> Coroutine[Any, Any, FileObject]: +) -> Coroutine[object, object, FileObject]: """ Async: Delete file @@ -1003,7 +1003,7 @@ def file_content_streaming( logging_obj: LiteLLMLoggingObj | None, _is_async: bool, client: Any | None, -) -> FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult]: +) -> FileContentStreamingResult | Coroutine[object, object, FileContentStreamingResult]: if logging_obj is not None: logging_obj.model = model or "" logging_obj.model_call_details["model"] = model or "" @@ -1028,8 +1028,8 @@ def file_content_streaming( headers=response.headers, ) - response: FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult] = FileContentStreamingResult( - stream_iterator=iter(()), headers={} + response: FileContentStreamingResult | Coroutine[object, object, FileContentStreamingResult] = ( + FileContentStreamingResult(stream_iterator=iter(()), headers={}) ) if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: openai_creds: Final = get_openai_credentials( diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index f2f88ea55a8..9829ef4e18f 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -47,6 +47,8 @@ import os import threading import time import uuid +from collections.abc import Mapping, Sequence +from datetime import datetime from typing import Any, Final import litellm @@ -408,8 +410,8 @@ class NewRelicLogger(CustomLogger): def _get_duration( self, kwargs: dict, - start_time: Any, - end_time: Any, + start_time: datetime | float | None, + end_time: datetime | float | None, standard_logging_object: StandardLoggingPayload | None = None, ) -> float | None: """ @@ -438,7 +440,7 @@ class NewRelicLogger(CustomLogger): self, kwargs: dict, standard_logging_object: StandardLoggingPayload | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Extract request parameters like temperature and max_tokens, preferring StandardLoggingPayload.model_parameters. @@ -450,7 +452,7 @@ class NewRelicLogger(CustomLogger): else: source_params = kwargs.get("optional_params") or {} - params: Final = {} + params: Final[dict[str, object]] = {} temperature: Final = source_params.get("temperature") if temperature is not None: @@ -502,7 +504,7 @@ class NewRelicLogger(CustomLogger): response_model: str, vendor: str, standard_logging_object: StandardLoggingPayload | None = None, - ) -> list[dict[str, Any]]: + ) -> Sequence[Mapping[str, object]]: """ Extract all messages (request + response) with sequence numbers and timestamps. @@ -512,7 +514,7 @@ class NewRelicLogger(CustomLogger): Adds timestamps from StandardLoggingPayload (preferred) or kwargs if available (converted to epoch milliseconds). """ - messages: Final = [] + messages: Final[list[dict[str, object]]] = [] sequence = 0 # Extract timestamps, preferring StandardLoggingPayload @@ -544,7 +546,7 @@ class NewRelicLogger(CustomLogger): else: request_messages = kwargs.get("messages") or [] for msg in request_messages: - message_data = { + message_data: dict[str, object] = { "role": msg.get("role") or "user", "sequence": sequence, "response.model": response_model, @@ -599,11 +601,11 @@ class NewRelicLogger(CustomLogger): num_messages: int, usage: dict[str, int], duration: float | None = None, - request_params: dict[str, Any] | None = None, + request_params: Mapping[str, object] | None = None, ): """Record LlmChatCompletionSummary event to New Relic.""" try: - event_data: Final = { + event_data: Final[dict[str, object]] = { "id": request_id, "request_id": request_id, "request.model": request_model, @@ -647,7 +649,7 @@ class NewRelicLogger(CustomLogger): request_id: str, llm_response_id: str, trace_id: str | None, - messages: list[dict[str, Any]], + messages: Sequence[Mapping[str, object]], ): """Record LlmChatCompletionMessage events to New Relic. @@ -666,7 +668,7 @@ class NewRelicLogger(CustomLogger): for message in messages: sequence = message["sequence"] - event_data = { + event_data: dict[str, object] = { "id": f"{llm_response_id}-{sequence}", "request_id": request_id, "completion_id": request_id, diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index e8f3b305139..d4e7fcb577e 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1,7 +1,7 @@ import os import threading from collections import OrderedDict -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime @@ -166,7 +166,7 @@ class OTELMetricAttributeFilter: exclude_list: list[str] | None = None -def _build_metric_attribute_filter(value: Any) -> OTELMetricAttributeFilter: +def _build_metric_attribute_filter(value: object) -> OTELMetricAttributeFilter: if isinstance(value, OTELMetricAttributeFilter): return value if not isinstance(value, dict): @@ -205,7 +205,7 @@ def _resolve_metric_attribute_filter( ) -def _normalize_team_metadata_keys(value: Any) -> list[str]: +def _normalize_team_metadata_keys(value: str | Iterable[object] | None) -> list[str]: """Coerce a team-metadata allowlist from a list or comma-separated string. config.yaml passes a YAML list; an env var passes a comma-separated string. @@ -1569,7 +1569,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): self.safe_set_attribute(span=span, key=RESPONSE_SERVICE_TIER_ATTRIBUTE, value=served_tier) @staticmethod - def _team_metadata_json(value: Any, allowed_keys: list[str]) -> str | None: + def _team_metadata_json(value: object, allowed_keys: list[str]) -> str | None: """JSON-serialize only the allowlisted sub-keys of a team's metadata. Returns ``None`` when nothing is allowlisted or no allowlisted key is @@ -3524,7 +3524,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): kwargs={"standard_logging_object": {"error_information": error_information}}, ) - def set_preprocessing_duration_attribute(self, span: Span | None, container: Any) -> None: + def set_preprocessing_duration_attribute(self, span: Span | None, container: object) -> None: """ Set ``litellm.preprocessing.duration_ms`` (proxy-receive -> first provider handoff) on the proxy SERVER span. ``litellm_received_at`` diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 975a9bd8639..3e75c9cbf93 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2607,7 +2607,7 @@ class PrometheusLogger(CustomLogger): for all successful requests (both streaming and non-streaming). """ - def _safe_get(self, obj: Any, key: str, default: object = None) -> Any: + def _safe_get(self, obj: object, key: str, default: object = None) -> Any: """Get value from dict or Pydantic model.""" if obj is None: return default @@ -4215,8 +4215,8 @@ class PrometheusLogger(CustomLogger): def _safe_duration_seconds( self, - start_time: Any, - end_time: Any, + start_time: object, + end_time: object, ) -> float | None: """ Compute the duration in seconds between two objects. diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index b083a796a00..97c6c90d2ba 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -6,12 +6,13 @@ Native provider tools (like Anthropic's web_search_20250305) are converted to this format for consistent interception and execution. """ +from collections.abc import Mapping from typing import Any, Final from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME -def get_litellm_web_search_tool() -> dict[str, Any]: +def get_litellm_web_search_tool() -> dict[str, object]: """ Get the standard LiteLLM web search tool definition. @@ -49,7 +50,7 @@ def get_litellm_web_search_tool() -> dict[str, Any]: } -def get_litellm_web_search_tool_openai() -> dict[str, Any]: +def get_litellm_web_search_tool_openai() -> dict[str, object]: """ Get the standard LiteLLM web search tool definition in OpenAI format. @@ -82,7 +83,7 @@ def get_litellm_web_search_tool_openai() -> dict[str, Any]: } -def get_litellm_web_search_tool_responses() -> dict[str, Any]: +def get_litellm_web_search_tool_responses() -> dict[str, object]: """ Get the standard LiteLLM web search tool definition in Responses API format. @@ -114,7 +115,7 @@ def get_litellm_web_search_tool_responses() -> dict[str, Any]: } -def is_web_search_tool_responses(tool: dict[str, Any]) -> bool: +def is_web_search_tool_responses(tool: Mapping[str, object]) -> bool: """ Check if a tool is a web search tool for the Responses API. @@ -195,7 +196,7 @@ def is_web_search_tool_chat_completion(tool: dict[str, Any]) -> bool: return False -def is_anthropic_native_web_search_tool(tool: dict[str, Any]) -> bool: +def is_anthropic_native_web_search_tool(tool: Mapping[str, object]) -> bool: """ Check if a tool is an Anthropic-native ``web_search_*`` tool. diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 199ab020559..fe4b6583c55 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -24,7 +24,7 @@ class WebSearchTransformation: @staticmethod def transform_request( - response: Any, + response: object, stream: bool, response_format: str = "anthropic", ) -> tuple[bool, list[dict]]: @@ -66,7 +66,7 @@ class WebSearchTransformation: @staticmethod def _detect_from_responses_response( - response: Any, + response: object, ) -> tuple[bool, list[dict]]: """Parse a Responses API response for ``litellm_web_search`` function calls. @@ -399,7 +399,7 @@ class WebSearchTransformation: def build_web_search_tool_result_block( tool_use_id: str, search_response: SearchResponse | None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Build an Anthropic-native ``web_search_tool_result`` content block. @@ -433,7 +433,7 @@ class WebSearchTransformation: emitted with an empty result list (signals "search ran, no results" rather than "search did not run"). """ - items: Final[list[dict[str, Any]]] = [] + items: Final[list[dict[str, object]]] = [] if search_response is not None: results: Final = getattr(search_response, "results", None) or [] for r in results: diff --git a/litellm/interactions/agents/http_handler.py b/litellm/interactions/agents/http_handler.py index 14000ffaffd..ec9df0fb488 100644 --- a/litellm/interactions/agents/http_handler.py +++ b/litellm/interactions/agents/http_handler.py @@ -6,7 +6,7 @@ Extends InteractionsHTTPHandler so that the shared HTTP infrastructure duplicated. BaseAgentsAPIConfig stays as pure transform code. """ -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from typing import Any, Final import httpx @@ -39,11 +39,11 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, extra_headers: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> AgentCreateResponse | Coroutine[Any, Any, AgentCreateResponse]: + ) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: if _is_async: return self.async_create_agent( agents_api_config=agents_api_config, @@ -94,7 +94,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, extra_headers: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentCreateResponse: @@ -145,7 +145,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> AgentListResponse | Coroutine[Any, Any, AgentListResponse]: + ) -> AgentListResponse | Coroutine[object, object, AgentListResponse]: if _is_async: return self.async_list_agents( agents_api_config=agents_api_config, @@ -220,7 +220,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> AgentCreateResponse | Coroutine[Any, Any, AgentCreateResponse]: + ) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: if _is_async: return self.async_get_agent( agents_api_config=agents_api_config, @@ -299,7 +299,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> AgentDeleteResult | Coroutine[Any, Any, AgentDeleteResult]: + ) -> AgentDeleteResult | Coroutine[object, object, AgentDeleteResult]: if _is_async: return self.async_delete_agent( agents_api_config=agents_api_config, @@ -378,7 +378,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> AgentVersionsResponse | Coroutine[Any, Any, AgentVersionsResponse]: + ) -> AgentVersionsResponse | Coroutine[object, object, AgentVersionsResponse]: if _is_async: return self.async_list_agent_versions( agents_api_config=agents_api_config, diff --git a/litellm/interactions/agents/main.py b/litellm/interactions/agents/main.py index b63bea42f4f..1ca28adf0a4 100644 --- a/litellm/interactions/agents/main.py +++ b/litellm/interactions/agents/main.py @@ -30,7 +30,7 @@ Usage: import asyncio import contextvars -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import Any, Final @@ -75,7 +75,7 @@ def _make_logging_obj( model: str, custom_llm_provider: str, call_type: str, - optional_params: dict[str, Any], + optional_params: dict[str, object], ) -> LiteLLMLoggingObj: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) @@ -102,7 +102,7 @@ async def acreate( base_environment: InteractionEnvironment | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, ) -> AgentCreateResponse: @@ -146,10 +146,10 @@ def create( base_environment: InteractionEnvironment | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, -) -> AgentCreateResponse | Coroutine[Any, Any, AgentCreateResponse]: +) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: """ Sync: Create a managed agent on the provider side. @@ -244,7 +244,7 @@ def list( extra_headers: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, -) -> AgentListResponse | Coroutine[Any, Any, AgentListResponse]: +) -> AgentListResponse | Coroutine[object, object, AgentListResponse]: """Sync: List all agents on the provider side.""" local_vars: Final = locals() custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" @@ -320,7 +320,7 @@ def get( extra_headers: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, -) -> AgentCreateResponse | Coroutine[Any, Any, AgentCreateResponse]: +) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: """Sync: Get a specific agent by name.""" local_vars: Final = locals() custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" @@ -397,7 +397,7 @@ def delete( extra_headers: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, -) -> AgentDeleteResult | Coroutine[Any, Any, AgentDeleteResult]: +) -> AgentDeleteResult | Coroutine[object, object, AgentDeleteResult]: """Sync: Delete a specific agent by name.""" local_vars: Final = locals() custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" @@ -474,7 +474,7 @@ def list_versions( extra_headers: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, -) -> AgentVersionsResponse | Coroutine[Any, Any, AgentVersionsResponse]: +) -> AgentVersionsResponse | Coroutine[object, object, AgentVersionsResponse]: """Sync: List versions of a specific agent.""" local_vars: Final = locals() custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 9657b444969..39ccc26c38c 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -34,8 +34,8 @@ class LiteLLMResponsesInteractionsConfig: model: str, input: InteractionInput | None, optional_params: InteractionsAPIOptionalRequestParams, - **kwargs, - ) -> dict[str, Any]: + **kwargs: object, + ) -> dict[str, object]: """ Transform an Interactions API request to a Responses API request. @@ -45,7 +45,7 @@ class LiteLLMResponsesInteractionsConfig: - tools -> tools (similar format) - generation_config -> temperature, top_p, etc. """ - responses_request: Final[dict[str, Any]] = { + responses_request: Final[dict[str, object]] = { "model": model, } @@ -201,15 +201,15 @@ class LiteLLMResponsesInteractionsConfig: - Extract usage """ # Extract text from outputs and build both `outputs` (legacy) and `steps` (new schema). - outputs: Final[list[dict[str, Any]]] = [] - steps: Final[list[dict[str, Any]]] = [] + outputs: Final[list[dict[str, object]]] = [] + steps: Final[list[dict[str, object]]] = [] if hasattr(responses_response, "output") and responses_response.output: for output_item in responses_response.output: # Use getattr with None default to safely access content content = getattr(output_item, "content", None) if content is not None: content_items = content if isinstance(content, list) else [content] - model_output_contents: list[dict[str, Any]] = [] + model_output_contents: list[dict[str, object]] = [] for content_item in content_items: # Check if content_item has text attribute text = getattr(content_item, "text", None) @@ -264,7 +264,7 @@ class LiteLLMResponsesInteractionsConfig: # Add usage if available # Map Responses API usage (input_tokens, output_tokens) to Interactions API spec format # (total_input_tokens, total_output_tokens) - usage: Final = getattr(responses_response, "usage", None) + usage: Final[object] = getattr(responses_response, "usage", None) if usage: interactions_response_dict["usage"] = { "total_input_tokens": getattr(usage, "input_tokens", 0), diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index a2c3d510fae..8a33e9b39c5 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -229,7 +229,7 @@ def create( ) -> ( InteractionsAPIResponse | Iterator[InteractionsAPIStreamingResponse] - | Coroutine[Any, Any, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]] + | Coroutine[object, object, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]] ): """ Sync: Create a new interaction using Google's Interactions API. @@ -406,7 +406,7 @@ def get( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> InteractionsAPIResponse | Coroutine[Any, Any, InteractionsAPIResponse]: +) -> InteractionsAPIResponse | Coroutine[object, object, InteractionsAPIResponse]: """Sync: Get an interaction by its ID.""" local_vars: Final = locals() custom_llm_provider = custom_llm_provider or "gemini" @@ -510,7 +510,7 @@ def delete( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> DeleteInteractionResult | Coroutine[Any, Any, DeleteInteractionResult]: +) -> DeleteInteractionResult | Coroutine[object, object, DeleteInteractionResult]: """Sync: Delete an interaction by its ID.""" local_vars: Final = locals() custom_llm_provider = custom_llm_provider or "gemini" @@ -612,7 +612,7 @@ def cancel( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> CancelInteractionResult | Coroutine[Any, Any, CancelInteractionResult]: +) -> CancelInteractionResult | Coroutine[object, object, CancelInteractionResult]: """Sync: Cancel an interaction by its ID.""" local_vars: Final = locals() custom_llm_provider = custom_llm_provider or "gemini" diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 1738e30d865..2cdcfe4879c 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -419,7 +419,7 @@ def safe_deep_copy(data): if litellm.safe_memory_mode is True: return data - litellm_parent_otel_span: Any | None = None + litellm_parent_otel_span: object | None = None # Step 1: Remove the litellm_parent_otel_span litellm_parent_otel_span = None if isinstance(data, dict): @@ -510,7 +510,7 @@ def independent_snapshot( } -def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: +def filter_exceptions_from_params(data: object, max_depth: int = 20) -> Any: """ Recursively filter out Exception objects and callable objects from dicts/lists. @@ -542,7 +542,7 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: return None if isinstance(data, dict): - result: Final[dict[str, Any]] = {} + result: Final[dict[str, object]] = {} for k, v in data.items(): # Skip exception and callable values if isinstance(v, Exception) or (callable(v) and not isinstance(v, type)): @@ -556,7 +556,7 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: continue return result elif isinstance(data, list): - result_list: Final[list[Any]] = [] + result_list: Final[list[object]] = [] for item in data: # Skip exception and callable items if isinstance(item, Exception) or (callable(item) and not isinstance(item, type)): @@ -624,7 +624,7 @@ def redact_nested_match_and_regex_keys( # Iterative traversal; `seen` guards against cyclic refs preserved by deepcopy. try: seen: Final[set] = set() - stack: Final[list[Any]] = [redacted] + stack: Final[list[object]] = [redacted] while stack: node = stack.pop() node_id = id(node) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index b53a2d36753..c83c266a17e 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -168,7 +168,7 @@ class ResponseMetadata: def update_response_metadata( - result: Any, + result: object, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict, diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ba59e3fa997..56c1d605700 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1708,8 +1708,8 @@ def _find_server_tool_result( def convert_to_anthropic_tool_invoke( tool_calls: list[ChatCompletionAssistantToolCall], - web_search_results: list[Any] | None = None, - tool_results: list[Any] | None = None, + web_search_results: Sequence[object] | None = None, + tool_results: Sequence[object] | None = None, ) -> list[AnthropicMessagesToolUseParam | dict[str, Any]]: """ OpenAI tool invokes: @@ -5349,7 +5349,7 @@ class NormalizedToolCall(TypedDict): arguments: dict[str, object] -def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> dict[str, object]: +def _parse_tool_call_arguments(raw: object, tool_name: str | None, context: str) -> dict[str, object]: # Anthropic's tool_use blocks already carry a parsed dict in "input"; # chat completions and the Responses API carry a JSON string that may be # truncated by the model, so route those through the repair-aware parser. diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index cc5879df56d..ca993d40708 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -4,7 +4,7 @@ import copy import json import traceback from collections import deque -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import ( TYPE_CHECKING, Any, @@ -418,7 +418,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): augmented["usage"] = augmented_usage return augmented - def _next_compaction_event(self) -> dict[str, Any] | None: + def _next_compaction_event(self) -> dict[str, object] | None: """Return the next compaction content-block SSE event, or ``None``. Anthropic delivers compaction as a single delta (no token-by-token @@ -457,7 +457,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): "delta": {"type": "compaction_delta", "content": summary_content}, } - stop_event: Final = { + stop_event: Final[dict[str, object]] = { "type": "content_block_stop", "index": compaction_index, } @@ -989,7 +989,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): self.current_content_block_index += 1 @staticmethod - def _delta_has_content(processed_chunk: dict[str, Any]) -> bool: + def _delta_has_content(processed_chunk: Mapping[str, object]) -> bool: """Return True if a translated chunk carries a non-empty ``content_block_delta`` payload. diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 3d62b8b4784..b62e55f30f3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -87,7 +87,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): Processes both `system` and `messages` content blocks. """ - def _sanitize(cache_control: Any) -> None: + def _sanitize(cache_control: object) -> None: if isinstance(cache_control, dict): cache_control.pop("scope", None) @@ -152,7 +152,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return system_param @staticmethod - def _as_system_content_blocks(value: Any) -> list: + def _as_system_content_blocks(value: object) -> list: if value is None: return [] if isinstance(value, list): @@ -162,7 +162,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return [value] @staticmethod - def _is_system_role_message(message: Any) -> bool: + def _is_system_role_message(message: object) -> bool: return isinstance(message, dict) and message.get("role") == "system" _CONVERTED_SYSTEM_NOTE: Final = ( diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 5e61d0a1dd9..6db8c6a6c9f 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -114,8 +115,8 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: """ Transform search request for Azure AI Search API @@ -162,7 +163,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): url: Final = f"{api_base}/indexes/{index_name}/docs/search?api-version=2024-07-01" # Build the request body for Azure AI Search with vector search - request_body: Final = { + request_body: Final[dict[str, object]] = { "search": "*", # Get all documents (filtered by vector similarity) "vectorQueries": [ { diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 220fcedb0f8..0334b7f267c 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -39,7 +39,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( user_api_key_dict: Any | None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. @@ -62,7 +62,7 @@ class BaseTranslation(ABC): return {} # Transform keys to be prefixed with 'user_api_key_' - transformed: Final = {} + transformed: Final[dict[str, object]] = {} for key, value in user_dict.items(): # Skip None values and internal fields if value is None or key.startswith("_"): @@ -155,7 +155,7 @@ class BaseTranslation(ABC): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: Sequence[Any] | None = None, + responses_so_far: Sequence[object] | None = None, ) -> Sequence[bytes] | None: """ Build the streaming chunks that deliver a guardrail block message and @@ -178,8 +178,8 @@ class BaseTranslation(ABC): def build_stream_error_items( self, exc: "HTTPException", - responses_so_far: Sequence[Any] | None = None, - ) -> Sequence[Any] | None: + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[object] | None: """ Build the stream items that surface a guardrail HTTPException (a block with the default exception-on-block config, or a failed scan) after the diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 66ee5f10679..6e27cc7024f 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -52,8 +52,8 @@ _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( def merge_bedrock_aws_request_params( - litellm_params: Mapping[str, Any], - optional_params: Mapping[str, Any], + litellm_params: Mapping[str, object], + optional_params: Mapping[str, object], ) -> dict[str, Any]: """Merge deployment and request parameters without allowing auth escalation. @@ -303,7 +303,7 @@ def normalize_json_schema_custom_types_to_object(schema: dict) -> None: Uses an explicit stack (not recursion) to satisfy recursive-function guards in CI. """ - stack: Final[list[Any]] = [schema] + stack: Final[list[object]] = [schema] seen: Final[set[int]] = set() while stack: node = stack.pop() @@ -901,7 +901,7 @@ def _get_bedrock_converse_strict_tools_flag(base_model: str) -> bool | None: return None -def normalize_bedrock_opus_output_config_effort(model: str, output_config: Any) -> None: +def normalize_bedrock_opus_output_config_effort(model: str, output_config: object) -> None: """ Normalize Anthropic ``output_config.effort`` values for Bedrock Opus ids. @@ -1424,6 +1424,11 @@ class BedrockEventStreamDecoderBase: return chunk.decode() +def _decoded_json_value(raw: str) -> object: + """Decode a JSON document into an opaque value for isinstance narrowing.""" + return json.loads(raw) + + def get_anthropic_beta_from_headers(headers: dict) -> list[str]: """ Extract anthropic-beta header values and convert them to a list. @@ -1451,7 +1456,7 @@ def get_anthropic_beta_from_headers(headers: dict) -> list[str]: anthropic_beta_header = anthropic_beta_header.strip() if anthropic_beta_header.startswith("[") and anthropic_beta_header.endswith("]"): try: - parsed: Final = json.loads(anthropic_beta_header) + parsed: Final = _decoded_json_value(anthropic_beta_header) if isinstance(parsed, list): return [str(beta).strip() for beta in parsed] except json.JSONDecodeError: @@ -1464,8 +1469,8 @@ def get_anthropic_beta_from_headers(headers: dict) -> list[str]: def resolve_s3_encryption_key_id( - litellm_params: Mapping[str, Any], - optional_params: Mapping[str, Any] | None = None, + litellm_params: Mapping[str, object], + optional_params: Mapping[str, object] | None = None, ) -> str | None: """ Resolve the SSE-KMS key configured for Bedrock batch/file S3 objects. diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index ba76c7e628c..18d47301ee5 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -47,7 +47,7 @@ def _nova_canvas_task_body( task_type: str | None, mask_prompt: str | None, out_painting_mode: str | None, -) -> dict[str, Any]: +) -> dict[str, object]: """Build InvokeModel body task section (without imageGenerationConfig).""" if task_type == "BACKGROUND_REMOVAL": return { @@ -60,7 +60,7 @@ def _nova_canvas_task_body( "OUTPAINTING requires either a mask image or a mask prompt. " "Pass mask= or maskPrompt= in the request." ) - out_params: Final[dict[str, Any]] = { + out_params: Final[dict[str, object]] = { "image": image_b64, "text": text, } @@ -79,7 +79,7 @@ def _nova_canvas_task_body( # Honour explicit IMAGE_VARIATION even when a mask is present (mask is ignored # for this task type; callers use INPAINTING when they want mask semantics). if task_type == "IMAGE_VARIATION": - var_params_explicit: Final[dict[str, Any]] = { + var_params_explicit: Final[dict[str, object]] = { "images": [image_b64], "text": text, } @@ -100,7 +100,7 @@ def _nova_canvas_task_body( "or omit taskType for automatic routing (mask → INPAINTING, else IMAGE_VARIATION)." ) if mask_b64 is not None or mask_prompt is not None or task_type == "INPAINTING": - in_params: Final[dict[str, Any]] = {"image": image_b64, "text": text} + in_params: Final[dict[str, object]] = {"image": image_b64, "text": text} if mask_prompt is not None: in_params["maskPrompt"] = mask_prompt elif mask_b64 is not None: @@ -114,7 +114,7 @@ def _nova_canvas_task_body( "See https://docs.aws.amazon.com/nova/latest/userguide/image-gen-req-resp-structure.html" ) return {"taskType": "INPAINTING", "inPaintingParams": in_params} - var_params: Final[dict[str, Any]] = { + var_params: Final[dict[str, object]] = { "images": [image_b64], "text": text, } @@ -250,9 +250,9 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): image_edit_optional_params: ImageEditOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, object]: supported: Final = set(self.get_supported_openai_params(model)) - mapped: Final[dict[str, Any]] = dict(image_edit_optional_params) + mapped: Final[dict[str, object]] = dict(image_edit_optional_params) _size: Final = mapped.pop("size", None) if _size is not None and isinstance(_size, str) and "x" in _size: w, h = _size.split("x", 1) @@ -327,7 +327,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): cfg_scale: Final = op.pop("cfgScale", None) seed: Final = op.pop("seed", None) - image_generation_config: Final[dict[str, Any]] = {} + image_generation_config: Final[dict[str, object]] = {} nested_igc: Final = op.pop("imageGenerationConfig", None) if isinstance(nested_igc, dict): image_generation_config.update(nested_igc) diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 2d72db0cdba..6940077391f 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -203,7 +203,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = f"{api_base}/{encoded_vector_store_id}/retrieve" - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "retrievalQuery": BedrockKBRetrievalQuery(text=query), } @@ -288,7 +288,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): data_source_id: Final = metadata.get("x-amz-bedrock-kb-data-source-id", "unknown") if metadata else "unknown" return f"bedrock-kb-document-{data_source_id}" - def _get_attributes_from_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]: + def _get_attributes_from_metadata(self, metadata: dict[str, object]) -> dict[str, object]: """ Extract all attributes from Bedrock KB metadata. Returns a copy of the metadata dictionary. diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 013053e5bd5..62b631a7671 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -84,7 +84,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): BFL-specific params are passed through directly. """ - optional_params: Final[dict[str, Any]] = {} + optional_params: Final[dict[str, object]] = {} # Pass through BFL-specific params bfl_params: Final = [ @@ -246,7 +246,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): b64_image: Final = base64.b64encode(image_bytes).decode("utf-8") # Build request body - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "prompt": prompt, "input_image": b64_image, } diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index 2de78242c43..cfdb55f2048 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -9,6 +9,7 @@ Proxies the Gemini v1beta Agents API: GET /v1beta/agents/{name}/versions list versions """ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -87,7 +88,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): def get_complete_url( self, api_base: str | None, - litellm_params: dict[str, Any], + litellm_params: Mapping[str, object], ) -> str: return f"{self._base_url(api_base)}/agents" @@ -132,9 +133,9 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): def transform_create_request( self, name: str, - litellm_params: dict[str, Any], - ) -> dict[str, Any]: - body: Final[dict[str, Any]] = {"name": name} + litellm_params: Mapping[str, object], + ) -> dict[str, object]: + body: Final[dict[str, object]] = {"name": name} for key in _GEMINI_AGENT_BODY_KEYS: value = litellm_params.get(key) if value is not None: @@ -174,10 +175,10 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): def transform_list_request( self, api_base: str | None, - litellm_params: dict[str, Any], - ) -> tuple[str, dict[str, Any]]: + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: url: Final = f"{self._base_url(api_base)}/agents" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if litellm_params.get("page_size"): params["pageSize"] = litellm_params["page_size"] if litellm_params.get("page_token"): @@ -207,8 +208,8 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): self, name: str, api_base: str | None, - litellm_params: dict[str, Any], - ) -> tuple[str, dict[str, Any]]: + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: url: Final = f"{self._base_url(api_base)}/agents/{name}" return url, {} @@ -236,7 +237,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): self, name: str, api_base: str | None, - litellm_params: dict[str, Any], + litellm_params: Mapping[str, object], ) -> str: return f"{self._base_url(api_base)}/agents/{name}" @@ -262,10 +263,10 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): self, name: str, api_base: str | None, - litellm_params: dict[str, Any], - ) -> tuple[str, dict[str, Any]]: + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: url: Final = f"{self._base_url(api_base)}/agents/{name}/versions" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if litellm_params.get("page_size"): params["pageSize"] = litellm_params["page_size"] if litellm_params.get("page_token"): diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 34f0cd854c4..0e96b3577fd 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -122,8 +123,8 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: """ Transform search request for Azure AI Search API @@ -165,7 +166,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): url: Final = f"{api_base}/v2/vectordb/entities/search" # Build the request body for Azure AI Search with vector search - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "collectionName": index_name, "data": [query_vector], "annsField": "book_intro_vector", diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index 2263a98551e..f8926df1f3f 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -5,6 +5,7 @@ Maps OpenAI TTS spec to MiniMax TTS API (WebSocket-based HTTP API) Reference: https://platform.minimax.io/docs """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -86,8 +87,8 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): def _resolve_voice_id( self, - voice: str | dict[str, Any] | None, - params: dict[str, Any], + voice: str | Mapping[str, object] | None, + params: dict[str, object], ) -> str: """ Determine the MiniMax voice_id based on provided voice input or parameters. @@ -127,7 +128,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): """ Map OpenAI parameters to MiniMax TTS parameters """ - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Work on a copy so we don't mutate the caller's dictionary params: Final = dict(optional_params) if optional_params else {} @@ -242,7 +243,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): # Output format: 'url' or 'hex' (default is 'hex') output_format: Final = params.pop("output_format", "hex") - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "model": model, "text": input, "stream": False, # HTTP endpoint doesn't support streaming diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 88f04c59e01..9f596505f91 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -117,16 +117,16 @@ class OpenAICountTokensConfig: def transform_request_to_count_tokens( self, model: str, - input: str | list[Any], + input: str | Sequence[object], tools: list[dict[str, Any]] | None = None, instructions: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform request to OpenAI Responses API token counting format. The Responses API uses `input` (not `messages`) and `instructions` (not `system`). """ - request: Final[dict[str, Any]] = { + request: Final[dict[str, object]] = { "model": model, "input": input, } @@ -145,7 +145,7 @@ class OpenAICountTokensConfig: "Authorization": f"Bearer {api_key}", } - def validate_request(self, model: str, input: str | list[Any]) -> None: + def validate_request(self, model: str, input: str | Sequence[object]) -> None: if not model: raise ValueError("model parameter is required") @@ -155,18 +155,18 @@ class OpenAICountTokensConfig: @staticmethod def _transform_tools_for_responses_api( tools: list[dict[str, Any]], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform OpenAI chat tools format to Responses API tools format. Chat format: {"type": "function", "function": {"name": "...", "parameters": {...}}} Responses format: {"type": "function", "name": "...", "parameters": {...}} """ - transformed: Final = [] + transformed: Final[list[dict[str, object]]] = [] for tool in tools: if tool.get("type") == "function" and "function" in tool: func = tool["function"] - item: dict[str, Any] = { + item: dict[str, object] = { "type": "function", "name": func.get("name", ""), "description": func.get("description", ""), @@ -191,7 +191,7 @@ class OpenAICountTokensConfig: (input_items, instructions) tuple where instructions is extracted from system/developer messages. """ - input_items: Final[list[dict[str, Any]]] = [] + input_items: Final[list[dict[str, object]]] = [] instructions_parts: Final[list[str]] = [] for msg in messages: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1530c154e93..1db28193d10 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -110,7 +110,7 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] -def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: +def _next_stream_sequence_number(responses_so_far: Sequence[object] | None) -> int: sequence_numbers: Final = ( item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) for item in reversed(responses_so_far or ()) @@ -337,7 +337,7 @@ class OpenAIResponsesHandler(BaseTranslation): def _extract_input_text_and_images( self, - message: Any, + message: Mapping[str, object], msg_idx: int, texts_to_check: list[str], images_to_check: list[str], @@ -661,8 +661,8 @@ class OpenAIResponsesHandler(BaseTranslation): def build_stream_error_items( self, exc: "HTTPException", - responses_so_far: Sequence[Any] | None = None, - ) -> Sequence[Any] | None: + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[object] | None: from litellm.proxy.common_request_processing import ( serialize_http_exception_detail, ) diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index f1b6dcb330a..94dc30f41e5 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -571,8 +571,8 @@ class OpenAIVideoConfig(BaseVideoConfig): def _add_image_to_files( self, - files_list: list[tuple[str, Any]], - image: Any, + files_list: list[tuple[str, FileTypes]], + image: FileContent, field_name: str, ) -> None: """Add an image to the files list with appropriate content type""" diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index e3a2bf34854..b01c25aad0c 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -152,7 +152,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[dict, RequestFiles]: - content_parts: Final[list[dict[str, Any]]] = [] + content_parts: Final[list[dict[str, object]]] = [] # Add source image(s) as base64 data URLs if image is not None: @@ -174,7 +174,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): if prompt: content_parts.append({"type": "text", "text": prompt}) - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "model": model, "messages": [ { diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 6573ca827f0..f07acf2f728 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -127,7 +127,7 @@ class PassThroughEndpointHandler(BaseTranslation): async def process_output_response( self, - response: Any, + response: object, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Any | None = None, @@ -236,7 +236,7 @@ class LlmPassthroughRouteHandler(BaseTranslation): async def process_output_response( self, - response: Any, + response: object, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Any | None = None, diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index cb29a598d30..a911fa62719 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -13,7 +13,7 @@ This module decodes them into float arrays for OpenAI-compatible responses. import base64 import struct -from typing import Any, Final +from typing import Final import httpx @@ -117,7 +117,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): } @staticmethod - def _decode_base64_embedding(embedding_value: Any) -> list[float]: + def _decode_base64_embedding(embedding_value: object) -> object: """ Decode a Perplexity embedding into a list of floats. @@ -154,7 +154,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): model_response.object = raw_response_json.get("object", "list") raw_data: Final = raw_response_json.get("data", []) - decoded_data: Final[list[dict[str, Any]]] = [] + decoded_data: Final[list[dict[str, object]]] = [] for item in raw_data: decoded_item = dict(item) decoded_item["embedding"] = self._decode_base64_embedding(item.get("embedding")) diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index 282cb7a92a7..38a06a37f7e 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -91,7 +92,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """RAGFlow vector stores are management-only, search is not supported.""" raise NotImplementedError("RAGFlow vector stores support dataset management only, not search/retrieval") @@ -121,7 +122,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): raise ValueError("name is required for RAGFlow dataset creation") # Build request body - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "name": name, } diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index df9b1f8c66a..7ecc5e8ff3d 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -2,7 +2,7 @@ import json import traceback from collections.abc import Coroutine from datetime import datetime -from typing import Any, Final, Literal +from typing import Final, Literal import httpx @@ -207,7 +207,7 @@ class VertexFineTuningAPI(VertexLLM): timeout: float | httpx.Timeout, kwargs: dict | None = None, original_hyperparameters: dict | None = {}, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: verbose_logger.debug("creating fine tuning job, args= %s", create_fine_tuning_job_data) _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 94bca9460dd..bf5f95d3f38 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -430,7 +430,7 @@ def _clear_oauth_state_cookie(response: Response, request: Request, state: str) ) -def _get_validated_client_redirect_uri(request: Request, state_data: dict[str, Any]) -> str: +def _get_validated_client_redirect_uri(request: Request, state_data: Mapping[str, object]) -> str: """Return a trusted (same-origin, loopback, or ops-allowlisted) client redirect URI from OAuth state. """ @@ -469,7 +469,7 @@ def _resolve_oauth2_server_for_root_endpoints( return None -def _normalize_for_token_comparison(value: Any) -> str: +def _normalize_for_token_comparison(value: object) -> str: """Stringify ``value`` for token-rule comparison. Booleans are lower-cased so Python's ``True`` / ``False`` line up with @@ -481,8 +481,8 @@ def _normalize_for_token_comparison(value: Any) -> str: def _validate_token_response( - token_response: dict[str, Any], - validation_rules: dict[str, Any], + token_response: Mapping[str, object], + validation_rules: Mapping[str, object], server_id: str, ) -> None: """Raise HTTPException 403 if any validation rule doesn't match the token response. @@ -496,10 +496,10 @@ def _validate_token_response( responses of ``{"verified": true}``. """ for key, expected in validation_rules.items(): - actual: Any = token_response.get(key) + actual: object | None = token_response.get(key) # Try dot-notation traversal when top-level lookup returns None if actual is None and "." in key: - obj: Any = token_response + obj: object = token_response for part in key.split("."): if isinstance(obj, dict): obj = obj.get(part) diff --git a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py index ce7e963f55f..bbd1c9aaf1e 100644 --- a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py +++ b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py @@ -9,7 +9,7 @@ MCP Spec Reference: https://modelcontextprotocol.io/specification/2025-11-25/client/elicitation """ -from typing import TYPE_CHECKING, Any, Final, Union +from typing import TYPE_CHECKING, Final, Protocol, Union from litellm._logging import verbose_logger @@ -37,11 +37,21 @@ except ImportError: MCP_ELICITATION_AVAILABLE = False +class _DownstreamElicitSession(Protocol): + """The downstream MCP client session methods this module relays elicitation requests through.""" + + async def elicit_url(self, message: str, url: str, elicitation_id: str) -> "ElicitResult": ... + + async def elicit_form(self, message: str, requestedSchema: dict[str, object]) -> "ElicitResult": ... + + async def elicit(self, message: str, requestedSchema: dict[str, object]) -> "ElicitResult": ... + + async def handle_elicitation_request( - context: Any, + context: object, params: "ElicitRequestParams", - downstream_session: Any | None = None, - downstream_capabilities: Any | None = None, + downstream_session: _DownstreamElicitSession | None = None, + downstream_capabilities: object = None, ) -> Union["ElicitResult", "ErrorData"]: """ Handle an MCP elicitation/create request from an upstream MCP server. @@ -94,8 +104,8 @@ async def handle_elicitation_request( async def _relay_elicitation_to_downstream( params: "ElicitRequestParams", - downstream_session: Any, - downstream_capabilities: Any | None = None, + downstream_session: _DownstreamElicitSession, + downstream_capabilities: object = None, ) -> Union["ElicitResult", "ErrorData"]: """ Relay an elicitation request to the downstream MCP client. @@ -111,17 +121,17 @@ async def _relay_elicitation_to_downstream( mode: Final = getattr(params, "mode", "form") # Check if the downstream client supports the requested mode if downstream_capabilities is not None: - elicit_caps: Final = getattr(downstream_capabilities, "elicitation", None) + elicit_caps: Final[object] = getattr(downstream_capabilities, "elicitation", None) if elicit_caps is None: verbose_logger.info("MCP elicitation: downstream client does not support elicitation") return ElicitResult(action="decline") if mode == "url": - url_cap: Final = getattr(elicit_caps, "url", None) + url_cap: Final[object] = getattr(elicit_caps, "url", None) if url_cap is None: verbose_logger.info("MCP elicitation: downstream client does not support URL mode") return ElicitResult(action="decline") if mode == "form": - form_cap: Final = getattr(elicit_caps, "form", None) + form_cap: Final[object] = getattr(elicit_caps, "form", None) if form_cap is None: verbose_logger.info("MCP elicitation: downstream client does not support form mode") return ElicitResult(action="decline") @@ -135,14 +145,14 @@ async def _relay_elicitation_to_downstream( result = await downstream_session.elicit_url( message=params.message, url=params.url, - elicitation_id=getattr(params, "elicitationId", None), + elicitation_id=params.elicitationId, ) elif isinstance(params, ElicitRequestFormParams): # Form mode: relay structured form to client verbose_logger.info("MCP elicitation: relaying form mode to downstream") result = await downstream_session.elicit_form( message=params.message, - requestedSchema=getattr(params, "requestedSchema", None), + requestedSchema=params.requestedSchema, ) else: # Fallback for generic ElicitRequestParams — pass an empty schema diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 989b08b929a..57b60ff68a2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3477,7 +3477,7 @@ if MCP_AVAILABLE: is best-effort in that mode. """ - def _bytes_for_hash(value: Any) -> bytes | None: + def _bytes_for_hash(value: object) -> bytes | None: """Only hash str/bytes secrets; skip mocks and other unexpected types.""" if value is None: return None diff --git a/litellm/proxy/a2a/version_convert.py b/litellm/proxy/a2a/version_convert.py index 35587ee274c..bde2ff45a88 100644 --- a/litellm/proxy/a2a/version_convert.py +++ b/litellm/proxy/a2a/version_convert.py @@ -25,7 +25,6 @@ The two wire shapes: """ from collections.abc import Callable -from types import ModuleType from typing import Final, Literal from pydantic import BaseModel @@ -181,7 +180,7 @@ def _send_result_to(result: JsonDict, target: A2AVersion, request_id: RequestId) ) if target == "1.0": - compat_result: Final = _validate_message_or_task(result, types_v03) + compat_result: Final = _validate_message_or_task(result) response: Final = types_v03.SendMessageResponse( root=types_v03.SendMessageSuccessResponse( id=str(request_id) if request_id is not None else "", @@ -285,7 +284,7 @@ def _stream_result_to(result: JsonDict, target: A2AVersion, request_id: RequestI ) if target == "1.0": - event: Final = _validate_stream_event(result, types_v03) + event: Final = _validate_stream_event(result) wrapper: Final = types_v03.SendStreamingMessageSuccessResponse( id=str(request_id) if request_id is not None else "", result=event, # pyright: ignore[reportArgumentType] @@ -318,13 +317,17 @@ def _convert_agent_card(card: JsonDict, target: A2AVersion) -> JsonDict: return MessageToDict(core, preserving_proto_field_name=False) -def _validate_message_or_task(result: JsonDict, types_v03: ModuleType) -> BaseModel: +def _validate_message_or_task(result: JsonDict) -> BaseModel: + from a2a.compat.v0_3.conversions import types_v03 + if result.get("kind") == "task": return types_v03.Task.model_validate(result) return types_v03.Message.model_validate(result) -def _validate_stream_event(result: JsonDict, types_v03: ModuleType) -> BaseModel: +def _validate_stream_event(result: JsonDict) -> BaseModel: + from a2a.compat.v0_3.conversions import types_v03 + kind: Final = result.get("kind") if kind == "task": return types_v03.Task.model_validate(result) diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index cc165504113..4c83a7b799a 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -17,8 +17,8 @@ from ... import Client @dataclass class ModelYamlInfo: model_name: str - model_params: dict[str, Any] - model_info: dict[str, Any] + model_params: dict[str, object] + model_info: dict[str, object] model_id: str access_groups: list[str] provider: str diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py index e36307ae2df..f6ce96d8777 100644 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -13,14 +13,14 @@ pattern: global spend, feature flags, config, or other shared read-through data. import asyncio import time from collections.abc import Awaitable, Callable -from typing import Any, Final, Protocol, TypeVar +from typing import Final, Protocol, TypeVar from litellm._logging import verbose_proxy_logger T = TypeVar("T") -class AsyncCacheProtocol(Protocol): +class AsyncCacheProtocol(Protocol[T]): """Protocol for cache backends used by EventDrivenCacheCoordinator. Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params @@ -30,18 +30,18 @@ class AsyncCacheProtocol(Protocol): async def async_get_cache( self, key: str, - parent_otel_span: Any = None, + parent_otel_span: object = None, local_only: bool = False, - **kwargs: Any, - ) -> Any: ... + **kwargs: object, + ) -> T | None: ... async def async_set_cache( self, key: str, - value: Any, + value: T, local_only: bool = False, - **kwargs: Any, - ) -> Any: ... + **kwargs: object, + ) -> object: ... class EventDrivenCacheCoordinator: @@ -64,11 +64,11 @@ class EventDrivenCacheCoordinator: self._query_in_progress = False self._log_prefix = log_prefix - async def _get_cached(self, cache_key: str, cache: AsyncCacheProtocol) -> Any | None: + async def _get_cached(self, cache_key: str, cache: AsyncCacheProtocol[T]) -> T | None: """Return value from cache if present, else None.""" return await cache.async_get_cache(key=cache_key) - def _log_cache_hit(self, value: T) -> None: + def _log_cache_hit(self, value: object) -> None: if self._log_prefix: verbose_proxy_logger.debug("%s Cache hit, value: %s", self._log_prefix, value) @@ -98,7 +98,7 @@ class EventDrivenCacheCoordinator: self, event: asyncio.Event, cache_key: str, - cache: AsyncCacheProtocol, + cache: AsyncCacheProtocol[T], ) -> T | None: """Wait for loader to finish, then read from cache.""" await event.wait() @@ -118,7 +118,7 @@ class EventDrivenCacheCoordinator: async def _load_and_cache( self, cache_key: str, - cache: AsyncCacheProtocol, + cache: AsyncCacheProtocol[T], load_fn: Callable[[], Awaitable[T]], ) -> T | None: """Double-check cache, run load_fn, set cache, return value. Caller must call _signal_done in finally.""" @@ -163,7 +163,7 @@ class EventDrivenCacheCoordinator: async def get_or_load( self, cache_key: str, - cache: AsyncCacheProtocol, + cache: AsyncCacheProtocol[T], load_fn: Callable[[], Awaitable[T]], ) -> T | None: """ diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 96621b08ba1..2b730c450fb 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -1,6 +1,6 @@ import json import re -from collections.abc import Collection +from collections.abc import Collection, Mapping from typing import Any, Final import orjson @@ -186,7 +186,7 @@ def _safe_get_request_headers(request: Request | None) -> dict: if request is None: return {} state: Final = getattr(request, "state", None) - cached: Final = getattr(state, "_cached_headers", None) + cached: Final[object] = getattr(state, "_cached_headers", None) if isinstance(cached, dict): return cached if cached is not None: @@ -344,7 +344,9 @@ async def get_request_body(request: Request) -> dict[str, Any]: return {} -def extract_nested_form_metadata(form_data: dict[str, Any], prefix: str = "litellm_metadata[") -> dict[str, Any]: +def extract_nested_form_metadata( + form_data: Mapping[str, object], prefix: str = "litellm_metadata[" +) -> dict[str, object]: """ Extract nested metadata from form data with bracket notation. @@ -382,7 +384,7 @@ def extract_nested_form_metadata(form_data: dict[str, Any], prefix: str = "litel if not form_data: return {} - metadata: Final[dict[str, Any]] = {} + metadata: Final[dict[str, object]] = {} for key, value in form_data.items(): # Skip keys that don't start with the prefix @@ -430,7 +432,7 @@ def extract_nested_form_metadata(form_data: dict[str, Any], prefix: str = "litel return metadata -def get_tags_from_request_body(request_body: dict) -> list[str]: +def get_tags_from_request_body(request_body: Mapping[str, object]) -> list[str]: """ Extract tags from request body metadata. @@ -447,12 +449,12 @@ def get_tags_from_request_body(request_body: dict) -> list[str]: if isinstance(metadata, str): from litellm.litellm_core_utils.safe_json_loads import safe_json_loads - parsed: Final = safe_json_loads(metadata) + parsed: Final[object] = safe_json_loads(metadata) metadata = parsed if isinstance(parsed, dict) else {} elif not isinstance(metadata, dict): metadata = {} - tags_in_metadata: Final[Any] = metadata.get("tags", []) - tags_in_request_body: Final[Any] = request_body.get("tags", []) + tags_in_metadata: Final[object] = metadata.get("tags", []) + tags_in_request_body: Final[object] = request_body.get("tags", []) combined_tags: Final[list[str]] = [] ###################################### diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index aaee1d3e264..892ff9771cf 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -7,7 +7,7 @@ FastAPI route handlers for ALL container file endpoints. import json from pathlib import Path -from typing import Any, Final +from typing import Final from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse @@ -194,7 +194,7 @@ async def _process_binary_request( user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "file_id": file_id, **( await get_container_forwarding_params( @@ -374,7 +374,7 @@ async def _process_request( ) query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "query_params": query_params, **path_params, } diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 4bd007769b8..2190ae55fd2 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -503,7 +503,7 @@ class PrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, - http_client: Any | None = None, + http_client: object | None = None, *, expected_generation: int | None = None, ) -> bool: @@ -541,7 +541,7 @@ class PrismaWrapper: async def _recreate_prisma_client_locked( self, new_db_url: str, - http_client: Any | None = None, + http_client: object | None = None, *, expected_generation: int | None = None, ) -> bool: diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index ae92adcb1ee..c6e3f8ce34c 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -54,7 +54,7 @@ def _part_text(part: Mapping[str, object]) -> str | None: return None -def _iter_text_parts_in_content(content: Any) -> Iterator[str]: +def _iter_text_parts_in_content(content: object) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or multimodal list). Non-text parts (images, audio, …) are skipped.""" if isinstance(content, str): @@ -75,13 +75,13 @@ def _iter_text_parts_in_content(content: Any) -> Iterator[str]: yield text -def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: +def _coerce_input_to_messages(input_value: object) -> list[dict[str, object]]: """Coerce a Responses-API ``data["input"]`` value into chat-style messages.""" if isinstance(input_value, str): return [{"role": "user", "content": input_value}] if not isinstance(input_value, list): return [] - messages: Final[list[dict[str, Any]]] = [] + messages: Final[list[dict[str, object]]] = [] for item in input_value: if isinstance(item, str): messages.append({"role": "user", "content": item}) @@ -110,7 +110,7 @@ def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: return messages -def _iter_inspection_messages(data: dict[str, Any]) -> Iterator[dict[str, Any]]: +def _iter_inspection_messages(data: Mapping[str, object]) -> Iterator[object]: """Yield every message-like dict, walking ``messages`` AND ``input``.""" messages: Final = data.get("messages") if isinstance(messages, list): @@ -118,7 +118,7 @@ def _iter_inspection_messages(data: dict[str, Any]) -> Iterator[dict[str, Any]]: yield from _coerce_input_to_messages(data.get("input")) -def iter_message_text(data: dict[str, Any]) -> Iterator[str]: +def iter_message_text(data: Mapping[str, object]) -> Iterator[str]: """Yield every text fragment from ``messages`` AND ``input``. Walks every role (user, assistant, system, …) — guardrails inspect @@ -139,7 +139,7 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: """ visited = 0 - def _rewrite_content(content: Any) -> Any: + def _rewrite_content(content: object) -> object: nonlocal visited if isinstance(content, str): if content: @@ -147,7 +147,7 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: return visit(content) return content if isinstance(content, list): - new_parts: Final[list[Any]] = [] + new_parts: Final[list[object]] = [] for part in content: if isinstance(part, str) and part: visited += 1 @@ -218,7 +218,7 @@ def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: list[d data["input"] = "\n".join(text_parts) -def has_non_string_content(data: dict[str, Any]) -> bool: +def has_non_string_content(data: Mapping[str, object]) -> bool: """Return True if any inspected content is not a plain string. Used by hooks whose mask/redact path operates on string offsets and diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index f834426d619..d82944c44ed 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -139,7 +139,7 @@ class QualifireGuardrail(CustomGuardrail): ] ) - def _convert_messages_to_api_format(self, messages: list[AllMessageValues]) -> list[dict[str, Any]]: + def _convert_messages_to_api_format(self, messages: list[AllMessageValues]) -> list[dict[str, object]]: """ Convert LiteLLM messages to Qualifire API format. Supports tool calls for tool_selection_quality_check. @@ -167,7 +167,7 @@ class QualifireGuardrail(CustomGuardrail): text_parts.append(part) content = "\n".join(text_parts) - api_message: dict[str, Any] = { + api_message: dict[str, object] = { "role": role, "content": content if isinstance(content, str) else str(content), } @@ -205,7 +205,7 @@ class QualifireGuardrail(CustomGuardrail): return api_messages - def _convert_tools_to_api_format(self, tools: list[Any] | None) -> list[dict[str, Any]] | None: + def _convert_tools_to_api_format(self, tools: list[object] | None) -> list[dict[str, object]] | None: """ Convert OpenAI-format tools to Qualifire API format. @@ -264,13 +264,13 @@ class QualifireGuardrail(CustomGuardrail): def _build_evaluate_payload( self, - api_messages: list[dict[str, Any]], + api_messages: list[dict[str, object]], output: str | None, assertions: list[str] | None, - available_tools: list[dict[str, Any]] | None, - ) -> dict[str, Any]: + available_tools: list[dict[str, object]] | None, + ) -> dict[str, object]: """Build payload dictionary for the /api/evaluation/evaluate endpoint.""" - payload: Final[dict[str, Any]] = {"messages": api_messages} + payload: Final[dict[str, object]] = {"messages": api_messages} if output is not None: payload["output"] = output @@ -305,7 +305,7 @@ class QualifireGuardrail(CustomGuardrail): messages: list[AllMessageValues], output: str | None, dynamic_params: dict[str, Any], - available_tools: list[Any] | None = None, + available_tools: list[object] | None = None, ) -> None: """ Core Qualifire check logic - shared between hooks. diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 3865ba4ed0e..07340e95835 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -97,7 +97,7 @@ class SingulrGuardrail(CustomGuardrail): request_data: dict[str, Any], inputs: GenericGuardrailAPIInputs, input_type: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: if not request_data: texts: Final = inputs.get("texts", []) @@ -138,7 +138,7 @@ class SingulrGuardrail(CustomGuardrail): if value ) - async def _call_api(self, payload: dict[str, Any]) -> SingulrGuardrailResponse | None: + async def _call_api(self, payload: dict[str, object]) -> SingulrGuardrailResponse | None: endpoint: Final = f"{self.singulr_api_base}{_GUARD_ENDPOINT}" verbose_proxy_logger.debug("Singulr: %s", endpoint) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 46b00829b74..267087817d0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -878,7 +878,7 @@ class UnifiedLLMGuardrails(CustomLogger): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[object], request_data: dict, guardrail_to_apply: CustomGuardrail | None = None, buffer_until_moderated_default: bool = False, diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index c2f5dbb4032..fe658a13c24 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -9,7 +9,7 @@ import copy import json import traceback from datetime import datetime, timezone -from typing import Annotated, Any, Final +from typing import Annotated, Final from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -62,7 +62,7 @@ def _validate_team_callback(data: "AddTeamCallback") -> None: raise _callback_config_error(error) -def _redact_callback_secrets(metadata: Any) -> Any: +def _redact_callback_secrets(metadata: object) -> object: """Strip secret values out of a team-metadata snapshot before audit logging. Both ``team_metadata["logging"]`` (list of ``AddTeamCallback`` dicts) and @@ -176,8 +176,8 @@ def _log_audit_task_exception(task: "asyncio.Task[None]") -> None: async def _emit_team_callback_audit_log( *, team_id: str, - before_metadata: Any, - after_metadata: Any, + before_metadata: object, + after_metadata: object, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None, ) -> None: diff --git a/litellm/proxy/openai_evals_endpoints/endpoints.py b/litellm/proxy/openai_evals_endpoints/endpoints.py index 25d73e0dc1b..abfbed5f822 100644 --- a/litellm/proxy/openai_evals_endpoints/endpoints.py +++ b/litellm/proxy/openai_evals_endpoints/endpoints.py @@ -35,7 +35,7 @@ async def create_eval( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Create a new evaluation. @@ -131,7 +131,7 @@ async def list_evals( order_by: str | None = None, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ List evaluations with pagination. @@ -228,7 +228,7 @@ async def get_eval( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Get a specific evaluation by ID. @@ -316,7 +316,7 @@ async def update_eval( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Update an evaluation. @@ -406,7 +406,7 @@ async def delete_eval( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Delete an evaluation. @@ -494,7 +494,7 @@ async def cancel_eval( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Cancel a running evaluation. @@ -587,7 +587,7 @@ async def create_run( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Create a new run for an evaluation. @@ -690,7 +690,7 @@ async def list_runs( order: str | None = None, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ List all runs for an evaluation with pagination. @@ -780,7 +780,7 @@ async def get_run( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Get a specific run by ID. @@ -867,7 +867,7 @@ async def cancel_run( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Cancel a running run. @@ -956,7 +956,7 @@ async def delete_run( request: Request, custom_llm_provider: str | None = "openai", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> object: """ Delete a run. diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index 67fe25160ec..061a852b701 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -6,6 +6,7 @@ Configuration structure: - policy_attachments: Define WHERE policies apply (teams, keys, models) """ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from litellm._logging import verbose_proxy_logger @@ -25,8 +26,8 @@ _reset_color_code: Final = "\033[0m" def _print_policies_on_startup( - policies_config: dict[str, Any], - policy_attachments_config: list[dict[str, Any]] | None = None, + policies_config: Mapping[str, Mapping[str, object]], + policy_attachments_config: Sequence[Mapping[str, object]] | None = None, ) -> None: """ Print loaded policies to console on startup (similar to model list). diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 183a03cc13c..d5cf1249fbf 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -87,7 +87,7 @@ def _get_embedding_config_cache() -> InMemoryCache: return _embedding_config_cache -def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> Any: +def _redact_sensitive_litellm_params(litellm_params: object, _depth: int = 0) -> Any: """ Replace credential-bearing values in ``litellm_params`` with ``REDACTED_BY_LITELM`` while preserving non-secret keys (``api_base``, @@ -119,7 +119,7 @@ def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> An return json.dumps(_redact_sensitive_litellm_params(parsed, _depth + 1)) if not isinstance(litellm_params, dict): return litellm_params - out: Final[dict[str, Any]] = {} + out: Final[dict[str, object]] = {} for k, v in litellm_params.items(): if _LITELLM_PARAMS_MASKER.is_sensitive_key(k): out[k] = REDACTED_BY_LITELM_STRING diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index fa563d5a678..73a0159fc9f 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -7,7 +7,7 @@ so this implementation skips the embedding step and directly uploads files. from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Final, cast from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -83,7 +83,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): """ vector_store_id = self.vector_store_config.get("vector_store_id") - vector_store_config: Final = cast(dict[str, Any], self.vector_store_config) + vector_store_config: Final = self.vector_store_config # Get API credentials api_key: Final = cast(str | None, vector_store_config.get("api_key")) or GeminiModelInfo.get_api_key() @@ -228,7 +228,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): url: Final = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore" # Build request body with chunking config and metadata if provided - request_body: Final[dict[str, Any]] = {"displayName": filename} + request_body: Final[dict[str, object]] = {"displayName": filename} # Add chunking configuration if provided chunking_strategy: Final = self.chunking_strategy @@ -244,7 +244,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): # Add custom metadata if provided in vector_store_config custom_metadata: Final = cast( - list[dict[str, Any]] | None, + list[dict[str, object]] | None, self.vector_store_config.get("custom_metadata"), ) if custom_metadata: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d4b9f4e8cce..aa229270800 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -4,7 +4,7 @@ import asyncio import os from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast import litellm from litellm.constants import ( @@ -41,6 +41,9 @@ from ..llms.vertex_ai.vertex_llm_base import VertexBase from ..llms.xai.realtime.handler import XAIRealtime from ..utils import client as wrapper_client +if TYPE_CHECKING: + from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig + azure_realtime: Final = AzureOpenAIRealtime() openai_realtime: Final = OpenAIRealtime() bedrock_realtime: Final = BedrockRealtime() @@ -50,7 +53,7 @@ base_llm_http_handler = BaseLLMHTTPHandler() _EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) -def _with_resolved_session_model(session: dict[str, Any], model_name: str) -> dict[str, Any]: +def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: if "model" not in session: return session return {**session, "model": model_name} @@ -70,7 +73,7 @@ def _get_realtime_http_provider_config( dynamic_api_base: str | None, dynamic_api_key: str | None, litellm_params: GenericLiteLLMParams, -) -> tuple[Any, str, str]: +) -> tuple["BaseRealtimeHTTPConfig | None", str, str]: """ Return (provider_config, resolved_api_base, resolved_api_key) for the realtime HTTP endpoints (client_secrets / realtime_calls). diff --git a/litellm/repositories/config_repository.py b/litellm/repositories/config_repository.py index 76b9a3a5809..2e8e760db07 100644 --- a/litellm/repositories/config_repository.py +++ b/litellm/repositories/config_repository.py @@ -17,6 +17,11 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +def _decoded_json(raw: str) -> object: + """Decode a JSON-encoded config row value into an opaque object.""" + return json.loads(raw) + + class _ConfigRow(Protocol): @property def param_name(self) -> str: ... @@ -48,7 +53,7 @@ class _PrismaHandle(Protocol): class ConfigParam: """Simple wrapper for config parameter from DB.""" - def __init__(self, param_name: str, param_value: Any): + def __init__(self, param_name: str, param_value: object): self.param_name = param_name self.param_value = param_value @@ -85,12 +90,12 @@ class ConfigRepository: record: Final = await self._config_table.find_unique(where={"param_name": param_name}) if record is None: return None - param_value = record.param_value + param_value: object = record.param_value if isinstance(param_value, str): - param_value = json.loads(param_value) + param_value = _decoded_json(param_value) return ConfigParam(param_name=param_name, param_value=param_value) - async def set_param(self, param_name: str, param_value: Any) -> ConfigParam: + async def set_param(self, param_name: str, param_value: object) -> ConfigParam: """Set a config parameter in the database.""" value_json: Final = json.dumps(param_value) if not isinstance(param_value, str) else param_value await self._config_table.upsert( @@ -115,9 +120,9 @@ class ConfigRepository: records: Final = await self._config_table.find_many() result: Final[dict[str, object]] = {} for record in records: - param_value = record.param_value + param_value: object = record.param_value if isinstance(param_value, str): - param_value = json.loads(param_value) + param_value = _decoded_json(param_value) result[record.param_name] = param_value return result diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index b59ce6e3621..709910753f2 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -13,7 +13,8 @@ from __future__ import annotations import hashlib import json -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger @@ -26,6 +27,9 @@ from litellm.router_strategy.adaptive_router.config import ( ) from litellm.router_strategy.adaptive_router.signals import Turn +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + # Identity fields hashed into a derived session key so the same conversation # from the same caller produces a stable key, while different keys/teams/users # stay segregated even if they happen to send identical first messages. @@ -100,8 +104,8 @@ def _last_user_content(messages: list[dict[str, Any]] | None) -> str | None: def _recent_tool_results( - messages: list[dict[str, Any]] | None, -) -> list[dict[str, Any]]: + messages: Sequence[Mapping[str, object]] | None, +) -> list[dict[str, object]]: """Extract the current turn's tool result payloads from the request messages. Tool results are `role == "tool"` messages that sit at the tail of the @@ -115,7 +119,7 @@ def _recent_tool_results( """ if not messages: return [] - results: Final[list[dict[str, Any]]] = [] + results: Final[list[dict[str, object]]] = [] for msg in reversed(messages): if not isinstance(msg, dict): break @@ -154,7 +158,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple: raw_tool_calls = getattr(msg, "tool_calls", None) if raw_tool_calls is None and isinstance(msg, dict): raw_tool_calls = msg.get("tool_calls") - tool_calls: Final[list[dict[str, Any]]] = [] + tool_calls: Final[list[dict[str, object]]] = [] for tc in raw_tool_calls or []: if isinstance(tc, dict): tool_calls.append(tc) @@ -174,11 +178,11 @@ class AdaptiveRouterPostCallHook(CustomLogger): async def async_post_call_response_headers_hook( self, - data: dict[str, Any], - user_api_key_dict: Any, - response: Any, + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + response: object, request_headers: dict[str, str] | None = None, - litellm_call_info: dict[str, Any] | None = None, + litellm_call_info: dict[str, object] | None = None, ) -> dict[str, str] | None: """ Surface the chosen logical model as the `x-litellm-adaptive-router-model` @@ -209,7 +213,7 @@ class AdaptiveRouterPostCallHook(CustomLogger): async def _record( self, kwargs: dict[str, Any], - response_obj: Any, + response_obj: object, response_status: int, ) -> None: try: diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index 91e4cad4d27..e99d473d455 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -16,6 +16,7 @@ then cheapest `model_info.input_cost_per_token`). """ import math +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional from litellm._logging import verbose_router_logger @@ -98,7 +99,7 @@ class QualityRouter(CustomLogger): self._tier_to_models_cache = self._build_tier_index() return self._tier_to_models_cache - def _get_routing_preferences(self, deployment: Any) -> dict[str, Any] | None: + def _get_routing_preferences(self, deployment: object) -> dict[str, Any] | None: """ Extract litellm_routing_preferences from a deployment, handling both dict-shaped and Pydantic-object-shaped deployments. @@ -119,7 +120,7 @@ class QualityRouter(CustomLogger): return model_info.get("litellm_routing_preferences") return getattr(model_info, "litellm_routing_preferences", None) - def _get_deployment_input_cost(self, deployment: Any) -> float | None: + def _get_deployment_input_cost(self, deployment: object) -> float | None: """ Extract `input_cost_per_token` from a deployment's model_info. @@ -144,7 +145,7 @@ class QualityRouter(CustomLogger): except (TypeError, ValueError): return None - def _get_deployment_model_name(self, deployment: Any) -> str | None: + def _get_deployment_model_name(self, deployment: object) -> str | None: """Extract `model_name` from a dict- or object-shaped deployment.""" if isinstance(deployment, dict): return deployment.get("model_name") @@ -304,8 +305,8 @@ class QualityRouter(CustomLogger): def _stash_decision( self, - request_kwargs: dict[str, Any] | None, - decision: dict[str, Any], + request_kwargs: dict[str, object] | None, + decision: Mapping[str, object], ) -> None: """ Stash the routing decision in request_kwargs.metadata so the Router can diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3d37ca216a7..2bcac84ec19 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -1,6 +1,6 @@ import hashlib import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, Final @@ -39,7 +39,7 @@ _REQUEST_SCOPED_STATUS_CODES: Final = frozenset((404,)) def _trigger_cooldown_for_failed_deployment( litellm_router: LitellmRouter, - kwargs: Mapping[str, Any], + kwargs: Mapping[str, object], exception: Exception, ) -> None: """ @@ -218,7 +218,7 @@ PRE_ROUTING_SELECTED_MODEL_KEY: Final = "pre_routing_selected_model" _ROUTER_METADATA_BUCKETS: Final = ("metadata", "litellm_metadata") -def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selected_model: str) -> None: +def record_pre_routing_selection(request_kwargs: Mapping[str, object] | None, selected_model: str) -> None: """ Remember which model a pre-routing hook picked, so fallback lookup can key off it. @@ -257,14 +257,14 @@ def clear_pre_routing_selection(request_kwargs: Mapping[str, object] | None) -> del bucket[PRE_ROUTING_SELECTED_MODEL_KEY] -def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: +def get_pre_routing_selection(kwargs: Mapping[str, object]) -> str | None: """The model a pre-routing hook selected for this request, if one did.""" buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS) selections: Final = (bucket.get(PRE_ROUTING_SELECTED_MODEL_KEY) for bucket in buckets if isinstance(bucket, dict)) return next((selected for selected in selections if isinstance(selected, str) and selected), None) -def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: +def fallback_lookup_groups(kwargs: Mapping[str, object], model_group: str | None) -> tuple[str, ...]: """ Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, and the requested group still resolves when no tier-keyed chain exists, so configs keyed @@ -413,7 +413,7 @@ def creates_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: async def run_async_fallback( - *args: tuple[Any], + *args: object, litellm_router: LitellmRouter, fallback_model_group: list[str], original_model_group: str, @@ -630,5 +630,5 @@ def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool: return False -def run_non_standard_fallback_format(fallbacks: list[str] | list[dict[str, Any]], model_group: str): +def run_non_standard_fallback_format(fallbacks: Sequence[str] | Sequence[Mapping[str, object]], model_group: str): pass diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 48b1f24ae8a..01d42627001 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -526,8 +526,8 @@ async def async_io_token_pre_call_check( def io_token_reconcile_success( dual_cache: DualCache, - kwargs: Any, - response_obj: Any, + kwargs: Mapping[str, object] | None, + response_obj: object, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj @@ -577,8 +577,8 @@ def io_token_reconcile_success( async def async_io_token_reconcile_success( dual_cache: DualCache, - kwargs: Any, - response_obj: Any, + kwargs: Mapping[str, object] | None, + response_obj: object, *, parent_otel_span: Span | None = None, ) -> None: @@ -638,7 +638,7 @@ async def async_io_token_reconcile_success( def io_token_refund_failure( dual_cache: DualCache, - kwargs: Any, + kwargs: Mapping[str, object] | None, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) @@ -689,7 +689,7 @@ def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping async def async_io_token_refund_failure( dual_cache: DualCache, - kwargs: Any, + kwargs: Mapping[str, object] | None, *, parent_otel_span: Span | None = None, ) -> None: diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index ab5ef5853c9..309894957ea 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -10,10 +10,19 @@ import traceback from collections.abc import Callable from functools import partial from types import MappingProxyType -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_router_logger +if TYPE_CHECKING: + from litellm.types.router import SearchToolTypedDict + + +class _SearchToolsRouter(Protocol): + """The one router attribute the search-tool helpers read and replace.""" + + search_tools: "list[SearchToolTypedDict]" + class SearchAPIRouter: """ @@ -45,7 +54,7 @@ class SearchAPIRouter: return resolved_api_key, resolved_api_base @staticmethod - async def update_router_search_tools(router_instance: Any, search_tools: list): + async def update_router_search_tools(router_instance: _SearchToolsRouter, search_tools: list): """ Update the router with search tools from the database. @@ -83,7 +92,7 @@ class SearchAPIRouter: @staticmethod def get_matching_search_tools( - router_instance: Any, + router_instance: _SearchToolsRouter, search_tool_name: str, ) -> list: """ @@ -175,7 +184,7 @@ class SearchAPIRouter: @staticmethod async def async_search_with_fallbacks_helper( - router_instance: Any, + router_instance: _SearchToolsRouter, model: str, original_generic_function: Callable, **kwargs, diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 2c7f1f8389d..e86c8e7c919 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -266,7 +266,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): """ from litellm._uuid import uuid - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "Name": secret_name, "SecretString": secret_value, "ClientRequestToken": str(uuid.uuid4()), @@ -415,7 +415,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): """ from litellm._uuid import uuid - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "SecretId": secret_name, "SecretString": secret_value, "ClientRequestToken": str(uuid.uuid4()), diff --git a/litellm/skills/main.py b/litellm/skills/main.py index ae1ce150368..9d2ed524ce5 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -5,7 +5,7 @@ Provides create, list, get, and delete operations for skills import asyncio import contextvars -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import Any, Final @@ -35,7 +35,7 @@ DEFAULT_ANTHROPIC_API_BASE: Final = "https://api.anthropic.com/v1" _litellm_skills_handler = None -def _get_user_api_key_auth_from_kwargs(kwargs: dict[str, Any]) -> Any | None: +def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object]) -> Any | None: for metadata_key in ("metadata", "litellm_metadata"): metadata = kwargs.get(metadata_key) if isinstance(metadata, dict) and metadata.get("user_api_key_auth") is not None: @@ -44,7 +44,7 @@ def _get_user_api_key_auth_from_kwargs(kwargs: dict[str, Any]) -> Any | None: def _get_skill_request_metadata( - kwargs: dict[str, Any], + kwargs: Mapping[str, object], extra_body: dict[str, Any] | None, ) -> dict[str, Any] | None: if extra_body and isinstance(extra_body.get("metadata"), dict): @@ -73,7 +73,7 @@ async def acreate_skill( files: list[Any] | None = None, display_title: str | None = None, extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -136,12 +136,12 @@ def create_skill( files: list[Any] | None = None, display_title: str | None = None, extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> Skill | Coroutine[Any, Any, Skill]: +) -> Skill | Coroutine[object, object, Skill]: """ Create a new skill @@ -330,7 +330,7 @@ def list_skills( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> ListSkillsResponse | Coroutine[Any, Any, ListSkillsResponse]: +) -> ListSkillsResponse | Coroutine[object, object, ListSkillsResponse]: """ List all skills @@ -444,7 +444,7 @@ def list_skills( async def aget_skill( skill_id: str, extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -501,11 +501,11 @@ async def aget_skill( def get_skill( skill_id: str, extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> Skill | Coroutine[Any, Any, Skill]: +) -> Skill | Coroutine[object, object, Skill]: """ Get a skill by ID @@ -608,7 +608,7 @@ def get_skill( async def adelete_skill( skill_id: str, extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -665,11 +665,11 @@ async def adelete_skill( def delete_skill( skill_id: str, extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> DeleteSkillResponse | Coroutine[Any, Any, DeleteSkillResponse]: +) -> DeleteSkillResponse | Coroutine[object, object, DeleteSkillResponse]: """ Delete a skill by ID diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 474c652ff3a..8bb0235ea2a 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from enum import Enum @@ -128,31 +129,31 @@ class VertexSearchDataStoreExtraBody(TypedDict, total=False): pageToken: str offset: int oneBoxPageSize: int - pageCategories: list[str] - imageQuery: dict[str, Any] + pageCategories: Sequence[str] + imageQuery: Mapping[str, object] filter: str canonicalFilter: str orderBy: str - userInfo: dict[str, Any] + userInfo: Mapping[str, object] languageCode: str - facetSpecs: list[dict[str, Any]] - boostSpec: dict[str, Any] - params: dict[str, Any] - queryExpansionSpec: dict[str, Any] - spellCorrectionSpec: dict[str, Any] + facetSpecs: Sequence[Mapping[str, object]] + boostSpec: Mapping[str, object] + params: Mapping[str, object] + queryExpansionSpec: Mapping[str, object] + spellCorrectionSpec: Mapping[str, object] userPseudoId: str - contentSearchSpec: dict[str, Any] + contentSearchSpec: Mapping[str, object] rankingExpression: str rankingExpressionBackend: str safeSearch: bool - userLabels: dict[str, str] - naturalLanguageQueryUnderstandingSpec: dict[str, Any] - searchAsYouTypeSpec: dict[str, Any] - displaySpec: dict[str, Any] - crowdingSpecs: list[dict[str, Any]] + userLabels: Mapping[str, str] + naturalLanguageQueryUnderstandingSpec: Mapping[str, object] + searchAsYouTypeSpec: Mapping[str, object] + displaySpec: Mapping[str, object] + crowdingSpecs: Sequence[Mapping[str, object]] relevanceThreshold: str - relevanceScoreSpec: dict[str, Any] - customRankingParams: dict[str, Any] + relevanceScoreSpec: Mapping[str, object] + customRankingParams: Mapping[str, object] class VertexSearchEngineExtraBody(VertexSearchDataStoreExtraBody, total=False): @@ -166,7 +167,7 @@ class VertexSearchEngineExtraBody(VertexSearchDataStoreExtraBody, total=False): (per-store scoping/filtering) and ``numResultsPerDataStore``. """ - dataStoreSpecs: list[dict[str, Any]] + dataStoreSpecs: Sequence[Mapping[str, object]] numResultsPerDataStore: int @@ -256,7 +257,7 @@ class IndexCreateLiteLLMParams(BaseModel): class IndexCreateRequest(BaseModel): index_name: str litellm_params: IndexCreateLiteLLMParams - index_info: dict[str, Any] | None = None + index_info: dict[str, object] | None = None class BaseVectorStoreAuthCredentials(TypedDict, total=False): @@ -270,7 +271,7 @@ class LiteLLM_ManagedVectorStoreIndex(BaseModel): id: str index_name: str litellm_params: IndexCreateLiteLLMParams - index_info: dict[str, Any] | None = None + index_info: dict[str, object] | None = None created_at: datetime | None = None created_by: str | None = None updated_at: datetime | None = None diff --git a/litellm/vector_store_files/main.py b/litellm/vector_store_files/main.py index 7af8dc7d435..5bc3c8f1525 100644 --- a/litellm/vector_store_files/main.py +++ b/litellm/vector_store_files/main.py @@ -39,7 +39,7 @@ def _ensure_provider(custom_llm_provider: str | None) -> str: def _prepare_registry_credentials( *, vector_store_id: str, - kwargs: dict[str, Any], + kwargs: dict[str, object], ) -> None: if litellm.vector_store_registry is None: return @@ -116,7 +116,7 @@ def create( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> VectorStoreFileObject | Coroutine[Any, Any, VectorStoreFileObject]: +) -> VectorStoreFileObject | Coroutine[object, object, VectorStoreFileObject]: local_vars: Final = locals() try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") @@ -245,7 +245,7 @@ def list( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> VectorStoreFileListResponse | Coroutine[Any, Any, VectorStoreFileListResponse]: +) -> VectorStoreFileListResponse | Coroutine[object, object, VectorStoreFileListResponse]: local_vars: Final = locals() try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") @@ -355,7 +355,7 @@ def retrieve( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> VectorStoreFileObject | Coroutine[Any, Any, VectorStoreFileObject]: +) -> VectorStoreFileObject | Coroutine[object, object, VectorStoreFileObject]: local_vars: Final = locals() try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") @@ -463,7 +463,7 @@ def retrieve_content( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> VectorStoreFileContentResponse | Coroutine[Any, Any, VectorStoreFileContentResponse]: +) -> VectorStoreFileContentResponse | Coroutine[object, object, VectorStoreFileContentResponse]: local_vars: Final = locals() try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") @@ -577,7 +577,7 @@ def update( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> VectorStoreFileObject | Coroutine[Any, Any, VectorStoreFileObject]: +) -> VectorStoreFileObject | Coroutine[object, object, VectorStoreFileObject]: local_vars: Final = locals() try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") @@ -692,7 +692,7 @@ def delete( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> VectorStoreFileDeleteResponse | Coroutine[Any, Any, VectorStoreFileDeleteResponse]: +) -> VectorStoreFileDeleteResponse | Coroutine[object, object, VectorStoreFileDeleteResponse]: local_vars: Final = locals() try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") From 362fb4cffe71d2df1b8e3d335a0d6550bdcf4724 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:35:01 +0000 Subject: [PATCH 02/38] refactor(typing): replace Any with proven types in 42 more backend files --- .../providers/bedrock_agentcore/handler.py | 12 ++--- litellm/a2a_protocol/streaming_iterator.py | 4 +- litellm/a2a_protocol/utils.py | 7 +-- litellm/caching/caching_handler.py | 4 +- litellm/experimental_mcp_client/client.py | 12 ++++- litellm/files/main.py | 3 +- litellm/integrations/arize/_utils.py | 23 +++++++-- .../focus/destinations/s3_destination.py | 50 +++++++++++-------- litellm/integrations/prometheus.py | 13 +++-- litellm/interactions/http_handler.py | 30 +++++------ .../messages/fake_stream_iterator.py | 32 ++++++------ litellm/llms/bedrock/chat/invoke_handler.py | 12 +++-- litellm/llms/cohere/embed/transformation.py | 13 +++-- .../llms/dashscope/rerank/transformation.py | 4 +- .../llms/dataforseo/search/transformation.py | 4 +- .../text_to_speech/transformation.py | 22 ++++---- .../fireworks_ai/rerank/transformation.py | 4 +- litellm/llms/gemini/count_tokens/handler.py | 4 +- litellm/llms/gigachat/file_handler.py | 12 ++++- litellm/llms/huggingface/embedding/handler.py | 16 ++++-- .../minimax/text_to_speech/transformation.py | 2 +- .../openai/vector_stores/transformation.py | 2 +- .../guardrail_translation/handler.py | 11 ++-- .../text_to_speech/transformation.py | 17 ++++--- litellm/proxy/client/cli/commands/auth.py | 14 +++--- litellm/proxy/client/cli/commands/models.py | 22 ++++++-- litellm/proxy/client/cli/commands/users.py | 13 +++-- litellm/proxy/client/http_client.py | 5 +- litellm/proxy/client/models.py | 9 ++-- .../container_endpoints/handler_factory.py | 6 +-- .../cato_networks/cato_networks.py | 8 +-- .../guardrail_hooks/dynamoai/dynamoai.py | 8 +-- .../guardrail_hooks/singulr/singulr.py | 7 ++- .../tool_policy/tool_policy_guardrail.py | 9 +++- litellm/proxy/guardrails/usage_endpoints.py | 13 +++-- .../usage_endpoints/ai_usage_chat.py | 36 +++++++------ litellm/realtime_api/main.py | 8 +-- litellm/responses/utils.py | 7 ++- .../adaptive_router/signals.py | 13 ++--- litellm/router_utils/cooldown_handlers.py | 8 +-- litellm/skills/main.py | 16 +++--- litellm/vector_store_files/main.py | 34 ++++++------- 42 files changed, 336 insertions(+), 213 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index db57072ca38..a4e6fa50901 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -6,8 +6,8 @@ completion bridge that would otherwise strip the envelope. """ import json -from collections.abc import AsyncIterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Mapping +from typing import Any, Final from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( @@ -28,7 +28,7 @@ class BedrockAgentCoreA2AHandler: @staticmethod async def handle_non_streaming( request_id: str, - params: dict[str, Any], + params: Mapping[str, object], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, ) -> dict[str, Any]: @@ -56,7 +56,7 @@ class BedrockAgentCoreA2AHandler: verbose_logger.info("BedrockAgentCore A2A: Sending non-streaming request to %s", url) client: Final = get_async_httpx_client( - llm_provider=cast(Any, httpxSpecialProvider.A2AProvider), + llm_provider=httpxSpecialProvider.A2AProvider, ) response: Final = await client.post( url, @@ -74,7 +74,7 @@ class BedrockAgentCoreA2AHandler: @staticmethod async def handle_streaming( request_id: str, - params: dict[str, Any], + params: Mapping[str, object], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, ) -> AsyncIterator[dict[str, Any]]: @@ -103,7 +103,7 @@ class BedrockAgentCoreA2AHandler: verbose_logger.info("BedrockAgentCore A2A: Sending streaming request to %s", url) client: Final = get_async_httpx_client( - llm_provider=cast(Any, httpxSpecialProvider.A2AProvider), + llm_provider=httpxSpecialProvider.A2AProvider, ) response: Final = await client.post( url, diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 413691f233d..67db8e905e3 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -148,9 +148,9 @@ class A2AStreamingIterator: except Exception as e: verbose_logger.debug("Error in A2A streaming completion handler: %s", e) - def _build_logging_result(self, usage: litellm.Usage) -> dict[str, Any]: + def _build_logging_result(self, usage: litellm.Usage) -> dict[str, object]: """Build a result dict for logging.""" - result: Final[dict[str, Any]] = { + result: Final[dict[str, object]] = { "id": getattr(self.request, "id", "unknown"), "jsonrpc": "2.0", "usage": (usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)), diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index f2e61f66105..5ffca68130b 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -2,6 +2,7 @@ Utility functions for A2A protocol. """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import litellm @@ -46,7 +47,7 @@ class A2ARequestUtils: return " ".join(text_parts) @staticmethod - def extract_text_from_response(response_dict: dict[str, Any]) -> str: + def extract_text_from_response(response_dict: Mapping[str, object]) -> str: """ Extract text content from A2A response result. @@ -109,7 +110,7 @@ class A2ARequestUtils: @staticmethod def calculate_usage_from_request_response( request: "SendMessageRequest | SendStreamingMessageRequest", - response_dict: dict[str, Any], + response_dict: Mapping[str, object], ) -> tuple[int, int, int]: """ Calculate token usage from A2A request and response. @@ -145,5 +146,5 @@ def extract_text_from_a2a_message(message: Any) -> str: return A2ARequestUtils.extract_text_from_message(message) -def extract_text_from_a2a_response(response_dict: dict[str, Any]) -> str: +def extract_text_from_a2a_response(response_dict: Mapping[str, object]) -> str: return A2ARequestUtils.extract_text_from_response(response_dict) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 8fe60876b4e..0de88eacaa5 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -672,7 +672,7 @@ class LLMCachingHandler: def _async_log_cache_hit_on_callbacks( self, logging_obj: LiteLLMLoggingObj, - cached_result: Any, + cached_result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: bool, @@ -1184,7 +1184,7 @@ class LLMCachingHandler: logging_obj: LiteLLMLoggingObj, model: str, kwargs: dict[str, Any], - cached_result: Any, + cached_result: object, is_async: bool, is_embedding: bool = False, custom_llm_provider: str | None = None, diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 34af6fcffba..f40941d62cc 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -5,7 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 import os -from collections.abc import Awaitable, Callable, Generator, Sequence +from collections.abc import Awaitable, Callable, Generator from contextlib import AbstractAsyncContextManager from datetime import timedelta from functools import partial @@ -13,11 +13,19 @@ from importlib import metadata from typing import Any, Final, Protocol, TypeAlias, TypeVar import httpx +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client +from mcp.shared.message import SessionMessage +from typing_extensions import Unpack -_TransportContext: TypeAlias = AbstractAsyncContextManager[Sequence[Any]] +_TransportStreams: TypeAlias = tuple[ + MemoryObjectReceiveStream[SessionMessage | Exception], + MemoryObjectSendStream[SessionMessage], + Unpack[tuple[object, ...]], +] +_TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams] class _StreamableHttpClientFactory(Protocol): diff --git a/litellm/files/main.py b/litellm/files/main.py index e769a0a0508..19da77b7364 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -14,6 +14,7 @@ from functools import partial from typing import Any, Final, Literal, cast import httpx +from openai import AsyncOpenAI, OpenAI # Type aliases for provider parameters FileCreateProvider = Literal[ @@ -1002,7 +1003,7 @@ def file_content_streaming( timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj | None, _is_async: bool, - client: Any | None, + client: OpenAI | AsyncOpenAI | None, ) -> FileContentStreamingResult | Coroutine[object, object, FileContentStreamingResult]: if logging_obj is not None: logging_obj.model = model or "" diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index e7e1ab538d5..5a5324eae5e 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -2,7 +2,7 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final -from typing_extensions import override +from typing_extensions import ReadOnly, TypedDict, override from litellm._logging import verbose_logger from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import ( @@ -492,12 +492,12 @@ def _sanitize_optional_params(optional_params: dict | None) -> dict: return optional_params -def _set_metadata_attributes(span: "Span", metadata: Any | None, span_attrs) -> None: +def _set_metadata_attributes(span: "Span", metadata: object | None, span_attrs) -> None: if metadata is not None: safe_set_attribute(span, span_attrs.METADATA, safe_dumps(metadata)) -def _extract_metadata_tools(metadata: Any | None) -> list | None: +def _extract_metadata_tools(metadata: object | None) -> list | None: if not isinstance(metadata, dict): return None llm_obj: Final = metadata.get("llm") @@ -670,7 +670,22 @@ def _get_tool_calls(message) -> list | None: return tool_calls if isinstance(tool_calls, list) and tool_calls else None -def _normalize_tool_call(raw_tc) -> dict[str, Any] | None: +class _NormalizedToolCallFunction(TypedDict): + """The ``function`` sub-object of a normalized tool call.""" + + name: ReadOnly[object] + arguments: ReadOnly[object] + + +class _NormalizedToolCall(TypedDict): + """A tool call reduced to the stable shape the OpenInference emitters read.""" + + id: ReadOnly[object] + type: ReadOnly[object] + function: ReadOnly[_NormalizedToolCallFunction] + + +def _normalize_tool_call(raw_tc) -> _NormalizedToolCall | None: """Normalize a single tool_call (dict or Pydantic) into a stable shape: {"id": str|None, "type": str, "function": {"name": str|None, "arguments": str|None}} diff --git a/litellm/integrations/focus/destinations/s3_destination.py b/litellm/integrations/focus/destinations/s3_destination.py index d6530b889d9..661cf1933ff 100644 --- a/litellm/integrations/focus/destinations/s3_destination.py +++ b/litellm/integrations/focus/destinations/s3_destination.py @@ -3,14 +3,26 @@ from __future__ import annotations import asyncio +from collections.abc import Mapping from datetime import timezone -from typing import Any, Final +from typing import Final, TypedDict import boto3 +from typing_extensions import ReadOnly from .base import FocusDestination, FocusTimeWindow +class _S3ClientKwargs(TypedDict, total=False): + """Optional boto3 client arguments the destination config may supply.""" + + region_name: ReadOnly[str] + endpoint_url: ReadOnly[str] + aws_access_key_id: ReadOnly[str] + aws_secret_access_key: ReadOnly[str] + aws_session_token: ReadOnly[str] + + class FocusS3Destination(FocusDestination): """Handles uploading serialized exports to S3 buckets.""" @@ -18,7 +30,7 @@ class FocusS3Destination(FocusDestination): self, *, prefix: str, - config: dict[str, Any] | None = None, + config: Mapping[str, str] | None = None, ) -> None: config = config or {} bucket_name: Final = config.get("bucket_name") @@ -47,25 +59,23 @@ class FocusS3Destination(FocusDestination): key_prefix: Final = "/".join(filter(None, parts)) return f"{key_prefix}/{filename}" if key_prefix else filename + def _client_kwargs(self) -> _S3ClientKwargs: + """Collect the boto3 client arguments the destination config provides.""" + region: Final = self.config.get("region_name") + endpoint: Final = self.config.get("endpoint_url") + key_id: Final = self.config.get("aws_access_key_id") + secret: Final = self.config.get("aws_secret_access_key") + token: Final = self.config.get("aws_session_token") + return { + **(_S3ClientKwargs(region_name=region) if region else _S3ClientKwargs()), + **(_S3ClientKwargs(endpoint_url=endpoint) if endpoint else _S3ClientKwargs()), + **(_S3ClientKwargs(aws_access_key_id=key_id) if key_id else _S3ClientKwargs()), + **(_S3ClientKwargs(aws_secret_access_key=secret) if secret else _S3ClientKwargs()), + **(_S3ClientKwargs(aws_session_token=token) if token else _S3ClientKwargs()), + } + def _upload(self, content: bytes, object_key: str) -> None: - client_kwargs: Final[dict[str, Any]] = {} - region_name: Final = self.config.get("region_name") - if region_name: - client_kwargs["region_name"] = region_name - endpoint_url: Final = self.config.get("endpoint_url") - if endpoint_url: - client_kwargs["endpoint_url"] = endpoint_url - - session_kwargs: Final[dict[str, Any]] = {} - for key in ( - "aws_access_key_id", - "aws_secret_access_key", - "aws_session_token", - ): - if self.config.get(key): - session_kwargs[key] = self.config[key] - - s3_client: Final = boto3.client("s3", **client_kwargs, **session_kwargs) + s3_client: Final = boto3.client("s3", **self._client_kwargs()) s3_client.put_object( Bucket=self.bucket_name, Key=object_key, diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 3e75c9cbf93..6766d246894 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -10,7 +10,7 @@ import sys from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import replace from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast from pydantic import BaseModel @@ -142,6 +142,9 @@ class _ExcludedLabelMetric: return self._metric.labels(*kept_values) if kept_values else self._metric +_MetricLike: TypeAlias = "NoOpMetric | _ExcludedLabelMetric | MetricWrapperBase" + + def _get_budget_metrics_per_request_timeout() -> float: raw: Final = os.getenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT") if raw is None: @@ -1652,7 +1655,7 @@ class PrometheusLogger(CustomLogger): cache_creation_detail_tokens: Final = PrometheusLogger._resolve_cache_write_tokens(prompt_details) - detail_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]]] = [ + detail_metrics: Final[list[tuple[_MetricLike, DEFINED_PROMETHEUS_METRICS, object]]] = [ ( self.litellm_input_cached_tokens_metric, "litellm_input_cached_tokens_metric", @@ -1705,7 +1708,7 @@ class PrometheusLogger(CustomLogger): if not isinstance(usage_object, dict): return - media_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]]] = [ + media_metrics: Final[list[tuple[_MetricLike, DEFINED_PROMETHEUS_METRICS, object]]] = [ ( self.litellm_video_duration_seconds_metric, "litellm_video_duration_seconds_metric", @@ -1727,7 +1730,7 @@ class PrometheusLogger(CustomLogger): def _inc_sparse_usage_counters( self, - counters_with_values: Sequence[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]], + counters_with_values: Sequence[tuple[_MetricLike, DEFINED_PROMETHEUS_METRICS, object]], enum_values: UserAPIKeyLabelValues, label_context: PrometheusLabelFactoryContext | None = None, ) -> None: @@ -2623,7 +2626,7 @@ class PrometheusLogger(CustomLogger): """ standard_logging_payload: Final = request_kwargs.get("standard_logging_object", {}) or {} _litellm_params: Final = request_kwargs.get("litellm_params", {}) or {} - _metadata_raw: Final = self._safe_get(standard_logging_payload, "metadata") or {} + _metadata_raw: Final[object] = self._safe_get(standard_logging_payload, "metadata") or {} if isinstance(_metadata_raw, dict): _metadata = _metadata_raw else: diff --git a/litellm/interactions/http_handler.py b/litellm/interactions/http_handler.py index 044c171653c..17ec4a3398d 100644 --- a/litellm/interactions/http_handler.py +++ b/litellm/interactions/http_handler.py @@ -4,7 +4,7 @@ HTTP Handler for Interactions API requests. This module handles the HTTP communication for the Google Interactions API. """ -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from typing import Any, Final import httpx @@ -96,8 +96,8 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): model: str | None = None, agent: str | None = None, input: InteractionInput | None = None, - extra_headers: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -105,7 +105,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): ) -> ( InteractionsAPIResponse | Iterator[InteractionsAPIStreamingResponse] - | Coroutine[Any, Any, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]] + | Coroutine[object, object, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]] ): """ Create a new interaction (synchronous or async based on _is_async flag). @@ -211,8 +211,8 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): model: str | None = None, agent: str | None = None, input: InteractionInput | None = None, - extra_headers: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, stream: bool | None = None, @@ -345,11 +345,11 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> InteractionsAPIResponse | Coroutine[Any, Any, InteractionsAPIResponse]: + ) -> InteractionsAPIResponse | Coroutine[object, object, InteractionsAPIResponse]: """Get an interaction by ID.""" if _is_async: return self.async_get_interaction( @@ -407,7 +407,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> InteractionsAPIResponse: @@ -464,11 +464,11 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> DeleteInteractionResult | Coroutine[Any, Any, DeleteInteractionResult]: + ) -> DeleteInteractionResult | Coroutine[object, object, DeleteInteractionResult]: """Delete an interaction by ID.""" if _is_async: return self.async_delete_interaction( @@ -527,7 +527,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> DeleteInteractionResult: @@ -585,11 +585,11 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, - ) -> CancelInteractionResult | Coroutine[Any, Any, CancelInteractionResult]: + ) -> CancelInteractionResult | Coroutine[object, object, CancelInteractionResult]: """Cancel an interaction by ID.""" if _is_async: return self.async_cancel_interaction( @@ -648,7 +648,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> CancelInteractionResult: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py index 14f1b7697cf..d0fed3225af 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py @@ -9,6 +9,7 @@ the LLM doesn't make a tool call, and we need to return a stream to the user. """ import json +from collections.abc import Mapping from typing import Any, Final, cast from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -38,7 +39,7 @@ class FakeAnthropicMessagesStreamIterator: self.chunks = self._create_streaming_chunks() self.current_index = 0 - def _create_content_block_chunks(self, block_dict: dict[str, Any], index: int) -> list[bytes]: + def _create_content_block_chunks(self, block_dict: Mapping[str, object], index: int) -> list[bytes]: """Build SSE chunks for a single content block.""" chunks: Final = [] block_type: Final = block_dict.get("type") @@ -133,14 +134,14 @@ class FakeAnthropicMessagesStreamIterator: response_dict: Final = cast(dict[str, Any], self.response) # 1. message_start event - usage: Final = response_dict.get("usage", {}) + usage: Final = self.response.get("usage") message_start: Final = { "type": "message_start", "message": { - "id": response_dict.get("id"), + "id": self.response.get("id"), "type": "message", - "role": response_dict.get("role", "assistant"), - "model": response_dict.get("model"), + "role": self.response.get("role", "assistant"), + "model": self.response.get("model"), "content": [], "stop_reason": None, "stop_sequence": None, @@ -161,21 +162,24 @@ class FakeAnthropicMessagesStreamIterator: # 5. message_delta event (with final usage and stop_reason) # Include cache usage fields so clients that only read message_delta # (like Claude Code's SDK) see the full input token breakdown. - delta_usage: Final[dict[str, Any]] = { + delta_usage: Final[dict[str, int]] = { "output_tokens": usage.get("output_tokens", 0) if usage else 0, } if usage: - if usage.get("input_tokens") is not None: - delta_usage["input_tokens"] = usage["input_tokens"] - if usage.get("cache_creation_input_tokens") is not None: - delta_usage["cache_creation_input_tokens"] = usage["cache_creation_input_tokens"] - if usage.get("cache_read_input_tokens") is not None: - delta_usage["cache_read_input_tokens"] = usage["cache_read_input_tokens"] + input_tokens: Final = usage.get("input_tokens") + if input_tokens is not None: + delta_usage["input_tokens"] = input_tokens + cache_creation_input_tokens: Final = usage.get("cache_creation_input_tokens") + if cache_creation_input_tokens is not None: + delta_usage["cache_creation_input_tokens"] = cache_creation_input_tokens + cache_read_input_tokens: Final = usage.get("cache_read_input_tokens") + if cache_read_input_tokens is not None: + delta_usage["cache_read_input_tokens"] = cache_read_input_tokens message_delta: Final = { "type": "message_delta", "delta": { - "stop_reason": response_dict.get("stop_reason"), - "stop_sequence": response_dict.get("stop_sequence"), + "stop_reason": self.response.get("stop_reason"), + "stop_sequence": self.response.get("stop_sequence"), }, "usage": delta_usage, } diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index fc34e403beb..c39c88240c5 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -163,7 +163,7 @@ async def make_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -) -> tuple[Any, httpx.Headers]: +) -> "tuple[MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict], httpx.Headers]": try: if client is None: client = get_async_httpx_client( @@ -199,7 +199,9 @@ async def make_call( messages=messages, encoding=litellm.encoding, ) - completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict] = ( + MockResponseIterator(model_response=model_response, json_mode=json_mode) + ) elif bedrock_invoke_provider == "anthropic": decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( model=model, @@ -248,7 +250,7 @@ def make_sync_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -) -> tuple[Any, httpx.Headers]: +) -> "tuple[MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict], httpx.Headers]": try: if client is None: client = _get_httpx_client( @@ -283,7 +285,9 @@ def make_sync_call( messages=messages, encoding=litellm.encoding, ) - completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict] = ( + MockResponseIterator(model_response=model_response, json_mode=json_mode) + ) elif bedrock_invoke_provider == "anthropic": decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( model=model, diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py index eb3f65bec94..bac899c4142 100644 --- a/litellm/llms/cohere/embed/transformation.py +++ b/litellm/llms/cohere/embed/transformation.py @@ -10,7 +10,8 @@ Convers Docs - https://docs.cohere.com/v2/reference/embed """ -from typing import Any, Final, cast +from collections.abc import Sized +from typing import Final, Protocol, cast import httpx @@ -30,6 +31,12 @@ from litellm.utils import is_base64_encoded from ..common_utils import CohereError +class _SupportsEncode(Protocol): + """Tokenizer handle: the embedding usage path only encodes text to measure its token length.""" + + def encode(self, text: str, /) -> Sized: ... + + class CohereEmbeddingConfig(BaseEmbeddingConfig): """ Reference: https://docs.cohere.com/v2/reference/embed @@ -133,7 +140,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): ), ) - def _calculate_usage(self, input: list[str], encoding: Any, meta: dict) -> Usage: + def _calculate_usage(self, input: list[str], encoding: _SupportsEncode, meta: dict) -> Usage: input_tokens = 0 text_tokens: Final[int | None] = meta.get("billed_units", {}).get("input_tokens") @@ -169,7 +176,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): data: dict | CohereEmbeddingRequest, model_response: EmbeddingResponse, model: str, - encoding: Any, + encoding: _SupportsEncode, input: list, ) -> EmbeddingResponse: response_json: Final = response.json() diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 3dd3996b2ee..490757a0948 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -148,7 +148,7 @@ class DashScopeRerankConfig(BaseRerankConfig): if "documents" not in optional_rerank_params: raise ValueError("documents is required for DashScope rerank") - request: Final[dict[str, Any]] = { + request: Final[dict[str, object]] = { "model": model, "query": optional_rerank_params["query"], "documents": optional_rerank_params["documents"], @@ -209,7 +209,7 @@ class DashScopeRerankConfig(BaseRerankConfig): # which already matches LiteLLM's RerankResponseDocument shape. transformed_results: Final[list[dict]] = [] for r in results: - item: dict[str, Any] = { + item: dict[str, object] = { "index": r["index"], "relevance_score": r["relevance_score"], } diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index eedffd844ef..fcd4ae70645 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -4,7 +4,7 @@ Calls DataForSEO SERP API to search the web. DataForSEO API Reference: https://docs.dataforseo.com/v3/serp/google/organic/live/advanced/?bash """ -from typing import Any, Final, Literal +from typing import Final, Literal import httpx @@ -126,7 +126,7 @@ class DataForSEOSearchConfig(BaseSearchConfig): List[Dict]: Request body for DataForSEO API (array of task objects as required by API) """ # DataForSEO expects an array of task objects - task: Final[dict[str, Any]] = {} + task: Final[dict[str, object]] = {} # Convert query to string if it's a list if isinstance(query, list): diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index 3439f4872c3..3cf9a983efe 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -80,8 +80,8 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): def _resolve_voice_id( self, - voice: str | dict[str, Any] | None, - params: dict[str, Any], + voice: str | dict[str, object] | None, + params: dict[str, object], ) -> str: """ Determine the ElevenLabs voice_id based on provided voice input or parameters. @@ -115,17 +115,17 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): optional_params: dict, voice: str | dict | None = None, drop_params: bool = False, - kwargs: dict[str, Any] | None = None, + kwargs: dict[str, object] | None = None, ) -> tuple[str | None, dict]: """ Map OpenAI parameters to ElevenLabs TTS parameters """ - mapped_params: Final[dict[str, Any]] = {} - query_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} + query_params: Final[dict[str, object]] = {} # Work on a copy so we don't mutate the caller's dictionary params: Final = dict(optional_params) if optional_params else {} - passthrough_kwargs: Final[dict[str, Any]] = kwargs if kwargs is not None else {} + passthrough_kwargs: Final[dict[str, object]] = kwargs if kwargs is not None else {} # Extract voice identifier mapped_voice: Final = self._resolve_voice_id(voice, params) @@ -205,7 +205,7 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): params: Final = dict(optional_params) if optional_params else {} extra_body: Final = params.pop("extra_body", None) - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "text": input, "model_id": model, } @@ -229,10 +229,10 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): def _add_elevenlabs_specific_params( self, mapped_voice: str, - query_params: dict[str, Any], - mapped_params: dict[str, Any], - kwargs: dict[str, Any] | None, - remaining_params: dict[str, Any], + query_params: dict[str, object], + mapped_params: dict[str, object], + kwargs: dict[str, object] | None, + remaining_params: dict[str, object], ) -> None: if kwargs is None: kwargs = {} diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 8ef2c9acccb..e142622aa1b 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -67,11 +67,11 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): max_chunks_per_doc: int | None = None, max_tokens_per_doc: int | None = None, instruction: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Map Cohere rerank params to Fireworks AI rerank params """ - params: Final[dict[str, Any]] = { + params: Final[dict[str, object]] = { "query": query, "documents": documents, } diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index 1920cd698f5..cb2be2c860e 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -58,9 +58,9 @@ class GoogleAIStudioTokenCounter: self, api_base: str | None = None, api_key: str | None = None, - headers: dict[str, Any] | None = None, + headers: dict[str, object] | None = None, model: str = "", - litellm_params: dict[str, Any] | None = None, + litellm_params: dict[str, object] | None = None, ) -> tuple[dict[str, Any], str]: """ Returns a Tuple of headers and url for the Google Gen AI Studio countTokens endpoint. diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index 163e944f124..359553e144f 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -50,13 +50,21 @@ def _parse_data_url(data_url: str) -> tuple[bytes, str, str] | None: return content_bytes, content_type, ext +def _content_type_or_default(headers: Mapping[str, str]) -> str: + """Return the response's ``content-type`` header, falling back to ``image/jpeg`` when absent.""" + try: + return headers["content-type"] + except KeyError: + return "image/jpeg" + + def _download_image_sync(url: str) -> tuple[bytes, str, str]: """Download image from URL synchronously.""" client: Final = _get_httpx_client(params={"ssl_verify": False}) response: Final = client.get(url) response.raise_for_status() - content_type: Final = response.headers.get("content-type", "image/jpeg") + content_type: Final = _content_type_or_default(response.headers) ext: Final = content_type.split("/")[-1].split(";")[0] or "jpg" return response.content, content_type, ext @@ -71,7 +79,7 @@ async def _download_image_async(url: str) -> tuple[bytes, str, str]: response: Final = await client.get(url) response.raise_for_status() - content_type: Final = response.headers.get("content-type", "image/jpeg") + content_type: Final = _content_type_or_default(response.headers) ext: Final = content_type.split("/")[-1].split(";")[0] or "jpg" return response.content, content_type, ext diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 12c070b3461..57d1357ee46 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -1,7 +1,7 @@ import json import os -from collections.abc import Callable -from typing import Any, Final, Literal, get_args +from collections.abc import Sequence +from typing import Final, Literal, Protocol, get_args import httpx @@ -29,6 +29,12 @@ hf_tasks_embeddings: Final = ( ) +class _SupportsTokenEncode(Protocol): + """Token encoder handle. Only ``encode`` is ever called on it here.""" + + def encode(self, text: str, *, disallowed_special: tuple[str, ...]) -> Sequence[int]: ... + + def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> str | None: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): @@ -173,7 +179,7 @@ class HuggingFaceEmbedding(BaseLLM): model_response: EmbeddingResponse, model: str, input: list, - encoding: Any, + encoding: _SupportsTokenEncode, ) -> EmbeddingResponse: output_data: Final = [] if "similarities" in embeddings: @@ -234,7 +240,7 @@ class HuggingFaceEmbedding(BaseLLM): api_base: str, api_key: str | None, headers: dict, - encoding: Callable, + encoding: _SupportsTokenEncode, client: AsyncHTTPHandler | None = None, ): ## TRANSFORMATION ## @@ -294,7 +300,7 @@ class HuggingFaceEmbedding(BaseLLM): optional_params: dict, litellm_params: dict, logging_obj: LiteLLMLoggingObj, - encoding: Callable, + encoding: _SupportsTokenEncode, api_key: str | None = None, api_base: str | None = None, timeout: float | httpx.Timeout = httpx.Timeout(None), diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index f8926df1f3f..e38a8a2c3a3 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -123,7 +123,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): optional_params: dict, voice: str | dict | None = None, drop_params: bool = False, - kwargs: dict[str, Any] | None = None, + kwargs: Mapping[str, object] | None = None, ) -> tuple[str | None, dict]: """ Map OpenAI parameters to MiniMax TTS parameters diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index f6c093f2e2a..125e5168c69 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -98,7 +98,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = f"{api_base}/{encoded_vector_store_id}/search" diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index f07acf2f728..1f295a6e656 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -6,6 +6,7 @@ It uses the field targeting configuration from litellm_logging_obj to extract specific fields for guardrail processing. """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional from litellm._logging import verbose_proxy_logger @@ -89,7 +90,7 @@ class PassThroughEndpointHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, - ) -> Any: + ) -> Mapping[str, object]: """ Process input by applying guardrails to targeted fields or full payload. """ @@ -130,9 +131,9 @@ class PassThroughEndpointHandler(BaseTranslation): response: object, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, - ) -> Any: + ) -> object: """ Process output response by applying guardrails to targeted fields. @@ -239,9 +240,9 @@ class LlmPassthroughRouteHandler(BaseTranslation): response: object, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, - ) -> Any: + ) -> object: provider: Final = (request_data or {}).get("custom_llm_provider") handler_cls: Final = _get_provider_handlers().get(provider or "") if handler_cls is None: diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 332f892ae6b..d7b4ad22a01 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -29,6 +29,7 @@ from litellm.types.llms.vertex_ai_text_to_speech import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.openai import HttpxBinaryResponseContent else: LiteLLMLoggingObj = Any @@ -131,19 +132,19 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): model: str, input: str, voice: str | dict | None, - optional_params: dict, - litellm_params_dict: dict, + optional_params: dict[str, object], + litellm_params_dict: dict[str, object], logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout, - extra_headers: dict[str, Any] | None, - base_llm_http_handler: Any, + extra_headers: dict[str, object] | None, + base_llm_http_handler: "BaseLLMHTTPHandler", aspeech: bool, api_base: str | None, api_key: str | None, - **kwargs: Any, + **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[object, object, "HttpxBinaryResponseContent"], ]: """ Dispatch method to handle Vertex AI TTS requests @@ -227,7 +228,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Returns: Tuple of (mapped_voice_str, mapped_params) """ - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} ########################################################## # Map voice using helper @@ -428,7 +429,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): speakingRate=speaking_rate, ) - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "input": dict(vertex_input), "voice": dict(vertex_voice), "audioConfig": dict(vertex_audio_config), diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 2fad9f933c1..4da31c82b57 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -1,8 +1,8 @@ import sys import time import webbrowser -from collections.abc import Callable, Mapping -from typing import Any, Final +from collections.abc import Callable, Mapping, Sequence +from typing import Any, Final, TypeVar from urllib.parse import urlencode import click @@ -112,6 +112,8 @@ class CliAuthResult(TypedDict): team_id: str | None +_TeamMapping: Final = TypeVar("_TeamMapping", bound=Mapping[str, object]) + KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" KEYRING_ENABLE_HINT: Final = "keyring --enable (or unset PYTHON_KEYRING_BACKEND)" @@ -353,7 +355,7 @@ def get_key_input(): return None -def display_interactive_team_selection(teams: list[dict[str, Any]], selected_index: int = 0) -> None: +def display_interactive_team_selection(teams: Sequence[Mapping[str, Any]], selected_index: int = 0) -> None: """Display teams with one highlighted for selection""" console: Final = Console() @@ -391,7 +393,7 @@ def display_interactive_team_selection(teams: list[dict[str, Any]], selected_ind console.print(f" Budget: [dim]{budget_str}[/dim]\n") -def prompt_team_selection(teams: list[dict[str, Any]]) -> dict[str, Any] | None: +def prompt_team_selection(teams: Sequence[_TeamMapping]) -> _TeamMapping | None: """Interactive team selection with arrow keys""" if not teams: return None @@ -441,8 +443,8 @@ def prompt_team_selection(teams: list[dict[str, Any]]) -> dict[str, Any] | None: def prompt_team_selection_fallback( - teams: list[dict[str, Any]], -) -> dict[str, Any] | None: + teams: Sequence[_TeamMapping], +) -> _TeamMapping | None: """Fallback team selection for non-interactive environments""" if not teams: return None diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index 4c83a7b799a..f2b38c6eab4 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -1,17 +1,32 @@ # stdlib imports import re from collections import defaultdict +from collections.abc import Callable from dataclasses import dataclass from datetime import datetime -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal # third party imports import click import rich import yaml +from typing_extensions import NotRequired, ReadOnly, TypedDict # local imports from ... import Client +from ._cli_context import cli_context_values + +if TYPE_CHECKING: + from rich.console import JustifyMethod + + +class _ModelInfoColumnConfig(TypedDict): + """Rendering config for one column of the ``models info`` table.""" + + header: ReadOnly[str] + style: ReadOnly[str] + justify: NotRequired[ReadOnly["JustifyMethod"]] + get_value: ReadOnly[Callable[..., str]] @dataclass @@ -84,7 +99,8 @@ def format_cost_per_1k_tokens(cost: float | None) -> str: def create_client(ctx: click.Context) -> Client: """Helper function to create a client from context.""" - return Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + return Client(base_url=context["base_url"], api_key=context["api_key"]) @click.group() @@ -216,7 +232,7 @@ def get_models_info(ctx: click.Context, output_format: Literal["table", "json"], table: Final = rich.table.Table(title="Models Information") # Define all possible columns with their configurations - column_configs: Final[dict[str, dict[str, Any]]] = { + column_configs: Final[dict[str, _ModelInfoColumnConfig]] = { "public_model": { "header": "Public Model", "style": "cyan", diff --git a/litellm/proxy/client/cli/commands/users.py b/litellm/proxy/client/cli/commands/users.py index 2cfba5ec357..a5ebefd1c4a 100644 --- a/litellm/proxy/client/cli/commands/users.py +++ b/litellm/proxy/client/cli/commands/users.py @@ -4,6 +4,7 @@ import click import rich from ... import UsersManagementClient +from ._cli_context import cli_context_values @click.group() @@ -15,7 +16,8 @@ def users(): @click.pass_context def list_users(ctx: click.Context): """List all users""" - client: Final = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = UsersManagementClient(base_url=context["base_url"], api_key=context["api_key"]) users = client.list_users() if isinstance(users, dict) and "users" in users: users = users["users"] @@ -46,7 +48,8 @@ def list_users(ctx: click.Context): @click.pass_context def get_user(ctx: click.Context, user_id: str): """Get information about a specific user""" - client: Final = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = UsersManagementClient(base_url=context["base_url"], api_key=context["api_key"]) result: Final = client.get_user(user_id=user_id) rich.print_json(data=result) @@ -60,7 +63,8 @@ def get_user(ctx: click.Context, user_id: str): @click.pass_context def create_user(ctx: click.Context, email, role, alias, team, max_budget): """Create a new user""" - client: Final = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = UsersManagementClient(base_url=context["base_url"], api_key=context["api_key"]) user_data: Final = { "user_email": email, "user_role": role, @@ -80,6 +84,7 @@ def create_user(ctx: click.Context, email, role, alias, team, max_budget): @click.pass_context def delete_user(ctx: click.Context, user_ids): """Delete one or more users by user_id""" - client: Final = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = UsersManagementClient(base_url=context["base_url"], api_key=context["api_key"]) result: Final = client.delete_user(list(user_ids)) rich.print_json(data=result) diff --git a/litellm/proxy/client/http_client.py b/litellm/proxy/client/http_client.py index 18344f267b9..aa0b986b1ad 100644 --- a/litellm/proxy/client/http_client.py +++ b/litellm/proxy/client/http_client.py @@ -1,5 +1,6 @@ """HTTP client for making requests to the LiteLLM proxy server.""" +from collections.abc import Mapping from typing import Any, Final import requests @@ -25,8 +26,8 @@ class HTTPClient: method: str, uri: str, *, - data: dict[str, Any] | list | bytes | None = None, - json: dict[str, Any] | list | None = None, + data: Mapping[str, object] | list | bytes | None = None, + json: Mapping[str, object] | list | None = None, headers: dict[str, str] | None = None, **kwargs: Any, ) -> Any: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 4b16087e15b..10626f95e49 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -1,4 +1,5 @@ import builtins +from collections.abc import Mapping from typing import Any, Final import requests @@ -68,8 +69,8 @@ class ModelsManagementClient: def new( self, model_name: str, - model_params: dict[str, Any], - model_info: dict[str, Any] | None = None, + model_params: Mapping[str, object], + model_info: Mapping[str, object] | None = None, return_request: bool = False, ) -> dict[str, Any] | requests.Request: """ @@ -245,8 +246,8 @@ class ModelsManagementClient: def update( self, model_id: str, - model_params: dict[str, Any], - model_info: dict[str, Any] | None = None, + model_params: Mapping[str, object], + model_info: Mapping[str, object] | None = None, return_request: bool = False, ) -> dict[str, Any] | requests.Request: """ diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 892ff9771cf..95642bc74bc 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -207,7 +207,7 @@ async def _process_binary_request( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - content: Final = await processor.base_process_llm_request( + content: Final[object] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -268,7 +268,7 @@ async def _process_multipart_upload_request( user_api_key_dict: UserAPIKeyAuth, route_type: str, container_id: str, -): +) -> object: """Process multipart file upload requests.""" from litellm.proxy.common_utils.http_parsing_utils import ( convert_upload_files_to_file_data, @@ -357,7 +357,7 @@ async def _process_request( user_api_key_dict: UserAPIKeyAuth, route_type: str, path_params: dict[str, str], -): +) -> object: """Common request processing logic.""" from litellm.proxy.proxy_server import ( general_settings, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 958f84e18de..9c635128510 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -299,7 +299,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return data if action_type == "monitor_action": verbose_proxy_logger.info("Cato: monitor action") - elif action_type == "block_action": + elif action_type == "block_action" and required_action is not None: self._handle_block_action(res.get("analysis_result", {}), required_action) elif action_type == "anonymize_action": return self._anonymize_request(res, data) @@ -310,7 +310,7 @@ class CatoNetworksGuardrail(CustomGuardrail): def _handle_block_action( self, analysis_result: _CatoAnalysisResult, - required_action: Any, + required_action: _CatoRequiredAction, ) -> None: detection_message: Final = required_action.get("detection_message", None) verbose_proxy_logger.info( @@ -410,7 +410,7 @@ class CatoNetworksGuardrail(CustomGuardrail): res: Final[_CatoAnalyzeResponse] = response.json() required_action: Final = res.get("required_action") action_type: Final = required_action and required_action.get("action_type", None) - if action_type and action_type == "block_action": + if action_type == "block_action" and required_action is not None: self._handle_block_action_on_output(res.get("analysis_result", {}), required_action) redacted_chat: Final = res.get("redacted_chat", None) @@ -425,7 +425,7 @@ class CatoNetworksGuardrail(CustomGuardrail): def _handle_block_action_on_output( self, analysis_result: _CatoAnalysisResult, - required_action: Any, + required_action: _CatoRequiredAction, ) -> None: detection_message: Final = required_action.get("detection_message", None) verbose_proxy_logger.info( diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index 694db182fe7..bc419b359c1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -6,7 +6,7 @@ # +-------------------------------------------------------------+ import os -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterable from datetime import datetime from typing import Any, Final @@ -188,7 +188,7 @@ class DynamoAIGuardrails(CustomGuardrail): applied_policies: Final = response.get("appliedPolicies", []) violations_detected: Final[list[str]] = [] - violation_details: Final[dict[str, Any]] = {} + violation_details: Final[dict[str, object]] = {} # For now, only handle BLOCK action if final_action == "BLOCK": @@ -404,7 +404,7 @@ class DynamoAIGuardrails(CustomGuardrail): # to avoid sending empty content to DynamoAI (e.g., during tool calls) if isinstance(response, litellm.ModelResponse): has_text_content = False - dynamoai_messages: Final[list[dict[str, Any]]] = [] + dynamoai_messages: Final[list[dict[str, str]]] = [] for choice in response.choices: if isinstance(choice, litellm.Choices): @@ -446,7 +446,7 @@ class DynamoAIGuardrails(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[ModelResponseStream], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 07340e95835..5109f09d9c2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -4,6 +4,7 @@ from urllib.parse import urlparse import httpx import pydantic +from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException @@ -34,6 +35,10 @@ _GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm" _DEFAULT_TIMEOUT: Final = 30.0 +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + class SingulrGuardrail(CustomGuardrail): def __init__( self, @@ -43,7 +48,7 @@ class SingulrGuardrail(CustomGuardrail): singulr_guardrail_id: str | None = None, block_on_error: bool | None = None, timeout: float | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.singulr_api_key = singulr_api_key or os.environ.get("SINGULR_API_KEY") self.singulr_api_base = (singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).rstrip( diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index cd983801c34..f54e4bc30f7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -20,9 +20,10 @@ Configuration in proxy config YAML: mode: post_call """ -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException +from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -39,6 +40,10 @@ if TYPE_CHECKING: GUARDRAIL_NAME: Final = "tool_policy" +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + def _get_request_object_permission_ids( request_data: dict, ) -> tuple[str | None, str | None]: @@ -106,7 +111,7 @@ class ToolPolicyGuardrail(CustomGuardrail): ToolPolicyRegistry (synced from DB). """ - def __init__(self, **kwargs: Any) -> None: + def __init__(self, **kwargs: Unpack[_CustomGuardrailOptions]) -> None: if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ GuardrailEventHooks.pre_call, diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 7a0edbddca8..9490eda9d47 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -355,6 +355,11 @@ def _to_dict(value: object) -> dict[str, Any]: return {} +def _field_str(mapping: Mapping[str, object], key: str, default: str) -> str: + """Stringify `mapping[key]`, falling back to `default` when the key is absent.""" + return str(mapping.get(key, default)) + + def _get_guardrail_attrs(g: "_DbOrConfigGuardrail") -> tuple[Any, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" gid: Final = _get_guardrail_field(g, "guardrail_id") @@ -383,9 +388,9 @@ def _guardrail_overview_rows( req, blocked = a["requests"], a["blocked"] fail_rate = (100.0 * blocked / req) if req else 0.0 litellm_params = _to_dict(_get_guardrail_field(g, "litellm_params")) - provider = str(litellm_params.get("guardrail", "Unknown")) + provider = _field_str(litellm_params, "guardrail", "Unknown") guardrail_info = _to_dict(_get_guardrail_field(g, "guardrail_info")) - gtype = str(guardrail_info.get("type", "Guardrail")) + gtype = _field_str(guardrail_info, "type", "Guardrail") prev_fail = 0.0 for k in lookup_keys: if k in prev_agg: @@ -624,8 +629,8 @@ async def guardrails_usage_detail( return UsageDetailResponse( guardrail_id=guardrail_id, guardrail_name=_guardrail_name or guardrail_id, - type=str(guardrail_info.get("type", "Guardrail")), - provider=str(litellm_params.get("guardrail", "Unknown")), + type=_field_str(guardrail_info, "type", "Guardrail"), + provider=_field_str(litellm_params, "guardrail", "Unknown"), requestsEvaluated=requests, failRate=round(fail_rate, 1), avgScore=None, diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 9d5ddda017a..4fcc798f93c 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -6,7 +6,7 @@ usage/spend data by querying the aggregated daily activity endpoints. import json from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import date -from typing import Any, Final, Literal, Protocol, cast, overload +from typing import Any, Final, Literal, NamedTuple, Protocol, cast, overload from typing_extensions import ReadOnly, TypedDict @@ -82,6 +82,15 @@ class _DayDump(TypedDict, total=False): breakdown: ReadOnly[Mapping[str, Mapping[str, _EntityEntry]]] +class _EntityTotal(NamedTuple): + """Running per-entity totals accumulated while summarising a usage dump.""" + + alias: str + spend: float + requests: float + tokens: float + + class _UsageDump(Protocol): @overload def get(self, key: Literal["metadata"], default: Mapping[str, float], /) -> Mapping[str, float]: ... @@ -241,7 +250,7 @@ def _parse_csv_ids(raw: str | None) -> list[str] | None: async def _query_activity( table_name: str, entity_id_field: str, - entity_id: Any | None, + entity_id: str | list[str] | None, start_date: str, end_date: str, *, @@ -382,23 +391,22 @@ def _summarise_entity_data(data: _UsageDump, entity_label: str) -> str: if not results: return f"No {entity_label} usage data found for the given date range." - totals: Final[dict[str, dict[str, Any]]] = {} + totals: Final[dict[str, _EntityTotal]] = {} for day in results: for eid, entry in day.get("breakdown", {}).get("entities", {}).items(): - if eid not in totals: - alias = entry.get("metadata", {}).get("alias", eid) - totals[eid] = {"alias": alias, "spend": 0.0, "requests": 0, "tokens": 0} + previous = totals.get(eid) m = entry.get("metrics", {}) - totals[eid]["spend"] += m.get("spend", 0) - totals[eid]["requests"] += m.get("api_requests", 0) - totals[eid]["tokens"] += m.get("total_tokens", 0) + totals[eid] = _EntityTotal( + alias=previous.alias if previous is not None else entry.get("metadata", {}).get("alias", eid), + spend=(previous.spend if previous is not None else 0.0) + m.get("spend", 0), + requests=(previous.requests if previous is not None else 0) + m.get("api_requests", 0), + tokens=(previous.tokens if previous is not None else 0) + m.get("total_tokens", 0), + ) lines: Final = [f"{entity_label} Usage ({len(totals)} {entity_label.lower()}s):", ""] - for eid, d in sorted(totals.items(), key=lambda x: -x[1]["spend"]): - label = d["alias"] if d["alias"] != eid else eid - lines.append( - f"- {label} (ID: {eid}): ${d['spend']:.4f} | {int(d['requests'])} reqs | {int(d['tokens'])} tokens" - ) + for eid, d in sorted(totals.items(), key=lambda x: -x[1].spend): + label = d.alias if d.alias != eid else eid + lines.append(f"- {label} (ID: {eid}): ${d.spend:.4f} | {int(d.requests)} reqs | {int(d.tokens)} tokens") return "\n".join(lines) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index aa229270800..f6c872d3b92 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -42,6 +42,8 @@ from ..llms.xai.realtime.handler import XAIRealtime from ..utils import client as wrapper_client if TYPE_CHECKING: + from fastapi import WebSocket + from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig azure_realtime: Final = AzureOpenAIRealtime() @@ -332,12 +334,12 @@ async def _resolve_vertex_access_token_bounded( @wrapper_client async def _arealtime( model: str, - websocket: Any, # fastapi websocket + websocket: "WebSocket", # fastapi websocket api_base: str | None = None, api_key: str | None = None, api_version: str | None = None, azure_ad_token: str | None = None, - client: Any | None = None, + client: object | None = None, timeout: float | None = None, query_params: RealtimeQueryParams | None = None, **kwargs, @@ -574,7 +576,7 @@ _TRANSCRIPTION_QUERY_PARAMS: Final[RealtimeQueryParams] = {"intent": "transcript def _azure_realtime_health_protocol( - model: str, realtime_protocol: str | None, model_params: Mapping[str, Any] + model: str, realtime_protocol: str | None, model_params: Mapping[str, object] ) -> tuple[str, RealtimeQueryParams | None]: query_params: Final = _TRANSCRIPTION_QUERY_PARAMS if _is_transcription_only_realtime_model(model, "azure") else None configured_raw: Final = ( diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 39675faf735..3ca7b0503bf 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,7 +1,7 @@ import base64 import re from collections.abc import Iterable, Mapping, Sequence -from typing import Any, Final, Optional, Union, cast, get_type_hints, overload +from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload from pydantic import BaseModel from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire payloads without a runtime conversion @@ -59,6 +59,9 @@ def _as_input_text_part(part: object) -> object: return part +_RequestInputT: Final = TypeVar("_RequestInputT") + + class ResponsesAPIRequestUtils: """Helper utils for constructing ResponseAPI requests""" @@ -502,7 +505,7 @@ class ResponsesAPIRequestUtils: return response @staticmethod - def _restore_encrypted_content_item_ids_in_input(request_input: object) -> Any: + def _restore_encrypted_content_item_ids_in_input(request_input: _RequestInputT) -> _RequestInputT: """Decode litellm-encoded item IDs in request input back to original IDs. Called before forwarding the request to the upstream provider so the diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index 310a7717b38..7b69714aad9 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -13,6 +13,7 @@ bounded list of recent tool call signatures. from __future__ import annotations import re +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Any, Final @@ -92,7 +93,7 @@ class Turn: user_content: str | None = None assistant_content: str | None = None tool_calls: list[dict[str, Any]] = field(default_factory=list) - tool_results: list[dict[str, Any]] = field(default_factory=list) + tool_results: Sequence[Mapping[str, object]] = field(default_factory=list[Mapping[str, object]]) response_status: int | None = None @@ -104,7 +105,7 @@ _TOKEN_RE: Final = re.compile(r"[A-Za-z0-9]+") def _tokens(text: str | None) -> set[str]: if not text: return set() - return {t.lower() for t in _TOKEN_RE.findall(text)} + return {match.group(0).lower() for match in _TOKEN_RE.finditer(text)} def _jaccard(a: set[str], b: set[str]) -> float: @@ -160,7 +161,7 @@ def _detect_satisfaction(curr_user: str | None) -> bool: return any(p.search(curr_user) for p in _SATISFACTION_PATTERNS) -def _detect_failure(tool_results: list[dict[str, Any]]) -> bool: +def _detect_failure(tool_results: Sequence[Mapping[str, object]]) -> bool: """Any tool result explicitly flagged as an error. We do NOT treat empty content as failure — many tools legitimately return @@ -209,7 +210,7 @@ _EXHAUSTION_KEYWORDS: Final = ( ) -def _detect_exhaustion(status: int | None, tool_results: list[dict[str, Any]]) -> bool: +def _detect_exhaustion(status: int | None, tool_results: Sequence[Mapping[str, object]]) -> bool: if status is not None and status in _EXHAUSTION_STATUSES: return True for r in tool_results: @@ -222,7 +223,7 @@ def _detect_exhaustion(status: int | None, tool_results: list[dict[str, Any]]) - def detect_user_feedback( previous_user_content: str | None, current_user_content: str | None, - tool_results: list[dict[str, Any]], + tool_results: Sequence[Mapping[str, object]], allow_satisfaction: bool, ) -> SignalDelta: return SignalDelta( @@ -238,7 +239,7 @@ def detect_response_signals( current_assistant_content: str | None, tool_call_history: list[str], tool_calls: list[dict[str, Any]], - tool_results: list[dict[str, Any]], + tool_results: Sequence[Mapping[str, object]], response_status: int | None, ) -> SignalDelta: return SignalDelta( diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 86d9bb5c3ed..4534fa114b3 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -259,7 +259,7 @@ def _should_run_cooldown_logic( litellm_router_instance: LitellmRouter, deployment: str | None, exception_status: str | int, - original_exception: Any, + original_exception: Exception, time_to_cooldown: float | None = None, ) -> bool: """ @@ -318,7 +318,7 @@ def _should_cooldown_deployment( litellm_router_instance: LitellmRouter, deployment: str, exception_status: str | int, - original_exception: Any, + original_exception: Exception, requested_model_group: str | None = None, ) -> bool: """ @@ -412,7 +412,7 @@ def _should_cooldown_deployment( def _set_cooldown_deployments( litellm_router_instance: LitellmRouter, - original_exception: Any, + original_exception: Exception, exception_status: str | int, deployment: str | None = None, time_to_cooldown: float | None = None, @@ -547,7 +547,7 @@ def _get_cooldown_deployments(litellm_router_instance: LitellmRouter, parent_ote def should_cooldown_based_on_allowed_fails_policy( litellm_router_instance: LitellmRouter, deployment: str, - original_exception: Any, + original_exception: Exception, allowed_fails_override: int | None = None, cooldown_time_override: float | None = None, cache_key_suffix: str | None = None, diff --git a/litellm/skills/main.py b/litellm/skills/main.py index 9d2ed524ce5..002419dbad4 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -72,7 +72,7 @@ def _get_litellm_skills_handler(): async def acreate_skill( files: list[Any] | None = None, display_title: str | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, @@ -135,7 +135,7 @@ async def acreate_skill( def create_skill( files: list[Any] | None = None, display_title: str | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, @@ -262,7 +262,7 @@ async def alist_skills( limit: int | None = None, page: str | None = None, source: str | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -325,7 +325,7 @@ def list_skills( limit: int | None = None, page: str | None = None, source: str | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -443,7 +443,7 @@ def list_skills( @client async def aget_skill( skill_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -500,7 +500,7 @@ async def aget_skill( @client def get_skill( skill_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -607,7 +607,7 @@ def get_skill( @client async def adelete_skill( skill_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -664,7 +664,7 @@ async def adelete_skill( @client def delete_skill( skill_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, diff --git a/litellm/vector_store_files/main.py b/litellm/vector_store_files/main.py index 5bc3c8f1525..3b6f1de3c7a 100644 --- a/litellm/vector_store_files/main.py +++ b/litellm/vector_store_files/main.py @@ -2,7 +2,7 @@ import asyncio import contextvars -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import Any, Final @@ -57,9 +57,9 @@ async def acreate( vector_store_id: str, file_id: str, attributes: VectorStoreFileAttributes | None = None, - chunking_strategy: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + chunking_strategy: Mapping[str, object] | None = None, + extra_headers: dict[str, str] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -109,9 +109,9 @@ def create( vector_store_id: str, file_id: str, attributes: VectorStoreFileAttributes | None = None, - chunking_strategy: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + chunking_strategy: Mapping[str, object] | None = None, + extra_headers: dict[str, str] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -187,7 +187,7 @@ async def alist( filter: str | None = None, limit: int | None = None, order: str | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -240,7 +240,7 @@ def list( filter: str | None = None, limit: int | None = None, order: str | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_query: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -308,7 +308,7 @@ async def aretrieve( *, vector_store_id: str, file_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -351,7 +351,7 @@ def retrieve( *, vector_store_id: str, file_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -417,7 +417,7 @@ async def aretrieve_content( *, vector_store_id: str, file_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -459,7 +459,7 @@ def retrieve_content( *, vector_store_id: str, file_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -526,7 +526,7 @@ async def aupdate( vector_store_id: str, file_id: str, attributes: VectorStoreFileAttributes, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -572,7 +572,7 @@ def update( vector_store_id: str, file_id: str, attributes: VectorStoreFileAttributes, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, @@ -646,7 +646,7 @@ async def adelete( *, vector_store_id: str, file_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -688,7 +688,7 @@ def delete( *, vector_store_id: str, file_id: str, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, From 459858829e8a01df74c1aee1372f8447967fc43c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:08:42 +0000 Subject: [PATCH 03/38] refactor(typing): replace Any with proven types in 89 more backend files --- .../proxy/hooks/managed_vector_stores.py | 12 +++-- litellm/_redis_credential_provider.py | 17 ++++++- litellm/_service_logger.py | 46 +++++++++++++++---- litellm/a2a_protocol/card_resolver.py | 3 +- .../watsonx_orchestrate/transformation.py | 14 +++--- litellm/assistants/utils.py | 45 ++++++++++-------- litellm/batches/batch_utils.py | 12 ++--- litellm/compression/compress.py | 12 ++--- litellm/containers/endpoint_factory.py | 14 +++--- litellm/exceptions.py | 4 +- litellm/fine_tuning/main.py | 14 +++--- litellm/images/utils.py | 3 +- .../datadog/datadog_cost_management.py | 5 +- .../dotprompt/dotprompt_manager.py | 7 +-- litellm/integrations/focus/focus_logger.py | 4 +- .../generic_prompt_manager.py | 3 +- litellm/integrations/humanloop.py | 6 +-- .../opentelemetry_utils/gen_ai_semconv.py | 6 +-- .../opik_payload_builder/payload_builders.py | 8 ++-- litellm/integrations/weave/weave_otel.py | 5 +- .../dot_notation_indexing.py | 17 ++++--- .../json_validation_rule.py | 6 +-- litellm/litellm_core_utils/logging_utils.py | 2 +- litellm/litellm_core_utils/safe_json_dumps.py | 2 +- litellm/llms/a2a/common_utils.py | 3 +- .../messages/interceptors/advisor.py | 10 ++-- .../responses_adapters/streaming_iterator.py | 10 ++-- .../llms/anthropic/files/transformation.py | 4 +- .../text_to_speech/transformation.py | 13 +++--- litellm/llms/azure/realtime/handler.py | 14 ++++-- .../llms/azure/responses/transformation.py | 6 +-- .../anthropic/count_tokens/token_counter.py | 8 ++-- .../llms/base_llm/agents/transformation.py | 24 +++++----- .../base_llm/guardrail_translation/utils.py | 12 ++--- .../vector_store_files/transformation.py | 23 +++++----- litellm/llms/bedrock/base_aws_llm.py | 22 +++++++-- .../llms/custom_httpx/container_handler.py | 6 +-- .../gemini/google_genai/transformation.py | 6 +-- litellm/llms/jina_ai/rerank/transformation.py | 6 +-- litellm/llms/litellm_proxy/skills/handler.py | 30 ++++++------ .../litellm_proxy/skills/sandbox_executor.py | 39 ++++++++++++++-- litellm/llms/openai/fine_tuning/handler.py | 17 +++---- .../llms/openai/image_variations/handler.py | 4 +- litellm/llms/openai/realtime/handler.py | 9 ++-- .../responses/count_tokens/token_counter.py | 8 ++-- .../vector_store_files/transformation.py | 25 +++++----- litellm/llms/predibase/chat/transformation.py | 25 ++++++++-- .../audio_transcription/transformation.py | 6 +-- .../llms/vertex_ai/rag_engine/ingestion.py | 8 ++-- .../llms/vertex_ai/videos/transformation.py | 4 +- .../embedding/transformation_multimodal.py | 6 +-- litellm/llms/voyage/rerank/transformation.py | 4 +- litellm/llms/xai/chat/transformation.py | 2 +- litellm/llms/xai/responses/transformation.py | 14 +++--- litellm/proxy/caching_routes.py | 10 ++-- litellm/proxy/client/chat.py | 4 +- litellm/proxy/client/cli/commands/agents.py | 10 ++-- litellm/proxy/client/keys.py | 17 +++---- .../proxy/common_utils/performance_utils.py | 20 ++++++-- .../proxy/container_endpoints/endpoints.py | 8 ++-- litellm/proxy/db/exception_handler.py | 4 +- .../guardrails/guardrail_hooks/azure/base.py | 4 +- .../guardrail_hooks/azure/text_moderation.py | 18 ++++---- .../block_code_execution/__init__.py | 6 +-- .../guardrail_hooks/custom_code/primitives.py | 6 +-- .../generic_guardrail_api.py | 10 ++-- .../model_armor/model_armor.py | 5 +- .../guardrails/guardrail_hooks/noma/noma.py | 4 +- .../guardrail_hooks/pangea/pangea.py | 6 +-- .../panw_prisma_airs/panw_prisma_airs.py | 2 +- litellm/proxy/guardrails/usage_tracking.py | 16 ++++--- .../shared_health_check_manager.py | 9 ++-- litellm/proxy/hooks/batch_rate_limiter.py | 8 ++-- .../proxy/hooks/key_management_event_hooks.py | 6 +-- .../management_endpoints/common_utils.py | 7 +-- litellm/proxy/realtime_endpoints/endpoints.py | 11 +++-- .../proxy/response_polling/polling_handler.py | 2 +- litellm/proxy/vector_store_endpoints/utils.py | 7 +-- litellm/rag/ingestion/bedrock_ingestion.py | 10 ++-- litellm/repositories/table_repositories.py | 2 +- litellm/router_strategy/lowest_latency.py | 2 +- .../encrypted_content_affinity_check.py | 21 ++++++--- litellm/router_utils/prompt_caching_cache.py | 4 +- .../custom_secret_manager_loader.py | 4 +- litellm/types/containers/main.py | 42 +++++++++-------- litellm/types/llms/oci.py | 28 +++++------ litellm/types/llms/openai_evals.py | 33 ++++++------- .../proxy/management_endpoints/scim_v2.py | 12 ++--- litellm/types/videos/main.py | 19 ++++---- 89 files changed, 594 insertions(+), 418 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index 254d816039c..3b8c19f0097 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -16,6 +16,7 @@ from litellm.llms.base_llm.managed_resources.utils import ( is_base64_encoded_unified_id, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import LLMResponseTypes from litellm.types.vector_stores import ( VectorStoreCreateOptionalRequestParams, VectorStoreCreateResponse, @@ -24,6 +25,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.caching.caching import DualCache from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient @@ -156,7 +158,7 @@ class _PROXY_LiteLLMManagedVectorStores( # Create vector store for each model # Convert TypedDict to Dict[str, Any] for base class compatibility - request_data_dict: Dict[str, Any] = dict(create_request) + request_data_dict: Dict[str, object] = dict(create_request) responses = await self.create_resource_for_each_model( llm_router=llm_router, request_data=request_data_dict, @@ -209,7 +211,7 @@ class _PROXY_LiteLLMManagedVectorStores( limit: Optional[int] = None, after: Optional[str] = None, order: Optional[str] = None, - ) -> Dict[str, Any]: + ) -> Dict[str, object]: """ List vector stores created by a user. @@ -301,7 +303,7 @@ class _PROXY_LiteLLMManagedVectorStores( async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - cache: Any, + cache: "DualCache", data: Dict, call_type: str, ) -> Union[Exception, str, Dict, None]: @@ -403,8 +405,8 @@ class _PROXY_LiteLLMManagedVectorStores( self, data: Dict, user_api_key_dict: UserAPIKeyAuth, - response: Any, - ) -> Any: + response: LLMResponseTypes, + ) -> LLMResponseTypes: """ Post-call hook to transform responses. diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index 98fa62629a8..ba0398789a6 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -1,7 +1,7 @@ import asyncio import threading import time -from typing import Any, Final +from typing import Final, Protocol from redis.credentials import CredentialProvider @@ -18,6 +18,19 @@ _token_cache: Final[dict[str, tuple[str, float]]] = {} _token_cache_lock: Final = threading.Lock() +class AzureAccessToken(Protocol): + """The ``azure.core.credentials.AccessToken`` shape this module reads.""" + + @property + def token(self) -> str: ... + + +class AzureCredential(Protocol): + """The ``azure-identity`` credential surface this module calls.""" + + def get_token(self, *scopes: str) -> AzureAccessToken: ... + + def _generate_gcp_iam_access_token(service_account: str) -> str: """ Generate GCP IAM access token for Redis authentication. @@ -115,7 +128,7 @@ class AzureADCredentialProvider(CredentialProvider): fail authentication after the initial token expired (~1 hour TTL). """ - def __init__(self, credential: Any, username: str | None = None) -> None: + def __init__(self, credential: AzureCredential, username: str | None = None) -> None: self._credential = credential self._username = username diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 42a86763b6d..703aa197a63 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import litellm from litellm._logging import verbose_logger @@ -24,7 +24,30 @@ else: UserAPIKeyAuth = Any -def _get_otel_v2_class() -> type | None: +class _ServiceSpanLogger(Protocol): + """The OTel logger surface this module drives: the two service-span hooks it calls.""" + + async def async_service_success_hook( + self, + payload: ServiceLoggerPayload, + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: datetime | float | None = None, + event_metadata: dict | None = None, + ) -> None: ... + + async def async_service_failure_hook( + self, + payload: ServiceLoggerPayload, + error: str | None = "", + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: datetime | float | None = None, + event_metadata: dict | None = None, + ) -> None: ... + + +def _get_otel_v2_class() -> type[_ServiceSpanLogger] | None: """Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent. Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry @@ -54,7 +77,7 @@ class ServiceLogging(CustomLogger): if "prometheus_system" in litellm.service_callback: self.prometheusServicesLogger = PrometheusServicesLogger() - def _resolve_otel_service_logger(self, callback: Any) -> Any | None: + def _resolve_otel_service_logger(self, callback: object) -> _ServiceSpanLogger | None: """Resolve the OTel logger (legacy or V2) to emit a service span on. Returns the logger instance whose ``async_service_*_hook`` should fire for @@ -69,18 +92,21 @@ class ServiceLogging(CustomLogger): """ otel_v2_cls: Final = _get_otel_v2_class() - def _is_otel_logger(obj: Any) -> bool: + def _as_otel_logger(obj: object) -> _ServiceSpanLogger | None: if isinstance(obj, OpenTelemetry): - return True - return otel_v2_cls is not None and isinstance(obj, otel_v2_cls) + return obj + if otel_v2_cls is not None and isinstance(obj, otel_v2_cls): + return obj + return None - if _is_otel_logger(callback): - return callback + resolved_callback: Final = _as_otel_logger(callback) + if resolved_callback is not None: + return resolved_callback if callback == "otel": from litellm.proxy.proxy_server import open_telemetry_logger - if open_telemetry_logger is not None and _is_otel_logger(open_telemetry_logger): - return open_telemetry_logger + if open_telemetry_logger is not None: + return _as_otel_logger(open_telemetry_logger) return None def service_success_hook( diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 25f2e1a9a0d..b663e3085fb 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -4,6 +4,7 @@ Custom A2A Card Resolver for LiteLLM. Extends the A2A SDK's card resolver to support multiple well-known paths. """ +from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -152,7 +153,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): async def get_agent_card( self, relative_card_path: str | None = None, - http_kwargs: dict[str, Any] | None = None, + http_kwargs: Mapping[str, object] | None = None, ) -> "AgentCard": """ Fetch the agent card, trying multiple well-known paths. diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py index 3748d8043cc..57c4a4677d0 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -8,7 +8,7 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model: """ import asyncio -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import Any, Final from uuid import uuid4 @@ -51,9 +51,9 @@ class WatsonxOrchestrateTransformation: wxo_agent_id: str, text: str, thread_id: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build the WXO POST /v1/orchestrate/runs request body.""" - body: Final[dict[str, Any]] = { + body: Final[dict[str, object]] = { "agent_id": wxo_agent_id, "message": { "role": "user", @@ -70,7 +70,7 @@ class WatsonxOrchestrateTransformation: return body @staticmethod - def extract_text_from_wxo_result(result: Any) -> str: + def extract_text_from_wxo_result(result: object) -> str: """ Extract response text from a WXO run result. @@ -103,7 +103,7 @@ class WatsonxOrchestrateTransformation: return "" @staticmethod - def extract_text_from_a2a_message_response(a2a_response: dict[str, Any]) -> str: + def extract_text_from_a2a_message_response(a2a_response: Mapping[str, object]) -> str: result: Final = a2a_response.get("result") if not isinstance(result, dict): verbose_logger.warning("WXO: A2A response missing result object") @@ -119,7 +119,7 @@ class WatsonxOrchestrateTransformation: return "" @staticmethod - def build_a2a_message_response(request_id: str, text: str) -> dict[str, Any]: + def build_a2a_message_response(request_id: str, text: str) -> dict[str, object]: """ Build a standard A2A non-streaming SendMessageResponse (kind=message). """ @@ -140,7 +140,7 @@ class WatsonxOrchestrateTransformation: request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Emit standard A2A streaming events from a completed text response. diff --git a/litellm/assistants/utils.py b/litellm/assistants/utils.py index e41cff8419a..a2841e3ff93 100644 --- a/litellm/assistants/utils.py +++ b/litellm/assistants/utils.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from typing import Final import litellm @@ -10,20 +11,22 @@ def get_optional_params_add_message( role: str | None, content: str | list[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None, attachments: list[Attachment] | None, - metadata: dict | None, + metadata: Mapping[str, object] | None, custom_llm_provider: str, - **kwargs, -): + **kwargs: object, +) -> dict[str, object]: """ Azure doesn't support 'attachments' for creating a message Reference - https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-messages?tabs=python#create-message """ - passed_params: Final = locals() - custom_llm_provider = passed_params.pop("custom_llm_provider") - special_params: Final = passed_params.pop("kwargs") - for k, v in special_params.items(): - passed_params[k] = v + passed_params: Final[Mapping[str, object]] = { + "role": role, + "content": content, + "attachments": attachments, + "metadata": metadata, + **kwargs, + } default_params: Final = { "role": None, @@ -33,10 +36,10 @@ def get_optional_params_add_message( } non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])} - optional_params = {} + optional_params: dict[str, object] = {} ## raise exception if non-default value passed for non-openai/azure embedding calls - def _check_valid_arg(supported_params): + def _check_valid_arg(supported_params: Sequence[str]) -> Mapping[str, object] | None: if len(non_default_params.keys()) > 0: keys: Final = list(non_default_params.keys()) for k in keys: @@ -71,14 +74,18 @@ def get_optional_params_image_gen( style: str | None = None, user: str | None = None, custom_llm_provider: str | None = None, - **kwargs, -): + **kwargs: object, +) -> dict[str, object]: # retrieve all parameters passed to the function - passed_params: Final = locals() - custom_llm_provider = passed_params.pop("custom_llm_provider") - special_params: Final = passed_params.pop("kwargs") - for k, v in special_params.items(): - passed_params[k] = v + passed_params: Final[Mapping[str, object]] = { + "n": n, + "quality": quality, + "response_format": response_format, + "size": size, + "style": style, + "user": user, + **kwargs, + } default_params: Final = { "n": None, @@ -90,10 +97,10 @@ def get_optional_params_image_gen( } non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])} - optional_params = {} + optional_params: dict[str, object] = {} ## raise exception if non-default value passed for non-openai/azure embedding calls - def _check_valid_arg(supported_params): + def _check_valid_arg(supported_params: Sequence[str]) -> Mapping[str, object] | None: if len(non_default_params.keys()) > 0: keys: Final = list(non_default_params.keys()) for k in keys: diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 3831f57a10d..97be5f77d79 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -160,7 +160,7 @@ def _classify_output_line_stats( def _safe_output_line_stats( - entry: Mapping[str, Any], + entry: Mapping[str, object], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, model_info: ModelInfo | None, @@ -182,7 +182,7 @@ def _safe_output_line_stats( def _compute_output_line_stats( - entry: Mapping[str, Any], + entry: Mapping[str, object], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, model_info: ModelInfo | None, @@ -213,7 +213,7 @@ def _compute_output_line_stats( def _output_line_cost( - response_body: Mapping[str, Any], + response_body: Mapping[str, object], usage: Usage, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, @@ -556,7 +556,7 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]: def _parse_batch_output_line(line: bytes) -> dict | None: try: - parsed: Final = json.loads(line) + parsed: Final[object] = json.loads(line) except ValueError as e: verbose_logger.warning("skipping malformed batch output line: %s", str(e)) return None @@ -601,7 +601,7 @@ def _count_entry_tokens( return 0 -def _count_prompt_or_input_tokens(model: str, value: Any) -> int: +def _count_prompt_or_input_tokens(model: str, value: object) -> int: """Token-count a ``prompt`` / ``input`` field that the OpenAI batch schema allows in four shapes: @@ -680,7 +680,7 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[st def _get_response_from_batch_job_output_file( batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai" -) -> Mapping[str, Any]: +) -> Mapping[str, object]: """ Get the response from the batch job output file """ diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index f844b3a3d7f..c646baf9d9e 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -66,7 +66,7 @@ def _build_retrieval_tools(keys: list[str], call_type: str) -> list[dict]: return cast(list[dict], anthropic_tools) -def _content_to_text(content: Any) -> str: +def _content_to_text(content: object) -> str: """ Convert OpenAI/Anthropic message content blocks to plain text. @@ -78,7 +78,7 @@ def _content_to_text(content: Any) -> str: Implemented iteratively (stack-based) to avoid unbounded recursion. """ parts: Final[list[str]] = [] - stack: Final[list[Any]] = [content] + stack: Final[list[object]] = [content] while stack: item = stack.pop() if isinstance(item, str): @@ -111,7 +111,7 @@ def _normalize_messages_for_compression( f"Unsupported call_type={call_type!r} for compression. Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}." ) - original_messages: Final[list[dict[str, Any]]] = [dict(m) for m in messages] + original_messages: Final[list[dict[str, object]]] = [dict(m) for m in messages] normalized_messages: Final[list[dict]] = [] for msg in original_messages: @@ -132,7 +132,7 @@ def _extract_last_user_message(messages: list[dict]) -> str: return "" -def _extract_tool_use_ids(content: Any) -> list[str]: +def _extract_tool_use_ids(content: object) -> list[str]: if not isinstance(content, list): return [] tool_use_ids: Final[list[str]] = [] @@ -147,7 +147,7 @@ def _extract_tool_use_ids(content: Any) -> list[str]: return tool_use_ids -def _extract_tool_result_ids(content: Any) -> set[str]: +def _extract_tool_result_ids(content: object) -> set[str]: if not isinstance(content, list): return set() tool_result_ids: Final[set[str]] = set() @@ -337,7 +337,7 @@ def compress( compression_trigger: int = 200_000, compression_target: int | None = None, embedding_model: str | None = None, - embedding_model_params: dict[str, Any] | None = None, + embedding_model_params: Mapping[str, object] | None = None, compression_cache: DualCache | None = None, ) -> CompressedResult: """ diff --git a/litellm/containers/endpoint_factory.py b/litellm/containers/endpoint_factory.py index 09bc7eda41f..25fc223cde0 100644 --- a/litellm/containers/endpoint_factory.py +++ b/litellm/containers/endpoint_factory.py @@ -11,7 +11,7 @@ import json from collections.abc import Callable from functools import partial from pathlib import Path -from typing import Any, Final, Literal +from typing import Final, Literal import litellm from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT @@ -56,9 +56,9 @@ def create_sync_endpoint_function(endpoint_config: dict) -> Callable: def endpoint_func( timeout: int = 600, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, **kwargs, ): local_vars: Final = locals() @@ -145,9 +145,9 @@ def create_async_endpoint_function( async def async_endpoint_func( timeout: int = 600, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, **kwargs, ): local_vars: Final = locals() diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 286f7528896..16202321709 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -85,7 +85,7 @@ _RATE_LIMIT_CATEGORY_VALUES: Final = frozenset(c.value for c in RateLimitErrorCa _RATE_LIMIT_TYPE_VALUES: Final = frozenset(t.value for t in RateLimitType) -def validate_rate_limit_category(value: Any) -> str | None: +def validate_rate_limit_category(value: object) -> str | None: """Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`. Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus @@ -100,7 +100,7 @@ def validate_rate_limit_category(value: Any) -> str | None: return None -def validate_rate_limit_type(value: Any) -> str | None: +def validate_rate_limit_type(value: object) -> str | None: """Return ``value`` only if it matches a known :class:`RateLimitType`. See :func:`validate_rate_limit_category` for the rationale. diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index 48bb4cc6380..38be0666008 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -11,7 +11,7 @@ https://platform.openai.com/docs/api-reference/fine-tuning import asyncio import contextvars import os -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import Any, Final, Literal @@ -37,8 +37,8 @@ vertex_fine_tuning_apis_instance: Final = VertexFineTuningAPI() def _prepare_azure_extra_body( extra_body: dict[str, Any] | None, - kwargs: dict[str, Any], - azure_specific_hyperparams: dict[str, Any], + kwargs: Mapping[str, object], + azure_specific_hyperparams: Mapping[str, object], ) -> dict[str, Any]: """ Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters. @@ -138,7 +138,7 @@ def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, v def _resolve_fine_tuning_timeout( - timeout: Any, + timeout: float | str | httpx.Timeout | None, custom_llm_provider: str, ) -> float | httpx.Timeout: """Normalise a raw timeout value to a float (seconds) or httpx.Timeout for fine-tuning calls.""" @@ -163,7 +163,7 @@ def create_fine_tuning_job( extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, **kwargs, -) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: +) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: """ Creates a fine-tuning job which begins the process of creating a new model from a given dataset. @@ -375,7 +375,7 @@ def cancel_fine_tuning_job( extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, **kwargs, -) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: +) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: """ Immediately cancel a fine-tune job. @@ -682,7 +682,7 @@ def retrieve_fine_tuning_job( extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, **kwargs, -) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: +) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: """ Get info about a fine-tuning job. """ diff --git a/litellm/images/utils.py b/litellm/images/utils.py index 2f080d88de4..49b70870de6 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from io import BufferedReader, BytesIO from typing import Any, Final, cast, get_type_hints @@ -61,7 +62,7 @@ class ImageEditRequestUtils: @staticmethod def get_requested_image_edit_optional_param( - params: dict[str, Any], + params: Mapping[str, object], ) -> ImageEditOptionalRequestParams: """ Filter parameters to only include those defined in ImageEditOptionalRequestParams. diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 7255c9c761c..538dd95abdd 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -1,6 +1,7 @@ import asyncio import os import time +from collections.abc import Mapping from datetime import datetime from typing import Any, Final, cast @@ -181,7 +182,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): # cast because StandardLoggingMetadata is a TypedDict; we iterate it # as a generic mapping below. - metadata: Final[dict[str, Any]] = cast(dict[str, Any], log.get("metadata") or {}) + metadata: Final[Mapping[str, object]] = cast(dict[str, Any], log.get("metadata") or {}) # Backwards-compat: team/user/model_group preserved regardless of allowlist. if metadata.get("user_api_key_alias"): @@ -233,7 +234,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): tags[key] = normalize_datadog_tag_value(value) @staticmethod - def _add_tag(tags: dict[str, str], key: str, value: Any) -> None: + def _add_tag(tags: dict[str, str], key: str, value: object) -> None: if value: tags[key] = str(value) diff --git a/litellm/integrations/dotprompt/dotprompt_manager.py b/litellm/integrations/dotprompt/dotprompt_manager.py index f1ef011cdb7..c646dbf4e2e 100644 --- a/litellm/integrations/dotprompt/dotprompt_manager.py +++ b/litellm/integrations/dotprompt/dotprompt_manager.py @@ -4,6 +4,7 @@ Builds on top of PromptManagementBase to provide .prompt file support. """ import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from litellm.integrations.custom_prompt_management import CustomPromptManagement @@ -347,14 +348,14 @@ class DotpromptManager(CustomPromptManagement): metadata: Final = json_data.get("metadata", {}) self.prompt_manager.add_prompt(prompt_id, content, metadata) - def load_prompts_from_json(self, prompts_data: dict[str, dict[str, Any]]) -> None: + def load_prompts_from_json(self, prompts_data: dict[str, dict[str, object]]) -> None: """Load multiple prompts from JSON data.""" self.prompt_manager.load_prompts_from_json_data(prompts_data) - def get_prompts_as_json(self) -> dict[str, dict[str, Any]]: + def get_prompts_as_json(self) -> dict[str, dict[str, object]]: """Get all prompts in JSON format.""" return self.prompt_manager.get_all_prompts_as_json() - def convert_prompt_file_to_json(self, file_path: str) -> dict[str, Any]: + def convert_prompt_file_to_json(self, file_path: str) -> Mapping[str, object]: """Convert a .prompt file to JSON format.""" return self.prompt_manager.prompt_file_to_json(file_path) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 74ef6f70a65..c9b47835948 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -102,7 +102,7 @@ class FocusLogger(CustomLogger): # No time bounds → export all available data await self._export_all(limit=limit) - async def dry_run_export_usage_data(self, limit: int | None = DEFAULT_DRY_RUN_LIMIT) -> dict[str, Any]: + async def dry_run_export_usage_data(self, limit: int | None = DEFAULT_DRY_RUN_LIMIT) -> dict[str, object]: """Return transformed data without uploading.""" engine: Final = self._ensure_engine() return await engine.dry_run_export_usage_data(limit=limit) @@ -153,7 +153,7 @@ class FocusLogger(CustomLogger): **trigger_kwargs, ) - def _build_scheduler_trigger(self) -> dict[str, Any]: + def _build_scheduler_trigger(self) -> dict[str, str | int]: """Return scheduler configuration for the selected frequency.""" if self.frequency == "interval": seconds: Final = self.interval_seconds or 60 diff --git a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py index bed3bdb58d1..77d315d0cee 100644 --- a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py +++ b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py @@ -4,6 +4,7 @@ Fetches prompts from any API that implements the /beta/litellm_prompt_management """ import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -349,7 +350,7 @@ class GenericPromptManager(CustomPromptManagement): def _apply_variables( self, prompt_client: PromptManagementClient, - variables: dict[str, Any], + variables: Mapping[str, object], ) -> PromptManagementClient: """ Apply variables to the prompt template. diff --git a/litellm/integrations/humanloop.py b/litellm/integrations/humanloop.py index 405854b0ce9..9e52ccd3c02 100644 --- a/litellm/integrations/humanloop.py +++ b/litellm/integrations/humanloop.py @@ -4,7 +4,7 @@ Humanloop integration https://humanloop.com/ """ -from typing import Any, Final, cast +from typing import Final, cast import httpx from typing_extensions import TypedDict @@ -24,7 +24,7 @@ class PromptManagementClient(TypedDict): prompt_id: str prompt_template: list[AllMessageValues] model: str | None - optional_params: dict[str, Any] | None + optional_params: dict[str, object] | None class HumanLoopPromptManager(DualCache): @@ -36,7 +36,7 @@ class HumanLoopPromptManager(DualCache): return cast(PromptManagementClient | None, self.get_cache(key=humanloop_prompt_id)) def _compile_prompt_helper( - self, prompt_template: list[AllMessageValues], prompt_variables: dict[str, Any] + self, prompt_template: list[AllMessageValues], prompt_variables: dict[str, object] ) -> list[AllMessageValues]: """ Helper function to compile the prompt by substituting variables in the template. diff --git a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py index 0e58cf67795..b5eedc42fe9 100644 --- a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py +++ b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py @@ -117,7 +117,7 @@ class OTELGenAISemconvMixin: if TYPE_CHECKING: config: "OpenTelemetryConfig" - def safe_set_attribute(self, span: Span, key: str, value: Any) -> None: ... + def safe_set_attribute(self, span: Span, key: str, value: object) -> None: ... def _capture_in_event(self) -> bool: ... @@ -195,13 +195,13 @@ class OTELGenAISemconvMixin: if value: self.safe_set_attribute(span=span, key=semconv_key, value=value) - def _build_inference_details_attrs(self, kwargs: dict, response_obj: dict, provider: str) -> dict[str, Any]: + def _build_inference_details_attrs(self, kwargs: dict, response_obj: dict, provider: str) -> dict[str, str]: """Build the attribute payload for the inference-details event. Always includes provider/operation; input/output messages are added only when content capture is enabled and non-empty. Mixin-internal. """ - attrs: Final[dict[str, Any]] = { + attrs: Final[dict[str, str]] = { "event_name": _INFERENCE_DETAILS_EVENT_NAME, "gen_ai.provider.name": provider, "gen_ai.operation.name": self._gen_ai_operation_name(kwargs), diff --git a/litellm/integrations/opik/opik_payload_builder/payload_builders.py b/litellm/integrations/opik/opik_payload_builder/payload_builders.py index 855b84ba4c8..3aaf5bfc162 100644 --- a/litellm/integrations/opik/opik_payload_builder/payload_builders.py +++ b/litellm/integrations/opik/opik_payload_builder/payload_builders.py @@ -15,8 +15,8 @@ def build_trace_payload( response_obj: dict[str, Any], start_time: datetime, end_time: datetime, - input_data: Any, - output_data: Any, + input_data: object, + output_data: object, metadata: dict[str, object], tags: list[str], thread_id: str | None, @@ -45,8 +45,8 @@ def build_span_payload( response_obj: dict[str, Any], start_time: datetime, end_time: datetime, - input_data: Any, - output_data: Any, + input_data: object, + output_data: object, metadata: dict[str, object], tags: list[str], usage: dict[str, int], diff --git a/litellm/integrations/weave/weave_otel.py b/litellm/integrations/weave/weave_otel.py index 1fc53d14a54..f2cc64a9ba2 100644 --- a/litellm/integrations/weave/weave_otel.py +++ b/litellm/integrations/weave/weave_otel.py @@ -3,6 +3,7 @@ from __future__ import annotations import base64 import json import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from opentelemetry.trace import Status, StatusCode @@ -59,7 +60,7 @@ class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): safe_set_attribute(span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt)) -def _set_weave_specific_attributes(span: Span, kwargs: dict[str, Any], response_obj: Any): +def _set_weave_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj: Any): """ Sets Weave-specific metadata attributes onto the OTEL span. @@ -169,7 +170,7 @@ def get_weave_otel_config() -> WeaveOtelConfig: ) -def set_weave_otel_attributes(span: Span, kwargs: dict[str, Any], response_obj: Any): +def set_weave_otel_attributes(span: Span, kwargs: Mapping[str, object], response_obj: object): """ Sets OpenTelemetry span attributes for Weave observability. Uses the same attribute setting logic as other OTEL integrations for consistency. diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 80a27007329..1dac67ecbf6 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -23,12 +23,13 @@ Used by JWT Auth to get the user role from the token, and by additional_drop_params to remove nested fields from optional parameters. """ +from collections.abc import Mapping from typing import Any, Final, TypeVar T = TypeVar("T") -def get_nested_value(data: dict[str, Any], key_path: str, default: T | None = None) -> T | None: +def get_nested_value(data: Mapping[str, object], key_path: str, default: T | None = None) -> T | None: """ Retrieves a value from a nested dictionary using dot notation. @@ -107,7 +108,7 @@ def _parse_path_segments(path: str) -> list: def _delete_nested_value_custom( - data: dict[str, Any] | list[Any], + data: dict[str, object] | list[object], segments: list, segment_index: int = 0, ) -> None: @@ -168,13 +169,15 @@ def _delete_nested_value_custom( if segment in data: next_segment: Final = segments[segment_index + 1] if segment_index + 1 < len(segments) else None + child: Final = data[segment] + # If next segment is array notation, current field should be list if next_segment and (next_segment.startswith("[")): - if isinstance(data[segment], list): - _delete_nested_value_custom(data[segment], segments, segment_index + 1) + if isinstance(child, list): + _delete_nested_value_custom(child, segments, segment_index + 1) # Otherwise navigate into dict - elif isinstance(data[segment], dict): - _delete_nested_value_custom(data[segment], segments, segment_index + 1) + elif isinstance(child, dict): + _delete_nested_value_custom(child, segments, segment_index + 1) def delete_nested_value( @@ -182,7 +185,7 @@ def delete_nested_value( path: str, depth: int = 0, max_depth: int = 20, -) -> dict[str, Any]: +) -> dict[str, object]: """ Delete a field from nested data using JSONPath notation. diff --git a/litellm/litellm_core_utils/json_validation_rule.py b/litellm/litellm_core_utils/json_validation_rule.py index 12f952d1d69..9fd4c03ac9e 100644 --- a/litellm/litellm_core_utils/json_validation_rule.py +++ b/litellm/litellm_core_utils/json_validation_rule.py @@ -5,10 +5,10 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH def normalize_json_schema_types( - schema: dict[str, Any] | list[Any] | Any, + schema: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, -) -> dict[str, Any] | list[Any] | Any: +) -> object: """ Normalize JSON schema types from uppercase to lowercase format. @@ -47,7 +47,7 @@ def normalize_json_schema_types( return [normalize_json_schema_types(item, depth + 1, max_depth) for item in schema] if isinstance(schema, dict): - normalized_schema: Final[dict[str, Any]] = {} + normalized_schema: Final[dict[str, object]] = {} for key, value in schema.items(): if key == "type" and isinstance(value, str) and value in type_mapping: diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 91c8ba36b26..f3b1b29a9ad 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -184,7 +184,7 @@ def _get_parent_otel_span_from_logging_obj( def convert_litellm_response_object_to_str( - response_obj: Any | LiteLLMModelResponse, + response_obj: object, ) -> str | None: """ Get the string of the response object from LiteLLM diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index a1b71593dda..5b99e8cba98 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -30,7 +30,7 @@ def safe_dumps( def _transform(key: str | None, value: str) -> str: return value if value_transform is None else value_transform(key, value) - def _serialize(obj: Any, seen: set, depth: int, key: str | None = None) -> Any: + def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> Any: # Check for maximum depth. if depth > max_depth: return "MaxDepthExceeded" diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 178b4c47a0f..57eadfe36d2 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -2,6 +2,7 @@ Common utilities for A2A (Agent-to-Agent) Protocol """ +from collections.abc import Mapping from typing import Any, Final from pydantic import BaseModel @@ -91,7 +92,7 @@ def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_d return " ".join(text_parts) -def extract_text_from_a2a_response(response_dict: dict[str, Any], max_depth: int = 10) -> str: +def extract_text_from_a2a_response(response_dict: Mapping[str, object], max_depth: int = 10) -> str: """ Extract text content from A2A response result. diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py index 701211049db..4a6b65bb2b1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py @@ -266,7 +266,7 @@ def _make_synthetic_advisor_tool() -> dict: } -def _find_advisor_tool_use(response: Any) -> dict | None: +def _find_advisor_tool_use(response: object) -> dict | None: """Return the first tool_use block with name='advisor', or None.""" content: Final = response.get("content") if isinstance(response, dict) else [] if not isinstance(content, list): @@ -277,7 +277,7 @@ def _find_advisor_tool_use(response: Any) -> dict | None: return None -def _extract_response_text(response: Any) -> str: +def _extract_response_text(response: object) -> str: """Extract concatenated text from all text blocks in a response.""" content: Final = response.get("content") if isinstance(response, dict) else [] if not isinstance(content, list): @@ -291,7 +291,7 @@ _PROVIDER_SPECIFIC_KEYS: Final = frozenset({"provider_specific_fields"}) def _build_advisor_context( messages: list[dict], - executor_response: Any, + executor_response: object, advisor_use_block: dict, ) -> list[dict]: """ @@ -327,7 +327,7 @@ def _build_advisor_context( def _inject_advisor_turn( messages: list[dict], - executor_response: Any, + executor_response: object, advisor_use_block: dict, advisor_text: str, ) -> list[dict]: @@ -355,7 +355,7 @@ def _inject_advisor_turn( def _inject_max_uses_error( messages: list[dict], - executor_response: Any, + executor_response: object, advisor_use_block: dict, ) -> list[dict]: """ diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 292d2622c7f..b9ab350f221 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -42,9 +42,9 @@ class AnthropicResponsesStreamWrapper: self._pending_tool_ids: dict[str, str] = {} # item_id -> call_id / name accumulator self._sent_message_start = False self._sent_message_stop = False - self._chunk_queue: deque = deque() + self._chunk_queue: deque[dict[str, object]] = deque() - def _make_message_start(self) -> dict[str, Any]: + def _make_message_start(self) -> dict[str, object]: return { "type": "message_start", "message": { @@ -68,7 +68,7 @@ class AnthropicResponsesStreamWrapper: self._current_block_index += 1 return self._current_block_index - def _open_block(self, item_id: str | None, content_block: Mapping[str, Any]) -> int: + def _open_block(self, item_id: str | None, content_block: Mapping[str, object]) -> int: block_idx = self._next_block_index() if item_id: self._item_id_to_block_index[item_id] = block_idx @@ -81,7 +81,7 @@ class AnthropicResponsesStreamWrapper: ) return block_idx - def _process_event(self, event: Any) -> None: + def _process_event(self, event: object) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) if event_type is None and isinstance(event, dict): @@ -247,7 +247,7 @@ class AnthropicResponsesStreamWrapper: def __aiter__(self) -> "AnthropicResponsesStreamWrapper": return self - async def __anext__(self) -> dict[str, Any]: + async def __anext__(self) -> dict[str, object]: # Return any queued chunks first if self._chunk_queue: return self._chunk_queue.popleft() diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index fe7f57d7a13..7b5ab78af8d 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -14,7 +14,7 @@ Anthropic Files API endpoints: import calendar import time -from typing import Any, Final, cast +from typing import Final, cast import httpx from openai.types.file_deleted import FileDeleted @@ -226,7 +226,7 @@ class AnthropicFilesConfig(BaseFilesConfig): ) -> tuple[str, dict]: api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE url: Final = f"{api_base.rstrip('/')}/v1/files" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, str]] = {} if purpose: params["purpose"] = purpose return url, params diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 8f96f80d15e..133e40dc1ab 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -20,6 +20,7 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.openai import HttpxBinaryResponseContent else: LiteLLMLoggingObj = Any @@ -75,15 +76,15 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): litellm_params_dict: dict, logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout, - extra_headers: dict[str, Any] | None, - base_llm_http_handler: Any, + extra_headers: dict[str, object] | None, + base_llm_http_handler: "BaseLLMHTTPHandler", aspeech: bool, api_base: str | None, api_key: str | None, - **kwargs: Any, + **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[object, object, "HttpxBinaryResponseContent"], ]: """ Dispatch method to handle AWS Polly TTS requests @@ -251,7 +252,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): def _sign_polly_request( self, - request_body: dict[str, Any], + request_body: dict[str, object], endpoint_url: str, litellm_params: dict, ) -> tuple[dict[str, str], str]: @@ -337,7 +338,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): engine: Final = optional_params.get("engine", self.DEFAULT_ENGINE) # Build request body - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "Engine": engine, "OutputFormat": output_format, "Text": input, diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 88492ef996e..e9913f0108d 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -6,7 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final, cast +from typing import Any, Final, Protocol, cast from litellm._logging import _redact_string, verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -31,6 +31,12 @@ async def forward_messages(client_ws: Any, backend_ws: Any): pass +class _ProxyClientWebSocket(Protocol): + """Client-facing websocket handle: this path only closes it after a failed handshake.""" + + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + + class AzureOpenAIRealtime(AzureChatCompletion): @staticmethod def get_auth_headers(api_key: str | None, azure_ad_token: str | None) -> Mapping[str, str]: @@ -104,17 +110,17 @@ class AzureOpenAIRealtime(AzureChatCompletion): async def async_realtime( self, model: str, - websocket: Any, + websocket: _ProxyClientWebSocket, logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, api_version: str | None = None, azure_ad_token: str | None = None, - client: Any | None = None, + client: object | None = None, timeout: float | None = None, realtime_protocol: str | None = None, query_params: RealtimeQueryParams | None = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, litellm_metadata: dict | None = None, ): import websockets diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 0dd5e87e4ba..2a59cabfaf0 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -96,7 +96,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Then filter out status from message items if isinstance(validated_input, list): - filtered_input: Final[list[Any]] = [] + filtered_input: Final[list[object]] = [] for item in validated_input: if isinstance(item, dict) and item.get("type") == "message": # Filter out status field from message items @@ -123,7 +123,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): if "tools" in response_api_optional_request_params and isinstance( response_api_optional_request_params["tools"], list ): - new_tools: Final[list[dict[str, Any]]] = [] + new_tools: Final[list[dict[str, object]]] = [] for tool in response_api_optional_request_params["tools"]: if isinstance(tool, dict) and "function" in tool: new_tool: dict[str, Any] = deepcopy(tool) @@ -291,7 +291,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): url: Final = self._construct_url_for_response_id_in_path( api_base=api_base, response_id=response_id, path_suffix="/input_items" ) - params: Final[dict[str, Any]] = {} + params: Final[dict[str, str | int]] = {} if after is not None: params["after"] = after if before is not None: diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 955090b9b90..c31d2c427bb 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -28,12 +28,12 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, + contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, object]] | None = None, + system: object | None = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. diff --git a/litellm/llms/base_llm/agents/transformation.py b/litellm/llms/base_llm/agents/transformation.py index 970639939f1..9d139b289c4 100644 --- a/litellm/llms/base_llm/agents/transformation.py +++ b/litellm/llms/base_llm/agents/transformation.py @@ -10,7 +10,7 @@ InteractionsHTTPHandler). """ from abc import ABC, abstractmethod -from typing import Any +from collections.abc import Mapping import httpx @@ -35,7 +35,7 @@ class BaseAgentsAPIConfig(ABC): def get_complete_url( self, api_base: str | None, - litellm_params: dict[str, Any], + litellm_params: Mapping[str, object], ) -> str: """Return the full URL for POST /agents (create).""" @@ -43,7 +43,7 @@ class BaseAgentsAPIConfig(ABC): def validate_environment( self, headers: dict[str, str], - litellm_params: dict[str, Any], + litellm_params: dict[str, object], ) -> dict[str, str]: """Validate credentials and return auth headers.""" @@ -51,8 +51,8 @@ class BaseAgentsAPIConfig(ABC): def transform_create_request( self, name: str, - litellm_params: dict[str, Any], - ) -> dict[str, Any]: + litellm_params: Mapping[str, object], + ) -> dict[str, object]: """Map name + litellm_params to the provider's create-agent body.""" @abstractmethod @@ -71,8 +71,8 @@ class BaseAgentsAPIConfig(ABC): def transform_list_request( self, api_base: str | None, - litellm_params: dict[str, Any], - ) -> tuple[str, dict[str, Any]]: + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: """Return (url, query_params) for GET /agents.""" @abstractmethod @@ -91,8 +91,8 @@ class BaseAgentsAPIConfig(ABC): self, name: str, api_base: str | None, - litellm_params: dict[str, Any], - ) -> tuple[str, dict[str, Any]]: + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: """Return (url, query_params) for GET /agents/{name}.""" @abstractmethod @@ -112,7 +112,7 @@ class BaseAgentsAPIConfig(ABC): self, name: str, api_base: str | None, - litellm_params: dict[str, Any], + litellm_params: Mapping[str, object], ) -> str: """Return the URL for DELETE /agents/{name}.""" @@ -133,8 +133,8 @@ class BaseAgentsAPIConfig(ABC): self, name: str, api_base: str | None, - litellm_params: dict[str, Any], - ) -> tuple[str, dict[str, Any]]: + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: """Return (url, query_params) for GET /agents/{name}/versions.""" @abstractmethod diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 9b6f9c47105..a67ca9bffa8 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -2,13 +2,13 @@ from __future__ import annotations import json from collections.abc import Callable, Iterator, Sequence -from typing import Any, Final, TypeVar +from typing import Final, TypeVar from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage -def _anthropic_stream_chunk_events(item: Any) -> list[dict]: +def _anthropic_stream_chunk_events(item: object) -> list[dict]: if isinstance(item, dict): return [item] if isinstance(item, bytes): @@ -36,7 +36,7 @@ def _anthropic_stream_chunk_events(item: Any) -> list[dict]: return events -def _usage_from_anthropic_stream_chunks(original_response: list[Any]) -> AnthropicUsage | None: +def _usage_from_anthropic_stream_chunks(original_response: Sequence[object]) -> AnthropicUsage | None: input_tokens = 0 output_tokens = 0 found_usage = False @@ -79,7 +79,7 @@ def _usage_tokens(usage_obj: object, key: str, fallback_key: str) -> int: return int(getattr(usage_obj, key, getattr(usage_obj, fallback_key, 0)) or 0) -def blocked_response_usage(original_response: Any | None) -> AnthropicUsage: +def blocked_response_usage(original_response: object) -> AnthropicUsage: """ Token usage for a synthetic guardrail-blocked response. @@ -179,7 +179,7 @@ def blocked_responses_stream_usage(original_response: object) -> ResponseAPIUsag return blocked_responses_api_usage(completed) -def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool: +def effective_skip_system_message_for_guardrail(guardrail_to_apply: object) -> bool: per: Final = getattr(guardrail_to_apply, "skip_system_message_in_guardrail", None) if per is not None: return bool(per) @@ -188,7 +188,7 @@ def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool return bool(getattr(litellm, "skip_system_message_in_guardrail", False)) -def effective_skip_tool_message_for_guardrail(guardrail_to_apply: Any) -> bool: +def effective_skip_tool_message_for_guardrail(guardrail_to_apply: object) -> bool: per: Final = getattr(guardrail_to_apply, "skip_tool_message_in_guardrail", None) if per is not None: return bool(per) diff --git a/litellm/llms/base_llm/vector_store_files/transformation.py b/litellm/llms/base_llm/vector_store_files/transformation.py index 74aa283113c..9fb4d3e9dac 100644 --- a/litellm/llms/base_llm/vector_store_files/transformation.py +++ b/litellm/llms/base_llm/vector_store_files/transformation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -43,10 +44,10 @@ class BaseVectorStoreFilesConfig(ABC): self, *, operation: str, - non_default_params: dict[str, Any], - optional_params: dict[str, Any], + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], drop_params: bool, - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """Map non-default OpenAI params to provider-specific params.""" return optional_params @@ -87,7 +88,7 @@ class BaseVectorStoreFilesConfig(ABC): vector_store_id: str, create_request: VectorStoreFileCreateRequest, api_base: str, - ) -> tuple[str, dict[str, Any]]: ... + ) -> tuple[str, dict[str, object]]: ... @abstractmethod def transform_create_vector_store_file_response( @@ -103,7 +104,7 @@ class BaseVectorStoreFilesConfig(ABC): vector_store_id: str, query_params: VectorStoreFileListQueryParams, api_base: str, - ) -> tuple[str, dict[str, Any]]: ... + ) -> tuple[str, dict[str, object]]: ... @abstractmethod def transform_list_vector_store_files_response( @@ -119,7 +120,7 @@ class BaseVectorStoreFilesConfig(ABC): vector_store_id: str, file_id: str, api_base: str, - ) -> tuple[str, dict[str, Any]]: ... + ) -> tuple[str, dict[str, object]]: ... @abstractmethod def transform_retrieve_vector_store_file_response( @@ -135,7 +136,7 @@ class BaseVectorStoreFilesConfig(ABC): vector_store_id: str, file_id: str, api_base: str, - ) -> tuple[str, dict[str, Any]]: ... + ) -> tuple[str, dict[str, object]]: ... @abstractmethod def transform_retrieve_vector_store_file_content_response( @@ -152,7 +153,7 @@ class BaseVectorStoreFilesConfig(ABC): file_id: str, update_request: VectorStoreFileUpdateRequest, api_base: str, - ) -> tuple[str, dict[str, Any]]: ... + ) -> tuple[str, dict[str, object]]: ... @abstractmethod def transform_update_vector_store_file_response( @@ -168,7 +169,7 @@ class BaseVectorStoreFilesConfig(ABC): vector_store_id: str, file_id: str, api_base: str, - ) -> tuple[str, dict[str, Any]]: ... + ) -> tuple[str, dict[str, object]]: ... @abstractmethod def transform_delete_vector_store_file_response( @@ -196,8 +197,8 @@ class BaseVectorStoreFilesConfig(ABC): self, *, headers: dict[str, str], - optional_params: dict[str, Any], - request_data: dict[str, Any], + optional_params: Mapping[str, object], + request_data: Mapping[str, object], api_base: str, api_key: str | None = None, ) -> tuple[dict[str, str], bytes | None]: diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1e634ced29b..ec98a7a2c8f 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -4,7 +4,7 @@ import json import os import re import urllib.parse -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime from threading import Lock from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args @@ -125,7 +125,7 @@ class BaseAWSLLM: return get_ssl_verify(ssl_verify=ssl_verify) - def get_cache_key(self, credential_args: dict[str, str | None]) -> str: + def get_cache_key(self, credential_args: Mapping[str, str | bool | None]) -> str: """ Generate a unique cache key based on the credential arguments. """ @@ -135,8 +135,8 @@ class BaseAWSLLM: def _get_or_set_cached_credentials( self, - credential_args: dict[str, str | None], - credential_fetcher: Callable[[], tuple[Any, int | None]], + credential_args: Mapping[str, str | bool | None], + credential_fetcher: Callable[[], tuple[Credentials, int | None]], ) -> Any: """ Read-through IAM cache on the process-wide ``DualCache``. @@ -271,7 +271,19 @@ class BaseAWSLLM: aws_external_id, ) - args: Final = {k: v for k, v in locals().items() if k.startswith("aws_") or k == "ssl_verify"} + args: Final = { + "aws_access_key_id": aws_access_key_id, + "aws_secret_access_key": aws_secret_access_key, + "aws_session_token": aws_session_token, + "aws_region_name": aws_region_name, + "aws_session_name": aws_session_name, + "aws_profile_name": aws_profile_name, + "aws_role_name": aws_role_name, + "aws_web_identity_token": aws_web_identity_token, + "aws_sts_endpoint": aws_sts_endpoint, + "aws_external_id": aws_external_id, + "ssl_verify": ssl_verify, + } ######################################################### # Handle diff boto3 auth flows diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index dd20a8c2ed4..d4f1f0a4f1b 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -141,7 +141,7 @@ def _build_query_params( def _error_message_from_response(response: httpx.Response) -> str: try: - body: Final = response.json() + body: Final[object] = response.json() except ValueError: return response.text @@ -330,7 +330,7 @@ class GenericContainerHandler: timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, **kwargs: object, - ) -> Any: + ) -> ContainerEndpointResponse: """Synchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) if not endpoint_config: @@ -410,7 +410,7 @@ class GenericContainerHandler: timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, **kwargs: object, - ) -> Any: + ) -> ContainerEndpointResponse: """Asynchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) if not endpoint_config: diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index d220742b92b..1189af2d6a3 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -117,7 +117,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): _snake_to_camel, ) - _generate_content_config_dict: Final[dict[str, Any]] = {} + _generate_content_config_dict: Final[dict[str, object]] = {} supported_google_genai_params: Final = self.get_supported_generate_content_optional_params(model) # Create a set with both camelCase and snake_case versions for faster lookup supported_params_set: Final = set(supported_google_genai_params) @@ -175,7 +175,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): def _get_common_auth_components( self, litellm_params: dict, - ) -> tuple[Any, str | None, str | None]: + ) -> tuple[str | None, str | None, str | None]: """ Get common authentication components used by both sync and async methods. @@ -193,7 +193,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): auth_header: str | None, vertex_project: str | None, vertex_location: str | None, - vertex_credentials: Any, + vertex_credentials: str | None, stream: bool, api_base: str | None, litellm_params: dict, diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 199599d6b9c..a8f3388d092 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -6,8 +6,8 @@ Why separate file? Make it easy to see how transformation works Docs - https://jina.ai/reranker """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final from httpx import URL, Response @@ -39,7 +39,7 @@ class JinaAIRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: Sequence[str | Mapping[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 9d365572e6f..73f6ed23092 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -6,7 +6,7 @@ Used by the transformation layer and skills injection hook. """ import uuid -from typing import Any, Final +from typing import Final from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -76,7 +76,7 @@ class LiteLLMSkillsHandler: # this module FastAPI-free per the project layering rule. raise ValueError("Unable to record skill ownership: caller has no identity scope.") - skill_data: Final[dict[str, Any]] = { + skill_data: Final[dict[str, object]] = { "skill_id": skill_id, "display_title": data.display_title, "description": data.description, @@ -115,22 +115,24 @@ class LiteLLMSkillsHandler: verbose_logger.debug("LiteLLMSkillsHandler: Listing skills with limit=%s, offset=%s", limit, offset) - find_many_kwargs: Final[dict[str, Any]] = { - "take": limit, - "skip": offset, - "order": {"created_at": "desc"}, - } - if user_api_key_dict is not None and not is_proxy_admin(user_api_key_dict): - owner_scopes: Final = get_resource_owner_scopes(user_api_key_dict) - if not owner_scopes: - return [] - find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} + owner_scopes: Final = ( + get_resource_owner_scopes(user_api_key_dict) + if user_api_key_dict is not None and not is_proxy_admin(user_api_key_dict) + else None + ) + if owner_scopes is not None and not owner_scopes: + return [] - skills: Final = await SkillsRepository(prisma_client).table.find_many(**find_many_kwargs) + skills: Final = await SkillsRepository(prisma_client).table.find_many( + take=limit, + skip=offset, + order={"created_at": "desc"}, + where={"created_by": {"in": owner_scopes}} if owner_scopes else None, + ) return [_prisma_skill_to_litellm(s) for s in skills] @staticmethod - async def _load_skill(skill_id: str) -> Any | None: + async def _load_skill(skill_id: str) -> object | None: """Cache-first read of the Prisma skill row. Owner-scope filtering happens on the cached row, so the cache is per-skill not per-caller. """ diff --git a/litellm/llms/litellm_proxy/skills/sandbox_executor.py b/litellm/llms/litellm_proxy/skills/sandbox_executor.py index cdf4f8511e2..3e38dd81905 100644 --- a/litellm/llms/litellm_proxy/skills/sandbox_executor.py +++ b/litellm/llms/litellm_proxy/skills/sandbox_executor.py @@ -7,11 +7,40 @@ Supports Docker, Podman, and Kubernetes backends. import base64 import os -from typing import Any, Final +from typing import Any, Final, Protocol, TypedDict + +from typing_extensions import ReadOnly from litellm._logging import verbose_logger +class _SandboxRunResult(Protocol): + """Result of running code inside an llm-sandbox session.""" + + @property + def exit_code(self) -> int: ... + + @property + def stdout(self) -> str | None: ... + + +class _SandboxSession(Protocol): + """The subset of an llm-sandbox session used while collecting generated files.""" + + def run(self, code: str, /) -> _SandboxRunResult: ... + + def copy_from_runtime(self, src: str, dest: str, /) -> object: ... + + +class _GeneratedFile(TypedDict): + """A file produced inside the sandbox and carried back out as base64.""" + + name: ReadOnly[str] + path: ReadOnly[str] + content_base64: ReadOnly[str] + mime_type: ReadOnly[str] + + class SkillsSandboxExecutor: """ Executes skill code in llm-sandbox Docker container. @@ -77,7 +106,7 @@ class SkillsSandboxExecutor: try: # Create sandbox session - session_kwargs: Final[dict[str, Any]] = { + session_kwargs: Final[dict[str, object]] = { "lang": "python", "verbose": False, } @@ -197,9 +226,9 @@ sys.path.insert(0, '/sandbox') def _collect_generated_files( self, - session: Any, + session: _SandboxSession, original_files: dict[str, bytes], - ) -> list[dict[str, Any]]: + ) -> list[_GeneratedFile]: """ Collect files generated during execution. @@ -213,7 +242,7 @@ sys.path.insert(0, '/sandbox') Returns: List of generated files with base64 content """ - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[_GeneratedFile]] = [] try: import tempfile diff --git a/litellm/llms/openai/fine_tuning/handler.py b/litellm/llms/openai/fine_tuning/handler.py index 7fb99d61475..1ff5909a103 100644 --- a/litellm/llms/openai/fine_tuning/handler.py +++ b/litellm/llms/openai/fine_tuning/handler.py @@ -1,13 +1,14 @@ -from collections.abc import Coroutine -from typing import Any, Final, cast +from collections.abc import Coroutine, Mapping +from typing import Final, cast import httpx from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from openai.types.fine_tuning import FineTuningJob from litellm._logging import verbose_logger from litellm.types.utils import LiteLLMFineTuningJob -_AZURE_STATUS_MAP: Final = { +_AZURE_STATUS_MAP: Final[Mapping[object, str]] = { "pending": "queued", "notRunning": "queued", "running": "running", @@ -20,7 +21,7 @@ _AZURE_STATUS_MAP: Final = { # because LiteLLMFineTuningJob schema has no intermediate cancellation state. -def _normalize_fine_tuning_job_dict(data: dict[str, Any], is_azure: bool = False) -> dict[str, Any]: +def _normalize_fine_tuning_job_dict(data: dict[str, object], is_azure: bool = False) -> dict[str, object]: """ Normalize Azure OpenAI FineTuningJob response to match OpenAI schema. @@ -47,7 +48,7 @@ def _normalize_fine_tuning_job_dict(data: dict[str, Any], is_azure: bool = False return normalized -def _litellm_fine_tuning_job_from_response(response: Any, is_azure: bool = False) -> LiteLLMFineTuningJob: +def _litellm_fine_tuning_job_from_response(response: FineTuningJob, is_azure: bool = False) -> LiteLLMFineTuningJob: return LiteLLMFineTuningJob(**_normalize_fine_tuning_job_dict(response.model_dump(), is_azure=is_azure)) @@ -111,7 +112,7 @@ class OpenAIFineTuningAPI: max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -159,7 +160,7 @@ class OpenAIFineTuningAPI: max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -258,7 +259,7 @@ class OpenAIFineTuningAPI: max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, - ) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: + ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index dba1e9d01d3..bc02d274f24 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -104,7 +104,7 @@ class OpenAIImageVariationsHandler: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text: Final = getattr(e, "text", str(e)) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) @@ -221,7 +221,7 @@ class OpenAIImageVariationsHandler: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text: Final = getattr(e, "text", str(e)) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 0343f22e7d1..e3ecbac1a53 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -4,6 +4,7 @@ This file contains the calling OpenAI's `/v1/realtime` endpoint. This requires websockets, and is currently only supported on LiteLLM Proxy. """ +import ssl from typing import Any, Final, cast from litellm._logging import _redact_string, verbose_logger @@ -56,7 +57,7 @@ class OpenAIRealtime(OpenAIChatCompletion): headers["OpenAI-Beta"] = "realtime=v1" return headers - def _get_ssl_config(self, url: str) -> Any: + def _get_ssl_config(self, url: str) -> bool | str | ssl.SSLContext | None: """ Get SSL configuration for WebSocket connection. Override this in subclasses to customize SSL behavior. @@ -111,12 +112,12 @@ class OpenAIRealtime(OpenAIChatCompletion): logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, - client: Any | None = None, + client: object | None = None, timeout: float | None = None, query_params: RealtimeQueryParams | None = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, litellm_metadata: dict | None = None, - **kwargs: Any, + **kwargs: object, ): import websockets from websockets.asyncio.client import ClientConnection diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 64018df8b7a..c05d943b7cf 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -32,12 +32,12 @@ class OpenAITokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, + contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, object]] | None = None, + system: object = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index 8a2064f1823..8519b5f4bc4 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -1,4 +1,5 @@ -from typing import Any, Final, cast +from collections.abc import Mapping +from typing import Final, cast import httpx @@ -22,7 +23,7 @@ from litellm.types.vector_store_files import ( from litellm.utils import add_openai_metadata -def _clean_dict(source: dict[str, Any]) -> dict[str, Any]: +def _clean_dict(source: Mapping[str, object]) -> dict[str, object]: return {k: v for k, v in source.items() if v is not None} @@ -30,7 +31,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): ASSISTANTS_HEADER_KEY = "OpenAI-Beta" ASSISTANTS_HEADER_VALUE = "assistants=v2" - def get_auth_credentials(self, litellm_params: dict[str, Any]) -> VectorStoreFileAuthCredentials: + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> VectorStoreFileAuthCredentials: api_key: Final = litellm_params.get("api_key") if api_key is None: raise ValueError("api_key is required") @@ -82,7 +83,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): *, api_base: str | None, vector_store_id: str, - litellm_params: dict[str, Any], + litellm_params: Mapping[str, object], ) -> str: base_url = ( api_base @@ -101,8 +102,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): vector_store_id: str, create_request: VectorStoreFileCreateRequest, api_base: str, - ) -> tuple[str, dict[str, Any]]: - payload: Final[dict[str, Any]] = _clean_dict(dict(create_request)) + ) -> tuple[str, dict[str, object]]: + payload: Final[dict[str, object]] = _clean_dict(dict(create_request)) attributes: Final = payload.get("attributes") if isinstance(attributes, dict): filtered_attributes: Final = add_openai_metadata(attributes) @@ -133,7 +134,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): vector_store_id: str, query_params: VectorStoreFileListQueryParams, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: params: Final = _clean_dict(dict(query_params)) return api_base, params @@ -157,7 +158,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): vector_store_id: str, file_id: str, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") return f"{api_base}/{encoded_file_id}", {} @@ -181,7 +182,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): vector_store_id: str, file_id: str, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") return f"{api_base}/{encoded_file_id}/content", {} @@ -206,8 +207,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): file_id: str, update_request: VectorStoreFileUpdateRequest, api_base: str, - ) -> tuple[str, dict[str, Any]]: - payload: Final[dict[str, Any]] = dict(update_request) + ) -> tuple[str, dict[str, object]]: + payload: Final[dict[str, object]] = dict(update_request) attributes: Final = payload.get("attributes") if isinstance(attributes, dict): filtered_attributes: Final = add_openai_metadata(attributes) @@ -238,7 +239,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): vector_store_id: str, file_id: str, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") return f"{api_base}/{encoded_file_id}", {} diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 3265537d1aa..0ebac5185d7 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -18,6 +18,8 @@ from litellm.types.utils import Choices, Message, ModelResponse, Usage from ..common_utils import PredibaseError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -64,8 +66,23 @@ class PredibaseConfig(BaseConfig): typical_p: float | None = None, watermark: bool | None = None, ) -> None: - locals_: Final = locals().copy() - for key, value in locals_.items(): + locals_: Final = ( + ("best_of", best_of), + ("decoder_input_details", decoder_input_details), + ("details", details), + ("max_new_tokens", max_new_tokens), + ("repetition_penalty", repetition_penalty), + ("return_full_text", return_full_text), + ("seed", seed), + ("stop", stop), + ("temperature", temperature), + ("top_k", top_k), + ("top_p", top_p), + ("truncate", truncate), + ("typical_p", typical_p), + ("watermark", watermark), + ) + for key, value in locals_: if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -133,7 +150,7 @@ class PredibaseConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -217,7 +234,7 @@ class PredibaseConfig(BaseConfig): # Keep usage calculation non-blocking if token counting fails. pass output_text: Final = model_response["choices"][0]["message"].get("content", "") - if output_text is not None and len(output_text) > 0: + if encoding is not None and output_text is not None and len(output_text) > 0: completion_tokens = 0 try: completion_tokens = len(encoding.encode(model_response["choices"][0]["message"].get("content", ""))) diff --git a/litellm/llms/soniox/audio_transcription/transformation.py b/litellm/llms/soniox/audio_transcription/transformation.py index 318444cffec..8507ae73305 100644 --- a/litellm/llms/soniox/audio_transcription/transformation.py +++ b/litellm/llms/soniox/audio_transcription/transformation.py @@ -152,7 +152,7 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): and for filling in `file_id`/`audio_url`. This method exists so the config can be exercised in isolation by unit tests. """ - body: Final[dict[str, Any]] = {"model": model} + body: Final[dict[str, object]] = {"model": model} for key in SONIOX_PASSTHROUGH_PARAMS: value = optional_params.get(key) @@ -247,9 +247,9 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # For verbose_json, include word-level timing from tokens. if response_format == "verbose_json" and tokens: - words: Final[list[dict[str, Any]]] = [] + words: Final[list[dict[str, object]]] = [] for token in tokens: - word_entry: dict[str, Any] = {"word": token.get("text", "")} + word_entry: dict[str, object] = {"word": token.get("text", "")} if token.get("start_ms") is not None: word_entry["start"] = float(token["start_ms"]) / 1000.0 if token.get("end_ms") is not None: diff --git a/litellm/llms/vertex_ai/rag_engine/ingestion.py b/litellm/llms/vertex_ai/rag_engine/ingestion.py index 06e525a90ff..d9916209a14 100644 --- a/litellm/llms/vertex_ai/rag_engine/ingestion.py +++ b/litellm/llms/vertex_ai/rag_engine/ingestion.py @@ -14,7 +14,7 @@ Key differences from OpenAI: from __future__ import annotations import os -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm import get_secret_str from litellm._logging import verbose_logger @@ -26,12 +26,12 @@ if TYPE_CHECKING: from litellm.types.rag import RAGIngestOptions -def _get_str_or_none(value: Any) -> str | None: +def _get_str_or_none(value: object) -> str | None: """Cast config value to Optional[str].""" return str(value) if value is not None else None -def _get_int(value: Any, default: int) -> int: +def _get_int(value: str | float | None, default: int) -> int: """Cast config value to int with default.""" if value is None: return default @@ -205,7 +205,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): ) verbose_logger.info("Import started asynchronously") - def _build_transformation_config(self) -> Any: + def _build_transformation_config(self) -> object: """ Build Vertex AI TransformationConfig from unified chunking_strategy. diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index e6e3c2739c1..c66ad8e38b0 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -265,7 +265,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict # Ensure litellm_params is a dict for type checking - params_dict: Final[dict[str, Any]] = cast(dict[str, Any], litellm_params) if litellm_params is not None else {} + params_dict: Final[dict[str, object]] = ( + cast(dict[str, object], litellm_params) if litellm_params is not None else {} + ) vertex_project: Final = VertexBase.safe_get_vertex_ai_project(litellm_params=params_dict) vertex_credentials: Final = VertexBase.safe_get_vertex_ai_credentials(litellm_params=params_dict) diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py index 4bbb537804c..814d5ab7eb0 100644 --- a/litellm/llms/voyage/embedding/transformation_multimodal.py +++ b/litellm/llms/voyage/embedding/transformation_multimodal.py @@ -6,7 +6,7 @@ containing content blocks, unlike standard Voyage embeddings which use /v1/embeddings and a string/list `input` field. """ -from typing import Any, Final +from typing import Final import httpx @@ -98,7 +98,7 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): ) return {"Authorization": f"Bearer {api_key}"} - def _normalize_content_item(self, item: dict[str, Any]) -> dict[str, Any]: + def _normalize_content_item(self, item: dict[str, object]) -> dict[str, object]: item_type: Final = item.get("type") if item_type == "image_url": image_url = item.get("image_url") @@ -115,7 +115,7 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): return {"type": "image_url", "image_url": image_url} return item - def _normalize_input_item(self, item: Any) -> dict[str, Any]: + def _normalize_input_item(self, item: object) -> object: if isinstance(item, str): return {"content": [{"type": "text", "text": item}]} if isinstance(item, dict) and "content" in item: diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index ee330c92f1a..fea8452d934 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -43,7 +43,7 @@ class VoyageRerankConfig(BaseRerankConfig): instruction: str | None = None, ) -> dict: # Voyage AI uses 'top_k' instead of 'top_n' - optional_params: Final[dict[str, Any]] = {"query": query, "documents": documents} + optional_params: Final[dict[str, object]] = {"query": query, "documents": documents} if top_n is not None: optional_params["top_k"] = top_n if return_documents is not None: @@ -109,7 +109,7 @@ class VoyageRerankConfig(BaseRerankConfig): # Transform to LiteLLM format transformed_results: Final = [] for result in _results: - transformed_result: dict[str, Any] = { + transformed_result: dict[str, object] = { "index": result["index"], "relevance_score": result["relevance_score"], } diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index ae5849812bf..4590bdd5aa3 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -196,7 +196,7 @@ class XAIChatConfig(OpenAIGPTConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "XAIChatCompletionStreamingHandler": return XAIChatCompletionStreamingHandler( streaming_response=streaming_response, sync_stream=sync_stream, diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index d79e7d4c146..36ae15e1df3 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,4 +1,5 @@ -from typing import Any, Final +from collections.abc import Mapping +from typing import Final import litellm from litellm._logging import verbose_logger @@ -8,7 +9,6 @@ from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfi from litellm.llms.xai.common_utils import XAIModelInfo from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams -from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders @@ -44,7 +44,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return supported_params - def _transform_web_search_tool(self, tool: dict[str, Any]) -> XAIWebSearchTool | dict[str, Any]: + def _transform_web_search_tool(self, tool: Mapping[str, object]) -> Mapping[str, object]: """ Transform web_search tool to XAI format. @@ -55,7 +55,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): XAI does NOT support search_context_size (OpenAI-specific). """ - xai_tool: Final[dict[str, Any]] = {"type": "web_search"} + xai_tool: Final[dict[str, object]] = {"type": "web_search"} # Remove search_context_size if present (not supported by XAI) if "search_context_size" in tool: @@ -83,7 +83,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return xai_tool - def _transform_x_search_tool(self, tool: dict[str, Any]) -> XAIXSearchTool | dict[str, Any]: + def _transform_x_search_tool(self, tool: Mapping[str, object]) -> Mapping[str, object]: """ Transform x_search tool to XAI format. @@ -95,7 +95,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): - enable_image_understanding - enable_video_understanding """ - xai_tool: Final[dict[str, Any]] = {"type": "x_search"} + xai_tool: Final[dict[str, object]] = {"type": "x_search"} # Handle allowed_x_handles if "allowed_x_handles" in tool: @@ -157,7 +157,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): if not isinstance(tools_list, list): tools_list = [tools_list] - transformed_tools: Final[list[Any]] = [] + transformed_tools: Final[list[object]] = [] for tool in tools_list: if isinstance(tool, dict): tool_type = tool.get("type") diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index 16acd95af9c..eccbf75667d 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request @@ -19,7 +19,7 @@ router: Final = APIRouter( ) -def _extract_cache_params() -> dict[str, Any]: +def _extract_cache_params() -> dict[str, object]: """ Safely extracts and cleans cache parameters. @@ -56,8 +56,8 @@ async def cache_ping(): """ Endpoint for checking if cache can be pinged """ - litellm_cache_params: dict[str, Any] = {} - cleaned_cache_params: dict[str, Any] = {} + litellm_cache_params: dict[str, object] = {} + cleaned_cache_params: dict[str, object] = {} if litellm.cache is None: raise ProxyException( message=safe_dumps( @@ -162,7 +162,7 @@ async def cache_delete(request: Request): ) -def _get_redis_client_info(cache_instance) -> tuple[list, int]: +def _get_redis_client_info(cache_instance: RedisCache) -> tuple[list[object], int]: """ Helper function to safely get Redis client list information. diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index bd4d0df3ed0..a330d057490 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -73,7 +73,7 @@ class ChatClient: url: Final = f"{self._base_url}/chat/completions" # Build request data with required fields - data: Final[dict[str, Any]] = {"model": model, "messages": messages} + data: Final[dict[str, object]] = {"model": model, "messages": messages} # Add optional parameters if provided if temperature is not None: @@ -143,7 +143,7 @@ class ChatClient: url: Final = f"{self._base_url}/chat/completions" # Build request data with required fields - data: Final[dict[str, Any]] = {"model": model, "messages": messages, "stream": True} + data: Final[dict[str, object]] = {"model": model, "messages": messages, "stream": True} # Add optional parameters if provided if temperature is not None: diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index c591cbabee1..5cf0bd6f89f 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -8,7 +8,7 @@ from typing import Final import click import requests -from .auth import context_secret_vault, get_stored_api_key, login +from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login from .cmd_quoting import quote_for_cmd ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" @@ -289,8 +289,9 @@ def _is_interactive() -> bool: def resolve_api_key(ctx: click.Context) -> str: - base_url: Final = ctx.obj["base_url"] - api_key = ctx.obj.get("api_key") + ctx_obj: Final[CliContextObj] = ctx.obj + base_url: Final = ctx_obj["base_url"] + api_key = ctx_obj.get("api_key") if api_key: return api_key @@ -312,7 +313,8 @@ _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: - base_url: Final = ctx.obj["base_url"] + ctx_obj: Final[CliContextObj] = ctx.obj + base_url: Final = ctx_obj["base_url"] started_interactive: Final = _is_interactive() api_key: Final = resolve_api_key(ctx) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index fe100c5f676..028b338f412 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -1,4 +1,5 @@ import builtins +from collections.abc import Mapping from typing import Any, Final import requests @@ -72,7 +73,7 @@ class KeysManagementClient: requests.exceptions.RequestException: If the request fails with any other error """ url: Final = f"{self._base_url}/key/list" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, int | str]] = {} # Add optional query parameters if page is not None: @@ -119,9 +120,9 @@ class KeysManagementClient: team_id: str | None = None, user_id: str | None = None, budget_id: str | None = None, - config: dict[str, Any] | None = None, + config: Mapping[str, object] | None = None, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Generate an API key based on the provided data. @@ -149,7 +150,7 @@ class KeysManagementClient: """ url: Final = f"{self._base_url}/key/generate" - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} if models is not None: data["models"] = models if aliases is not None: @@ -189,7 +190,7 @@ class KeysManagementClient: keys: builtins.list[str] | None = None, key_aliases: builtins.list[str] | None = None, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Delete existing keys @@ -238,7 +239,7 @@ class KeysManagementClient: key_alias: str | None = None, team_id: str | None = None, user_id: str | None = None, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Update an existing API key's parameters. @@ -261,7 +262,7 @@ class KeysManagementClient: """ url: Final = f"{self._base_url}/key/update" - data: Final[dict[str, Any]] = {"key": key} + data: Final[dict[str, object]] = {"key": key} if key_alias is not None: data["key_alias"] = key_alias @@ -288,7 +289,7 @@ class KeysManagementClient: except Exception: raise Exception(f"Error updating key: {response_text}") - def info(self, key: str, return_request: bool = False) -> dict[str, Any] | requests.Request: + def info(self, key: str, return_request: bool = False) -> dict[str, object] | requests.Request: """ Get information about API keys. diff --git a/litellm/proxy/common_utils/performance_utils.py b/litellm/proxy/common_utils/performance_utils.py index 5d5334f2177..0b79599e8f6 100644 --- a/litellm/proxy/common_utils/performance_utils.py +++ b/litellm/proxy/common_utils/performance_utils.py @@ -15,10 +15,24 @@ import inspect import threading from collections.abc import Callable from pathlib import Path as PathLib -from typing import Any, Final +from types import ModuleType +from typing import Final, Protocol, TextIO from litellm._logging import verbose_proxy_logger + +class _LineProfiler(Protocol): + """The line_profiler.LineProfiler surface this module drives.""" + + def __call__(self, func: Callable[..., object]) -> Callable[..., object]: ... + + def add_function(self, func: Callable[..., object]) -> object: ... + + def dump_stats(self, filename: str) -> object: ... + + def print_stats(self, stream: TextIO) -> object: ... + + # Global profiling state _profile_lock: Final = threading.Lock() _profiler = None @@ -27,7 +41,7 @@ _sample_counter = 0 _sample_counter_lock: Final = threading.Lock() # Global line_profiler state -_line_profiler: Any | None = None +_line_profiler: _LineProfiler | None = None _line_profiler_lock: Final = threading.Lock() _wrapped_functions: Final[dict[str, Callable]] = {} # Store original functions @@ -157,7 +171,7 @@ def enable_line_profiler() -> None: verbose_proxy_logger.info("Line profiler enabled") -def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool: +def wrap_function_with_line_profiler(module: ModuleType, function_name: str) -> bool: """Dynamically wrap a function with line_profiler. Args: diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 4a088140725..85ef469ee69 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -1,6 +1,6 @@ #### Container Endpoints ##### -from typing import Any, Final +from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import ORJSONResponse @@ -208,7 +208,7 @@ async def list_containers( # Read query parameters query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = {"query_params": query_params} + data: Final[dict[str, object]] = {"query_params": query_params} # Extract custom_llm_provider using priority chain custom_llm_provider: Final = ( @@ -312,7 +312,7 @@ async def retrieve_container( ) # Include container_id in request data - data: Final[dict[str, Any]] = {"container_id": container_id} + data: Final[dict[str, object]] = {"container_id": container_id} # Extract custom_llm_provider using priority chain custom_llm_provider = ( @@ -417,7 +417,7 @@ async def delete_container( ) # Include container_id in request data - data: Final[dict[str, Any]] = {"container_id": container_id} + data: Final[dict[str, object]] = {"container_id": container_id} # Extract custom_llm_provider using priority chain custom_llm_provider = ( diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 5502543b926..f7362d3b809 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -204,7 +204,7 @@ class PrismaDBExceptionHandler: if isinstance(e, prisma.errors.PrismaError): return False - tb = getattr(e, "__traceback__", None) + tb = e.__traceback__ if hasattr(e, "__traceback__") else None while tb is not None: if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"): return True @@ -318,7 +318,7 @@ _DEFAULT_RECONNECT_TIMEOUT_SECONDS: Final = 2.0 _DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS: Final = 0.1 -def _coerce_timeout(value: Any, fallback: float) -> float: +def _coerce_timeout(value: object, fallback: float) -> float: """Return `value` if it is a real int/float, else `fallback`. Guards against tests that mock `prisma_client` and leave the timeout slots as MagicMock instances.""" diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 4d17c6edb31..42f0220cc4d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -45,7 +45,7 @@ class AzureGuardrailBase: self.api_base = api_base self.api_version: str = kwargs.get("api_version") or "2024-09-01" - async def _post_to_content_safety(self, endpoint_path: str, request_body: dict[str, Any]) -> dict[str, Any]: + async def _post_to_content_safety(self, endpoint_path: str, request_body: dict[str, object]) -> dict[str, Any]: """POST to an Azure Content Safety endpoint with standard auth headers. Args: @@ -94,7 +94,7 @@ class AzureGuardrailBase: # Tokenize into alternating non-whitespace and whitespace runs so # that original newlines, tabs, and multiple spaces are preserved # within each chunk. - tokens: Final = re.findall(r"\S+|\s+", text) + tokens: Final = [match.group(0) for match in re.finditer(r"\S+|\s+", text)] chunks: Final[list[str]] = [] current_chunk = "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 07e435c675b..0dca8be3307 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -3,7 +3,7 @@ Azure Text Moderation Native Guardrail Integrationfor LiteLLM """ -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Union, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast from fastapi import HTTPException @@ -14,18 +14,18 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs, LLMResponseTypes from .base import AzureGuardrailBase if TYPE_CHECKING: + from litellm.caching.caching import DualCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import ( AzureTextModerationGuardrailResponse, ) from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel - from litellm.types.utils import EmbeddingResponse, ImageResponse, ModelResponse class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardrail): @@ -219,10 +219,10 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr async def async_pre_call_hook( self, user_api_key_dict: "UserAPIKeyAuth", - cache: Any, + cache: "DualCache", data: dict[str, Any], call_type: CallTypesLiteral, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """ Pre-call hook to scan user prompts before sending to LLM. @@ -251,8 +251,8 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr self, data: dict, user_api_key_dict: "UserAPIKeyAuth", - response: Union[Any, "ModelResponse", "EmbeddingResponse", "ImageResponse"], - ) -> Any: + response: LLMResponseTypes, + ) -> LLMResponseTypes: from litellm.types.utils import Choices, ModelResponse if isinstance(response, ModelResponse) and response.choices: @@ -267,7 +267,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr ) return response - async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> Any: + async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> str: try: if response is not None and len(response) > 0: await self.async_make_request( @@ -281,7 +281,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr return f"data: {error_returned}\n\n" -def _message_content_to_text(content: Any) -> str: +def _message_content_to_text(content: object) -> str: if isinstance(content, str): return content if isinstance(content, list): diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py index 5feeafe8e95..64770fdf0f7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -1,6 +1,6 @@ """Block Code Execution guardrail: blocks or masks fenced code blocks by language.""" -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, cast from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations @@ -20,8 +20,8 @@ def _get_param( litellm_params: "LitellmParams", guardrail: "Guardrail", key: str, - default: Any = None, -) -> Any: + default: object = None, +) -> object: """Get a param from litellm_params, with fallback to raw guardrail litellm_params (for extra fields not on LitellmParams).""" value: Final = getattr(litellm_params, key, default) if value is not None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 53da8aeed42..24801aa2df1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -8,7 +8,7 @@ and provide safe, sandboxed functionality for common guardrail operations. import json import re from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Final from urllib.parse import urlparse import httpx @@ -16,7 +16,7 @@ from pydantic import JsonValue from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider # ============================================================================= @@ -508,7 +508,7 @@ async def http_request( async def _execute_http_request( - client: Any, + client: AsyncHTTPHandler, method: str, url: str, headers: dict[str, str] | None, 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 e3cf645ceaf..d8296003ae9 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 @@ -7,6 +7,7 @@ import fnmatch import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional import httpx @@ -23,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, @@ -73,7 +75,7 @@ def _header_value_allowed( def _sanitize_inbound_headers( - headers: Any, + headers: object, extra_allowlist: set[str] | None = None, ) -> dict[str, str] | None: """ @@ -175,7 +177,7 @@ class GenericGuardrailAPI(CustomGuardrail): headers: dict[str, Any] | None = None, api_base: str | None = None, api_key: str | None = None, - additional_provider_specific_params: dict[str, Any] | None = None, + additional_provider_specific_params: Mapping[str, object] | None = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", fail_on_error: bool | None = True, extra_headers: list | None = None, @@ -318,8 +320,8 @@ class GenericGuardrailAPI(CustomGuardrail): self, *, texts: list, - images: Any, - tools: Any, + images: list[str] | None, + tools: list[ChatCompletionToolParam] | None, guardrail_response: GenericGuardrailAPIResponse, ) -> GenericGuardrailAPIInputs: # Action is NONE or no modifications needed diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index d187b5b12e9..5d11c3643cd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -33,6 +33,7 @@ from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( ) from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.utils import ( CallTypes, CallTypesLiteral, @@ -97,7 +98,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): template_id: str | None = None, project_id: str | None = None, location: str | None = None, - credentials: Any | None = None, + credentials: VERTEX_CREDENTIALS_TYPES | None = None, api_endpoint: str | None = None, sanitize_error_detail: "bool | None" = True, **kwargs, @@ -147,7 +148,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else: return {"modelResponseData": {"text": content}} - def _extract_content_from_response(self, response: Any | ModelResponse) -> str: + def _extract_content_from_response(self, response: object) -> str: """ Extract text content from model response. diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 385e7d61dee..7ef0a9f73f3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -72,7 +72,7 @@ class NomaBlockedMessage(HTTPException): }, ) - def _is_result_true(self, result_obj: dict[str, Any] | None) -> bool: + def _is_result_true(self, result_obj: dict[str, object] | None) -> bool: """ Check if a result object has a "result" field that is True. @@ -454,7 +454,7 @@ class NomaGuardrail(CustomGuardrail): return False - def _is_result_true(self, result_obj: dict[str, Any] | None) -> bool: + def _is_result_true(self, result_obj: dict[str, object] | None) -> bool: """ Check if a result object has a "result" field that is True. diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 3d5d87e4d17..5acf837cf84 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -35,14 +35,14 @@ class PangeaGuardrailMissingSecrets(Exception): class _TextCompletionRequest: - def __init__(self, body): + def __init__(self, body: dict[str, object]) -> None: self.body = body def get_messages(self) -> list[dict]: return [{"role": "user", "content": self.body["prompt"]}] # This mutates the original dict, but we'll still return it anyways - def update_original_body(self, prompt_messages: list[dict]) -> Any: + def update_original_body(self, prompt_messages: list[dict]) -> dict[str, object]: assert len(prompt_messages) == 1 self.body["prompt"] = prompt_messages[0]["content"] return self.body @@ -159,7 +159,7 @@ class PangeaHandler(CustomGuardrail): call_type: str, ): transformer = None - messages: Any = None + messages: object = None if call_type == "text_completion" or call_type == "atext_completion": transformer = _TextCompletionRequest(data) messages = transformer.get_messages() diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 5f07f529e7a..b73d3adb99e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -721,7 +721,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): } } - def _record_scan_id(self, request_data: dict[str, Any], scan_result: Mapping[str, object]) -> None: + def _record_scan_id(self, request_data: dict[str, object], scan_result: Mapping[str, object]) -> None: """Surface the AIRS scan id on the response, so allowed calls are auditable too.""" scan_id: Final = scan_result.get("scan_id") add_guardrail_scan_id(request_data=request_data, scan_id=str(scan_id) if scan_id else None) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index b8ae09afc00..4f1d2380520 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -172,7 +172,7 @@ def _guardrail_status_to_action(status: str | None) -> str: return "passed" -def _parse_guardrail_info_from_payload(payload: Mapping[str, Any]) -> Sequence[Mapping[str, Any]]: +def _parse_guardrail_info_from_payload(payload: Mapping[str, object]) -> Sequence[Mapping[str, Any]]: """Extract guardrail_information from spend log payload metadata.""" meta = payload.get("metadata") if not meta: @@ -197,7 +197,7 @@ def _date_str(dt: datetime) -> str: return dt.astimezone(timezone.utc).strftime("%Y-%m-%d") -def _parse_payload_start_time(payload: Mapping[str, Any]) -> datetime | None: +def _parse_payload_start_time(payload: Mapping[str, object]) -> datetime | None: start_time: Final = payload.get("startTime") if isinstance(start_time, datetime): return start_time @@ -209,7 +209,9 @@ def _parse_payload_start_time(payload: Mapping[str, Any]) -> datetime | None: return None -def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Iterator[tuple[_UsageUnitKey, int]]: +def _iter_usage_unit_increments( + logs_to_process: Sequence[Mapping[str, object]], +) -> Iterator[tuple[_UsageUnitKey, int]]: for payload in logs_to_process: start_time = _parse_payload_start_time(payload) if not payload.get("request_id") or start_time is None: @@ -227,7 +229,7 @@ def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> yield _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)), units -def _sum_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Mapping[_UsageUnitKey, int]: +def _sum_usage_unit_increments(logs_to_process: Sequence[Mapping[str, object]]) -> Mapping[_UsageUnitKey, int]: ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=itemgetter(0)) return MappingProxyType( {key: sum(units for _, units in group) for key, group in groupby(ordered, key=itemgetter(0))} @@ -284,7 +286,7 @@ async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg async def process_spend_logs_guardrail_usage( prisma_client: PrismaClient, - logs_to_process: list[dict[str, Any]], + logs_to_process: Sequence[Mapping[str, object]], sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, pending: PendingRollups = _PENDING_ROLLUPS, ) -> None: @@ -295,7 +297,7 @@ async def process_spend_logs_guardrail_usage( if not logs_to_process: return # Aggregate daily metrics by (guardrail_id, date). Latency/score metrics dropped. - daily_guardrail: Final[dict[_MetricsKey, dict[str, Any]]] = defaultdict( + daily_guardrail: Final[dict[_MetricsKey, dict[str, int]]] = defaultdict( lambda: { "requests_evaluated": 0, "passed_count": 0, @@ -303,7 +305,7 @@ async def process_spend_logs_guardrail_usage( "flagged_count": 0, } ) - index_rows: Final[list[dict[str, Any]]] = [] + index_rows: Final[list[dict[str, object]]] = [] for payload in logs_to_process: request_id = payload.get("request_id") diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index f12cee4b636..79d54df97ae 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -1,6 +1,7 @@ import asyncio import json import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger @@ -143,8 +144,8 @@ class SharedHealthCheckManager: async def cache_health_check_results( self, - healthy_endpoints: list[dict[str, Any]], - unhealthy_endpoints: list[dict[str, Any]], + healthy_endpoints: Sequence[Mapping[str, object]], + unhealthy_endpoints: Sequence[Mapping[str, object]], ) -> None: """ Cache health check results in Redis. @@ -336,14 +337,14 @@ class SharedHealthCheckManager: verbose_proxy_logger.error("Error checking health check lock status: %s", str(e)) return False - async def get_health_check_status(self) -> dict[str, Any]: + async def get_health_check_status(self) -> dict[str, object]: """ Get the current status of health check coordination. Returns: Dict containing status information """ - status: Final = { + status: Final[dict[str, object]] = { "pod_id": self.pod_id, "redis_available": self.redis_cache is not None, "lock_ttl": self.lock_ttl, diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 5b814ad28fd..dcd34a1d9cb 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -62,6 +62,7 @@ from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.caching.caching import DualCache from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitDescriptor as _RateLimitDescriptor, ) @@ -73,8 +74,9 @@ if TYPE_CHECKING: ) from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.router import Router as _Router + from litellm.types.llms.openai import HttpxBinaryResponseContent - Span = _Span | Any + Span = _Span InternalUsageCache = _InternalUsageCache Router = _Router ParallelRequestLimiter = _ParallelRequestLimiter @@ -1011,7 +1013,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, file_id: str, user_api_key_dict: UserAPIKeyAuth, - ) -> Any: + ) -> "HttpxBinaryResponseContent": """ Fetch file content from managed files hook. @@ -1062,7 +1064,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - cache: Any, + cache: "DualCache", data: dict, call_type: str, ) -> Exception | str | dict | None: diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 88803d6442d..cdaa6d5a81c 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -1,7 +1,7 @@ import asyncio import json from datetime import datetime, timezone -from typing import Any, Final +from typing import Final import litellm from litellm._logging import verbose_proxy_logger @@ -89,8 +89,8 @@ class KeyManagementEventHooks: @staticmethod async def async_key_updated_hook( data: UpdateKeyRequest, - existing_key_row: Any, - response: Any, + existing_key_row: LiteLLM_VerificationToken, + response: object, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None = None, ): diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 2241884faf1..b5773e3e884 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,4 +1,5 @@ import math +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional, Union from fastapi import HTTPException, status @@ -438,7 +439,7 @@ _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: Final = ( ) -def _is_set_budget_value(value: Any) -> bool: +def _is_set_budget_value(value: object) -> bool: if value is None: return False if isinstance(value, list) and len(value) == 0: @@ -446,7 +447,7 @@ def _is_set_budget_value(value: Any) -> bool: return True -def _has_meaningful_budget_limit(budget_values: dict[str, Any]) -> bool: +def _has_meaningful_budget_limit(budget_values: Mapping[str, object]) -> bool: """A budget is meaningful if at least one limit is actually set; an empty list (no model restriction) and None both count as unset.""" return any(_is_set_budget_value(budget_values.get(field)) for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS) @@ -590,7 +591,7 @@ def _update_metadata_field(updated_kv: dict, field_name: str) -> None: updated_kv["metadata"] = {field_name: _value} -def _has_non_empty_value(value: Any) -> bool: +def _has_non_empty_value(value: object) -> bool: """Check if a value has real content (not None, not empty list, not blank string).""" if value is None: return False diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index 7f9cd251a8a..f41bf4dbd93 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -2,7 +2,7 @@ import json import time -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -24,6 +24,9 @@ from litellm.types.realtime import ( RealtimeTranscriptionSessionResponse, ) +if TYPE_CHECKING: + from litellm.router import Router + router: Final = APIRouter() _REALTIME_TOKEN_VERSION: Final = "realtime_v1" @@ -38,7 +41,7 @@ def _coerce_realtime_session_type(session_type: str | None) -> str: return "realtime" -def _append_model_candidate(candidates: list[str], model: Any) -> None: +def _append_model_candidate(candidates: list[str], model: object) -> None: if isinstance(model, str) and model and model not in candidates: candidates.append(model) @@ -116,7 +119,7 @@ async def _prepare_client_secret_session( req: RealtimeClientSecretRequest, user_api_key_dict: UserAPIKeyAuth, llm_model_list: list | None, - llm_router: Any, + llm_router: "Router | None", ) -> tuple[str, dict | None, str]: session_type: Final = _coerce_realtime_session_type(req.session.type if req.session else None) session_data: Final[dict | None] = req.session.model_dump(exclude_none=True) if req.session else None @@ -171,7 +174,7 @@ def _encode_realtime_token_payload( Encode metadata with the upstream ephemeral key so /realtime/calls can route without requiring model as a query param. """ - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, str | int | None]] = { "v": _REALTIME_TOKEN_VERSION, "ephemeral_key": ephemeral_key, "model_id": model_id, diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 3dfb67efb50..fe7fa79a3d9 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -89,7 +89,7 @@ class ResponsePollingHandler: error: dict | None = None, incomplete_details: dict | None = None, reasoning: dict | None = None, - tool_choice: Any | None = None, + tool_choice: object | None = None, tools: list | None = None, output: list | None = None, # Additional ResponsesAPIResponse fields diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 93f1510bf22..f224e02db32 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,6 +1,7 @@ import json import re -from typing import Any, Final, Literal +from collections.abc import Mapping +from typing import Final, Literal from fastapi import HTTPException, Request @@ -291,8 +292,8 @@ def _does_endpoint_match(endpoint_path: str, request_path: str) -> bool: def check_vector_store_permission( index_name: str, permission: str, - key_metadata: dict[str, Any] | None, - team_metadata: dict[str, Any] | None, + key_metadata: Mapping[str, object] | None, + team_metadata: Mapping[str, object] | None, ) -> bool: """ Check if a specific permission is allowed for a given vector store index. diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index 3d7056f8176..e412b7bcc8d 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -14,7 +14,7 @@ from __future__ import annotations import asyncio import json import uuid -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_arn_prefix @@ -26,12 +26,12 @@ if TYPE_CHECKING: from litellm.types.rag import RAGIngestOptions -def _get_str_or_none(value: Any) -> str | None: +def _get_str_or_none(value: object) -> str | None: """Cast config value to Optional[str].""" return str(value) if value is not None else None -def _get_int(value: Any, default: int) -> int: +def _get_int(value: str | float | None, default: int) -> int: """Cast config value to int with default.""" if value is None: return default @@ -122,7 +122,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): self._config_initialized = False # Track resources we create (for cleanup if needed) - self._created_resources: dict[str, Any] = {} + self._created_resources: dict[str, object] = {} async def _ensure_config_initialized(self): """Lazily initialize KB config - either detect from existing or create new.""" @@ -233,7 +233,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): verbose_logger.debug("Creating S3 bucket: %s", bucket_name) - create_params: Final[dict[str, Any]] = {"Bucket": bucket_name} + create_params: Final[dict[str, object]] = {"Bucket": bucket_name} if self.aws_region_name != "us-east-1": create_params["CreateBucketConfiguration"] = {"LocationConstraint": self.aws_region_name} diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 18cf884f267..b67d9e87831 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -21,7 +21,7 @@ class PrismaTableRepository(Generic[RowT_co]): table_name: str - def __init__(self, prisma_client: Any): + def __init__(self, prisma_client: object): self._prisma_client = prisma_client @property diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index eebe81ebba1..a1b67eaeaf9 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -14,7 +14,7 @@ from litellm.types.utils import LiteLLMPydanticObjectBase if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 3d35751394e..b623e31ce06 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,7 +37,7 @@ Safe to enable globally: """ import time -from typing import TYPE_CHECKING, Any, Final, Optional, cast +from typing import TYPE_CHECKING, Final, Optional, Protocol, cast import httpx @@ -51,11 +51,20 @@ from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues +from litellm.types.router import Deployment if TYPE_CHECKING: from litellm.router import Router +class _SupportsActiveCooldowns(Protocol): + """Cooldown-cache handle: this check only reads back the currently active cooldowns.""" + + async def async_get_active_cooldowns( + self, model_ids: list[str], parent_otel_span: Span | None + ) -> list[tuple[str, CooldownCacheValue]]: ... + + class EncryptedContentAffinityCheck(CustomLogger): """ Routes follow-up Responses API requests to the deployment that produced @@ -99,7 +108,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ) @staticmethod - def _extract_model_id_from_input(request_input: Any) -> str | None: + def _extract_model_id_from_input(request_input: object) -> str | None: """ Scan ``input`` items for litellm-encoded encrypted-content markers and return the ``model_id`` embedded in the first one found. @@ -151,7 +160,7 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _encryption_boundary_key( - litellm_params: Any, + litellm_params: object, ) -> tuple | None: """ ``(api_base, api_key)`` pair identifying an Azure resource. Two @@ -179,7 +188,7 @@ class EncryptedContentAffinityCheck(CustomLogger): self, healthy_deployments: list[dict], model_id: str, - ) -> tuple[list[dict], Any]: + ) -> tuple[list[dict], Deployment | None]: """ Deployments in ``healthy_deployments`` sharing the originating deployment's ``(api_base, api_key)``, alongside the originating @@ -289,7 +298,7 @@ class EncryptedContentAffinityCheck(CustomLogger): self, model: str, model_id: str, - originating: Any, + originating: Deployment | None, parent_otel_span: Span | None, ) -> Exception: # Public error messages intentionally omit the originating ``model_id`` so @@ -347,7 +356,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ) -> CooldownCacheValue | None: if self.router is None: return None - cooldown_cache: Final = getattr(self.router, "cooldown_cache", None) + cooldown_cache: Final[_SupportsActiveCooldowns | None] = getattr(self.router, "cooldown_cache", None) if cooldown_cache is None: return None try: diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 817c008fad3..39708e168f5 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: from litellm.router import Router litellm_router = Router - Span = _Span | Any + Span = _Span else: Span = Any litellm_router = Any @@ -34,7 +34,7 @@ class PromptCachingCache: self.in_memory_cache = InMemoryCache() @staticmethod - def serialize_object(obj: Any) -> Any: + def serialize_object(obj: Any) -> object: """Helper function to serialize Pydantic objects, dictionaries, or fallback to string.""" if hasattr(obj, "dict"): # If the object is a Pydantic model, use its `dict()` method diff --git a/litellm/secret_managers/custom_secret_manager_loader.py b/litellm/secret_managers/custom_secret_manager_loader.py index 14144b7230f..32c162ffc11 100644 --- a/litellm/secret_managers/custom_secret_manager_loader.py +++ b/litellm/secret_managers/custom_secret_manager_loader.py @@ -38,7 +38,9 @@ def load_custom_secret_manager(config_file_path: str | None = None) -> None: "CustomSecretManagerException - key_management_settings is required with custom_secret_manager field" ) - custom_secret_manager_path: Final = getattr(litellm._key_management_settings, "custom_secret_manager", None) + custom_secret_manager_path: Final[str | None] = getattr( + litellm._key_management_settings, "custom_secret_manager", None + ) if not custom_secret_manager_path: raise ValueError( diff --git a/litellm/types/containers/main.py b/litellm/types/containers/main.py index 27cbf437430..6a339fd2eac 100644 --- a/litellm/types/containers/main.py +++ b/litellm/types/containers/main.py @@ -1,7 +1,9 @@ +import builtins +from collections.abc import Mapping from typing import Any, Literal from pydantic import BaseModel -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class ExpiresAfter(BaseModel): @@ -23,15 +25,15 @@ class ContainerObject(BaseModel): name: str | None = None _hidden_params: dict[str, Any] = {} - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default: builtins.object = None) -> builtins.object: # Custom .get() method to access attributes with a default value if the attribute doesn't exist return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str) -> builtins.object: # Allow dictionary-style access to attributes return getattr(self, key) @@ -50,13 +52,13 @@ class DeleteContainerResult(BaseModel): object: Literal["container.deleted"] deleted: bool - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default: builtins.object = None) -> builtins.object: return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str) -> builtins.object: return getattr(self, key) def json(self, **kwargs): @@ -75,13 +77,13 @@ class ContainerListResponse(BaseModel): last_id: str | None = None has_more: bool - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default: builtins.object = None) -> builtins.object: return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str) -> builtins.object: return getattr(self, key) def json(self, **kwargs): @@ -98,7 +100,7 @@ class ContainerCreateOptionalRequestParams(TypedDict, total=False): Params here: https://platform.openai.com/docs/api-reference/containers/create """ - expires_after: dict[str, Any] | None # ExpiresAfter object + expires_after: ReadOnly[Mapping[str, object] | None] # ExpiresAfter object file_ids: list[str] | None extra_headers: dict[str, str] | None extra_body: dict[str, str] | None @@ -140,13 +142,13 @@ class ContainerFileObject(BaseModel): source: str _hidden_params: dict[str, Any] = {} - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default: builtins.object = None) -> builtins.object: return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str) -> builtins.object: return getattr(self, key) def json(self, **kwargs): @@ -165,13 +167,13 @@ class ContainerFileListResponse(BaseModel): last_id: str | None = None has_more: bool - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default: builtins.object = None) -> builtins.object: return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str) -> builtins.object: return getattr(self, key) def json(self, **kwargs): @@ -189,13 +191,13 @@ class DeleteContainerFileResponse(BaseModel): object: Literal["container.file.deleted", "container_file.deleted"] deleted: bool - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default: builtins.object = None) -> builtins.object: return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str) -> builtins.object: return getattr(self, key) def json(self, **kwargs): diff --git a/litellm/types/llms/oci.py b/litellm/types/llms/oci.py index ff56d3d183b..d87e4231337 100644 --- a/litellm/types/llms/oci.py +++ b/litellm/types/llms/oci.py @@ -1,7 +1,7 @@ from __future__ import annotations from enum import Enum -from typing import Any, Literal +from typing import Literal from pydantic import BaseModel, SerializeAsAny @@ -105,9 +105,9 @@ class OCIChatRequestPayload(BaseModel): # Honoured by GPT-5 family, Gemini 2.5, Grok reasoning variants, # Cohere Command-A-Reasoning. Ignored by non-reasoning models. reasoningEffort: str | None = None - responseFormat: dict[str, Any] | None = None - toolChoice: str | dict[str, Any] | None = None - logitBias: dict[str, Any] | None = None + responseFormat: dict[str, object] | None = None + toolChoice: str | dict[str, object] | None = None + logitBias: dict[str, object] | None = None logProbs: int | None = None @@ -163,7 +163,7 @@ class OCIResponseChoice(BaseModel): # reasoning phase without producing any visible content. message: OCIMessage | None = None finishReason: str | None = None - logprobs: dict[str, Any] | None = None + logprobs: dict[str, object] | None = None class OCIChatResponse(BaseModel): @@ -275,7 +275,7 @@ class CohereToolCall(BaseModel): """Tool call made by Cohere model.""" name: str - parameters: dict[str, Any] + parameters: dict[str, object] class CohereToolResult(BaseModel): @@ -286,7 +286,7 @@ class CohereToolResult(BaseModel): """ call: CohereToolCall - outputs: list[dict[str, Any]] + outputs: list[dict[str, object]] class CohereChatRequest(BaseModel): @@ -318,12 +318,12 @@ class CohereChatRequest(BaseModel): # OCI Cohere responseFormat is {"type": "TEXT" | "JSON_OBJECT", "schema"?: ...}; # there is no JSON_SCHEMA type. The shape is built in # OCIChatConfig._normalize_response_format. - responseFormat: dict[str, Any] | None = None + responseFormat: dict[str, object] | None = None preambleOverride: str | None = None - documents: list[dict[str, Any]] | None = None + documents: list[dict[str, object]] | None = None searchQueriesOnly: bool | None = None searchEntryPoint: str | None = None - grounding: dict[str, Any] | None = None + grounding: dict[str, object] | None = None isEcho: bool | None = None isSearchQueriesOnly: bool | None = None isRawPrompting: bool | None = None @@ -333,7 +333,7 @@ class CohereChatRequest(BaseModel): citationQuality: str | None = None maxInputTokens: int | None = None isStream: bool | None = None - streamOptions: dict[str, Any] | None = None + streamOptions: dict[str, object] | None = None class CohereUsage(BaseModel): @@ -342,8 +342,8 @@ class CohereUsage(BaseModel): promptTokens: int completionTokens: int totalTokens: int - promptTokensDetails: dict[str, Any] | None = None - completionTokensDetails: dict[str, Any] | None = None + promptTokensDetails: dict[str, object] | None = None + completionTokensDetails: dict[str, object] | None = None class CohereCitation(BaseModel): @@ -378,7 +378,7 @@ class CohereChatResponse(BaseModel): # Optional fields chatHistory: list[CohereMessage] | None = None citations: list[CohereCitation] | None = None - documents: list[dict[str, Any]] | None = None + documents: list[dict[str, object]] | None = None errorMessage: str | None = None isSearchRequired: bool | None = None prompt: str | None = None diff --git a/litellm/types/llms/openai_evals.py b/litellm/types/llms/openai_evals.py index c96ca515d60..519e3e82fff 100644 --- a/litellm/types/llms/openai_evals.py +++ b/litellm/types/llms/openai_evals.py @@ -2,7 +2,8 @@ Type definitions for OpenAI Evals API """ -from typing import Any, Literal +import builtins +from typing import Literal from pydantic import BaseModel from typing_extensions import Required, TypedDict @@ -15,7 +16,7 @@ class DataSourceConfigCustom(TypedDict, total=False): type: Required[Literal["custom"]] """Data source type - custom""" - item_schema: Required[dict[str, Any]] + item_schema: Required[dict[str, object]] """JSON schema describing the structure of each row in the dataset""" include_sample_schema: bool | None @@ -28,7 +29,7 @@ class DataSourceConfigLogs(TypedDict, total=False): type: Required[Literal["logs"]] """Data source type - logs""" - metadata: dict[str, Any] | None + metadata: dict[str, object] | None """Optional metadata for filtering logs""" @@ -38,7 +39,7 @@ class DataSourceConfigStoredCompletions(TypedDict, total=False): type: Required[Literal["stored_completions"]] """Data source type - stored_completions (deprecated)""" - metadata: dict[str, Any] | None + metadata: dict[str, object] | None """Optional metadata for filtering stored completions""" @@ -93,7 +94,7 @@ class CreateEvalRequest(TypedDict, total=False): testing_criteria: Required[list[GraderConfig]] """List of graders for all eval runs""" - metadata: dict[str, Any] | None + metadata: dict[str, object] | None """Set of 16 key-value pairs that can be attached to an object (max 64 char keys, 512 char values)""" @@ -103,7 +104,7 @@ class UpdateEvalRequest(TypedDict, total=False): name: str | None """Updated name""" - metadata: dict[str, Any] | None + metadata: dict[str, object] | None """Updated metadata""" @@ -145,13 +146,13 @@ class Eval(BaseModel): name: str | None = None """The name of the evaluation""" - data_source_config: dict[str, Any] + data_source_config: dict[str, builtins.object] """Configuration for the data source""" - testing_criteria: list[dict[str, Any]] + testing_criteria: list[dict[str, builtins.object]] """List of graders for the evaluation""" - metadata: dict[str, Any] | None = None + metadata: dict[str, builtins.object] | None = None """Additional metadata""" @@ -227,7 +228,7 @@ class DataSourceInlineConfig(TypedDict, total=False): type: Required[Literal["inline"]] """Data source type - inline""" - samples: Required[list[dict[str, Any]]] + samples: Required[list[dict[str, object]]] """List of inline samples to use for the run""" @@ -259,13 +260,13 @@ class CompletionConfig(TypedDict, total=False): class CreateRunRequest(TypedDict, total=False): """Request parameters for creating a run""" - data_source: Required[dict[str, Any]] + data_source: Required[dict[str, object]] """Data source configuration for the run (can be jsonl, completions, or responses type)""" name: str | None """Optional name for the run""" - metadata: dict[str, Any] | None + metadata: dict[str, object] | None """Optional metadata for the run""" @@ -330,7 +331,7 @@ class Run(BaseModel): status: Literal["queued", "running", "completed", "failed", "cancelled"] """Current status of the run""" - data_source: dict[str, Any] + data_source: dict[str, builtins.object] """Data source configuration used for the run""" eval_id: str @@ -348,7 +349,7 @@ class Run(BaseModel): model: str | None = None """Model used for the run, if any""" - per_model_usage: Any | None = None + per_model_usage: builtins.object | None = None """Model usage details per model, if available""" per_testing_criteria_results: list[PerTestingCriteriaResult] | None = None @@ -363,10 +364,10 @@ class Run(BaseModel): shared_with_openai: bool | None = None """Whether run is shared with OpenAI""" - metadata: dict[str, Any] | None = None + metadata: dict[str, builtins.object] | None = None """Additional metadata""" - error: dict[str, Any] | None = None + error: dict[str, builtins.object] | None = None """Error details if the run failed""" diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index 7825684cfe5..61fd5c36b16 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -36,7 +36,7 @@ class SCIMResource(BaseModel): schemas: list[str] id: str | None = None externalId: str | None = None - meta: dict[str, Any] | None = None + meta: dict[str, object] | None = None class SCIMUserName(BaseModel): @@ -119,7 +119,7 @@ class SCIMUser(SCIMResource): ) @model_serializer(mode="wrap") - def _omit_absent_optional_blocks(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: + def _omit_absent_optional_blocks(self, handler: SerializerFunctionWrapHandler) -> dict[str, object]: dumped: Final = handler(self) if self.enterprise_user is None: dumped.pop(SCIM_ENTERPRISE_USER_SCHEMA, None) @@ -169,7 +169,7 @@ class SCIMListResponse(BaseModel): class SCIMPatchOperation(BaseModel): op: str path: str | None = None - value: Any | None = None + value: object | None = None @field_validator("op", mode="before") @classmethod @@ -203,7 +203,7 @@ class SCIMServiceProviderConfig(BaseModel): changePassword: SCIMFeature = SCIMFeature(supported=False) sort: SCIMFeature = SCIMFeature(supported=False) etag: SCIMFeature = SCIMFeature(supported=False) - authenticationSchemes: list[dict[str, Any]] | None = None + authenticationSchemes: list[dict[str, object]] | None = None meta: dict[str, Any] | None = None @@ -231,7 +231,7 @@ class SCIMResourceType(BaseModel): schema_: str # "schema" is a reserved name in Pydantic context schemaExtensions: list[SCIMSchemaExtension] | None = None - meta: dict[str, Any] | None = None + meta: dict[str, object] | None = None def model_dump(self, **kwargs): d: Final = super().model_dump(**kwargs) @@ -266,4 +266,4 @@ class SCIMSchema(BaseModel): name: str description: str | None = None attributes: list[SCIMSchemaAttribute] = [] - meta: dict[str, Any] | None = None + meta: dict[str, object] | None = None diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 99b08f6caf6..f4369fd95af 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -1,3 +1,4 @@ +import builtins from typing import Any, Literal from openai.types.audio.transcription_create_params import FileTypes @@ -14,14 +15,14 @@ class VideoObject(BaseModel): created_at: int | None = None completed_at: int | None = None expires_at: int | None = None - error: dict[str, Any] | None = None + error: dict[str, builtins.object] | None = None progress: int | None = None remixed_from_video_id: str | None = None seconds: str | None = None size: str | None = None model: str | None = None usage: dict[str, Any] | None = None - _hidden_params: dict[str, Any] = {} + _hidden_params: dict[str, builtins.object] = {} def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator @@ -31,7 +32,7 @@ class VideoObject(BaseModel): # Custom .get() method to access attributes with a default value if the attribute doesn't exist return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key) -> builtins.object: # Allow dictionary-style access to attributes return getattr(self, key) @@ -47,7 +48,7 @@ class VideoResponse(BaseModel): """Response object for video generation requests.""" data: list[VideoObject] - hidden_params: dict[str, Any] = {} + hidden_params: dict[str, object] = {} def __contains__(self, key) -> bool: return hasattr(self, key) @@ -55,7 +56,7 @@ class VideoResponse(BaseModel): def get(self, key, default=None): return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key) -> object: return getattr(self, key) def json(self, **kwargs): @@ -73,8 +74,8 @@ class VideoCreateOptionalRequestParams(TypedDict, total=False): """ input_reference: FileTypes | None # File reference for input image - image: Any | None # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object - parameters: dict[str, Any] | None # Provider-specific parameters block passed directly to the API + image: object | None # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object + parameters: dict[str, object] | None # Provider-specific parameters block passed directly to the API model: str | None resolution: ReadOnly[str | None] seconds: str | None @@ -110,7 +111,7 @@ class CharacterObject(BaseModel): object: Literal["character"] = "character" created_at: int name: str - _hidden_params: dict[str, Any] = {} + _hidden_params: dict[str, builtins.object] = {} def __contains__(self, key) -> bool: return hasattr(self, key) @@ -118,7 +119,7 @@ class CharacterObject(BaseModel): def get(self, key, default=None): return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key) -> builtins.object: return getattr(self, key) def json(self, **kwargs): From 81481bea955701601e3a937c86ed32e2a6070b35 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:21:19 +0000 Subject: [PATCH 04/38] chore(lint): ratchet down Any and strict-typing budgets --- basedpyright-code-budget.json | 24 ++++++++++++------------ ruff-strict-budget.json | 14 +++++++------- type-discipline-budget.json | 10 +++++----- 3 files changed, 24 insertions(+), 24 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..813347fdd1a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 14076 + "limit": 13434 }, "reportArgumentType": { - "limit": 2216 + "limit": 2208 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4128 + "limit": 3370 }, "reportFunctionMemberAccess": { "limit": 7 @@ -48,16 +48,16 @@ "limit": 34 }, "reportInvalidTypeVarUse": { - "limit": 2 + "limit": 1 }, "reportMatchNotExhaustive": { "limit": 0 }, "reportMissingParameterType": { - "limit": 5601 + "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15306 + "limit": 15303 }, "reportMissingTypeStubs": { "limit": 40 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 181 + "limit": 180 }, "reportTypedDictNotRequiredAccess": { "limit": 24 @@ -105,16 +105,16 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38350 + "limit": 38324 }, "reportUnknownParameterType": { - "limit": 19626 + "limit": 19590 }, "reportUnknownVariableType": { - "limit": 29890 + "limit": 29873 }, "reportUnnecessaryCast": { - "limit": 111 + "limit": 110 }, "reportUnnecessaryComparison": { "limit": 692 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 826 + "limit": 823 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 9b1cc977a64..4bdcf8997c9 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 2985 + "limit": 2957 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 809 + "limit": 806 }, "ANN201": { - "limit": 2001 + "limit": 1982 }, "ANN202": { - "limit": 835 + "limit": 831 }, "ANN204": { - "limit": 693 + "limit": 683 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 307 + "limit": 122 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1073 + "limit": 1036 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..2318e3391af 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22364 + "limit": 22221 }, "LIT002": { - "limit": 26777 + "limit": 26776 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1039 + "limit": 1036 }, "LIT007": { "limit": 0 @@ -30,9 +30,9 @@ "limit": 16507 }, "LIT011": { - "limit": 5535 + "limit": 5533 }, "LIT012": { - "limit": 4495 + "limit": 4494 } } From b2ed6eaa059a9295595cfb691d1c7b9bfedd1398 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:44:43 +0000 Subject: [PATCH 05/38] chore(lint): ratchet down Any and strict-typing budgets --- basedpyright-code-budget.json | 24 ++++++++++++------------ ruff-strict-budget.json | 14 +++++++------- type-discipline-budget.json | 10 +++++----- 3 files changed, 24 insertions(+), 24 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d094c98f5ec..030632a3102 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 14074 + "limit": 13432 }, "reportArgumentType": { - "limit": 2216 + "limit": 2208 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4125 + "limit": 3367 }, "reportFunctionMemberAccess": { "limit": 7 @@ -48,16 +48,16 @@ "limit": 34 }, "reportInvalidTypeVarUse": { - "limit": 2 + "limit": 1 }, "reportMatchNotExhaustive": { "limit": 0 }, "reportMissingParameterType": { - "limit": 5601 + "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15306 + "limit": 15303 }, "reportMissingTypeStubs": { "limit": 40 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 181 + "limit": 180 }, "reportTypedDictNotRequiredAccess": { "limit": 24 @@ -105,16 +105,16 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38350 + "limit": 38324 }, "reportUnknownParameterType": { - "limit": 19625 + "limit": 19589 }, "reportUnknownVariableType": { - "limit": 29877 + "limit": 29860 }, "reportUnnecessaryCast": { - "limit": 111 + "limit": 110 }, "reportUnnecessaryComparison": { "limit": 692 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 826 + "limit": 823 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ae91b711e13..3ffdce1b0e4 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 2985 + "limit": 2957 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 809 + "limit": 806 }, "ANN201": { - "limit": 2000 + "limit": 1981 }, "ANN202": { - "limit": 835 + "limit": 831 }, "ANN204": { - "limit": 693 + "limit": 683 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 307 + "limit": 122 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1073 + "limit": 1036 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 6273fbce595..4cd8fec5aae 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22358 + "limit": 22215 }, "LIT002": { - "limit": 26774 + "limit": 26773 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1039 + "limit": 1036 }, "LIT007": { "limit": 0 @@ -30,9 +30,9 @@ "limit": 16494 }, "LIT011": { - "limit": 5535 + "limit": 5533 }, "LIT012": { - "limit": 4495 + "limit": 4494 } } From b6c10d31e8f95240386f40e0ee93277e530c71e7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:33:11 +0000 Subject: [PATCH 06/38] chore(lint): ratchet down Any and strict-typing budgets --- basedpyright-code-budget.json | 24 ++++++++++++------------ ruff-strict-budget.json | 14 +++++++------- type-discipline-budget.json | 10 +++++----- 3 files changed, 24 insertions(+), 24 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 3f96531cf6f..9365c7daad5 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 14074 + "limit": 13431 }, "reportArgumentType": { - "limit": 2215 + "limit": 2207 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4124 + "limit": 3371 }, "reportFunctionMemberAccess": { "limit": 7 @@ -48,16 +48,16 @@ "limit": 34 }, "reportInvalidTypeVarUse": { - "limit": 2 + "limit": 1 }, "reportMatchNotExhaustive": { "limit": 0 }, "reportMissingParameterType": { - "limit": 5601 + "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15290 + "limit": 15287 }, "reportMissingTypeStubs": { "limit": 40 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 181 + "limit": 180 }, "reportTypedDictNotRequiredAccess": { "limit": 24 @@ -105,16 +105,16 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38332 + "limit": 38306 }, "reportUnknownParameterType": { - "limit": 19625 + "limit": 19589 }, "reportUnknownVariableType": { - "limit": 29861 + "limit": 29846 }, "reportUnnecessaryCast": { - "limit": 111 + "limit": 110 }, "reportUnnecessaryComparison": { "limit": 687 @@ -123,7 +123,7 @@ "limit": 4 }, "reportUnnecessaryIsInstance": { - "limit": 823 + "limit": 820 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 4fcf650a8bc..204ed2929ee 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 2985 + "limit": 2957 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 809 + "limit": 806 }, "ANN201": { - "limit": 2000 + "limit": 1981 }, "ANN202": { - "limit": 835 + "limit": 831 }, "ANN204": { - "limit": 693 + "limit": 683 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 304 + "limit": 119 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1071 + "limit": 1034 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3c4c7760c6..c9a1f28438e 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22334 + "limit": 22192 }, "LIT002": { - "limit": 26763 + "limit": 26762 }, "LIT003": { "limit": 261 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1038 + "limit": 1035 }, "LIT007": { "limit": 0 @@ -30,9 +30,9 @@ "limit": 16480 }, "LIT011": { - "limit": 5520 + "limit": 5518 }, "LIT012": { - "limit": 4489 + "limit": 4488 } } From e5c1133a7942a7875e00ba263bb6dabf64f6a4ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 06:53:19 +0000 Subject: [PATCH 07/38] chore(lint): re-ratchet lint budgets after merging staging --- basedpyright-code-budget.json | 24 ++++++++++++------------ ruff-strict-budget.json | 14 +++++++------- type-discipline-budget.json | 10 +++++----- 3 files changed, 24 insertions(+), 24 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 2967fc2a505..eb7f484901f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 14074 + "limit": 13431 }, "reportArgumentType": { - "limit": 2215 + "limit": 2207 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4124 + "limit": 3371 }, "reportFunctionMemberAccess": { "limit": 7 @@ -48,16 +48,16 @@ "limit": 34 }, "reportInvalidTypeVarUse": { - "limit": 2 + "limit": 1 }, "reportMatchNotExhaustive": { "limit": 0 }, "reportMissingParameterType": { - "limit": 5601 + "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15288 + "limit": 15285 }, "reportMissingTypeStubs": { "limit": 40 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 181 + "limit": 180 }, "reportTypedDictNotRequiredAccess": { "limit": 24 @@ -105,16 +105,16 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38324 + "limit": 38298 }, "reportUnknownParameterType": { - "limit": 19625 + "limit": 19589 }, "reportUnknownVariableType": { - "limit": 29861 + "limit": 29846 }, "reportUnnecessaryCast": { - "limit": 111 + "limit": 110 }, "reportUnnecessaryComparison": { "limit": 687 @@ -123,7 +123,7 @@ "limit": 4 }, "reportUnnecessaryIsInstance": { - "limit": 819 + "limit": 816 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index be2b30fc189..f9360a2308e 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 2984 + "limit": 2956 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 809 + "limit": 806 }, "ANN201": { - "limit": 2000 + "limit": 1981 }, "ANN202": { - "limit": 835 + "limit": 831 }, "ANN204": { - "limit": 693 + "limit": 683 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 304 + "limit": 119 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1071 + "limit": 1035 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index d5bf3883be4..071f3418101 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22330 + "limit": 22188 }, "LIT002": { - "limit": 26763 + "limit": 26762 }, "LIT003": { "limit": 261 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1038 + "limit": 1035 }, "LIT007": { "limit": 0 @@ -30,9 +30,9 @@ "limit": 16477 }, "LIT011": { - "limit": 5519 + "limit": 5517 }, "LIT012": { - "limit": 4489 + "limit": 4488 } } From 98a0cf306f213f511744502b22ed3f3a2a00d5bc Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 15:17:00 -0700 Subject: [PATCH 08/38] fix(shadow_eval): size the judge output cap for a judge that reasons The cap covers reasoning tokens as well as the verdict, and the models people pick as judges reason before answering whether the call asks them to or not: Anthropic's 5 family thinks adaptively and cannot be told not to, so the reasoning bills against max_tokens with nothing in the request to opt out. At 1500 the reasoning consumed the budget and the reply arrived empty or cut off mid-object, which the attempt recorded as an unparseable judge verdict rather than a result. Headroom costs nothing: max_tokens is a ceiling and only generated tokens bill, so the only movement is that judge calls which used to bill their full budget and return nothing now return a verdict. Deliberately not passing reasoning_effort to bound the reasoning instead: is_thinking_enabled treats any reasoning_effort as thinking-enabled, which drops the forced tool_choice that json_mode relies on and turns thinking on with a 1024-token floor for judges that were not reasoning at all. --- litellm/integrations/shadow_eval_logger.py | 10 ++-- .../integrations/test_shadow_eval_logger.py | 49 +++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 27da785331a..a1716c0954d 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,9 +60,13 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object; a tighter budget truncates the JSON -# mid-object and the attempt is lost to an error row. -JUDGE_MAX_OUTPUT_TOKENS: Final = 1500 +# The judge answers with a small JSON object, but the cap covers reasoning tokens too, +# and the models people pick as judges reason before answering whether or not the call +# asks them to (Anthropic's 5 family thinks adaptively and cannot be told not to). A +# budget sized for the JSON alone is spent on invisible reasoning instead, and the reply +# arrives empty or truncated mid-object, which the attempt records as an unparseable +# verdict. Headroom is free: max_tokens is a ceiling, and only generated tokens bill. +JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 5628d69de26..877677505d6 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -120,6 +120,27 @@ def _router( return router +def _reasoning_judge_router(reasoning_tokens, verdict='{"preference": "A", "confidence": 0.9}'): + """A router whose judge arm reasons before it answers, the way Anthropic's 5 family + does whether or not the call asks it to. Reasoning is billed against the caller's own + max_tokens and the reply is cut off at that cap, so a cap that does not clear the + reasoning budget yields a truncated verdict or no verdict at all. One character stands + in for one token, which is what makes the cap the thing under test.""" + router = MagicMock() + router.model_group_alias = {} + router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) + + async def acompletion(**kwargs): + if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: + kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} + return {"choices": [{"message": {"content": "shadow answer"}}]} + budget_for_the_answer = kwargs["max_tokens"] - reasoning_tokens + return {"choices": [{"message": {"content": verdict[: max(0, budget_for_the_answer)]}}]} + + router.acompletion = MagicMock(side_effect=acompletion) + return router + + def _spend_counter(store=None): """In-memory stand-in for the proxy's cross-pod spend counter: reads take the max of the counter and the caller's fallback, exactly like get_current_spend does for a key @@ -1134,6 +1155,34 @@ class TestShadowPipeline: assert row["shadow_cost"] == 0.007 assert logger._test_counter["spend:shadow_eval:job-1"] == 0.007 + async def test_the_judge_output_cap_leaves_room_for_a_reasoning_judge(self): + """The output cap covers reasoning tokens as well as the answer, and the models + people pick as judges reason before answering whether or not the call asks them to. + A cap sized for the verdict JSON alone is spent on reasoning instead and the reply + arrives empty, which the attempt records as an unparseable verdict rather than a + result. The judge here burns a reasoning budget typical of a thinking model on a + comparison task, so the cap has to clear it for the verdict to survive.""" + reasoning_tokens = 2000 + logger = _logger(router=_reasoning_judge_router(reasoning_tokens), prisma=(prisma := _prisma())) + + await logger._run_shadow_eval( + job=_job(), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + row = prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"] + assert row["outcome"] in ("real", "shadow", "tie"), row["error"] + assert row["error"] is None + async def test_a_pipeline_error_after_the_shadow_call_keeps_its_billed_cost(self, monkeypatch: pytest.MonkeyPatch): """An unexpected error between the billed shadow call and the attempt write must still record the shadow cost, or the per-key dollar gate undercounts forever.""" From a2f926eb8f7fac36a193a851b035eabb92b27373 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 15:55:19 -0700 Subject: [PATCH 09/38] fix(shadow_eval): correct the judge output cap's causal claim The prior commit claimed claude-sonnet-5 reasons invisibly by default and eats the judge's budget regardless of what the call asks for. Verified against a live proxy: with no thinking param (what _call_judge sends today), forced tool-choice json_mode, native structured output, and even an explicit thinking=adaptive, the model returned 0 reasoning tokens and a clean compact verdict every time, on prompts up to several thousand characters. The real mechanism only shows up with an elevated reasoning_effort or output_config.effort on the request, which happens when the judge_model deployment is configured with one, e.g. an admin pointing the judge at their best reasoning model. Reproduced directly: reasoning_effort=max, 300-token cap, real Anthropic reply came back finish_reason=length, content=None, 299 of 300 tokens spent on reasoning. Same request at 4096 returned a valid verdict. This is a narrower, verified claim than the one it replaces. --- litellm/integrations/shadow_eval_logger.py | 12 +++++----- .../integrations/test_shadow_eval_logger.py | 22 +++++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index a1716c0954d..b554c4bc668 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,12 +60,12 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object, but the cap covers reasoning tokens too, -# and the models people pick as judges reason before answering whether or not the call -# asks them to (Anthropic's 5 family thinks adaptively and cannot be told not to). A -# budget sized for the JSON alone is spent on invisible reasoning instead, and the reply -# arrives empty or truncated mid-object, which the attempt records as an unparseable -# verdict. Headroom is free: max_tokens is a ceiling, and only generated tokens bill. +# The judge answers with a small JSON object, but the cap covers reasoning tokens too. A +# judge_model deployment configured with an elevated reasoning_effort or thinking budget +# (a realistic pick: an admin's best reasoning model doubling as the judge) spends most or +# all of a tight cap on that reasoning, invisibly to this call, and the reply arrives empty +# or truncated mid-object, which the attempt records as an unparseable verdict. Headroom is +# free: max_tokens is a ceiling, and only generated tokens bill. JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 877677505d6..367ad758772 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -121,11 +121,11 @@ def _router( def _reasoning_judge_router(reasoning_tokens, verdict='{"preference": "A", "confidence": 0.9}'): - """A router whose judge arm reasons before it answers, the way Anthropic's 5 family - does whether or not the call asks it to. Reasoning is billed against the caller's own - max_tokens and the reply is cut off at that cap, so a cap that does not clear the - reasoning budget yields a truncated verdict or no verdict at all. One character stands - in for one token, which is what makes the cap the thing under test.""" + """A router whose judge arm reasons before it answers, the way a deployment carrying an + elevated reasoning_effort does. Reasoning is billed against the caller's own max_tokens + and the reply is cut off at that cap, so a cap that does not clear the reasoning budget + yields a truncated verdict or no verdict at all. One character stands in for one token, + which is what makes the cap the thing under test.""" router = MagicMock() router.model_group_alias = {} router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) @@ -1156,12 +1156,12 @@ class TestShadowPipeline: assert logger._test_counter["spend:shadow_eval:job-1"] == 0.007 async def test_the_judge_output_cap_leaves_room_for_a_reasoning_judge(self): - """The output cap covers reasoning tokens as well as the answer, and the models - people pick as judges reason before answering whether or not the call asks them to. - A cap sized for the verdict JSON alone is spent on reasoning instead and the reply - arrives empty, which the attempt records as an unparseable verdict rather than a - result. The judge here burns a reasoning budget typical of a thinking model on a - comparison task, so the cap has to clear it for the verdict to survive.""" + """The output cap covers reasoning tokens as well as the answer, and a judge_model + deployment carrying an elevated reasoning_effort spends that budget before it writes + anything. A cap sized for the verdict JSON alone goes entirely to reasoning and the + reply arrives empty, which the attempt records as an unparseable verdict rather than + a result. The judge here burns a reasoning budget a live claude-sonnet-5 call was + measured at, so the cap has to clear it for the verdict to survive.""" reasoning_tokens = 2000 logger = _logger(router=_reasoning_judge_router(reasoning_tokens), prisma=(prisma := _prisma())) From 2f5bfae1a61b0821b6af9eabb045522adfa7b28a Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:21:13 -0700 Subject: [PATCH 10/38] refactor(shadow_eval): tighten the judge cap comment and type the test helper --- litellm/integrations/shadow_eval_logger.py | 9 +++------ .../integrations/test_shadow_eval_logger.py | 10 +++++----- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index b554c4bc668..2c56ecb8721 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,12 +60,9 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object, but the cap covers reasoning tokens too. A -# judge_model deployment configured with an elevated reasoning_effort or thinking budget -# (a realistic pick: an admin's best reasoning model doubling as the judge) spends most or -# all of a tight cap on that reasoning, invisibly to this call, and the reply arrives empty -# or truncated mid-object, which the attempt records as an unparseable verdict. Headroom is -# free: max_tokens is a ceiling, and only generated tokens bill. +# The judge answers with a small JSON object, but the cap covers reasoning tokens too: a +# judge deployment carrying an elevated reasoning_effort spends a tight cap before it ever +# answers, and the truncated reply is recorded as an unparseable verdict. JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 367ad758772..9fcbd116f63 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -120,12 +120,12 @@ def _router( return router -def _reasoning_judge_router(reasoning_tokens, verdict='{"preference": "A", "confidence": 0.9}'): +def _reasoning_judge_router( + reasoning_tokens: int, verdict: str = '{"preference": "A", "confidence": 0.9}' +) -> MagicMock: """A router whose judge arm reasons before it answers, the way a deployment carrying an - elevated reasoning_effort does. Reasoning is billed against the caller's own max_tokens - and the reply is cut off at that cap, so a cap that does not clear the reasoning budget - yields a truncated verdict or no verdict at all. One character stands in for one token, - which is what makes the cap the thing under test.""" + elevated reasoning_effort does: reasoning bills against the caller's own max_tokens and + the reply is cut off at that cap. One character stands in for one token.""" router = MagicMock() router.model_group_alias = {} router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) From da5e38ce9c1fe2ba4854952eac0cc2a694e1c38b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:33:16 -0700 Subject: [PATCH 11/38] refactor(shadow_eval): state the cap's constraint without the rationale --- litellm/integrations/shadow_eval_logger.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 2c56ecb8721..fc82ebafe09 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,9 +60,8 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object, but the cap covers reasoning tokens too: a -# judge deployment carrying an elevated reasoning_effort spends a tight cap before it ever -# answers, and the truncated reply is recorded as an unparseable verdict. +# Covers the judge's reasoning tokens as well as its small JSON answer: a judge deployment +# carrying an elevated reasoning_effort spends a tight cap before it ever answers. JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 From 5980055d7eae3d1ca28286979c5bd264cd37af57 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:50:35 -0700 Subject: [PATCH 12/38] feat(shadow_eval): say which shape produced an unparseable judge verdict The parser message alone cannot separate a judge that answered with nothing from one truncated mid-object, and the two want opposite fixes. Records the reply's shape, never its text, since no attempt row carries sampled content. --- litellm/integrations/shadow_eval_logger.py | 20 +++- .../integrations/test_shadow_eval_logger.py | 94 +++++++++++++++++++ 2 files changed, 113 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index fc82ebafe09..fb75ef74db9 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -345,6 +345,22 @@ def _failure_detail(e: BaseException) -> str: return f"{type(e).__name__}{location}: {e}" +def _judge_reply_shape(response: object) -> str: + """How an unparseable judge reply was shaped. The parser's own message cannot separate a + judge that answered with nothing from one truncated mid-object, and those want opposite + fixes. Shape only, never the reply text: the judge quotes the sampled turns it compares, + and no attempt row carries sampled content today.""" + try: + choice: Final = response["choices"][0] # pyright: ignore[reportIndexIssue] # judge replies are subscriptable payloads + content: Final = choice["message"]["content"] + finish: Final = choice.get("finish_reason") or "unknown" + except (AttributeError, KeyError, IndexError, TypeError): + return "unreadable judge reply" + served: Final = str(getattr(response, "model", None) or "unknown") + body: Final = f"{len(str(content))} chars" if content else "no content" + return f"finish_reason={finish}, content={body}, model={served}" + + def _call_cost(response: object) -> float: """Price one eval-arm call with the figure the spend pipeline bills: the router client stamps _hidden_params.response_cost from the deployment's own pricing, which the public @@ -1139,7 +1155,9 @@ class ShadowEvalLogger(CustomLogger): verdict: Final = PairwiseVerdict.model_validate(parse_json_verdict(raw)) except Exception as e: # noqa: BLE001 # malformed verdicts become error rows verbose_logger.debug("shadow_eval: unparseable judge verdict: %s", e) - return _CallFailure(f"unparseable judge verdict: {e}", cost=_call_cost(response)) + return _CallFailure( + f"unparseable judge verdict: {e}; {_judge_reply_shape(response)}", cost=_call_cost(response) + ) return _JudgeVerdict( preference=_unmask_preference(verdict.preference, real_is_a), confidence=max(0.0, min(1.0, verdict.confidence)), diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 9fcbd116f63..dbc6d4ec915 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -141,6 +141,27 @@ def _reasoning_judge_router( return router +def _judge_reply_router(content: str | None, finish_reason: str = "stop", served_model: str = "judge-pick") -> MagicMock: + """A router whose judge arm returns a caller-shaped reply, so the shapes that all land + on the same parser error can be posed apart: no content at all, versus JSON cut off + mid-object.""" + router = MagicMock() + router.model_group_alias = {} + router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) + + async def acompletion(**kwargs): + if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: + kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} + return {"choices": [{"message": {"content": "shadow answer"}}]} + return ModelResponse( + model=served_model, + choices=[{"index": 0, "finish_reason": finish_reason, "message": {"role": "assistant", "content": content}}], + ) + + router.acompletion = MagicMock(side_effect=acompletion) + return router + + def _spend_counter(store=None): """In-memory stand-in for the proxy's cross-pod spend counter: reads take the max of the counter and the caller's fallback, exactly like get_current_spend does for a key @@ -1126,6 +1147,79 @@ class TestShadowPipeline: assert row["judge_cost"] == expected_cost assert row["shadow_cost"] == expected_shadow_cost + async def _judge_error(self, router: MagicMock, monkeypatch: pytest.MonkeyPatch) -> str: + import litellm as litellm_module + + monkeypatch.setattr(litellm_module, "completion_cost", lambda completion_response: 0.007) + prisma = _prisma() + await _logger(router=router, prisma=prisma)._run_shadow_eval( + job=_job(), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + return prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"]["error"] + + async def test_a_judge_that_answered_nothing_is_told_apart_from_one_cut_off( + self, monkeypatch: pytest.MonkeyPatch + ): + """Both land on the same parser message, and they want opposite fixes: a judge + returning no content points at the reply never being text, while one cut off + mid-object points at the output cap. The row has to say which.""" + truncated = '{"preference": "A", "confidence": 0.9, "reasoning": "' + answered_nothing = await self._judge_error(_judge_reply_router(None), monkeypatch) + cut_off = await self._judge_error( + _judge_reply_router(truncated, finish_reason="length"), monkeypatch + ) + + assert "content=no content" in answered_nothing + assert "finish_reason=stop" in answered_nothing + assert f"content={len(truncated)} chars" in cut_off + assert "finish_reason=length" in cut_off + + async def test_an_unparseable_verdict_names_the_model_that_served_it(self, monkeypatch: pytest.MonkeyPatch): + """A judge_model that fans out over deployments hides which one truncates: without + the served model the operator cannot tell a bad deployment from a bad cap.""" + error = await self._judge_error(_judge_reply_router(None, served_model="claude-sonnet-5"), monkeypatch) + + assert "model=claude-sonnet-5" in error + + async def test_a_diagnosed_verdict_error_stays_groupable(self, monkeypatch: pytest.MonkeyPatch): + """The customer groups attempt rows by error text. Every varying part has to sit + after the first semicolon or each row becomes its own group.""" + first = await self._judge_error(_judge_reply_router(None, served_model="model-a"), monkeypatch) + second = await self._judge_error(_judge_reply_router(None, served_model="model-b"), monkeypatch) + + assert first != second + assert first.split(";")[0] == second.split(";")[0] + + async def test_a_judge_reply_that_cannot_be_read_still_records_an_error(self, monkeypatch: pytest.MonkeyPatch): + """The shape reader runs inside the failure path: it must never raise a second time + and cost the row entirely.""" + router = MagicMock() + router.model_group_alias = {} + router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) + + async def acompletion(**kwargs): + if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: + kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} + return {"choices": [{"message": {"content": "shadow answer"}}]} + return {"choices": []} + + router.acompletion = MagicMock(side_effect=acompletion) + + error = await self._judge_error(router, monkeypatch) + + assert "unparseable judge verdict" in error + assert "unreadable judge reply" in error + async def test_an_empty_shadow_reply_still_bills_its_cost(self, monkeypatch: pytest.MonkeyPatch): """A shadow call that returns no extractable text has still billed; pricing it at zero would keep the dollar gate open while shadow calls keep charging the key.""" From a2d5215a4fd7e8eb9ff4d2112545bef798b0e86e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:14:44 -0700 Subject: [PATCH 13/38] fix(proxy): gate the OpenAI websocket passthrough behind an explicit opt-in --- litellm/proxy/_types.py | 4 + .../llm_passthrough_endpoints.py | 100 +++++- test-quality-budget.json | 4 +- .../test_openai_ws_passthrough_routes.py | 326 +++++++++++------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 5 files changed, 296 insertions(+), 143 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b33e2fe7ff6..f83011835fd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2638,6 +2638,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): default=None, description="Set-up pass-through endpoints for provider-specific endpoints. Docs - https://docs.litellm.ai/docs/proxy/pass_through", ) + enable_openai_websocket_passthrough: bool | None = Field( + default=None, + description="Serve the OpenAI pass-through WebSocket route, which relays frames to OpenAI under the proxy's own provider credential without reading them. Off by default.", + ) user_header_name: str | None = Field( None, description="[DEPRECATED] Use 'user_header_mappings' instead. When set, the header value is treated as the end user id unless overridden by user_header_mappings.", diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6b1d6405a6a..97d25e20939 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -14,8 +14,9 @@ import json import os import re from collections.abc import AsyncGenerator, Callable, Mapping +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, cast +from typing import TYPE_CHECKING, Annotated, Final, Protocol, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -2345,19 +2346,99 @@ def _key_has_model_restrictions(user_api_key_dict: UserAPIKeyAuth) -> bool: return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for model in scoped_models) +@dataclass(frozen=True, slots=True) +class _OpenAIWebsocketRefusal: + close_reason: str + message: str + + +_OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( + close_reason="OpenAI websocket passthrough is disabled", + message=( + "OpenAI websocket passthrough is disabled on this gateway. A proxy admin can turn it on by " + "setting general_settings.enable_openai_websocket_passthrough to true." + ), +) + +_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal( + close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough", + message=( + "Keys with model restrictions cannot use OpenAI websocket passthrough, because this route " + "relays frames to the provider without reading which model they ask for." + ), +) + + +def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool: + setting: Final = general_settings.get("enable_openai_websocket_passthrough") + if isinstance(setting, str): + return str_to_bool(setting) is True + return setting is True + + +def _openai_websocket_refusal( + user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object] +) -> _OpenAIWebsocketRefusal | None: + if not _is_openai_websocket_passthrough_enabled(general_settings): + return _OPENAI_WS_DISABLED_REFUSAL + if _key_has_model_restrictions(user_api_key_dict): + return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL + return None + + +class _OpenAIWebsocketRelay(Protocol): + async def __call__( + self, + *, + websocket: WebSocket, + target: str, + custom_headers: dict[str, str], # mutable-ok: the relay takes a plain dict of upstream headers + user_api_key_dict: UserAPIKeyAuth, + forward_headers: bool, + endpoint: str, + accept_websocket: bool, + ) -> None: ... + + +def _proxy_general_settings() -> Mapping[str, object]: + from litellm.proxy.proxy_server import general_settings + + return general_settings + + +def _openai_websocket_relay() -> _OpenAIWebsocketRelay: + return websocket_passthrough_request + + @router.websocket("/openai_passthrough/{endpoint:path}") @router.websocket("/openai/{endpoint:path}") async def openai_websocket_proxy_route( websocket: WebSocket, endpoint: str, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], + general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], + relay: Annotated[_OpenAIWebsocketRelay, Depends(_openai_websocket_relay)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" - if _key_has_model_restrictions(user_api_key_dict): - await websocket.close( - code=1008, - reason="Keys with model restrictions cannot use OpenAI websocket passthrough", + requested_subprotocols: Final = tuple( + protocol.strip() + for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if protocol.strip() + ) + negotiated_subprotocol: Final = requested_subprotocols[0] if requested_subprotocols else None + + refusal: Final = _openai_websocket_refusal(user_api_key_dict, general_settings) + if refusal is not None: + await websocket.accept(subprotocol=negotiated_subprotocol) + await websocket.send_text( + json.dumps( + { + "type": "error", + "error": {"type": "invalid_request_error", "message": refusal.message}, + } + ) ) + await websocket.close(code=1008, reason=refusal.close_reason) return base_target_url: Final = os.getenv("OPENAI_API_BASE") or "https://api.openai.com/" @@ -2393,14 +2474,9 @@ async def openai_websocket_proxy_route( "Authorization": f"Bearer {openai_api_key}" } - requested_subprotocols: Final = tuple( - protocol.strip() - for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") - if protocol.strip() - ) - await websocket.accept(subprotocol=requested_subprotocols[0] if requested_subprotocols else None) + await websocket.accept(subprotocol=negotiated_subprotocol) - await websocket_passthrough_request( + await relay( websocket=websocket, target=wss_target, custom_headers=custom_headers, diff --git a/test-quality-budget.json b/test-quality-budget.json index 7ca563d25af..3c12371f02f 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 741 + "limit": 737 }, "TQ003": { "limit": 62 @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11003 + "limit": 10993 } } diff --git a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py index b22e202d9e0..6578b75ace1 100644 --- a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py @@ -1,16 +1,35 @@ -"""OpenAI passthrough must register WebSocket catch-all routes (#36088).""" +"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals.""" -from unittest.mock import AsyncMock, MagicMock, patch +import json +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType, SimpleNamespace +from typing import Final +from unittest.mock import patch import pytest from starlette.routing import WebSocketRoute from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _OPENAI_WS_DISABLED_REFUSAL, + _OPENAI_WS_MODEL_RESTRICTED_REFUSAL, + _openai_websocket_refusal, openai_websocket_proxy_route, router, ) +ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) +DISABLED_SETTINGS: Final = ( + MappingProxyType({}), + MappingProxyType({"enable_openai_websocket_passthrough": False}), + MappingProxyType({"enable_openai_websocket_passthrough": "false"}), + MappingProxyType({"enable_openai_websocket_passthrough": None}), +) +GET_CREDENTIALS: Final = ( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" +) + def test_openai_websocket_passthrough_routes_registered(): ws_paths = {route.path for route in router.routes if isinstance(route, WebSocketRoute)} @@ -18,164 +37,213 @@ def test_openai_websocket_passthrough_routes_registered(): assert "/openai_passthrough/{endpoint:path}" in ws_paths -def _mock_websocket(path: str, query: str, headers: dict[str, str] | None = None) -> MagicMock: - websocket = MagicMock() - websocket.url.path = path - websocket.url.query = query - websocket.headers = headers or {} - websocket.accept = AsyncMock() - websocket.close = AsyncMock() - return websocket +class _FakeWebSocket: + def __init__(self, path: str, query: str, subprotocols: str | None = None) -> None: + self.url = SimpleNamespace(path=path, query=query) + self.headers = {"sec-websocket-protocol": subprotocols} if subprotocols else {} + self.accepts: list[str | None] = [] + self.sent: list[str] = [] + self.closed: tuple[int, str] | None = None + + async def accept(self, subprotocol: str | None = None) -> None: + self.accepts.append(subprotocol) + + async def send_text(self, data: str) -> None: + self.sent.append(data) + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.closed = (code, reason) + + def error_message(self) -> str: + assert len(self.sent) == 1 + frame = json.loads(self.sent[0]) + assert frame["type"] == "error" + return frame["error"]["message"] + + +@dataclass(frozen=True, slots=True) +class _RelayCall: + target: str + custom_headers: Mapping[str, str] + forward_headers: bool + endpoint: str + accept_websocket: bool + + +class _FakeRelay: + def __init__(self) -> None: + self.calls: list[_RelayCall] = [] + + async def __call__( + self, + *, + websocket: _FakeWebSocket, + target: str, + custom_headers: dict[str, str], + user_api_key_dict: UserAPIKeyAuth, + forward_headers: bool, + endpoint: str, + accept_websocket: bool, + ) -> None: + self.calls.append( + _RelayCall( + target=target, + custom_headers=MappingProxyType(dict(custom_headers)), + forward_headers=forward_headers, + endpoint=endpoint, + accept_websocket=accept_websocket, + ) + ) + + +async def _serve( + websocket: _FakeWebSocket, + endpoint: str, + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], +) -> _FakeRelay: + relay = _FakeRelay() + await openai_websocket_proxy_route( + websocket=websocket, + endpoint=endpoint, + user_api_key_dict=user_api_key_dict, + general_settings=general_settings, + relay=relay, + ) + return relay @pytest.mark.asyncio @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) -async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix): - websocket = _mock_websocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") +async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch): + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-provider", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._join_url_paths", - return_value="https://api.openai.com/v1/realtime", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=UserAPIKeyAuth(), + with patch(GET_CREDENTIALS, return_value="sk-provider"): + relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) + + assert relay.calls == [ + _RelayCall( + target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview", + custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}), + forward_headers=False, + endpoint=f"/{prefix}/v1/realtime", + accept_websocket=False, ) - - kwargs = mock_ws.await_args.kwargs - assert kwargs["target"] == "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" - assert kwargs["custom_headers"] == {"Authorization": "Bearer sk-provider"} - assert kwargs["forward_headers"] is False - assert kwargs["endpoint"] == f"/{prefix}/v1/realtime" - assert kwargs["accept_websocket"] is False - websocket.accept.assert_awaited_once_with(subprotocol=None) - websocket.close.assert_not_awaited() + ] + assert websocket.accepts == [None] + assert websocket.sent == [] + assert websocket.closed is None @pytest.mark.asyncio async def test_openai_websocket_accepts_first_client_subprotocol(): - websocket = _mock_websocket( + websocket = _FakeWebSocket( "/openai/v1/realtime", "model=gpt-4o-realtime-preview", - headers={ - "sec-websocket-protocol": "realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1" - }, + subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1", ) - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-provider", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=UserAPIKeyAuth(), - ) + with patch(GET_CREDENTIALS, return_value="sk-provider"): + relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) - websocket.accept.assert_awaited_once_with(subprotocol="realtime") - assert mock_ws.await_args.kwargs["accept_websocket"] is False - websocket.close.assert_not_awaited() + assert websocket.accepts == ["realtime"] + assert [call.accept_websocket for call in relay.calls] == [False] + assert websocket.closed is None @pytest.mark.asyncio async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing(): - websocket = _mock_websocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") + websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value=None, - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=UserAPIKeyAuth(), - ) + with patch(GET_CREDENTIALS, return_value=None): + relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) - websocket.close.assert_awaited_once() - assert websocket.close.await_args.kwargs["code"] == 1011 - websocket.accept.assert_not_awaited() - mock_ws.assert_not_awaited() + assert websocket.closed is not None + assert websocket.closed[0] == 1011 + assert "OPENAI_API_KEY" in websocket.closed[1] + assert websocket.accepts == [] + assert relay.calls == [] @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_api_key_dict", - [ - UserAPIKeyAuth(models=["gpt-4o"]), - UserAPIKeyAuth(team_models=["gpt-4o-realtime-preview"]), - UserAPIKeyAuth(models=["all-team-models"], team_models=["gpt-4o"]), - ], +@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) +@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) +async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings): + websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") + + relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) + + assert "enable_openai_websocket_passthrough" in websocket.error_message() + assert websocket.accepts == [None] + assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason) + assert relay.calls == [] + + +@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) +def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): + assert _openai_websocket_refusal(UserAPIKeyAuth(), general_settings) is _OPENAI_WS_DISABLED_REFUSAL + + +@pytest.mark.parametrize("value", [True, "true", "True"]) +def test_openai_websocket_refusal_is_none_for_truthy_settings(value): + settings = MappingProxyType({"enable_openai_websocket_passthrough": value}) + assert _openai_websocket_refusal(UserAPIKeyAuth(), settings) is None + + +@pytest.mark.asyncio +async def test_openai_websocket_refusal_echoes_requested_subprotocol(): + websocket = _FakeWebSocket( + "/openai_passthrough/v1/realtime", + "model=gpt-4o-realtime-preview", + subprotocols="realtime, openai-beta.realtime-v1", + ) + + relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), MappingProxyType({})) + + assert websocket.accepts == ["realtime"] + assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason) + assert relay.calls == [] + + +RESTRICTED_KEYS: Final = ( + UserAPIKeyAuth(models=["gpt-4o"]), + UserAPIKeyAuth(team_models=["gpt-4o-realtime-preview"]), + UserAPIKeyAuth(models=["all-team-models"], team_models=["gpt-4o"]), ) +UNRESTRICTED_KEYS: Final = ( + UserAPIKeyAuth(), + UserAPIKeyAuth(models=["all-proxy-models"]), + UserAPIKeyAuth(models=["*"]), + UserAPIKeyAuth(models=["all-team-models"], team_models=["all-proxy-models"]), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("user_api_key_dict", RESTRICTED_KEYS) async def test_openai_websocket_rejects_model_restricted_keys(user_api_key_dict): - websocket = _mock_websocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") + websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws: - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/realtime", - user_api_key_dict=user_api_key_dict, - ) + relay = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED) - websocket.close.assert_awaited_once() - assert websocket.close.await_args.kwargs["code"] == 1008 - websocket.accept.assert_not_awaited() - mock_ws.assert_not_awaited() + assert "model restrictions" in websocket.error_message() + assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason) + assert relay.calls == [] + + +@pytest.mark.parametrize("user_api_key_dict", RESTRICTED_KEYS) +def test_openai_websocket_refusal_prefers_disabled_over_model_restriction(user_api_key_dict): + assert _openai_websocket_refusal(user_api_key_dict, MappingProxyType({})) is _OPENAI_WS_DISABLED_REFUSAL @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_api_key_dict", - [ - UserAPIKeyAuth(), - UserAPIKeyAuth(models=["all-proxy-models"]), - UserAPIKeyAuth(models=["*"]), - UserAPIKeyAuth(models=["all-team-models"], team_models=["all-proxy-models"]), - ], -) +@pytest.mark.parametrize("user_api_key_dict", UNRESTRICTED_KEYS) async def test_openai_websocket_allows_unrestricted_keys(user_api_key_dict): - websocket = _mock_websocket("/openai/v1/responses", "") + websocket = _FakeWebSocket("/openai/v1/responses", "") - with ( - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", - return_value="sk-provider", - ), - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", - new_callable=AsyncMock, - ) as mock_ws, - ): - await openai_websocket_proxy_route( - websocket=websocket, - endpoint="v1/responses", - user_api_key_dict=user_api_key_dict, - ) + with patch(GET_CREDENTIALS, return_value="sk-provider"): + relay = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED) - mock_ws.assert_awaited_once() - websocket.close.assert_not_awaited() + assert len(relay.calls) == 1 + assert websocket.sent == [] + assert websocket.closed is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index de1b7fe699e..cda1a3834ce 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25731,6 +25731,11 @@ export interface components { * @description If True and SSO is configured (MICROSOFT_CLIENT_ID, GOOGLE_CLIENT_ID, GENERIC_CLIENT_ID, or SAML_IDP_METADATA_URL/XML), disables username/password login on /login, /v2/login, and /v3/login so SSO is the only way to reach the Admin UI. An admin locked out of the UI can still administer the proxy over the API with the master key; unset this setting and restart the proxy to restore UI username/password login. Default is False. */ disable_password_login_when_sso_enabled?: boolean | null; + /** + * Enable Openai Websocket Passthrough + * @description Serve the OpenAI pass-through WebSocket route, which relays frames to OpenAI under the proxy's own provider credential without reading them. Off by default. + */ + enable_openai_websocket_passthrough?: boolean | null; /** * Enable Public Model Hub * @description Public model hub for users to see what models they have access to, supported openai params, etc. From 2674934e45ffe7fc08d93be9684835e0297b1e59 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:28:53 -0700 Subject: [PATCH 14/38] feat(helm): render nodeSelector, tolerations, and affinity on the componentized chart migrations Job --- helm/litellm/templates/migrations-job.yaml | 12 ++++ helm/litellm/tests/migration_job_tests.yaml | 68 ++++++++++++++++++++- helm/litellm/values.yaml | 7 +++ 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/helm/litellm/templates/migrations-job.yaml b/helm/litellm/templates/migrations-job.yaml index 8d33081e72f..de1cc2b103b 100644 --- a/helm/litellm/templates/migrations-job.yaml +++ b/helm/litellm/templates/migrations-job.yaml @@ -77,4 +77,16 @@ spec: volumes: {{- toYaml . | nindent 8 }} {{- end }} + {{- with .Values.migrationJob.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.migrationJob.affinity }} + affinity: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.migrationJob.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} {{- end }} diff --git a/helm/litellm/tests/migration_job_tests.yaml b/helm/litellm/tests/migration_job_tests.yaml index c3f3083ece5..2ebb1b44926 100644 --- a/helm/litellm/tests/migration_job_tests.yaml +++ b/helm/litellm/tests/migration_job_tests.yaml @@ -1,4 +1,4 @@ -suite: test migrations Job ServiceAccount resolution and pod hardening +suite: test migrations Job ServiceAccount resolution, pod hardening, and scheduling templates: - migrations-job.yaml values: @@ -188,3 +188,69 @@ tests: asserts: - notExists: path: spec.activeDeadlineSeconds + + - it: renders no scheduling fields by default + asserts: + - isNull: + path: spec.template.spec.nodeSelector + - isNull: + path: spec.template.spec.tolerations + - isNull: + path: spec.template.spec.affinity + + - it: renders nodeSelector, tolerations, and affinity from the migrationJob values + set: + migrationJob.nodeSelector: + intent: no-csi-nodes + migrationJob.tolerations: + - key: intent + operator: Equal + value: no-csi-nodes + effect: NoSchedule + migrationJob.affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: intent + operator: In + values: + - no-csi-nodes + asserts: + - equal: + path: spec.template.spec.nodeSelector + value: + intent: no-csi-nodes + - equal: + path: spec.template.spec.tolerations + value: + - key: intent + operator: Equal + value: no-csi-nodes + effect: NoSchedule + - equal: + path: spec.template.spec.affinity + value: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: intent + operator: In + values: + - no-csi-nodes + + - it: does not inherit the gateway's scheduling values + set: + gateway.nodeSelector: + intent: no-csi-nodes + gateway.tolerations: + - key: intent + operator: Equal + value: no-csi-nodes + effect: NoSchedule + asserts: + - isNull: + path: spec.template.spec.nodeSelector + - isNull: + path: spec.template.spec.tolerations diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 461330ba491..6c9fb9440c7 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -152,6 +152,13 @@ migrationJob: # the writable scratch space a read-only root filesystem needs. volumes: [] volumeMounts: [] + # Scheduling for the Job pod, same shape as gateway.nodeSelector / + # gateway.tolerations / gateway.affinity. The Job does not inherit the other + # components' scheduling values: a migration usually needs a larger node + # than the gateway, so pin it here explicitly. + nodeSelector: {} + tolerations: [] + affinity: {} image: repository: ghcr.io/berriai/litellm-migrations tag: "" # defaults to .Chart.AppVersion From 7351911b533717155599bdfcbf09701aa1760fe5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:39:56 -0700 Subject: [PATCH 15/38] fix(proxy): refuse OpenAI websocket passthrough on every enforced model allowlist and propagate the DB opt-in --- litellm/proxy/auth/auth_checks.py | 56 ++ .../llm_passthrough_endpoints.py | 37 +- litellm/proxy/proxy_server.py | 5 + .../proxy/auth/test_auth_checks.py | 602 +++++++----------- .../test_openai_ws_passthrough_routes.py | 128 ++-- tests/test_litellm/proxy/test_proxy_server.py | 82 ++- 6 files changed, 461 insertions(+), 449 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f83f0303deb..98d334ce2cc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4155,6 +4155,62 @@ async def _granted_model_lists( ) +async def enforced_model_allowlists( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[Sequence[str], ...]: + """One model allowlist per level that ``common_checks`` enforces on a request from this identity.""" + key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) + if prisma_client is None: + return (key_models,) + team_object: Final = ( + None + if valid_token.team_id is None + else await get_team_object( + team_id=valid_token.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + user_object: Final = ( + None + if team_object is not None + else await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + ) + project_object: Final = ( + None + if valid_token.project_id is None + else await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + return ( + key_models, + team_object.models if team_object is not None else (), + await _team_member_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + user_object.models if user_object is not None else (), + project_object.models if project_object is not None else (), + ) + + async def collect_matched_model_access_groups( model: str | Sequence[str] | None, valid_token: UserAPIKeyAuth | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 97d25e20939..e92c949299c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -13,7 +13,7 @@ import inspect import json import os import re -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Protocol, cast @@ -36,6 +36,7 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import enforced_model_allowlists from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( @@ -2341,9 +2342,8 @@ _OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset( ) -def _key_has_model_restrictions(user_api_key_dict: UserAPIKeyAuth) -> bool: - scoped_models: Final = (*user_api_key_dict.models, *user_api_key_dict.team_models) - return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for model in scoped_models) +def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool: + return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist) @dataclass(frozen=True, slots=True) @@ -2376,12 +2376,18 @@ def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, obje return setting is True -def _openai_websocket_refusal( - user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object] +class _OpenAIWebsocketModelAllowlists(Protocol): + async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ... + + +async def _openai_websocket_refusal( + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], + model_allowlists: _OpenAIWebsocketModelAllowlists, ) -> _OpenAIWebsocketRefusal | None: if not _is_openai_websocket_passthrough_enabled(general_settings): return _OPENAI_WS_DISABLED_REFUSAL - if _key_has_model_restrictions(user_api_key_dict): + if _has_model_restrictions(await model_allowlists(user_api_key_dict)): return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL return None @@ -2410,6 +2416,20 @@ def _openai_websocket_relay() -> _OpenAIWebsocketRelay: return websocket_passthrough_request +def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: + return await enforced_model_allowlists( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return resolve + + @router.websocket("/openai_passthrough/{endpoint:path}") @router.websocket("/openai/{endpoint:path}") async def openai_websocket_proxy_route( @@ -2418,6 +2438,7 @@ async def openai_websocket_proxy_route( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], relay: Annotated[_OpenAIWebsocketRelay, Depends(_openai_websocket_relay)], + model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" requested_subprotocols: Final = tuple( @@ -2427,7 +2448,7 @@ async def openai_websocket_proxy_route( ) negotiated_subprotocol: Final = requested_subprotocols[0] if requested_subprotocols else None - refusal: Final = _openai_websocket_refusal(user_api_key_dict, general_settings) + refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) if refusal is not None: await websocket.accept(subprotocol=negotiated_subprotocol) await websocket.send_text( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5a39b8c610a..e56395a6169 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6793,6 +6793,11 @@ class ProxyConfig: else: general_settings["apply_user_budget_to_team_keys"] = db_value if db_value is None else bool(db_value) + if "enable_openai_websocket_passthrough" not in self._yaml_general_settings_keys: + general_settings["enable_openai_websocket_passthrough"] = _general_settings.get( + "enable_openai_websocket_passthrough" + ) + ## STORE MODEL IN DB ## if "store_model_in_db" in _general_settings: value = _general_settings["store_model_in_db"] diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index be83ca57e76..45e10948267 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -126,14 +126,10 @@ def invalid_sso_user_defined_values(): def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_values): """Test generating JWT token with valid user role""" - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") # Check that decrypted_token is not None before using json.loads assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -159,9 +155,7 @@ def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_value team_alias="test-team", ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -188,9 +182,7 @@ def test_get_cli_jwt_auth_token_carries_team_grants_not_user_allowlist( team_model_aliases={"team-fast": "gpt-4.1-mini"}, ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -207,9 +199,7 @@ def test_get_cli_jwt_auth_token_keeps_user_allowlist_when_no_team( """A session token with no team bound still carries the user's own allowlist.""" token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -222,12 +212,8 @@ def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( valid_sso_user_defined_values, ): """Test that Experimental UI token uses fixed 10-minute expiry (does not use LITELLM_UI_SESSION_DURATION).""" - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) @@ -244,43 +230,33 @@ def test_experimental_ui_token_ignores_litellm_ui_session_duration( Experimental UI intentionally uses fixed 10-min expiry. If this test fails, the constant was incorrectly wired to the experimental flow.""" # Default LITELLM_UI_SESSION_DURATION is "24h" - token must still expire in ~10 min - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) now = get_utc_datetime() # Must be ~10 min, NOT 24h. If LITELLM_UI_SESSION_DURATION were incorrectly used, this would fail. - assert expires <= now + timedelta( - minutes=11 - ), "Experimental UI must use 10-min expiry, not LITELLM_UI_SESSION_DURATION" + assert expires <= now + timedelta(minutes=11), ( + "Experimental UI must use 10-min expiry, not LITELLM_UI_SESSION_DURATION" + ) def test_get_experimental_ui_login_jwt_auth_token_invalid( invalid_sso_user_defined_values, ): """Test generating JWT token with missing user role""" - with pytest.raises(Exception, match='User role is required for experimental UI login') as exc_info: - ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - invalid_sso_user_defined_values - ) + with pytest.raises(Exception, match="User role is required for experimental UI login") as exc_info: + ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(invalid_sso_user_defined_values) assert str(exc_info.value) == "User role is required for experimental UI login" -def test_get_key_object_from_ui_hash_key_valid( - valid_sso_user_defined_values, monkeypatch -): +def test_get_key_object_from_ui_hash_key_valid(valid_sso_user_defined_values, monkeypatch): """Test getting key object from valid UI hash key""" monkeypatch.setenv("EXPERIMENTAL_UI_LOGIN", "True") # Generate a valid token - token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - valid_sso_user_defined_values - ) + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) # Get key object key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) @@ -309,9 +285,7 @@ def test_get_key_object_from_ui_hash_key_invalid(): ("project", ProxyErrorTypes.project_model_access_denied), ], ) -def test_can_object_call_model_denials_return_forbidden( - object_type, expected_error_type -): +def test_can_object_call_model_denials_return_forbidden(object_type, expected_error_type): with pytest.raises(ProxyException) as exc_info: _can_object_call_model( model="restricted-model", @@ -568,9 +542,7 @@ async def test_get_key_object_should_reconnect_once_on_db_connection_error(): @pytest.mark.asyncio async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_error(): mock_prisma_client = MagicMock() - mock_prisma_client.get_data = AsyncMock( - side_effect=httpx.ConnectError("db not reachable after outage") - ) + mock_prisma_client.get_data = AsyncMock(side_effect=httpx.ConnectError("db not reachable after outage")) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) mock_cache = MagicMock() @@ -613,9 +585,7 @@ class TestAuthCacheRedisWritePolicy: @pytest.mark.asyncio async def test_get_key_object_db_load_publishes_to_redis(self): mock_prisma_client = MagicMock() - mock_prisma_client.get_data = AsyncMock( - return_value=UserAPIKeyAuth(token="hashed-token-db") - ) + mock_prisma_client.get_data = AsyncMock(return_value=UserAPIKeyAuth(token="hashed-token-db")) fake_redis = _fake_redis_cache() cache = UserApiKeyCache() @@ -630,8 +600,7 @@ class TestAuthCacheRedisWritePolicy: assert key_obj.token == "hashed-token-db" fake_redis.async_set_cache.assert_awaited_once() assert ( - fake_redis.async_set_cache.await_args.kwargs.get("key") - or fake_redis.async_set_cache.await_args.args[0] + fake_redis.async_set_cache.await_args.kwargs.get("key") or fake_redis.async_set_cache.await_args.args[0] ) == "hashed-token-db" @@ -640,9 +609,7 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -664,9 +631,7 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values assert expires >= get_utc_datetime() + timedelta(hours=23, minutes=59) -def test_get_cli_jwt_auth_token_custom_expiration( - valid_sso_user_defined_values, monkeypatch -): +def test_get_cli_jwt_auth_token_custom_expiration(valid_sso_user_defined_values, monkeypatch): """Test generating CLI JWT token with custom expiration via environment variable""" import importlib @@ -681,14 +646,10 @@ def test_get_cli_jwt_auth_token_custom_expiration( # Also reload auth_checks to pick up the new constant value importlib.reload(auth_checks) - token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token( - valid_sso_user_defined_values - ) + token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -706,18 +667,12 @@ def test_get_cli_jwt_auth_token_unique_per_session(valid_sso_user_defined_values from litellm.constants import CLI_SESSION_KEY_PREFIX def _decode(token: str) -> dict: - decrypted = decrypt_value_helper( - token, key="ui_hash_key", exception_type="debug" - ) + decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted is not None return json.loads(decrypted) - first = _decode( - ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - ) - second = _decode( - ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - ) + first = _decode(ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values)) + second = _decode(ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values)) assert first["token"].startswith(f"{CLI_SESSION_KEY_PREFIX}-") assert second["token"].startswith(f"{CLI_SESSION_KEY_PREFIX}-") @@ -740,9 +695,7 @@ def test_get_cli_jwt_auth_token_applies_fallback_budget(valid_sso_user_defined_v def test_get_cli_jwt_auth_token_no_fallback_when_budget_provided( valid_sso_user_defined_values, ): - token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - valid_sso_user_defined_values, max_budget=None - ) + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values, max_budget=None) decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") assert decrypted is not None assert json.loads(decrypted).get("max_budget") is None @@ -945,9 +898,7 @@ async def test_get_user_object_upsert_includes_user_email(): mock_prisma_client.db.litellm_usertable.create.assert_called_once() creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"] - assert ( - "user_email" in creation_args - ), "user_email should be included when upserting a new user" + assert "user_email" in creation_args, "user_email should be included when upserting a new user" assert creation_args["user_email"] == "test@example.com" assert creation_args["user_id"] == "new_test_user" @@ -962,12 +913,8 @@ async def test_get_user_object_backfills_null_email_from_cache_hit(): was returned unchanged and the DB was never updated. """ cache = UserApiKeyCache() - existing = LiteLLM_UserTable( - user_id="jwt-user-1", user_email=None, user_role="internal_user" - ) - await cache.async_set_cache( - key="jwt-user-1", value=existing, model_type=LiteLLM_UserTable - ) + existing = LiteLLM_UserTable(user_id="jwt-user-1", user_email=None, user_role="internal_user") + await cache.async_set_cache(key="jwt-user-1", value=existing, model_type=LiteLLM_UserTable) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) @@ -996,9 +943,7 @@ async def test_get_user_object_backfills_null_email_from_cache_hit(): assert update_kwargs["where"] == {"user_id": "jwt-user-1", "user_email": None} assert update_kwargs["data"]["user_email"] == "jwt-user-1@example.com" - refreshed = await cache.async_get_cache( - key="jwt-user-1", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-1", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "jwt-user-1@example.com" @@ -1010,9 +955,7 @@ async def test_get_user_object_backfills_null_email_from_db_read(): backfilled from the JWT-provided email before it is cached and returned. """ cache = UserApiKeyCache() - db_row = LiteLLM_UserTable( - user_id="jwt-user-3", user_email=None, user_role="internal_user" - ) + db_row = LiteLLM_UserTable(user_id="jwt-user-3", user_email=None, user_role="internal_user") backfilled_row = LiteLLM_UserTable( user_id="jwt-user-3", user_email="jwt-user-3@example.com", @@ -1020,15 +963,11 @@ async def test_get_user_object_backfills_null_email_from_db_read(): ) mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=[db_row, backfilled_row] - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=[db_row, backfilled_row]) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) - with patch( - "litellm.proxy.auth.auth_checks._should_check_db", return_value=True - ): + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): result = await get_user_object( user_id="jwt-user-3", prisma_client=mock_prisma_client, @@ -1042,9 +981,7 @@ async def test_get_user_object_backfills_null_email_from_db_read(): assert result.user_email == "jwt-user-3@example.com" mock_prisma_client.db.litellm_usertable.update_many.assert_called_once() - refreshed = await cache.async_get_cache( - key="jwt-user-3", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-3", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "jwt-user-3@example.com" @@ -1062,9 +999,7 @@ async def test_get_user_object_does_not_overwrite_existing_email(): user_email="operator-set@example.com", user_role="internal_user", ) - await cache.async_set_cache( - key="jwt-user-2", value=existing, model_type=LiteLLM_UserTable - ) + await cache.async_set_cache(key="jwt-user-2", value=existing, model_type=LiteLLM_UserTable) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=0) @@ -1091,12 +1026,8 @@ async def test_get_user_object_backfill_race_prefers_db_email(): with the value the DB accepted, not this request's proposed email. """ cache = UserApiKeyCache() - existing = LiteLLM_UserTable( - user_id="jwt-user-4", user_email=None, user_role="internal_user" - ) - await cache.async_set_cache( - key="jwt-user-4", value=existing, model_type=LiteLLM_UserTable - ) + existing = LiteLLM_UserTable(user_id="jwt-user-4", user_email=None, user_role="internal_user") + await cache.async_set_cache(key="jwt-user-4", value=existing, model_type=LiteLLM_UserTable) winner_row = LiteLLM_UserTable( user_id="jwt-user-4", @@ -1105,9 +1036,7 @@ async def test_get_user_object_backfill_race_prefers_db_email(): ) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=winner_row - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=winner_row) result = await get_user_object( user_id="jwt-user-4", @@ -1121,9 +1050,7 @@ async def test_get_user_object_backfill_race_prefers_db_email(): assert result is not None assert result.user_email == "winner@example.com" - refreshed = await cache.async_get_cache( - key="jwt-user-4", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-4", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "winner@example.com" @@ -1138,12 +1065,8 @@ async def test_get_user_object_backfill_caches_persisted_email_not_proposed(): optimistically caching the proposed email would serve a stale value. """ cache = UserApiKeyCache() - existing = LiteLLM_UserTable( - user_id="jwt-user-5", user_email=None, user_role="internal_user" - ) - await cache.async_set_cache( - key="jwt-user-5", value=existing, model_type=LiteLLM_UserTable - ) + existing = LiteLLM_UserTable(user_id="jwt-user-5", user_email=None, user_role="internal_user") + await cache.async_set_cache(key="jwt-user-5", value=existing, model_type=LiteLLM_UserTable) persisted_row = LiteLLM_UserTable( user_id="jwt-user-5", @@ -1152,9 +1075,7 @@ async def test_get_user_object_backfill_caches_persisted_email_not_proposed(): ) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=persisted_row - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=persisted_row) result = await get_user_object( user_id="jwt-user-5", @@ -1168,9 +1089,7 @@ async def test_get_user_object_backfill_caches_persisted_email_not_proposed(): assert result is not None assert result.user_email == "admin-edited@example.com" - refreshed = await cache.async_get_cache( - key="jwt-user-5", model_type=LiteLLM_UserTable - ) + refreshed = await cache.async_get_cache(key="jwt-user-5", model_type=LiteLLM_UserTable) assert refreshed is not None assert refreshed.user_email == "admin-edited@example.com" @@ -1224,10 +1143,7 @@ async def test_get_user_object_upsert_routes_default_team_to_membership(monkeypa mock_add_to_team.assert_awaited_once() passed_teams = mock_add_to_team.await_args[1]["teams"] assert [team.team_id for team in passed_teams] == ["default-team"] - assert ( - mock_add_to_team.await_args[1]["user_api_key_dict"].user_role - == LitellmUserRoles.PROXY_ADMIN - ) + assert mock_add_to_team.await_args[1]["user_api_key_dict"].user_role == LitellmUserRoles.PROXY_ADMIN def test_log_budget_lookup_failure_dry_run(): @@ -1252,9 +1168,7 @@ def test_log_budget_lookup_failure_skips_user_not_found(): @pytest.mark.asyncio -@patch( - "litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock -) +@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock) async def test_get_team_db_check_calls_new_team_on_upsert(mock_new_team, monkeypatch): """ Test that _get_team_db_check correctly calls the `new_team` function @@ -1288,12 +1202,8 @@ async def test_get_team_db_check_calls_new_team_on_upsert(mock_new_team, monkeyp @pytest.mark.asyncio -@patch( - "litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock -) -async def test_get_team_db_check_does_not_call_new_team_if_exists( - mock_new_team, monkeypatch -): +@patch("litellm.proxy.management_endpoints.team_endpoints.new_team", new_callable=AsyncMock) +async def test_get_team_db_check_does_not_call_new_team_if_exists(mock_new_team, monkeypatch): """ Test that _get_team_db_check does NOT call the `new_team` function if the team already exists in the database. @@ -1327,9 +1237,7 @@ async def test_get_team_db_check_does_not_call_new_team_if_exists( (MagicMock(), MagicMock(), True), # No vector stores to run ], ) -async def test_vector_store_access_check_early_returns( - prisma_client, vector_store_registry, expected_result -): +async def test_vector_store_access_check_early_returns(prisma_client, vector_store_registry, expected_result): """Test vector_store_access_check returns True for early exit conditions""" request_body = {"messages": [{"role": "user", "content": "test"}]} @@ -1411,9 +1319,7 @@ async def test_vector_store_access_check_skips_db_lookup_when_no_vector_stores_r ), # Partial access ], ) -def test_can_object_call_vector_stores_scenarios( - object_permissions, vector_store_ids, should_raise, error_type -): +def test_can_object_call_vector_stores_scenarios(object_permissions, vector_store_ids, should_raise, error_type): """Test _can_object_call_vector_stores with various permission scenarios""" # Convert dict to object if not None if object_permissions is not None: @@ -1421,11 +1327,7 @@ def test_can_object_call_vector_stores_scenarios( mock_permissions.vector_stores = object_permissions["vector_stores"] object_permissions = mock_permissions - object_type = ( - "key" - if error_type == ProxyErrorTypes.key_vector_store_access_denied - else "team" - ) + object_type = "key" if error_type == ProxyErrorTypes.key_vector_store_access_denied else "team" if should_raise: with pytest.raises(ProxyException) as exc_info: @@ -1460,9 +1362,7 @@ async def test_vector_store_access_check_with_permissions(): mock_prisma_client = MagicMock() mock_permissions = MagicMock() mock_permissions.vector_stores = ["store-1", "store-2"] - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=mock_permissions - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=mock_permissions) mock_vector_store_registry = MagicMock() mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-1"] @@ -1508,14 +1408,10 @@ async def test_vector_store_access_check_with_team_permissions(): mock_prisma_client = MagicMock() team_permissions = MagicMock() team_permissions.vector_stores = ["team-store-allowed"] - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=team_permissions - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=team_permissions) mock_vector_store_registry = MagicMock() - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = [ - "team-store-allowed" - ] + mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-allowed"] with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), @@ -1529,9 +1425,7 @@ async def test_vector_store_access_check_with_team_permissions(): assert result is True - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = [ - "team-store-denied" - ] + mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-denied"] with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), @@ -2092,9 +1986,7 @@ async def test_get_tag_objects_batch(): mock_cache.async_set_cache = AsyncMock() # Mock DB to return all uncached tags in ONE query - mock_prisma.db.litellm_tagtable.find_many = AsyncMock( - return_value=[uncached_tag_1, uncached_tag_2, uncached_tag_3] - ) + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(return_value=[uncached_tag_1, uncached_tag_2, uncached_tag_3]) # Call batch fetch tag_objects = await get_tag_objects_batch( @@ -2196,9 +2088,7 @@ async def test_get_tag_objects_batch_never_queries_db_for_unregistered_tags(): from litellm.proxy.auth.auth_checks import get_tag_objects_batch mock_prisma = MagicMock() - mock_prisma.db.litellm_tagtable.find_many = AsyncMock( - return_value=[_tag_registry_row("some-other-tag")] - ) + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(return_value=[_tag_registry_row("some-other-tag")]) cache = UserApiKeyCache() first = await get_tag_objects_batch( @@ -2209,9 +2099,7 @@ async def test_get_tag_objects_batch_never_queries_db_for_unregistered_tags(): assert first == {} # The only query is the names-only registry fetch; the tag itself is never looked up. - mock_prisma.db.litellm_tagtable.find_many.assert_called_once_with( - take=TAG_REGISTRY_MAX_SIZE + 1 - ) + mock_prisma.db.litellm_tagtable.find_many.assert_called_once_with(take=TAG_REGISTRY_MAX_SIZE + 1) second = await get_tag_objects_batch( tag_names=["unregistered-tag"], @@ -2380,9 +2268,7 @@ async def test_get_tag_objects_batch_oversized_registry_falls_back_and_stops_ref """Past the cap the registry is unusable: keep the old per-tag path, but stop rebuilding it.""" from litellm.proxy.auth.auth_checks import get_tag_objects_batch - oversized = [ - _tag_registry_row(f"tag-{index}") for index in range(TAG_REGISTRY_MAX_SIZE + 1) - ] + oversized = [_tag_registry_row(f"tag-{index}") for index in range(TAG_REGISTRY_MAX_SIZE + 1)] async def fake_find_many(**kwargs): if "where" not in kwargs: @@ -2399,10 +2285,7 @@ async def test_get_tag_objects_batch_oversized_registry_falls_back_and_stops_ref user_api_key_cache=cache, ) assert list(first) == ["tag-a"] - assert ( - await cache.async_get_cache(key=tag_registry_cache_key()) - == TAG_REGISTRY_OVERFLOW_SENTINEL - ) + assert await cache.async_get_cache(key=tag_registry_cache_key()) == TAG_REGISTRY_OVERFLOW_SENTINEL second = await get_tag_objects_batch( tag_names=["tag-b"], @@ -2427,17 +2310,12 @@ async def test_tag_max_budget_check_still_enforces_registered_tag_over_budget(): async def fake_find_many(**kwargs): if "where" not in kwargs: return [_tag_registry_row("paid-tag")] - return [ - _tag_db_row(name, max_budget=1.0) - for name in kwargs["where"]["tag_name"]["in"] - ] + return [_tag_db_row(name, max_budget=1.0) for name in kwargs["where"]["tag_name"]["in"]] mock_prisma = MagicMock() mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:paid-tag": return 1.5 return fallback_spend @@ -2846,8 +2724,7 @@ def _pass_through_request() -> Request: LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) - def pass_through_endpoint(): - ... + def pass_through_endpoint(): ... setattr(pass_through_endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) return Request(scope={"type": "http", "headers": [], "endpoint": pass_through_endpoint}) @@ -2857,8 +2734,7 @@ def _builtin_request() -> Request: """A Request dispatched to a built-in (non-pass-through) handler, e.g. what a custom path colliding with a core route actually resolves to.""" - def chat_completions(): - ... + def chat_completions(): ... return Request(scope={"type": "http", "headers": [], "endpoint": chat_completions}) @@ -3023,9 +2899,7 @@ async def test_virtual_key_soft_budget_check_without_user_obj(): ], ) @pytest.mark.asyncio -async def test_virtual_key_soft_budget_check_scenarios( - spend, soft_budget, expect_alert -): +async def test_virtual_key_soft_budget_check_scenarios(spend, soft_budget, expect_alert): """Test _virtual_key_soft_budget_check with various spend and soft_budget scenarios""" alert_triggered = False @@ -3054,9 +2928,9 @@ async def test_virtual_key_soft_budget_check_scenarios( await asyncio.sleep(0.1) - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + ) @pytest.mark.asyncio @@ -3167,9 +3041,7 @@ async def test_virtual_key_max_budget_alert_check_without_user_obj(): ], ) @pytest.mark.asyncio -async def test_virtual_key_max_budget_alert_check_scenarios( - spend, max_budget, expect_alert -): +async def test_virtual_key_max_budget_alert_check_scenarios(spend, max_budget, expect_alert): """Test _virtual_key_max_budget_alert_check with various spend and max_budget scenarios""" alert_triggered = False @@ -3198,9 +3070,9 @@ async def test_virtual_key_max_budget_alert_check_scenarios( await asyncio.sleep(0.1) - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" + ) @pytest.mark.asyncio @@ -3459,9 +3331,7 @@ async def test_custom_auth_common_checks_opt_in(): "prisma_client": None, "user_api_key_cache": MagicMock(), "proxy_logging_obj": MagicMock(), - "general_settings": ( - {"custom_auth_run_common_checks": True} if flag else {} - ), + "general_settings": ({"custom_auth_run_common_checks": True} if flag else {}), "llm_router": None, "user_custom_auth": user_custom_auth, "litellm_proxy_admin_name": "admin", @@ -3533,9 +3403,7 @@ async def test_virtual_key_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:key:test-hashed-token": return 1.5 return fallback_spend @@ -3569,9 +3437,7 @@ async def test_virtual_key_budget_check_fallback_no_counter(): proxy_logging_obj.budget_alerts = AsyncMock() # get_current_spend returns fallback_spend when no counter exists - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): return fallback_spend with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): @@ -3601,9 +3467,7 @@ def _over_budget_token(**overrides) -> UserAPIKeyAuth: def _patched_spend(value: float): - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): return value return patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend) @@ -3664,9 +3528,7 @@ async def test_budget_throttle_decision_cleared_before_caching(): otherwise it would re-apply (and compound) on every subsequent request.""" from litellm.proxy.auth.auth_checks import _copy_user_api_key_auth_for_cache - valid_token = _over_budget_token( - tpm_limit=1000, rpm_limit=100, metadata={"throttle_on_budget_exceeded": True} - ) + valid_token = _over_budget_token(tpm_limit=1000, rpm_limit=100, metadata={"throttle_on_budget_exceeded": True}) valid_token.budget_throttle_pct = 0.1 cached = _copy_user_api_key_auth_for_cache(user_api_key_obj=valid_token) @@ -3762,9 +3624,7 @@ async def test_team_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team:test-team": return 1.5 return fallback_spend @@ -3791,9 +3651,7 @@ async def test_end_user_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:end_user:customer-1": return 1.5 return fallback_spend @@ -3821,9 +3679,7 @@ async def test_tag_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:paid-tag": return 1.5 return fallback_spend @@ -3873,9 +3729,7 @@ async def test_team_member_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 1.5 return fallback_spend @@ -3915,9 +3769,7 @@ class TestGuardrailModificationCheck: team_object = MagicMock() team_object.metadata = {} # no permission - return _guardrail_modification_check( - request_body=request_body, team_object=team_object - ) + return _guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): # no-op — should return silently @@ -3965,9 +3817,7 @@ class TestGuardrailModificationCheck: return_value=False, ): with pytest.raises(HTTPException) as exc: - self._call( - {"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}} - ) + self._call({"metadata": {"opted_out_global_guardrails": ["some_guardrail"]}}) assert exc.value.status_code == 403 @pytest.mark.parametrize( @@ -4101,18 +3951,12 @@ async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): fake_budget_row = MagicMock() fake_budget_row.max_budget = 50.0 - fake_budget_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": 50.0} - ) + fake_budget_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": 50.0}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_budget_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_budget_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 70.0 return fallback_spend @@ -4203,15 +4047,11 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau fake_budget_row.max_budget = 50.0 prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_budget_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_budget_row) mocked_spend = 70.0 - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return mocked_spend return fallback_spend @@ -4292,18 +4132,12 @@ async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): fake_default_row = MagicMock() fake_default_row.max_budget = 65.0 - fake_default_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": 65.0} - ) + fake_default_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": 65.0}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_default_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_default_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 500.0 return fallback_spend @@ -4361,18 +4195,12 @@ async def test_team_member_budget_check_null_clone_with_null_default_skips_enfor fake_default_row = MagicMock() fake_default_row.max_budget = None - fake_default_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": None} - ) + fake_default_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": None}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_default_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_default_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 1000.0 return fallback_spend @@ -4430,18 +4258,12 @@ async def test_team_member_budget_check_zero_team_default_treated_as_no_cap(): # Team default budget row with max_budget=0.0 (the regression trigger). fake_default_row = MagicMock() fake_default_row.max_budget = 0.0 - fake_default_row.dict = MagicMock( - return_value={"budget_id": "budget-default", "max_budget": 0.0} - ) + fake_default_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": 0.0}) prisma_client = MagicMock() - prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=fake_default_row - ) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=fake_default_row) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -4499,9 +4321,7 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): prisma_client = MagicMock() prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -4549,19 +4369,13 @@ def _patch_validation_helpers(monkeypatch, *, end_user=None, user=None, fuzzy=No """Stub out the DB helpers resolve_and_validate_end_user_id delegates to.""" from litellm.proxy.auth import auth_checks - monkeypatch.setattr( - auth_checks, "get_end_user_object", AsyncMock(return_value=end_user) - ) + monkeypatch.setattr(auth_checks, "get_end_user_object", AsyncMock(return_value=end_user)) monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=user)) - monkeypatch.setattr( - auth_checks, "_get_fuzzy_user_object", AsyncMock(return_value=fuzzy) - ) + monkeypatch.setattr(auth_checks, "_get_fuzzy_user_object", AsyncMock(return_value=fuzzy)) @pytest.mark.asyncio -async def test_resolve_end_user_returns_none_for_none_input( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_returns_none_for_none_input(_validate_flag_on, monkeypatch): from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id _patch_validation_helpers(monkeypatch) @@ -4596,9 +4410,7 @@ async def test_resolve_end_user_passes_through_when_flag_disabled(monkeypatch): @pytest.mark.asyncio -async def test_resolve_end_user_passes_through_when_no_prisma_client( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_passes_through_when_no_prisma_client(_validate_flag_on, monkeypatch): from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id _patch_validation_helpers(monkeypatch) @@ -4632,9 +4444,7 @@ async def test_resolve_end_user_matches_end_user_table(_validate_flag_on, monkey @pytest.mark.asyncio -async def test_resolve_end_user_matches_user_table_by_user_id( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_matches_user_table_by_user_id(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4652,9 +4462,7 @@ async def test_resolve_end_user_matches_user_table_by_user_id( @pytest.mark.asyncio -async def test_resolve_end_user_matches_user_table_by_email( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_matches_user_table_by_email(_validate_flag_on, monkeypatch): """Email-shaped ids route through get_user_object with user_email set. The fuzzy lookup must happen inside get_user_object so it shares the @@ -4682,9 +4490,7 @@ async def test_resolve_end_user_matches_user_table_by_email( @pytest.mark.asyncio -async def test_resolve_end_user_non_email_id_does_not_pass_user_email( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_non_email_id_does_not_pass_user_email(_validate_flag_on, monkeypatch): """Non-email ids skip the email fuzzy path to avoid a pointless DB hit.""" from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4703,9 +4509,7 @@ async def test_resolve_end_user_non_email_id_does_not_pass_user_email( @pytest.mark.asyncio -async def test_resolve_end_user_drops_codex_opaque_identifier( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_drops_codex_opaque_identifier(_validate_flag_on, monkeypatch): from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id _patch_validation_helpers(monkeypatch) # all helpers return None @@ -4727,9 +4531,7 @@ async def test_resolve_end_user_drops_codex_opaque_identifier( @pytest.mark.asyncio -async def test_resolve_end_user_preserves_id_when_default_budget_configured( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_preserves_id_when_default_budget_configured(_validate_flag_on, monkeypatch): """Don't drop unregistered ids when litellm.max_end_user_budget_id is set. The default end-user budget is applied downstream when the id is present @@ -4766,9 +4568,7 @@ async def test_resolve_end_user_drops_unknown_email(_validate_flag_on, monkeypat @pytest.mark.asyncio -async def test_resolve_end_user_uses_cached_valid_result( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_uses_cached_valid_result(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4788,9 +4588,7 @@ async def test_resolve_end_user_uses_cached_valid_result( @pytest.mark.asyncio -async def test_resolve_end_user_uses_cached_invalid_result( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_uses_cached_invalid_result(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4809,9 +4607,7 @@ async def test_resolve_end_user_uses_cached_invalid_result( @pytest.mark.asyncio -async def test_resolve_end_user_swallows_db_errors_and_returns_none( - _validate_flag_on, monkeypatch -): +async def test_resolve_end_user_swallows_db_errors_and_returns_none(_validate_flag_on, monkeypatch): from litellm.proxy.auth import auth_checks from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id @@ -4932,19 +4728,13 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): ) # (1) team_id-keyed write fires with the refreshed object - written_keys = [ - (c.kwargs.get("key") or c.args[0]) - for c in cache.async_set_cache.await_args_list - ] + written_keys = [(c.kwargs.get("key") or c.args[0]) for c in cache.async_set_cache.await_args_list] assert written_keys == ["team_id:team-1234"], ( "Only the team_id-keyed write should fire; the alias key must be " "deleted, NOT written. " f"Got writes: {written_keys}" ) - written_value = ( - cache.async_set_cache.await_args.kwargs.get("value") - or cache.async_set_cache.await_args.args[1] - ) + written_value = cache.async_set_cache.await_args.kwargs.get("value") or cache.async_set_cache.await_args.args[1] assert written_value is team_table # (2) team_alias-keyed entry is deleted in BOTH the in-memory cache @@ -4978,10 +4768,7 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): logging_obj2.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with( key="team_id:team-no-alias" ) - written_keys_aliasless = [ - (c.kwargs.get("key") or c.args[0]) - for c in cache2.async_set_cache.await_args_list - ] + written_keys_aliasless = [(c.kwargs.get("key") or c.args[0]) for c in cache2.async_set_cache.await_args_list] assert written_keys_aliasless == ["team_id:team-no-alias"] @@ -5061,9 +4848,7 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): await _cache_team_object( team_id=team_id, - team_table=LiteLLM_TeamTableCachedObj( - team_id=team_id, models=["model-a", "model-b"] - ), + team_table=LiteLLM_TeamTableCachedObj(team_id=team_id, models=["model-a", "model-b"]), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -5173,9 +4958,7 @@ async def test_cache_team_object_tolerates_cache_invalidation_failures(): cache.async_set_cache = AsyncMock() cache.delete_cache = MagicMock(side_effect=Exception("redis down")) logging_obj = MagicMock() - logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( - side_effect=Exception("redis down") - ) + logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(side_effect=Exception("redis down")) await _cache_team_object( team_id="team-cache-outage", @@ -5188,10 +4971,7 @@ async def test_cache_team_object_tolerates_cache_invalidation_failures(): proxy_logging_obj=logging_obj, ) - written_keys = [ - (c.kwargs.get("key") or c.args[0]) - for c in cache.async_set_cache.await_args_list - ] + written_keys = [(c.kwargs.get("key") or c.args[0]) for c in cache.async_set_cache.await_args_list] assert written_keys == ["team_id:team-cache-outage"] @@ -5467,8 +5247,9 @@ async def test_common_checks_budget_reads_run_concurrently(): probe = _BudgetSpendConcurrencyProbe(expected=4) - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", probe + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", probe), ): task = asyncio.create_task( common_checks( @@ -5536,8 +5317,9 @@ async def test_common_checks_budget_gather_raises_highest_priority_scope(): request=MagicMock(spec=Request), ) - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), ): # Both team and end-user over budget: team wins on priority. _spend_by_counter.team = 999.0 @@ -5572,8 +5354,9 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), ): with pytest.raises(litellm.BudgetExceededError) as over: await common_checks( @@ -5615,9 +5398,11 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): async def _no_membership(*args, **kwargs): return None - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter - ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), + patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), + ): result = await common_checks( request_body={"messages": [{"role": "user", "content": "hi"}]}, team_object=team, @@ -5656,9 +5441,11 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag async def _no_membership(*args, **kwargs): return None - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter - ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), + patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), + ): with pytest.raises(litellm.BudgetExceededError) as exc_info: await common_checks( request_body={"messages": [{"role": "user", "content": "hi"}]}, @@ -5689,8 +5476,9 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 - with patch("litellm.proxy.proxy_server.prisma_client", None), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), ): with pytest.raises(litellm.BudgetExceededError): await common_checks( @@ -5773,10 +5561,11 @@ async def test_budget_checks_only_run_on_llm_api_routes(scope, route, expect_blo request=MagicMock(spec=Request), ) - with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter - ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), patch( - "litellm.proxy.auth.auth_checks.get_org_object", _get_org + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), + patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), + patch("litellm.proxy.auth.auth_checks.get_org_object", _get_org), ): if expect_blocked: with pytest.raises(litellm.BudgetExceededError): @@ -6275,9 +6064,7 @@ async def test_get_end_user_object_token_budget_gate_keeps_fetching_unrestricted mock_prisma = MagicMock() mock_prisma.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_endusertable.find_unique = AsyncMock( - return_value=_end_user_db_row("eu-anon-1", spend=100.0) - ) + mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(return_value=_end_user_db_row("eu-anon-1", spend=100.0)) cache = UserApiKeyCache() result = await get_end_user_object( @@ -6959,9 +6746,7 @@ async def test_common_checks_ignores_non_llm_route_when_enabled(monkeypatch): monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) router = _router_with_priced_and_unpriced_models() - result = await _run_common_checks( - model="unpriced-group", llm_router=router, route="/model/new" - ) + result = await _run_common_checks(model="unpriced-group", llm_router=router, route="/model/new") assert result is True @@ -7077,11 +6862,15 @@ def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): roles = LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/model-a"]) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles) + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles + ) is True ) assert ( - allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles) + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles + ) is False ) @@ -7159,8 +6948,7 @@ async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_cle assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0 assert ( - real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") - == 0.0 + real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") == 0.0 ), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up" @@ -7356,9 +7144,9 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t ) assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0 - assert ( - remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 - ), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor" + assert remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0, ( + "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor" + ) @pytest.mark.asyncio @@ -7541,3 +7329,81 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name): names are just as valid as the alphanumeric ones.""" message = await _run_key_budget_check(key_name) assert f"Key=prod-key ({key_name}) Current cost" in message + + +class _UntouchedPrisma: + def __getattr__(self, name: str) -> object: + raise AssertionError(f"database reached through {name}") + + +@pytest.mark.asyncio +async def test_enforced_model_allowlists_reads_every_level_from_cache(): + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + ) + from litellm.proxy.auth.auth_checks import enforced_model_allowlists + from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_reservation_cache_key, + ) + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=cache) + await cache.async_set_cache( + key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"]) + ) + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"), + value=LiteLLM_TeamMembership( + user_id="user-fake", + team_id="team-fake", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]), + ), + ) + await cache.async_set_cache( + key="project_id:project-fake", + value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]), + ) + await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"])) + prisma_client = _UntouchedPrisma() + + team_scoped = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth( + token="hashed-fake", + models=["all-team-models"], + team_models=["gpt-4o", "gpt-4o-mini"], + user_id="user-fake", + team_id="team-fake", + project_id="project-fake", + ), + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) + personal = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) + without_database = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth(token="hashed-fake", models=["gpt-4o"], user_id="user-fake", team_id="team-fake"), + prisma_client=None, + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert [list(scope) for scope in team_scoped] == [ + ["gpt-4o", "gpt-4o-mini"], + ["gpt-4o"], + ["gpt-4o-mini"], + [], + ["gpt-4.1"], + ] + assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []] + assert [list(scope) for scope in without_database] == [["gpt-4o"]] diff --git a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py index 6578b75ace1..c96e7684e97 100644 --- a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py @@ -1,7 +1,7 @@ """OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals.""" import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final @@ -15,10 +15,13 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _OPENAI_WS_DISABLED_REFUSAL, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL, _openai_websocket_refusal, + _proxy_model_allowlists, openai_websocket_proxy_route, router, ) +Scopes = tuple[Sequence[str], ...] + ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) DISABLED_SETTINGS: Final = ( MappingProxyType({}), @@ -96,21 +99,39 @@ class _FakeRelay: ) +class _FakeModelAllowlists: + def __init__(self, scopes: Scopes) -> None: + self.scopes = scopes + self.calls: list[UserAPIKeyAuth] = [] + + async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes: + self.calls.append(valid_token) + return self.scopes + + +@dataclass(frozen=True, slots=True) +class _Served: + relay: _FakeRelay + allowlists: _FakeModelAllowlists + + async def _serve( websocket: _FakeWebSocket, endpoint: str, user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], -) -> _FakeRelay: - relay = _FakeRelay() + scopes: Scopes = (), +) -> _Served: + served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes)) await openai_websocket_proxy_route( websocket=websocket, endpoint=endpoint, user_api_key_dict=user_api_key_dict, general_settings=general_settings, - relay=relay, + relay=served.relay, + model_allowlists=served.allowlists, ) - return relay + return served @pytest.mark.asyncio @@ -120,9 +141,9 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") with patch(GET_CREDENTIALS, return_value="sk-provider"): - relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) - assert relay.calls == [ + assert served.relay.calls == [ _RelayCall( target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview", custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}), @@ -145,10 +166,10 @@ async def test_openai_websocket_accepts_first_client_subprotocol(): ) with patch(GET_CREDENTIALS, return_value="sk-provider"): - relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) assert websocket.accepts == ["realtime"] - assert [call.accept_websocket for call in relay.calls] == [False] + assert [call.accept_websocket for call in served.relay.calls] == [False] assert websocket.closed is None @@ -157,13 +178,13 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") with patch(GET_CREDENTIALS, return_value=None): - relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) assert websocket.closed is not None assert websocket.closed[0] == 1011 assert "OPENAI_API_KEY" in websocket.closed[1] assert websocket.accepts == [] - assert relay.calls == [] + assert served.relay.calls == [] @pytest.mark.asyncio @@ -172,23 +193,26 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings): websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") - relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) assert "enable_openai_websocket_passthrough" in websocket.error_message() assert websocket.accepts == [None] assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason) - assert relay.calls == [] + assert served.relay.calls == [] +@pytest.mark.asyncio @pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) -def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): - assert _openai_websocket_refusal(UserAPIKeyAuth(), general_settings) is _OPENAI_WS_DISABLED_REFUSAL +async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): + refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(())) + assert refusal is _OPENAI_WS_DISABLED_REFUSAL +@pytest.mark.asyncio @pytest.mark.parametrize("value", [True, "true", "True"]) -def test_openai_websocket_refusal_is_none_for_truthy_settings(value): +async def test_openai_websocket_refusal_is_none_for_truthy_settings(value): settings = MappingProxyType({"enable_openai_websocket_passthrough": value}) - assert _openai_websocket_refusal(UserAPIKeyAuth(), settings) is None + assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None @pytest.mark.asyncio @@ -199,51 +223,73 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol(): subprotocols="realtime, openai-beta.realtime-v1", ) - relay = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), MappingProxyType({})) + served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), MappingProxyType({})) assert websocket.accepts == ["realtime"] assert websocket.closed == (1008, _OPENAI_WS_DISABLED_REFUSAL.close_reason) - assert relay.calls == [] + assert served.relay.calls == [] -RESTRICTED_KEYS: Final = ( - UserAPIKeyAuth(models=["gpt-4o"]), - UserAPIKeyAuth(team_models=["gpt-4o-realtime-preview"]), - UserAPIKeyAuth(models=["all-team-models"], team_models=["gpt-4o"]), +RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( + (("gpt-4o",),), + ((), ("gpt-4o-realtime-preview",)), + (("all-team-models",), ("gpt-4o",)), + ((), ("all-proxy-models",), ("gpt-4o",)), + ((), (), (), ("gpt-4o",)), + (("*",), (), (), (), ("gpt-4o",)), ) -UNRESTRICTED_KEYS: Final = ( - UserAPIKeyAuth(), - UserAPIKeyAuth(models=["all-proxy-models"]), - UserAPIKeyAuth(models=["*"]), - UserAPIKeyAuth(models=["all-team-models"], team_models=["all-proxy-models"]), +UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( + (), + ((),), + (("all-proxy-models",),), + (("*",),), + (("all-team-models",), ("all-proxy-models",)), + ((), (), (), (), ()), + (("*",), ("all-proxy-models",), ("all-team-models",), (), ()), ) @pytest.mark.asyncio -@pytest.mark.parametrize("user_api_key_dict", RESTRICTED_KEYS) -async def test_openai_websocket_rejects_model_restricted_keys(user_api_key_dict): +@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) +async def test_openai_websocket_rejects_model_restricted_identities(scopes): websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") + user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake") - relay = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED) + served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes) assert "model restrictions" in websocket.error_message() assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason) - assert relay.calls == [] - - -@pytest.mark.parametrize("user_api_key_dict", RESTRICTED_KEYS) -def test_openai_websocket_refusal_prefers_disabled_over_model_restriction(user_api_key_dict): - assert _openai_websocket_refusal(user_api_key_dict, MappingProxyType({})) is _OPENAI_WS_DISABLED_REFUSAL + assert served.relay.calls == [] + assert served.allowlists.calls == [user_api_key_dict] @pytest.mark.asyncio -@pytest.mark.parametrize("user_api_key_dict", UNRESTRICTED_KEYS) -async def test_openai_websocket_allows_unrestricted_keys(user_api_key_dict): +@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) +async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes): + allowlists = _FakeModelAllowlists(scopes) + + refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists) + + assert refusal is _OPENAI_WS_DISABLED_REFUSAL + assert allowlists.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES) +async def test_openai_websocket_allows_unrestricted_identities(scopes): websocket = _FakeWebSocket("/openai/v1/responses", "") with patch(GET_CREDENTIALS, return_value="sk-provider"): - relay = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED) + served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes) - assert len(relay.calls) == 1 + assert len(served.relay.calls) == 1 assert websocket.sent == [] assert websocket.closed is None + + +@pytest.mark.asyncio +async def test_proxy_model_allowlists_reads_the_key_scope_without_a_database(): + with patch("litellm.proxy.proxy_server.prisma_client", None): + scopes = await _proxy_model_allowlists()(UserAPIKeyAuth(models=["gpt-4o"])) + + assert tuple(tuple(scope) for scope in scopes) == (("gpt-4o",),) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d91928a203e..34216f7e1b9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11568,14 +11568,10 @@ async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(): reset_at = datetime.now(timezone.utc) + timedelta(days=10) key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert len(enqueued) == 1 @@ -11594,14 +11590,10 @@ async def test_team_window_spend_row_is_enqueued(): reset_at = datetime.now(timezone.utc) + timedelta(days=3) team_obj = MagicMock() - team_obj.budget_limits = [ - {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} - ] + team_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: - await increment_spend_counters( - token=None, team_id="team-1", user_id=None, response_cost=1.5 - ) + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) enqueued = await _drain(queue) assert len(enqueued) == 1 @@ -11620,9 +11612,7 @@ async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved() reset_at = datetime.now(timezone.utc) + timedelta(days=10) key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}] reservation = { "entries": [ {"counter_key": "spend:key:hashed-token", "reserved": 1.0}, @@ -11660,9 +11650,7 @@ async def test_sliding_window_without_reset_at_is_not_enqueued(): key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0}] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert enqueued == [] @@ -11680,9 +11668,7 @@ async def test_each_configured_window_gets_its_own_row_enqueue(): ] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert sorted(item["window_duration"] for item in enqueued) == ["1d", "30d"] @@ -11697,9 +11683,7 @@ async def test_no_window_spend_row_enqueued_without_budget_limits(): key_obj.budget_limits = None with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) enqueued = await _drain(queue) assert enqueued == [] @@ -11713,9 +11697,7 @@ async def test_window_spend_row_carries_the_request_start_time(): reset_at = datetime.now(timezone.utc) + timedelta(days=10) key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: await increment_spend_counters( @@ -11736,9 +11718,7 @@ async def test_team_window_spend_row_carries_the_request_start_time(): reset_at = datetime.now(timezone.utc) + timedelta(days=3) team_obj = MagicMock() - team_obj.budget_limits = [ - {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} - ] + team_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}] with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: await increment_spend_counters( @@ -12050,7 +12030,6 @@ async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_re assert not GUARDRAIL_RECONCILE_LOCK.locked() - @pytest.mark.asyncio async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeypatch): from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY @@ -12094,7 +12073,9 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) assert served_content() == "Begin every reply with AHOY" - prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with HOWDY")]) + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[db_row("Begin every reply with HOWDY")] + ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) assert served_content() == "Begin every reply with HOWDY" @@ -12547,3 +12528,40 @@ def test_disabling_docs_does_not_disable_other_routes(monkeypatch): assert client.get("/redoc").status_code == 404 assert client.get("/health/liveliness").status_code == 200 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_general_settings, expected", + [ + ({"enable_openai_websocket_passthrough": True}, True), + ({"enable_openai_websocket_passthrough": False}, False), + ({}, None), + ], +) +async def test_update_general_settings_propagates_openai_websocket_passthrough(db_general_settings, expected): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch("litellm.proxy.proxy_server.general_settings", {"enable_openai_websocket_passthrough": True}): + await proxy_config._update_general_settings(db_general_settings=db_general_settings) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["enable_openai_websocket_passthrough"] is expected + + +@pytest.mark.asyncio +async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough(): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"enable_openai_websocket_passthrough"} + + with patch("litellm.proxy.proxy_server.general_settings", {"enable_openai_websocket_passthrough": False}): + await proxy_config._update_general_settings(db_general_settings={"enable_openai_websocket_passthrough": True}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["enable_openai_websocket_passthrough"] is False From 2a7fc8de01c2dca87de9cfc9514a91fddd7a43c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:05:35 -0700 Subject: [PATCH 16/38] fix(proxy): keep the token's team model list in the websocket passthrough gate without a database --- litellm/proxy/auth/auth_checks.py | 2 +- tests/test_litellm/proxy/auth/test_auth_checks.py | 10 ++++++++-- .../proxy/test_openai_ws_passthrough_routes.py | 9 ++++++--- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 98d334ce2cc..120a0bb29ea 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4164,7 +4164,7 @@ async def enforced_model_allowlists( """One model allowlist per level that ``common_checks`` enforces on a request from this identity.""" key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) if prisma_client is None: - return (key_models,) + return (key_models, tuple(valid_token.team_models or ())) team_object: Final = ( None if valid_token.team_id is None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 45e10948267..5242a99c54c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7392,7 +7392,13 @@ async def test_enforced_model_allowlists_reads_every_level_from_cache(): proxy_logging_obj=proxy_logging_obj, ) without_database = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", models=["gpt-4o"], user_id="user-fake", team_id="team-fake"), + valid_token=UserAPIKeyAuth( + token="hashed-fake", + models=["gpt-4o"], + team_models=["gpt-4o-mini"], + user_id="user-fake", + team_id="team-fake", + ), prisma_client=None, user_api_key_cache=cache, proxy_logging_obj=proxy_logging_obj, @@ -7406,4 +7412,4 @@ async def test_enforced_model_allowlists_reads_every_level_from_cache(): ["gpt-4.1"], ] assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []] - assert [list(scope) for scope in without_database] == [["gpt-4o"]] + assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]] diff --git a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py index c96e7684e97..7d79192b884 100644 --- a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py @@ -14,6 +14,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _OPENAI_WS_DISABLED_REFUSAL, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL, + _has_model_restrictions, _openai_websocket_refusal, _proxy_model_allowlists, openai_websocket_proxy_route, @@ -288,8 +289,10 @@ async def test_openai_websocket_allows_unrestricted_identities(scopes): @pytest.mark.asyncio -async def test_proxy_model_allowlists_reads_the_key_scope_without_a_database(): +async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database(): + token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"]) with patch("litellm.proxy.proxy_server.prisma_client", None): - scopes = await _proxy_model_allowlists()(UserAPIKeyAuth(models=["gpt-4o"])) + scopes = await _proxy_model_allowlists()(token) - assert tuple(tuple(scope) for scope in scopes) == (("gpt-4o",),) + assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",)) + assert _has_model_restrictions(scopes) From f846388bb1174b5cdb89a7f019825ad83b6d6e81 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:24:30 -0700 Subject: [PATCH 17/38] fix(proxy): treat a missing user row as unrestricted in the websocket passthrough gate --- litellm/proxy/auth/auth_checks.py | 23 ++++++++++++++--- .../proxy/auth/test_auth_checks.py | 25 +++++++++++++++++++ 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 120a0bb29ea..b0e22153401 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4155,6 +4155,24 @@ async def _granted_model_lists( ) +async def _user_object_or_none( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_UserTable | None: + try: + return await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + except ValueError: + return None + + async def enforced_model_allowlists( valid_token: UserAPIKeyAuth, prisma_client: PrismaClient | None, @@ -4178,11 +4196,10 @@ async def enforced_model_allowlists( user_object: Final = ( None if team_object is not None - else await get_user_object( - user_id=valid_token.user_id, + else await _user_object_or_none( + valid_token=valid_token, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - user_id_upsert=False, proxy_logging_obj=proxy_logging_obj, ) ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5242a99c54c..7351c981838 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7336,6 +7336,31 @@ class _UntouchedPrisma: raise AssertionError(f"database reached through {name}") +class _MissingUserPrisma: + class db: + class litellm_usertable: + @staticmethod + async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: + return None + + +@pytest.mark.asyncio +async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted(): + from litellm.proxy.auth.auth_checks import enforced_model_allowlists + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + scopes = await enforced_model_allowlists( + valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"), + prisma_client=_MissingUserPrisma(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert [list(scope) for scope in scopes] == [[], [], [], [], []] + + @pytest.mark.asyncio async def test_enforced_model_allowlists_reads_every_level_from_cache(): from litellm.proxy._types import ( From 6ee33df952ef9102f14960ed04c46e8f49900d66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:25:12 -0700 Subject: [PATCH 18/38] fix(realtime): relay the upstream websocket close to the client instead of hanging When the provider closes the realtime websocket (for example Vertex Live refusing the session with 1008 "Publisher model ... was not found"), the proxy swallowed the close and kept waiting on the client, so the client sat on an open socket with nothing coming back and the session was logged as a $0 success The backend relay now returns the upstream close, and bidirectional_forward sends the client an OpenAI-style error event naming the upstream code and reason, then closes the client socket with the same code (or 1011 when the upstream code is one a server may not send). A session the upstream refused before sending any frame is logged through the failure handlers instead of as a success --- litellm/litellm_core_utils/realtime_errors.py | 8 + .../litellm_core_utils/realtime_streaming.py | 186 ++++++++++++------ .../test_realtime_errors.py | 10 + .../test_realtime_streaming.py | 169 +++++++++++++++- 4 files changed, 303 insertions(+), 70 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_errors.py b/litellm/litellm_core_utils/realtime_errors.py index e1b957f4325..3c064728a66 100644 --- a/litellm/litellm_core_utils/realtime_errors.py +++ b/litellm/litellm_core_utils/realtime_errors.py @@ -29,3 +29,11 @@ def websocket_close_reason(message: str, fallback: str) -> str: if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES: return message return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") + + +def client_close_code(upstream_code: int) -> int: + from websockets.frames import EXTERNAL_CLOSE_CODES, CloseCode + + if upstream_code in EXTERNAL_CLOSE_CODES or 3000 <= upstream_code < 5000: + return upstream_code + return int(CloseCode.INTERNAL_ERROR) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 8479e108d17..746343026ed 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,7 +1,9 @@ import asyncio import json -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +import traceback +from collections.abc import Coroutine, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast from typing_extensions import ReadOnly @@ -19,9 +21,11 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging +from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from websockets.exceptions import ConnectionClosed from litellm.types.guardrails import GuardrailEventHooks @@ -30,8 +34,22 @@ else: CLIENT_CONNECTION_CLASS = Any -class _ClientWebSocketExceptions(Protocol): - ConnectionClosed: type[Exception] +@dataclass(frozen=True, slots=True) +class BackendClose: + code: int + reason: str + + @property + def message(self) -> str: + if not self.reason: + return f"upstream websocket closed with code {self.code}" + return f"upstream websocket closed with code {self.code}: {self.reason}" + + +def backend_close_from(error: "ConnectionClosed") -> BackendClose: + if error.rcvd is None: + return BackendClose(code=1006, reason=str(error)) + return BackendClose(code=error.rcvd.code, reason=error.rcvd.reason) class _ASGIScope(TypedDict, total=False): @@ -69,10 +87,13 @@ class _ScopedWebSocket(Protocol): class _ClientWebSocket(_ScopedWebSocket, Protocol): - exceptions: _ClientWebSocketExceptions - async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... + async def close(self, code: int = 1000, reason: str | None = None) -> None: ... + + +class _LoggingWorker(Protocol): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... def _decode_json_object(payload: str) -> Mapping[str, object]: @@ -108,11 +129,14 @@ class RealTimeStreaming: backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, + logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, ): self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj + self._logging_worker = logging_worker self.messages: list[OpenAIRealtimeEvents] = [] + self._backend_sent_frames: bool = False self.input_message: dict = {} self.input_messages: list[dict[str, str]] = [] self.session_tools: list[dict] = [] @@ -388,7 +412,7 @@ class RealTimeStreaming: # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) ) @@ -1035,60 +1059,84 @@ class RealTimeStreaming: return True return False - async def backend_to_client_send_messages(self): + async def _relay_backend_messages(self) -> NoReturn: + while True: + try: + raw_response = await self.backend_ws.recv(decode=False) + except TypeError: + raw_response = await self.backend_ws.recv() + self._backend_sent_frames = True + + if isinstance(raw_response, bytes): + try: + raw_response = raw_response.decode("utf-8") + except UnicodeDecodeError: + verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.") + continue + + if self.provider_config: + try: + await self._handle_provider_config_message(raw_response) + except Exception as e: + verbose_logger.exception("Error processing backend message, skipping: %s", e) + continue + else: + event = self._parse_backend_event(raw_response) + if event is None: + await self.websocket.send_text(raw_response) + continue + + if self._should_drop_event_from_client(event): + continue + + if await self._handle_raw_backend_message(event, raw_response): + continue + + event = self._normalize_event_for_ga_client(event) + self.store_message(event) + + if not self._client_wants_beta: + await self.websocket.send_text(json.dumps(event)) + continue + + translated = self._translate_event_to_beta(event) + if translated is None: + continue + await self.websocket.send_text(json.dumps(translated)) + + async def backend_to_client_send_messages(self) -> BackendClose: import websockets try: - while True: - try: - raw_response = await self.backend_ws.recv(decode=False) - except TypeError: - raw_response = await self.backend_ws.recv() - - if isinstance(raw_response, bytes): - try: - raw_response = raw_response.decode("utf-8") - except UnicodeDecodeError: - verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.") - continue - - if self.provider_config: - try: - await self._handle_provider_config_message(raw_response) - except Exception as e: - verbose_logger.exception("Error processing backend message, skipping: %s", e) - continue - else: - event = self._parse_backend_event(raw_response) - if event is None: - await self.websocket.send_text(raw_response) - continue - - if self._should_drop_event_from_client(event): - continue - - if await self._handle_raw_backend_message(event, raw_response): - continue - - event = self._normalize_event_for_ga_client(event) - self.store_message(event) - - if not self._client_wants_beta: - await self.websocket.send_text(json.dumps(event)) - continue - - translated = self._translate_event_to_beta(event) - if translated is None: - continue - await self.websocket.send_text(json.dumps(translated)) - + await self._relay_backend_messages() except websockets.exceptions.ConnectionClosed as e: verbose_logger.exception("Connection closed in backend to client send messages - %s", e) - except Exception as e: - verbose_logger.exception("Error in backend to client send messages: %s", e) - finally: + close: Final = backend_close_from(e) + self._flush_unbilled_transcription_usage() + if self._backend_refused_session(close): + await self.log_backend_refusal(e) + else: + await self.log_messages() + return close + except asyncio.CancelledError: self._flush_unbilled_transcription_usage() await self.log_messages() + raise + except Exception as e: + verbose_logger.exception("Error in backend to client send messages: %s", e) + self._flush_unbilled_transcription_usage() + await self.log_messages() + return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket") + + def _backend_refused_session(self, close: BackendClose) -> bool: + return close.code != 1000 and not self._backend_sent_frames and not self.messages + + async def log_backend_refusal(self, error: Exception) -> None: + if not self.logging_obj: + return + self._logging_worker.ensure_initialized_and_enqueue( + self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True) + ) @staticmethod def _detect_beta_header(websocket: _ScopedWebSocket) -> bool: @@ -1484,20 +1532,28 @@ class RealTimeStreaming: except Exception as e: verbose_logger.debug("Error in client ack messages: %s", e) - async def bidirectional_forward(self): + async def bidirectional_forward(self) -> None: forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages()) + client_task: Final = asyncio.create_task(self.client_ack_messages()) try: - await self.client_ack_messages() - except self.websocket.exceptions.ConnectionClosed: - verbose_logger.debug("Connection closed") - forward_task.cancel() + await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED) + if not client_task.done(): + await self._close_client(forward_task.result()) finally: - if not forward_task.done(): - forward_task.cancel() - try: - await forward_task - except asyncio.CancelledError: - pass + forward_task.cancel() + client_task.cancel() + await asyncio.gather(forward_task, client_task, return_exceptions=True) + + async def _close_client(self, close: BackendClose) -> None: + try: + if close.code != 1000: + await self.websocket.send_text(realtime_error_event(close.message, error_type="server_error")) + await self.websocket.close( + code=client_close_code(close.code), + reason=websocket_close_reason(close.reason, fallback=close.message), + ) + except Exception as e: # noqa: BLE001 # the client may already be gone; the session is over either way + verbose_logger.debug("Could not relay the upstream close to the client: %s", e) def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py index 494d16b0b9b..1d2cf905f4e 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py @@ -1,8 +1,10 @@ import json +import pytest from litellm.litellm_core_utils.realtime_errors import ( WEBSOCKET_CLOSE_REASON_MAX_BYTES, + client_close_code, realtime_error_event, websocket_close_reason, ) @@ -42,3 +44,11 @@ def test_websocket_close_reason_truncates_multibyte_message_by_bytes(): assert len(reason.encode("utf-8")) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES assert reason == "あ" * (WEBSOCKET_CLOSE_REASON_MAX_BYTES // 3) assert "�" not in reason + + +@pytest.mark.parametrize( + ("upstream_code", "expected"), + [(1000, 1000), (1008, 1008), (1011, 1011), (4001, 4001), (1005, 1011), (1006, 1011), (1015, 1011), (2999, 1011)], +) +def test_client_close_code_only_forwards_codes_a_server_may_send(upstream_code, expected): + assert client_close_code(upstream_code) == expected diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 52e88db753a..1e0456079e9 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -1,8 +1,13 @@ +import asyncio import json +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from websockets.exceptions import ConnectionClosed +from websockets.frames import Close import litellm @@ -2941,13 +2946,11 @@ async def test_log_messages_routes_async_logging_through_bounded_worker(): realtime turn leaves a suspended task pinning its response in memory -> an unbounded leak. Regression for that fix.""" logging_obj = MagicMock() - streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj) + mock_worker = MagicMock() + streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj, logging_worker=mock_worker) streaming.messages = [{"type": "session.created"}] - with ( - patch("litellm.litellm_core_utils.realtime_streaming.GLOBAL_LOGGING_WORKER") as mock_worker, - patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task, - ): + with patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task: await streaming.log_messages() mock_worker.ensure_initialized_and_enqueue.assert_called_once() @@ -3111,3 +3114,159 @@ async def test_session_close_flush_noop_without_unbilled_usage(): isinstance(message, dict) and message.get("type") == "conversation.item.input_audio_transcription.completed" for message in streaming.messages ) + + + +_UPSTREAM_REFUSAL: Final = "Publisher model `publishers/google/models/gemini-live-2.5-flash` was not found" + + +class _InlineLoggingWorker: + def __init__(self) -> None: + self.enqueued: tuple[Coroutine[object, object, None], ...] = () + + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + self.enqueued = (*self.enqueued, async_coroutine) + + async def drain(self) -> None: + for coroutine in self.enqueued: + await coroutine + + +class _RecordingLogging: + def __init__(self) -> None: + self.logged_sessions: tuple[tuple[dict, ...], ...] = () + self.logged_failures: tuple[Exception, ...] = () + + async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None: + self.logged_sessions = (*self.logged_sessions, tuple(result)) + + async def dispatch_failure_handlers( + self, exception: Exception, traceback_exception: str, prefer_async_handlers: bool = False + ) -> None: + self.logged_failures = (*self.logged_failures, exception) + + +@dataclass(frozen=True, slots=True) +class _RelaySession: + streaming: RealTimeStreaming + logging: _RecordingLogging + worker: _InlineLoggingWorker + + async def run(self) -> None: + await asyncio.wait_for(self.streaming.bidirectional_forward(), timeout=2) + await self.worker.drain() + + +async def _wait_forever() -> str: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +def _client_ws_that_never_sends() -> MagicMock: + client_ws: Final = MagicMock() + client_ws.headers = {} + client_ws.receive_text = AsyncMock(side_effect=_wait_forever) + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + return client_ws + + +def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock: + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=list(frames)) + return backend_ws + + +def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession: + logging: Final = _RecordingLogging() + worker: Final = _InlineLoggingWorker() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker + ) + return _RelaySession(streaming=streaming, logging=logging, worker=worker) + + +def _error_events_sent_to(client_ws: MagicMock) -> list[dict]: + events: Final = (json.loads(call.args[0]) for call in client_ws.send_text.await_args_list) + return [event for event in events if event.get("type") == "error"] + + +@pytest.mark.asyncio +async def test_bidirectional_forward_relays_upstream_policy_close_to_client(): + client_ws: Final = _client_ws_that_never_sends() + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert error_event["error"]["type"] == "server_error" + assert "1008" in error_event["error"]["message"] + assert _UPSTREAM_REFUSAL in error_event["error"]["message"] + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + + +@pytest.mark.asyncio +async def test_bidirectional_forward_maps_abnormal_upstream_close_to_internal_error(): + client_ws: Final = _client_ws_that_never_sends() + session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(None, None))) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert "1006" in error_event["error"]["message"] + client_ws.close.assert_awaited_once() + assert client_ws.close.await_args.kwargs["code"] == 1011 + + +@pytest.mark.asyncio +async def test_bidirectional_forward_relays_normal_upstream_close_without_error_event(): + client_ws: Final = _client_ws_that_never_sends() + session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(Close(1000, ""), None))) + + await session.run() + + assert _error_events_sent_to(client_ws) == [] + client_ws.close.assert_awaited_once() + assert client_ws.close.await_args.kwargs["code"] == 1000 + + +@pytest.mark.asyncio +async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success(): + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert session.logging.logged_sessions == () + + +@pytest.mark.asyncio +async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success(): + client_ws: Final = _client_ws_that_never_sends() + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close)) + + await session.run() + + (logged_session,) = session.logging.logged_sessions + assert [event["type"] for event in logged_session] == ["session.created"] + assert session.logging.logged_failures == () + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + + +@pytest.mark.asyncio +async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close(): + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client went away")) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=_wait_forever) + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + assert session.logging.logged_sessions == ((),) + assert session.logging.logged_failures == () + client_ws.close.assert_not_awaited() From 85d45fbb4b6f9299af75b6891af61664176ad69f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:33:30 -0700 Subject: [PATCH 19/38] fix(realtime): relay the upstream close even when a client message hit the closed socket first When the upstream closes while the proxy is forwarding a client message, the client loop ends before the backend relay sees the close, and the relay skipped closing the client because it read the client loop's exit as the client hanging up. The client loop now reports why it stopped, so a close observed on the backend send still reaches the client with the error event and the upstream close code --- .../litellm_core_utils/realtime_streaming.py | 19 ++++++++-- .../test_realtime_streaming.py | 37 +++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 746343026ed..530391c7b57 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -3,6 +3,7 @@ import json import traceback from collections.abc import Coroutine, Mapping, Sequence from dataclasses import dataclass +from enum import Enum, auto from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast from typing_extensions import ReadOnly @@ -46,6 +47,11 @@ class BackendClose: return f"upstream websocket closed with code {self.code}: {self.reason}" +class ClientLoopExit(Enum): + CLIENT_DISCONNECTED = auto() + BACKEND_CLOSED = auto() + + def backend_close_from(error: "ConnectionClosed") -> BackendClose: if error.rcvd is None: return BackendClose(code=1006, reason=str(error)) @@ -1291,7 +1297,9 @@ class RealTimeStreaming: item["content"] = new_content return item - async def client_ack_messages(self): + async def client_ack_messages(self) -> ClientLoopExit: + import websockets + client_event: _ClientEventFrame try: while True: @@ -1529,16 +1537,21 @@ class RealTimeStreaming: if guardrail_turn_detection_injected and sent: self._guardrail_turn_detection_update_sent = True + except websockets.exceptions.ConnectionClosed as e: + verbose_logger.debug("Backend closed while forwarding a client message: %s", e) + return ClientLoopExit.BACKEND_CLOSED except Exception as e: verbose_logger.debug("Error in client ack messages: %s", e) + return ClientLoopExit.CLIENT_DISCONNECTED async def bidirectional_forward(self) -> None: forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages()) client_task: Final = asyncio.create_task(self.client_ack_messages()) try: await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED) - if not client_task.done(): - await self._close_client(forward_task.result()) + if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED: + return + await self._close_client(await forward_task) finally: forward_task.cancel() client_task.cancel() diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 1e0456079e9..41b7557f6b2 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3134,9 +3134,13 @@ class _InlineLoggingWorker: class _RecordingLogging: def __init__(self) -> None: + self.model_call_details: dict[str, object] = {} self.logged_sessions: tuple[tuple[dict, ...], ...] = () self.logged_failures: tuple[Exception, ...] = () + def pre_call(self, input: str | dict, api_key: str) -> None: + return None + async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None: self.logged_sessions = (*self.logged_sessions, tuple(result)) @@ -3257,6 +3261,39 @@ async def test_upstream_close_after_relayed_events_still_logs_the_session_as_suc client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) +@pytest.mark.asyncio +async def test_upstream_closing_while_a_client_message_is_forwarded_still_reaches_the_client(): + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + backend_closed: Final = asyncio.Event() + client_messages: Final = iter((json.dumps({"type": "response.create"}),)) + + async def receive_text() -> str: + message = next(client_messages, None) + return message if message is not None else await _wait_forever() + + async def send_to_backend(_message: str) -> None: + backend_closed.set() + raise upstream_close + + async def recv_from_backend() -> bytes: + await backend_closed.wait() + raise upstream_close + + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = receive_text + backend_ws: Final = MagicMock() + backend_ws.send = send_to_backend + backend_ws.recv = recv_from_backend + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert _UPSTREAM_REFUSAL in error_event["error"]["message"] + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + assert session.logging.logged_failures == (upstream_close,) + + @pytest.mark.asyncio async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close(): client_ws: Final = _client_ws_that_never_sends() From a0b2e7fca6c0dd2bd22e81906eea41d2e2c87426 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:35:08 -0700 Subject: [PATCH 20/38] fix(proxy): only a provably missing user row counts as unrestricted in the websocket passthrough gate --- litellm/proxy/auth/auth_checks.py | 15 +++++++++--- .../proxy/auth/test_auth_checks.py | 24 +++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b0e22153401..dc693317de0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2352,6 +2352,13 @@ async def _backfill_null_user_email( return updated_row +class UserNotFoundError(ValueError): + """The user row is provably absent, as opposed to merely unreadable, so a caller that reads a missing row as no user-level limits can key on it without also swallowing a database that would not answer.""" + + def __init__(self, user_id: str) -> None: + super().__init__(f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call.") + + @log_db_metrics async def get_user_object( user_id: str | None, @@ -2457,7 +2464,7 @@ async def get_user_object( value=None, last_db_access_time=last_db_access_time, ) - raise Exception + raise UserNotFoundError(user_id=user_id) if response.organization_memberships is not None and len(response.organization_memberships) > 0: # dump each organization membership to type LiteLLM_OrganizationMembershipTable @@ -2493,7 +2500,9 @@ async def get_user_object( ) return _response - except Exception as e: # if user not in db + except UserNotFoundError: + raise + except Exception as e: _log_budget_lookup_failure("user", e) raise ValueError( f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {e}" @@ -4169,7 +4178,7 @@ async def _user_object_or_none( user_id_upsert=False, proxy_logging_obj=proxy_logging_obj, ) - except ValueError: + except UserNotFoundError: return None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 7351c981838..2284a05b2e9 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7361,6 +7361,30 @@ async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestrict assert [list(scope) for scope in scopes] == [[], [], [], [], []] +class _UnreachableUserPrisma: + class db: + class litellm_usertable: + @staticmethod + async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: + raise RuntimeError("database gone") + + +@pytest.mark.asyncio +async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup(): + from litellm.proxy.auth.auth_checks import enforced_model_allowlists + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + with pytest.raises(ValueError, match="database gone"): + await enforced_model_allowlists( + valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), + prisma_client=_UnreachableUserPrisma(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + @pytest.mark.asyncio async def test_enforced_model_allowlists_reads_every_level_from_cache(): from litellm.proxy._types import ( From a90328aa8c620935f78bbfcb0a6b49f1858fee97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:42:37 -0700 Subject: [PATCH 21/38] refactor(typing): drop the dead self guard in Predibase init and use a plain list factory --- litellm/llms/predibase/chat/transformation.py | 2 +- litellm/router_strategy/adaptive_router/signals.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 0ebac5185d7..2a63c489395 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -83,7 +83,7 @@ class PredibaseConfig(BaseConfig): ("watermark", watermark), ) for key, value in locals_: - if key != "self" and value is not None: + if value is not None: setattr(self.__class__, key, value) @classmethod diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index 7b69714aad9..c28613b54eb 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -93,7 +93,7 @@ class Turn: user_content: str | None = None assistant_content: str | None = None tool_calls: list[dict[str, Any]] = field(default_factory=list) - tool_results: Sequence[Mapping[str, object]] = field(default_factory=list[Mapping[str, object]]) + tool_results: Sequence[Mapping[str, object]] = field(default_factory=list) response_status: int | None = None From c27f1e348dd2f6191177e4b1016388bce1b161f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:47:31 -0700 Subject: [PATCH 22/38] fix(ui): name the object arguments at two new call sites to bring the inline-object lint budget back under its ceiling --- ui/litellm-dashboard/eslint-budgets.json | 2 +- .../src/components/add_model/ClassificationMethodConfig.tsx | 5 +++-- .../add_model/build_complexity_router_config.test.ts | 5 +++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index 3f9163d9028..b3c77e287fc 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -3,7 +3,7 @@ "no-console": { "max": 12, "target": 0 }, "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 }, - "local/no-large-inline-object-arg": { "max": 555, "target": 300 }, + "local/no-large-inline-object-arg": { "max": 554, "target": 300 }, "local/no-long-condition-chain": { "max": 265, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, "testing-library/no-node-access": { "max": 716, "target": 500 }, diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index cc66103fc86..188ef7f8cb5 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -315,12 +315,13 @@ const ClassificationMethodConfig: React.FC = ({ timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, classification_rubric: selectedRubric, }; - onChange({ + const nextValue: ComplexityRouterConfigValue = { ...value, ...(selectedRubric && { classifier_llm_config: rubricConfig }), classification_prompt: classificationPrompt, classification_examples: classificationExamples, - }); + }; + onChange(nextValue); }; const handleClassifierModelChange = (model: string) => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 1b5bb9e72eb..81f05a94a61 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -1170,12 +1170,13 @@ describe("buildComplexityRouterConfig stall escalation", () => { }); it("emits the toggle and both knobs when it is on", () => { - const config = buildComplexityRouterConfig({ + const params: BuildComplexityRouterConfigParams = { ...baseParams, stallEscalationEnabled: true, stallEscalationWindow: 8, stallEscalationRepeatThreshold: 4, - }); + }; + const config = buildComplexityRouterConfig(params); expect(config.stall_escalation_enabled).toBe(true); expect(config.stall_escalation_window).toBe(8); expect(config.stall_escalation_repeat_threshold).toBe(4); From da9dbdba961ce2981f9d82669ab9e3eb3a9d90a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:52:48 -0700 Subject: [PATCH 23/38] fix(realtime): treat any client receive failure as a client hangup client_ack_messages classified a websockets ConnectionClosed raised by the client socket as the backend closing, so bidirectional_forward kept waiting on the upstream instead of ending the session. Starlette clients raise WebSocketDisconnect, but the realtime test client in tests/llm_translation/realtime raises websockets.exceptions.ConnectionClosed, which hung test_openai_realtime_simple.py until the run was killed. Only the receive_text call now maps every exception to CLIENT_DISCONNECTED; the loop body keeps ConnectionClosed as BACKEND_CLOSED, since the backend socket is the only websockets socket touched there. --- litellm/litellm_core_utils/realtime_streaming.py | 11 ++++++++++- .../litellm_core_utils/test_realtime_streaming.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 530391c7b57..bb7fbd81146 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1297,13 +1297,22 @@ class RealTimeStreaming: item["content"] = new_content return item + async def _receive_client_message(self) -> str | None: + try: + return await self.websocket.receive_text() + except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone + verbose_logger.debug("Client disconnected: %s", e) + return None + async def client_ack_messages(self) -> ClientLoopExit: import websockets client_event: _ClientEventFrame try: while True: - message = await self.websocket.receive_text() + message = await self._receive_client_message() + if message is None: + return ClientLoopExit.CLIENT_DISCONNECTED ## GUARDRAIL: intercept conversation.item.create for text-based injection. guardrail_turn_detection_injected = False diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 41b7557f6b2..00addb613c2 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3307,3 +3307,18 @@ async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close( assert session.logging.logged_sessions == ((),) assert session.logging.logged_failures == () client_ws.close.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the_backend_closing(): + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(None, None)) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=_wait_forever) + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + assert session.logging.logged_sessions == ((),) + assert session.logging.logged_failures == () + client_ws.close.assert_not_awaited() From 14f8677bfcdc160d3b3b424dc84a9c1727734939 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:57:51 -0700 Subject: [PATCH 24/38] fix(realtime): mark realtime sessions async so failure hooks fire once The relay's failure dispatch runs the async handler and then the legacy sync failure_handler for the proxy's callable callbacks. The realtime logging object carried no async marker, so failure_handler treated the session as a sync SDK call and fired every CustomLogger's sync failure hook on top of the async one: Langfuse recorded two ERROR observations per refused session, and OpenTelemetry, MLflow, Braintrust, Literal AI, DeepEval and New Relic implement the same sync hook. Plant the _arealtime marker in litellm_params the way aanthropic_messages and agenerate_content already do, so both dispatchers classify the session async. --- litellm/litellm_core_utils/litellm_logging.py | 1 + litellm/realtime_api/main.py | 4 +-- .../test_litellm_logging.py | 30 +++++++++++++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index df579f6df5b..15585c64efb 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1820,6 +1820,7 @@ class Logging(LiteLLMLoggingBaseClass): and litellm_params.get(CallTypes.aanthropic_messages.value, False) is not True and litellm_params.get(CallTypes.agenerate_content.value, False) is not True and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True + and litellm_params.get(CallTypes.arealtime.value, False) is not True ) def _is_assembled_stream_success(self, result=None) -> bool: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 3862aec445f..9de91dfcaa5 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -27,7 +27,7 @@ from litellm.types.realtime import ( RealtimeTranscriptionSessionRequest, ) from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import LlmProviders +from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager from ..litellm_core_utils.get_litellm_params import get_litellm_params @@ -355,7 +355,7 @@ async def _arealtime( user: Final = kwargs.get("user", None) litellm_params: Final = GenericLiteLLMParams(**kwargs) - litellm_params_dict: Final = get_litellm_params(**kwargs) + litellm_params_dict: Final = {**get_litellm_params(**kwargs), CallTypes.arealtime.value: True} model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( model=model, diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f1de7390b5b..af75691eb10 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -995,6 +995,35 @@ async def test_anthropic_messages_marks_litellm_params_async(): litellm.callbacks = original_callbacks +@pytest.mark.asyncio +async def test_arealtime_marks_litellm_params_async(monkeypatch): + """LIT-6973: ``_arealtime`` must plant ``_arealtime`` in ``litellm_params`` so + ``_is_sync_litellm_request`` classifies the session async and a failed session + reaches a CustomLogger's failure hook once, through the async path only, even + though the sync ``failure_handler`` still runs ahead of the async one.""" + captured = {} + async_logged = asyncio.Event() + + class CaptureLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + captured["litellm_params"] = kwargs.get("litellm_params", {}) + async_logged.set() + + logger = CaptureLogger() + logger.log_failure_event = MagicMock() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + with pytest.raises(ValueError, match="Unsupported model"): + await litellm._arealtime(model="anthropic/claude-x", websocket=MagicMock()) + await asyncio.wait_for(async_logged.wait(), timeout=10) + logger.log_failure_event.assert_not_called() + assert captured["litellm_params"].get("_arealtime") is True + assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False + + @pytest.mark.asyncio async def test_agenerate_content_marks_litellm_params_async(): """LIT-4475: the async ``agenerate_content`` entrypoint must plant @@ -1180,6 +1209,7 @@ def test_is_sync_litellm_request(): assert LitellmLogging._is_sync_litellm_request({}) is True assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False + assert LitellmLogging._is_sync_litellm_request({"_arealtime": True}) is False assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False From 412c36bb8e0663fd27e6c635d94f35d7407eabd3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 21:24:36 -0700 Subject: [PATCH 25/38] fix(realtime): detect an upstream refusal from received frames, not the session log The refusal predicate also required the session log to be empty, but that log is not limited to upstream frames. With gemini_live_defer_setup the handler stores a synthetic session.created before the relay starts, and the transcription usage flush appends a usage event before the check runs, so an upstream policy close with no received frames was still logged as a $0 success. Key the check off the received-frames flag only --- .../litellm_core_utils/realtime_streaming.py | 2 +- .../test_realtime_streaming.py | 39 +++++++++++++++++-- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index bb7fbd81146..984934daaac 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1135,7 +1135,7 @@ class RealTimeStreaming: return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket") def _backend_refused_session(self, close: BackendClose) -> bool: - return close.code != 1000 and not self._backend_sent_frames and not self.messages + return close.code != 1000 and not self._backend_sent_frames async def log_backend_refusal(self, error: Exception) -> None: if not self.logging_obj: diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 00addb613c2..09757c35570 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3031,12 +3031,12 @@ async def test_session_close_flushes_unbilled_transcription_usage(): messages before log_messages runs, and never forwarded to the client.""" from typing import Final - from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage + from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage, RealtimeResponseTypedDict client_ws: Final = MagicMock() client_ws.send_text = AsyncMock() backend_ws: Final = MagicMock() - backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + backend_ws.recv = AsyncMock(side_effect=[b'{"serverContent": {}}', ConnectionClosed(None, None)]) logging_obj: Final = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() @@ -3048,7 +3048,24 @@ async def test_session_close_flushes_unbilled_transcription_usage(): "total_tokens": 171, "input_token_details": {"text_tokens": 0, "audio_tokens": 153}, } + transcript_frame: Final[RealtimeResponseTypedDict] = { + "response": { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "transcript": "ahoy", + "item_id": "item_1", + "content_index": 0, + }, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_conversation_id": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } provider_config: Final = MagicMock() + provider_config.transform_realtime_response = MagicMock(return_value=transcript_frame) provider_config.unbilled_usage_on_session_close = MagicMock(return_value=usage) streaming: Final = RealTimeStreaming( @@ -3080,7 +3097,9 @@ async def test_session_close_flushes_unbilled_transcription_usage(): ) assert len(flushed) == 1 assert flushed[0] in logged_snapshots[0] - assert not client_ws.send_text.called + forwarded: Final = tuple(json.loads(call.args[0]) for call in client_ws.send_text.await_args_list) + assert [event.get("transcript") for event in forwarded] == ["ahoy"] + assert all("usage" not in event for event in forwarded) @pytest.mark.asyncio @@ -3246,6 +3265,20 @@ async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success(): assert session.logging.logged_sessions == () +@pytest.mark.asyncio +async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_failure(): + """LIT-6973: deferred Gemini Live setup stores a synthetic ``session.created`` before + the relay starts. It is not an upstream frame, so a refusal after it is still a refusal.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + session.streaming.store_message(json.dumps({"type": "session.created", "session": {"id": "sess_synthetic"}})) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert session.logging.logged_sessions == () + + @pytest.mark.asyncio async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success(): client_ws: Final = _client_ws_that_never_sends() From 78ad88f52c0259e20da5c7a2b15ba3d42525fd29 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 4 Sep 2026 22:14:44 -0700 Subject: [PATCH 26/38] fix(responses): decode JSON-string tool schemas before sending to the provider (#39844) * fix(responses): decode JSON-string tool schemas before sending to the provider A caller that hands a tool schema over already JSON-encoded reached the Responses API with a string `parameters`, and the provider rejected the request with a 400 naming the routed model instead of the offending tool. Decode it at the one place every Responses request converges, and refuse anything that is neither an object nor a string encoding one. Collapses the duplicated input/tool sanitization block shared by the request and compact-request builders into a single owner, so the decode cannot be wired into one path and not the other. * test(responses): pin null tool schemas as accepted, and type the parametrized cases The Responses API serves `parameters: null` and an omitted schema alike, so neither may raise. Pin both against a future tightening, annotate the parametrized inputs, and trim the docstrings back to what the code does not already say. --- .../llms/openai/responses/transformation.py | 81 ++++++++++++---- .../test_openai_responses_transformation.py | 97 +++++++++++++++++++ type-discipline-budget.json | 6 +- 3 files changed, 164 insertions(+), 20 deletions(-) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index b97521b90c2..d7c2fcace09 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name @@ -205,29 +206,76 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): `remove_cache_control_flag_from_messages_and_tools`; mirror that here. """ - input = self._validate_input_param(input) - tools = response_api_optional_request_params.get("tools") - input, tools = self.remove_cache_control_flag_from_input_and_tools(model=model, input=input, tools=tools) - sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( - model=model, tools=tools, litellm_params=litellm_params + replay_safe_input, sanitized_tools = self._prepared_input_and_tools( + model=model, + input=input, + tools=response_api_optional_request_params.get("tools"), + litellm_params=litellm_params, ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools - replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input) final_request_params: Final = dict( ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) ) return final_request_params + def _prepared_input_and_tools( + self, + model: str, + input: str | ResponseInputParam, + tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, + litellm_params: GenericLiteLLMParams, + ) -> tuple[str | ResponseInputParam, Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None]: + validated_input: Final = self._validate_input_param(input) + stripped_input, stripped_tools = self.remove_cache_control_flag_from_input_and_tools( + model=model, input=validated_input, tools=tools + ) + object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools) + sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( + model=model, tools=object_schema_tools, litellm_params=litellm_params + ) + return self._drop_foreign_tool_call_item_ids(stripped_input), sanitized_tools + + def _tools_with_object_parameters( + self, model: str, tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None + ) -> Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None: + """Decode tool schemas handed over already JSON-encoded, which the Responses validator + rejects with a 400 naming the routed model rather than the tool. A null or absent schema + is left alone because the API accepts both.""" + if tools is None: + return None + decoded: Final = [ # mutable-ok: request tools are a JSON list + self._tool_with_object_parameters(model=model, index=index, tool=tool) for index, tool in enumerate(tools) + ] + return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", decoded) # cast-ok: dict spread keeps each tool's shape + + def _tool_with_object_parameters(self, model: str, index: int, tool: object) -> object: + if not isinstance(tool, dict) or tool.get("parameters") is None: + return tool + parameters: Final = tool["parameters"] + if isinstance(parameters, dict): + return tool + decoded: Final = safe_json_loads(parameters) if isinstance(parameters, str) else None + if isinstance(decoded, dict): + return {**tool, "parameters": decoded} # mutable-ok: request tools are JSON dicts + raise litellm.BadRequestError( + message=( + f"Invalid type for 'tools[{index}].parameters': expected an object, " + f"but got {type(parameters).__name__} instead." + ), + model=model, + llm_provider=self.custom_llm_provider, + ) + def remove_cache_control_flag_from_input_and_tools( self, model: str, # allows overrides to selectively run this input: str | ResponseInputParam, - tools: list[ALL_RESPONSES_API_TOOL_PARAMS] | None = None, + tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None = None, ) -> tuple[ str | ResponseInputParam, - list[ALL_RESPONSES_API_TOOL_PARAMS] | None, + Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, ]: """Sibling of `remove_cache_control_flag_from_messages_and_tools` on the chat path. Strips Anthropic-only `cache_control` markers from @@ -272,9 +320,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): def _flatten_tool_schema_combinators_for_openai( self, model: str, - tools: list[ALL_RESPONSES_API_TOOL_PARAMS] | None, # mutable-ok: request tools are a JSON list + tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, litellm_params: GenericLiteLLMParams, - ) -> list[ALL_RESPONSES_API_TOOL_PARAMS] | None: # mutable-ok: request tools are a JSON list + ) -> Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None: """Flatten top-level schema combinators only where OpenAI's validator rejects them. OpenAI-compatible backends reusing this config (and the ChatGPT backend @@ -293,7 +341,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): flattened: Final = [ # mutable-ok: request tools are a JSON list self._flattened_tool_or_passthrough(tool) for tool in tools ] - return cast("list[ALL_RESPONSES_API_TOOL_PARAMS]", flattened) # cast-ok: dict spread keeps each tool's shape + return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", flattened) # cast-ok: spread keeps each tool's shape @staticmethod def _flattened_tool_or_passthrough(tool: object) -> object: @@ -786,15 +834,14 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): compact_path: Final = parsed_url.path.rstrip("/") + "/compact" url: Final = str(parsed_url.copy_with(path=compact_path)) - input = self._validate_input_param(input) - tools = response_api_optional_request_params.get("tools") - input, tools = self.remove_cache_control_flag_from_input_and_tools(model=model, input=input, tools=tools) - sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( - model=model, tools=tools, litellm_params=litellm_params + replay_safe_input, sanitized_tools = self._prepared_input_and_tools( + model=model, + input=input, + tools=response_api_optional_request_params.get("tools"), + litellm_params=litellm_params, ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools - replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input) data: Final = dict( ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) ) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index 4ac072d0ca6..a5b748e391c 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -300,6 +300,89 @@ class TestOpenAIResponsesAPIConfig: assert result["input"][0]["id"] == "toolu_01Foreign" + @pytest.mark.parametrize( + "raw_parameters", + [ + '{"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}', + '{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}', + ], + ) + def test_transform_decodes_json_string_tool_parameters(self, raw_parameters: str): + """A JSON-encoded schema must reach the provider as an object.""" + result = self.config.transform_responses_api_request( + model=self.model, + input="weather in Paris", + response_api_optional_request_params={ + "tools": [{"type": "function", "name": "get_weather", "parameters": raw_parameters}] + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0]["parameters"] == { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + + def test_transform_decodes_json_string_tool_parameters_on_compact_request(self): + """The compact request path builds the same wire body, so it must decode too.""" + _url, data = self.config.transform_compact_response_api_request( + model=self.model, + input="weather in Paris", + response_api_optional_request_params={ + "tools": [{"type": "function", "name": "get_weather", "parameters": '{"type": "object"}'}] + }, + api_base="https://api.openai.com/v1/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["tools"][0]["parameters"] == {"type": "object"} + + @pytest.mark.parametrize("raw_parameters", ['"just a string"', "not json at all", "[1, 2, 3]", 42]) + def test_transform_rejects_tool_parameters_that_are_not_an_object(self, raw_parameters: object): + """Neither an object nor a string encoding one is a client error naming the tool index.""" + with pytest.raises(litellm.BadRequestError) as exc_info: + self.config.transform_responses_api_request( + model=self.model, + input="weather in Paris", + response_api_optional_request_params={ + "tools": [ + {"type": "web_search_preview"}, + {"type": "function", "name": "get_weather", "parameters": raw_parameters}, + ] + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "tools[1].parameters" in str(exc_info.value) + + def test_transform_leaves_object_null_and_absent_tool_parameters_untouched(self): + """The API accepts an object schema, an explicit null, an omitted schema and a built-in + tool, so decoding must forward all four unchanged rather than raising.""" + schema = {"type": "object", "properties": {"city": {"type": "string"}}} + tools = [ + {"type": "function", "name": "get_weather", "parameters": schema}, + {"type": "function", "name": "null_args", "parameters": None}, + {"type": "function", "name": "no_args"}, + {"type": "web_search_preview"}, + ] + + result = self.config.transform_responses_api_request( + model=self.model, + input="weather in Paris", + response_api_optional_request_params={"tools": tools}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0]["parameters"] == schema + assert result["tools"][1]["parameters"] is None + assert "parameters" not in result["tools"][2] + assert result["tools"][3] == {"type": "web_search_preview"} + def test_transform_compact_drops_foreign_tool_call_item_ids(self): """The compact request path replays input the same way, so it must apply the same id drop.""" @@ -864,6 +947,20 @@ class TestAzureResponsesAPIConfig: self.model = "gpt-4o" self.logging_obj = MagicMock() + def test_azure_decodes_json_string_tool_parameters(self): + """Azure reaches the same wire through `super()`, after un-nesting a chat-shaped tool.""" + result = self.config.transform_responses_api_request( + model=self.model, + input="weather in Paris", + response_api_optional_request_params={ + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": '{"type":"object"}'}}] + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0]["parameters"] == {"type": "object"} + def test_azure_get_complete_url_with_version_types(self): """Test Azure get_complete_url with different API version types""" base_url = "https://litellm8397336933.openai.azure.com" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 4a3ca612d41..481e3591ce9 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22183 + "limit": 22181 }, "LIT002": { "limit": 26745 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16468 + "limit": 16464 }, "LIT011": { - "limit": 5510 + "limit": 5506 }, "LIT012": { "limit": 4486 From 74613f9bd47d8e3068e6e2f1f519675ac15b7ab8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:52:26 -0700 Subject: [PATCH 27/38] fix(realtime): redact credentials from the relayed upstream close The handshake error path already runs client-facing error strings through _redact_string; the relay's _close_client did not, so a secret echoed in an upstream close reason could reach the client verbatim. Mirror the handshake path and scrub the close message and reason before relaying them. --- .../litellm_core_utils/realtime_streaming.py | 8 +++++--- .../test_realtime_streaming.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 984934daaac..b448cb7c9ff 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cas from typing_extensions import ReadOnly import litellm -from litellm._logging import verbose_logger +from litellm._logging import _redact_string, verbose_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.types.llms.openai import ( @@ -1567,12 +1567,14 @@ class RealTimeStreaming: await asyncio.gather(forward_task, client_task, return_exceptions=True) async def _close_client(self, close: BackendClose) -> None: + redacted_message: Final = _redact_string(close.message) + redacted_reason: Final = _redact_string(close.reason) try: if close.code != 1000: - await self.websocket.send_text(realtime_error_event(close.message, error_type="server_error")) + await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error")) await self.websocket.close( code=client_close_code(close.code), - reason=websocket_close_reason(close.reason, fallback=close.message), + reason=websocket_close_reason(redacted_reason, fallback=redacted_message), ) except Exception as e: # noqa: BLE001 # the client may already be gone; the session is over either way verbose_logger.debug("Could not relay the upstream close to the client: %s", e) diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 09757c35570..bcfacf16205 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3229,6 +3229,24 @@ async def test_bidirectional_forward_relays_upstream_policy_close_to_client(): client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) +@pytest.mark.asyncio +async def test_upstream_close_reason_with_a_secret_is_redacted_before_reaching_the_client(): + """LIT-6973: the relayed close mirrors the handshake path and scrubs credential + patterns, so an upstream error echoing a token never reaches the client verbatim.""" + secret: Final = "sk-live-abcdef0123456789abcdef0123" + client_ws: Final = _client_ws_that_never_sends() + upstream_close: Final = ConnectionClosed(Close(1008, f"auth failed for {secret}"), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert secret not in error_event["error"]["message"] + relayed_reason: Final = client_ws.close.await_args.kwargs["reason"] + assert secret not in relayed_reason + assert "REDACTED" in relayed_reason + + @pytest.mark.asyncio async def test_bidirectional_forward_maps_abnormal_upstream_close_to_internal_error(): client_ws: Final = _client_ws_that_never_sends() From fafd294878fa7d7de600d70ed58906e5c0c900d2 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 4 Sep 2026 23:52:33 -0700 Subject: [PATCH 28/38] fix(mcp): let config.yaml MCP servers pin server_id (#39286) * fix(mcp): let config.yaml MCP servers pin server_id A config-defined MCP server's id is a hash of server_name|url|transport| auth_type|alias, recomputed on every config load, so editing any of those fields mints a new id. Every key and team granted the old id via object_permission.mcp_servers keeps pointing at an id that no longer exists, and the server disappears from tools/list for them with nothing logged. load_servers_from_config now uses an explicit server_id from the server's config entry when present and falls back to the existing hash otherwise, so grants survive url/name/alias edits. Rejected at config load: a blank or non-string server_id, two entries claiming the same id, a pinned id already held by a database-backed server, and a pinned id that is another entry's server_name or alias (expand_permission_list matches ids before names, so that one would capture the other server's grants). Because the database registry loads after the config on startup, a database row that lands on a pinned config id is reported as a warning from the database reload instead, where it is decidable; the warning is latched on the shadowed set so the config-reload timer does not reprint it every interval. Deployments that do not set server_id keep the exact id they have today. * fix(mcp): close two more pinned-id capture paths A pinned server_id equal to an alias supplied through litellm_settings mcp_aliases was accepted, because the collision index only held the entry's own alias field. expand_permission_list matches ids before names, so grants written for the aliased server resolved to the pinning one. mcp_aliases keys whose target is a config server are now reserved the same way. A pinned server_id equal to a database-backed server's name, server_name or alias had the same effect against the database side, and could not be rejected at config load because the database registry is not loaded yet. The database reload now warns about it, latched like the existing shadow warning. * fix(mcp): reserve only the aliases the loader actually assigns Reserving every mcp_aliases key targeting a config server was too broad in two ways: the mapping is ignored when the entry sets its own alias, and only the first mapping for a server is ever applied. Both cases made a pinned server_id that could never have collided abort proxy startup. Reserve only the name load_servers_from_config will really assign. The database capture warning also fired for a database server whose own id is the config server_id. There the database row wins the id outright through get_registry precedence, so the shadow warning above it is the accurate one and the capture message contradicted it. Skip those rows. Also mark the two litellm-internal patches in the reload test helper, which the test-quality gate counts; the database reload has no other seam. * fix(mcp): match the loader's alias check exactly, is None not falsiness load_servers_from_config consults mcp_aliases only when the entry has no alias key at all, so an entry setting alias: "" gets no mapped alias. The collision index used falsiness and reserved the mapped name anyway, which failed startup on a pinned server_id that could never have collided with it. * fix(mcp): skip one identifier, not the whole database row A database row can shadow one config server_id by id and capture another by name at the same time. Skipping the entire row when its id shadowed a config entry dropped the second warning, leaving the operator with half a diagnosis. Skip only the identifier equal to the row's own id. * fix(mcp): reject conflicting self-pinned server ids * fix(mcp): validate config server names before building the identifier index The collision check reads every entry's body up front, so a malformed entry under an invalid name surfaced as an AttributeError instead of the name validation error the loader gave before this change. --- .../mcp_server/mcp_server_manager.py | 223 ++++++- .../mcp_server/test_mcp_server_manager.py | 563 ++++++++++++++++++ 2 files changed, 782 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bfc5f629faf..bcbcc6bc579 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,9 +13,19 @@ import json import os import re import time -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import ( + AsyncIterator, + Awaitable, + Callable, + Container, + Iterable, + Mapping, + MutableMapping, + Sequence, +) from contextlib import asynccontextmanager from dataclasses import dataclass, replace +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse @@ -307,6 +317,7 @@ class MCPServerConfig(TypedDict, total=False): :meth:`MCPServerManager.load_servers_from_config`. Every key is optional: YAML supplies whatever the admin wrote, and each read applies its own default.""" + server_id: ReadOnly[str] alias: str description: str mcp_info: MCPInfo @@ -400,6 +411,164 @@ def _blank_to_none(value: str | None) -> str | None: return value.strip() or None +def _pinned_config_server_id(raw_server_id: object, server_name: str) -> str | None: + """Return the ``server_id`` an admin pinned for this config.yaml server, or ``None`` when absent. + + Without a pin the id is derived by hashing ``server_name|url|transport|auth_type|alias``, so + editing any of those fields mints a new id and every ``object_permission.mcp_servers`` grant + holding the old one silently stops matching. A pinned id is used verbatim and survives those + edits. Blank and non-string values are rejected rather than silently falling back to the hash, + because a config that pins an id and still churns is the failure this field exists to prevent. + + Under ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX`` the tool prefix is derived from the server_id, so + pinning an id other than the one already in use renames every tool that server exposes. + """ + if raw_server_id is None: + return None + if not isinstance(raw_server_id, str) or not raw_server_id.strip(): + raise ValueError( + f"Invalid config for MCP server '{server_name}': server_id must be a non-empty string " + f"(got {raw_server_id!r})." + ) + return raw_server_id.strip() + + +def _first_mapped_alias(server_name: str, mcp_aliases: Mapping[str, str] | None) -> str | None: + """The ``mcp_aliases`` name ``load_servers_from_config`` will assign to this server, if any. + + Mirrors that loop, which takes the first mapping pointing at the server and stops. A later + mapping for the same server is never applied, so it stays free for another entry to pin. + """ + if mcp_aliases is None: + return None + return next( + (alias_name for alias_name, target_server_name in mcp_aliases.items() if target_server_name == server_name), + None, + ) + + +def _assigned_alias( + server_name: str, server_config: MCPServerConfig, mcp_aliases: Mapping[str, str] | None +) -> str | None: + """The alias ``load_servers_from_config`` will give this entry: its own, else the first mapping. + + ``is None``, not falsiness: the loader only consults the mapping when the key is absent, so an + entry that sets ``alias: ""`` gets no mapped alias and reserves nothing. + """ + alias: Final = server_config.get("alias") + return _first_mapped_alias(server_name, mcp_aliases) if alias is None else alias + + +def _validate_config_server_names(mcp_servers_config: Mapping[str, MCPServerConfig]) -> None: + """Reject bad server names before ``_config_identifier_owners`` reads any entry's body. + + The identifier index walks every entry up front, so without this pass a malformed entry under + a bad name would surface as an ``AttributeError`` from the index instead of the name error. + """ + for server_name in mcp_servers_config: + validate_mcp_server_name(server_name) + + +def _config_identifier_owners( + mcp_servers_config: Mapping[str, MCPServerConfig], + mcp_aliases: Mapping[str, str] | None, +) -> Mapping[str, frozenset[str]]: + """Map every server_name and alias in the config to the entries that own it. + + ``expand_permission_list`` resolves a grant against the registry keys before it falls back to + matching alias and server_name, so an id equal to another entry's name or alias captures that + entry's grants. Derived ids are hashes and never collide with a name, so this only matters once + an id is pinned. + + An alias is either set on the entry or mapped to it from ``litellm_settings.mcp_aliases``. Only + a name the loader below will really assign is reserved: the mapping is ignored for an entry that + sets its own ``alias``, and only the first mapping wins for one that does not, so reserving every + mapping would fail startup on a pin that was never going to collide. + + One identifier can have several owners when an entry's alias equals another entry's name. All of + them are kept: a grant naming that identifier resolves to every match while no id is pinned, and + a pin equal to it would narrow the grant to the pinning entry alone, even when that entry is one + of the owners. + """ + claims: Final = tuple( + (identifier, server_name) + for server_name, server_config in mcp_servers_config.items() + for identifier in (server_name, _assigned_alias(server_name, server_config, mcp_aliases)) + if identifier + ) + return MappingProxyType( + {identifier: frozenset(owner for claimed, owner in claims if claimed == identifier) for identifier, _ in claims} + ) + + +def _config_ids_capturing_db_identifiers( + config_server_ids: Container[str], + db_servers: Iterable[MCPServer], +) -> frozenset[str]: + """Config server ids that are a database-backed server's name, server_name or alias. + + ``expand_permission_list`` matches a grant against the registry keys before it matches names, so + such an id answers every grant written for the database server, and the database server itself + stops being reachable by name. The config load cannot catch this because the database registry + is not loaded yet, so it is reported from the reload that does have both halves. + + An identifier equal to the database server's own id is skipped: ``get_registry`` is + ``config_mcp_servers | registry``, so there the database server wins the id outright and the + shadow warning above is the accurate one. Reporting both would contradict. The skip is per + identifier rather than per server, so a row that shadows one config id and captures another + still reports the capture. + """ + return frozenset( + identifier + for server in db_servers + for identifier in (server.name, server.server_name, server.alias) + if identifier and identifier != server.server_id and identifier in config_server_ids + ) + + +def _reject_config_server_id_collision( + assigned_server_ids: Mapping[str, str], + server_id: str, + server_name: str, + pinned: bool, + db_backed_server_ids: Mapping[str, object], + identifier_owners: Mapping[str, frozenset[str]], +) -> None: + """Raise when ``server_id`` is already taken, either by an earlier config entry or by the database. + + Two config entries sharing an id would silently overwrite each other in ``config_mcp_servers``, + and an id already held by a database-backed server is hidden by it, because ``get_registry`` is + ``config_mcp_servers | registry`` and the right operand wins. A pinned id that is another + entry's server_name or alias captures that entry's permission grants the same way. Derived ids + cannot collide (the unique config key is part of the hash input), so all three only happen once + an id is pinned. + + Pinning an identifier this entry itself owns is allowed, because a grant naming it already + resolved here, but only when no other entry owns it too. An entry whose alias is this entry's + server_name shares the identifier, and pinning it would take that entry's grants. + """ + claimed_by = assigned_server_ids.get(server_id) + if claimed_by is not None: + raise ValueError( + f"Invalid config for MCP server '{server_name}': server_id '{server_id}' is already " + f"used by MCP server '{claimed_by}'. Each mcp_servers entry needs its own id." + ) + if pinned and server_id in db_backed_server_ids: + raise ValueError( + f"Invalid config for MCP server '{server_name}': server_id '{server_id}' belongs to a " + "database-backed MCP server. The database entry takes precedence over config.yaml, so " + "this server would never be reachable." + ) + other_owners: Final = identifier_owners.get(server_id, frozenset()) - frozenset((server_name,)) + if pinned and other_owners: + owner_names: Final = "', '".join(sorted(other_owners)) + raise ValueError( + f"Invalid config for MCP server '{server_name}': server_id '{server_id}' is the " + f"server_name or alias of MCP server '{owner_names}'. Permission entries naming " + f"'{server_id}' would resolve to '{server_name}' alone and no longer reach '{owner_names}'." + ) + + def _uses_issuer_anchor(manual_issuer: str | None, is_discovery_auth_type: bool) -> bool: """Whether the endpoints are authoritatively anchored to an admin-pinned issuer (RFC 8414 §3.3). @@ -1565,6 +1734,11 @@ class MCPServerManager: # empty result, or failure). Used to throttle re-probes for servers that do # not return instructions, and to apply a short cooldown after failures. self._upstream_initialize_instructions_probed_at: dict[str, float] = {} + # Last set of config server ids found shadowed by database rows. reload_servers_from_database + # runs on the config-reload timer, so this keeps a standing misconfiguration from re-logging + # the same warning every interval; a change in the set logs again. + self._warned_shadowed_config_server_ids: frozenset[str] = frozenset() + self._warned_capturing_config_server_ids: frozenset[str] = frozenset() self._oauth_discovery_on_startup = _mcp_oauth_discovery_on_startup_enabled() self._oauth_discovery_generation_counter = 0 self._oauth_discovery_slots: tuple[_OAuthDiscoverySlot, ...] = () @@ -1958,10 +2132,14 @@ class MCPServerManager: # Track which aliases have been used to ensure only first occurrence is used used_aliases: Final = set() + # server_id -> the config server_name that claimed it, so a pinned id cannot silently + # overwrite another server's entry in self.config_mcp_servers. + assigned_server_ids: MutableMapping[str, str] = {} # mutable-ok: per-load collision index + _validate_config_server_names(mcp_servers_config) + identifier_owners: Final = _config_identifier_owners(mcp_servers_config, mcp_aliases) for server_name, raw_server_config in mcp_servers_config.items(): server_config: MCPServerConfig = raw_server_config - validate_mcp_server_name(server_name) _mcp_info: MCPInfo = server_config.get("mcp_info", None) or {} # Preserve all custom fields from config while setting defaults for core fields mcp_info: MCPInfo = _mcp_info.copy() @@ -1994,14 +2172,24 @@ class MCPServerManager: name_for_prefix = get_server_prefix(temp_server) server_url = server_config.get("url", None) or "" - # Generate stable server ID based on parameters - server_id = self._generate_stable_server_id( + # An explicitly pinned server_id wins; otherwise derive one from the parameters. + pinned_server_id = _pinned_config_server_id(server_config.get("server_id"), server_name) + server_id = pinned_server_id or self._generate_stable_server_id( server_name=server_name, url=server_url, transport=server_config.get("transport", MCPTransport.http), auth_type=server_config.get("auth_type", None), alias=alias, ) + _reject_config_server_id_collision( + assigned_server_ids, + server_id, + server_name, + pinned=pinned_server_id is not None, + db_backed_server_ids=self.registry, + identifier_owners=identifier_owners, + ) + assigned_server_ids[server_id] = server_name _warn_on_server_name_fields( server_id=server_id, @@ -6123,6 +6311,33 @@ class MCPServerManager: verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) + # get_registry() is ``config_mcp_servers | registry``, so a database row sharing an id with a + # config.yaml server hides that server everywhere. Only reachable once an operator pins + # ``server_id`` in config.yaml; say so rather than letting the server disappear silently. + shadowed_config_server_ids: Final = frozenset(self.config_mcp_servers.keys() & registered_registry.keys()) + if shadowed_config_server_ids and shadowed_config_server_ids != self._warned_shadowed_config_server_ids: + verbose_logger.warning( + "config.yaml MCP server_id(s) %s are also database-backed MCP servers. The database " + "entry takes precedence, so the config.yaml server is unreachable. Give the config " + "entry a different server_id.", + ", ".join(sorted(shadowed_config_server_ids)), + ) + self._warned_shadowed_config_server_ids = shadowed_config_server_ids + + # The mirror image of the block above: a config server_id that is a database server's name + # answers that server's grants instead, because ids are matched before names. + capturing_config_server_ids: Final = _config_ids_capturing_db_identifiers( + self.config_mcp_servers.keys(), registered_registry.values() + ) + if capturing_config_server_ids and capturing_config_server_ids != self._warned_capturing_config_server_ids: + verbose_logger.warning( + "config.yaml MCP server_id(s) %s are the name or alias of a database-backed MCP " + "server. Permission entries naming them resolve to the config.yaml server, not the " + "database one. Give the config entry a different server_id.", + ", ".join(sorted(capturing_config_server_ids)), + ) + self._warned_capturing_config_server_ids = capturing_config_server_ids + await self._hydrate_config_servers_dcr_clients() def get_mcp_servers_from_ids(self, server_ids: list[str]) -> list[MCPServer]: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 02b1a19081a..9745e508703 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -11428,6 +11428,569 @@ class TestOpenApiHandlerRelaysUpstreamAuth: assert "upstream returned HTTP 503" in result.content[0].text +class TestConfigServerIdPinning: + """config.yaml servers may pin ``server_id`` so permission grants survive connection edits.""" + + @staticmethod + def _config(**overrides: object) -> dict[str, dict[str, object]]: + return { + "docs_server": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + **overrides, + } + } + + @pytest.mark.asyncio + async def test_derived_id_churns_when_connection_fields_change(self): + """The behavior the pin exists to escape: editing the url mints a brand-new id.""" + manager = MCPServerManager() + + await manager.load_servers_from_config(self._config()) + before = next(iter(manager.config_mcp_servers)) + + manager.config_mcp_servers.clear() + await manager.load_servers_from_config(self._config(url="https://prod.example.com/mcp")) + after = next(iter(manager.config_mcp_servers)) + + assert before != after + + @pytest.mark.asyncio + async def test_pinned_id_survives_url_transport_auth_and_alias_edits(self): + manager = MCPServerManager() + + await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) + assert list(manager.config_mcp_servers) == ["docs-prod-1"] + assert manager.config_mcp_servers["docs-prod-1"].server_id == "docs-prod-1" + + manager.config_mcp_servers.clear() + await manager.load_servers_from_config( + self._config( + server_id="docs-prod-1", + url="https://prod.example.com/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.bearer_token, + alias="docs", + ) + ) + + assert list(manager.config_mcp_servers) == ["docs-prod-1"] + assert manager.config_mcp_servers["docs-prod-1"].url == "https://prod.example.com/mcp" + + @pytest.mark.asyncio + async def test_absent_server_id_keeps_the_derived_hash(self): + manager = MCPServerManager() + + await manager.load_servers_from_config(self._config()) + + derived = manager._generate_stable_server_id( + server_name="docs_server", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=None, + alias=None, + ) + assert list(manager.config_mcp_servers) == [derived] + + @pytest.mark.asyncio + @pytest.mark.parametrize("bad_value", ["", " ", 123, True, ["docs-prod-1"]]) + async def test_blank_or_non_string_server_id_is_rejected(self, bad_value: Any): + manager = MCPServerManager() + + with pytest.raises(ValueError, match="server_id must be a non-empty string"): + await manager.load_servers_from_config(self._config(server_id=bad_value)) + + @pytest.mark.asyncio + async def test_two_servers_pinning_the_same_id_are_rejected(self): + manager = MCPServerManager() + config: Dict[str, Any] = { + "docs_server": {"url": "https://a.example.com/mcp", "server_id": "shared-id"}, + "wiki_server": {"url": "https://b.example.com/mcp", "server_id": "shared-id"}, + } + + with pytest.raises(ValueError, match="already used by MCP server 'docs_server'"): + await manager.load_servers_from_config(config) + + @pytest.mark.asyncio + async def test_pinned_id_colliding_with_a_derived_id_is_rejected(self): + """A pin that lands on another entry's derived hash collides just as hard.""" + manager = MCPServerManager() + derived = manager._generate_stable_server_id( + server_name="docs_server", + url="https://a.example.com/mcp", + transport=MCPTransport.http, + auth_type=None, + alias=None, + ) + config: Dict[str, Any] = { + "docs_server": {"url": "https://a.example.com/mcp", "transport": MCPTransport.http}, + "wiki_server": {"url": "https://b.example.com/mcp", "server_id": derived}, + } + + with pytest.raises(ValueError, match="already used by MCP server 'docs_server'"): + await manager.load_servers_from_config(config) + + @pytest.mark.asyncio + async def test_pinned_id_colliding_with_a_db_backed_server_is_rejected(self): + """get_registry() is ``config | registry``, so the db row would hide the config server. + + The registry is seeded by hand because on a real startup the config loads before the + database does, so this check only fires on a later reload. The startup ordering is covered + by ``test_db_row_arriving_on_a_pinned_config_id_warns``; the warning there is not redundant. + """ + manager = MCPServerManager() + manager.registry["db-uuid-1"] = MCPServer( + server_id="db-uuid-1", + name="db_server", + transport=MCPTransport.http, + url="https://db.example.com/mcp", + ) + + with pytest.raises(ValueError, match="belongs to a database-backed MCP server"): + await manager.load_servers_from_config(self._config(server_id="db-uuid-1")) + + @pytest.mark.asyncio + async def test_derived_id_matching_a_db_backed_server_is_not_rejected(self): + """Only a pinned id is an authoring error; a hash collision must not fail startup.""" + manager = MCPServerManager() + derived = manager._generate_stable_server_id( + server_name="docs_server", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=None, + alias=None, + ) + manager.registry[derived] = MCPServer( + server_id=derived, + name="db_server", + transport=MCPTransport.http, + url="https://db.example.com/mcp", + ) + + await manager.load_servers_from_config(self._config()) + + assert derived in manager.config_mcp_servers + + @pytest.mark.asyncio + async def test_pinned_id_is_stripped_of_surrounding_whitespace(self): + manager = MCPServerManager() + + await manager.load_servers_from_config(self._config(server_id=" docs-prod-1 ")) + + assert list(manager.config_mcp_servers) == ["docs-prod-1"] + + @staticmethod + async def _reload_with_db_server(manager: MCPServerManager, server_id: str, db_name: str = "db_server") -> None: + row = LiteLLM_MCPServerTable( + server_id=server_id, + server_name=db_name, + alias=db_name, + url="https://db.example.com/mcp", + transport=MCPTransport.http, + ) + raw_row = MagicMock() + raw_row.model_dump.return_value = row.model_dump() + repository = MagicMock() + repository.table.find_many = AsyncMock(return_value=[raw_row]) + built = MCPServer( + server_id=server_id, + name=db_name, + server_name=db_name, + url="https://db.example.com/mcp", + transport=MCPTransport.http, + ) + with ( + patch( # test-quality-ok: the db reload path has no seam but its own repository + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repository, + ), + patch( # test-quality-ok: same, the prisma client is fetched inside the reload + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch.object(manager, "build_mcp_server_from_table", new=AsyncMock(return_value=built)), + ): + await manager.reload_servers_from_database() + + @pytest.mark.asyncio + async def test_db_row_arriving_on_a_pinned_config_id_warns(self, caplog): + """The db row loads after config on startup, so the config server is hidden then, not at load.""" + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "docs-prod-1") + + assert any("docs-prod-1" in m and "database entry takes precedence" in m for m in caplog.messages) + assert manager.get_registry()["docs-prod-1"].url == "https://db.example.com/mcp" + + @pytest.mark.asyncio + async def test_db_row_with_a_distinct_id_does_not_warn(self, caplog): + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "db-uuid-1") + + assert all("database entry takes precedence" not in m for m in caplog.messages) + assert set(manager.get_registry()) == {"docs-prod-1", "db-uuid-1"} + + @pytest.mark.asyncio + async def test_pinned_id_matching_another_entrys_server_name_is_rejected(self): + """expand_permission_list resolves against registry keys first, so this steals the grants.""" + manager = MCPServerManager() + + with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): + await manager.load_servers_from_config( + { + "wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, + "docs_server": { + "server_id": "wiki_server", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + } + ) + + @pytest.mark.asyncio + async def test_pinned_id_matching_another_entrys_alias_is_rejected(self): + manager = MCPServerManager() + + with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): + await manager.load_servers_from_config( + { + "wiki_server": { + "alias": "wiki", + "url": "https://wiki.example.com/mcp", + "transport": MCPTransport.http, + }, + "docs_server": { + "server_id": "wiki", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + } + ) + + @pytest.mark.asyncio + async def test_pinning_a_servers_own_name_is_allowed(self): + """The most natural pin an operator writes; it resolves to the same server either way.""" + manager = MCPServerManager() + + await manager.load_servers_from_config(self._config(server_id="docs_server")) + + assert list(manager.config_mcp_servers) == ["docs_server"] + + @pytest.mark.asyncio + async def test_pinning_a_servers_own_alias_is_allowed(self): + manager = MCPServerManager() + + await manager.load_servers_from_config(self._config(alias="docs", server_id="docs")) + + assert list(manager.config_mcp_servers) == ["docs"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("aliasing_entry_first", [True, False]) + async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected(self, aliasing_entry_first: bool): + """A grant naming 'docs_server' reaches both servers unpinned; the pin would narrow it to one.""" + manager = MCPServerManager() + wiki = ( + "wiki_server", + {"alias": "docs_server", "url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, + ) + docs = ( + "docs_server", + {"server_id": "docs_server", "url": "https://example.com/mcp", "transport": MCPTransport.http}, + ) + + with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): + await manager.load_servers_from_config(dict((wiki, docs) if aliasing_entry_first else (docs, wiki))) + + @pytest.mark.asyncio + async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected(self): + manager = MCPServerManager() + + with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): + await manager.load_servers_from_config( + { + "wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, + "docs_server": { + "server_id": "docs_server", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + }, + mcp_aliases={"docs_server": "wiki_server"}, + ) + + @pytest.mark.asyncio + async def test_pinning_own_alias_shared_with_a_later_entry_is_rejected(self): + """Nothing rejects duplicate aliases, so the first entry's pin would answer the second's grants.""" + manager = MCPServerManager() + + with pytest.raises(ValueError, match="server_name or alias of MCP server 'docs_server'"): + await manager.load_servers_from_config( + { + "wiki_server": { + "alias": "shared", + "server_id": "shared", + "url": "https://wiki.example.com/mcp", + "transport": MCPTransport.http, + }, + "docs_server": { + "alias": "shared", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + } + ) + + @pytest.mark.asyncio + async def test_own_name_pin_resolves_grants_like_the_unpinned_name(self): + """The negative control: a sole-owner self-pin must keep loading and answer the same grants.""" + manager = MCPServerManager() + + await manager.load_servers_from_config( + { + "wiki_server": {"alias": "wiki", "url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, + "docs_server": { + "server_id": "docs_server", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + } + ) + wiki_id = next(sid for sid, server in manager.config_mcp_servers.items() if server.alias == "wiki") + + assert manager.expand_permission_list(["docs_server"]) == ["docs_server"] + assert manager.expand_permission_list(["wiki"]) == [wiki_id] + + @pytest.mark.asyncio + async def test_derived_id_is_not_checked_against_names(self): + """Unpinned configs must keep loading; only a pinned id can be an authoring error.""" + manager = MCPServerManager() + + await manager.load_servers_from_config( + { + "wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, + "docs_server": {"url": "https://example.com/mcp", "transport": MCPTransport.http}, + } + ) + + assert len(manager.config_mcp_servers) == 2 + + @pytest.mark.asyncio + async def test_shadow_warning_is_not_repeated_on_every_reload(self, caplog): + """reload_servers_from_database runs on the config-reload timer; one warning, not one a tick.""" + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "docs-prod-1") + first_round = [m for m in caplog.messages if "database entry takes precedence" in m] + await self._reload_with_db_server(manager, "docs-prod-1") + second_round = [m for m in caplog.messages if "database entry takes precedence" in m] + + assert len(first_round) == 1 + assert second_round == first_round + + @pytest.mark.asyncio + async def test_shadow_warning_fires_again_when_the_shadowed_set_changes(self, caplog): + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "docs-prod-1") + await self._reload_with_db_server(manager, "db-uuid-1") + await self._reload_with_db_server(manager, "docs-prod-1") + + assert len([m for m in caplog.messages if "database entry takes precedence" in m]) == 2 + + @pytest.mark.asyncio + async def test_pinned_id_matching_a_mapped_alias_is_rejected(self): + """An alias can also arrive from litellm_settings.mcp_aliases; it is reserved just the same.""" + manager = MCPServerManager() + + with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"): + await manager.load_servers_from_config( + { + "wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, + "docs_server": { + "server_id": "wiki", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + }, + {"wiki": "wiki_server"}, + ) + + @pytest.mark.asyncio + async def test_pinning_a_servers_own_mapped_alias_is_allowed(self): + manager = MCPServerManager() + + await manager.load_servers_from_config( + self._config(server_id="docs"), + {"docs": "docs_server"}, + ) + + assert list(manager.config_mcp_servers) == ["docs"] + + @pytest.mark.asyncio + async def test_mapped_alias_for_an_unknown_server_reserves_nothing(self): + """A dangling mcp_aliases entry is never applied, so it must not fail an unrelated pin.""" + manager = MCPServerManager() + + await manager.load_servers_from_config( + self._config(server_id="wiki"), + {"wiki": "a_server_that_does_not_exist"}, + ) + + assert list(manager.config_mcp_servers) == ["wiki"] + + @pytest.mark.asyncio + async def test_config_id_that_is_a_db_server_name_warns(self, caplog): + """The mirror of the shadow case: here the config entry captures the db server's grants.""" + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="db_server")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "db-uuid-1") + + assert any("db_server" in m and "name or alias of a database-backed" in m for m in caplog.messages) + + @pytest.mark.asyncio + async def test_capture_warning_is_not_repeated_on_every_reload(self, caplog): + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="db_server")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "db-uuid-1") + await self._reload_with_db_server(manager, "db-uuid-1") + + assert len([m for m in caplog.messages if "name or alias of a database-backed" in m]) == 1 + + @pytest.mark.asyncio + async def test_config_id_unrelated_to_db_names_does_not_warn(self, caplog): + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="docs-prod-1")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "db-uuid-1") + + assert all("name or alias of a database-backed" not in m for m in caplog.messages) + + @pytest.mark.asyncio + async def test_mapped_alias_for_a_server_with_its_own_alias_reserves_nothing(self): + """load_servers_from_config ignores the mapping when the entry sets alias, so it is free.""" + manager = MCPServerManager() + + await manager.load_servers_from_config( + { + "wiki_server": { + "alias": "wiki_prod", + "url": "https://wiki.example.com/mcp", + "transport": MCPTransport.http, + }, + "docs_server": { + "server_id": "wiki", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + }, + {"wiki": "wiki_server"}, + ) + + assert "wiki" in manager.config_mcp_servers + assert len(manager.config_mcp_servers) == 2 + + @pytest.mark.asyncio + async def test_only_the_first_mapped_alias_for_a_server_is_reserved(self): + """Only the first mapping is applied, so pinning the second one must still load.""" + manager = MCPServerManager() + + await manager.load_servers_from_config( + { + "wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http}, + "docs_server": { + "server_id": "wiki_two", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + }, + {"wiki_one": "wiki_server", "wiki_two": "wiki_server"}, + ) + + assert "wiki_two" in manager.config_mcp_servers + + @pytest.mark.asyncio + async def test_invalid_name_is_reported_before_any_entry_body_is_read(self): + """The identifier index walks every entry up front, so a bad name must still fail on the name.""" + with pytest.raises(Exception, match="Server name cannot contain"): + await MCPServerManager().load_servers_from_config({"my-server": None}) + + @pytest.mark.asyncio + async def test_a_shadowing_db_server_reports_only_the_shadow_warning(self, caplog): + """The db row wins the id outright, so the capture message would contradict the shadow one.""" + manager = MCPServerManager() + await manager.load_servers_from_config(self._config(server_id="db_server")) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "db_server") + + assert any("database entry takes precedence" in m for m in caplog.messages) + assert all("name or alias of a database-backed" not in m for m in caplog.messages) + assert manager.get_registry()["db_server"].url == "https://db.example.com/mcp" + + @pytest.mark.asyncio + async def test_an_explicitly_blank_alias_still_blocks_the_mapping(self): + """The loader only consults mcp_aliases when the key is absent, so a blank alias frees it.""" + manager = MCPServerManager() + + await manager.load_servers_from_config( + { + "wiki_server": { + "alias": "", + "url": "https://wiki.example.com/mcp", + "transport": MCPTransport.http, + }, + "docs_server": { + "server_id": "wiki", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + }, + {"wiki": "wiki_server"}, + ) + + assert "wiki" in manager.config_mcp_servers + assert manager.config_mcp_servers["wiki"].url == "https://example.com/mcp" + + @pytest.mark.asyncio + async def test_a_row_that_shadows_one_id_still_reports_capturing_another(self, caplog): + """Skipping is per identifier, not per row, so the second collision is not lost.""" + manager = MCPServerManager() + await manager.load_servers_from_config( + { + "docs_server": { + "server_id": "shadow_x", + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + }, + "wiki_server": { + "server_id": "capture_y", + "url": "https://wiki.example.com/mcp", + "transport": MCPTransport.http, + }, + } + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await self._reload_with_db_server(manager, "shadow_x", db_name="capture_y") + + assert any("shadow_x" in m and "database entry takes precedence" in m for m in caplog.messages) + assert any("capture_y" in m and "name or alias of a database-backed" in m for m in caplog.messages) + + class TestLitellmAdmissionKeyIsNeverTheSubjectToken: """The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it From f3cf5578989e558438b11c335b69702812d6b738 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 4 Sep 2026 23:57:53 -0700 Subject: [PATCH 29/38] feat(dashboard): configure classifier vision input (#39840) --- .../add_model/ClassificationMethodConfig.tsx | 10 ++- .../add_model/ClassifierVisionConfig.tsx | 79 ++++++++++++++++++ .../add_model/ComplexityRouterConfig.test.tsx | 81 +++++++++++++++++++ .../build_complexity_router_config.test.ts | 40 ++++++++- .../build_complexity_router_config.ts | 10 +-- .../edit_auto_router_modal.test.tsx | 57 +++++++++++++ 6 files changed, 268 insertions(+), 9 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/ClassifierVisionConfig.tsx diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index cc66103fc86..15cfff01766 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -15,6 +15,7 @@ import { RestrictedSection, restrictedBy } from "./TierRestrictions"; import HeuristicScoringConfig from "./HeuristicScoringConfig"; import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect"; import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig"; +import ClassifierVisionConfig from "./ClassifierVisionConfig"; import type { ReasoningEffort } from "./complexity_router_tiers"; import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults"; import { @@ -315,12 +316,13 @@ const ClassificationMethodConfig: React.FC = ({ timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, classification_rubric: selectedRubric, }; - onChange({ + const nextValue: ComplexityRouterConfigValue = { ...value, ...(selectedRubric && { classifier_llm_config: rubricConfig }), classification_prompt: classificationPrompt, classification_examples: classificationExamples, - }); + }; + onChange(nextValue); }; const handleClassifierModelChange = (model: string) => { @@ -577,6 +579,10 @@ const ClassificationMethodConfig: React.FC = ({ value={value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS }} onChange={(classifier_llm_config) => onChange({ ...value, classifier_llm_config })} /> + onChange({ ...value, classifier_llm_config })} + />
Classifier Prompt diff --git a/ui/litellm-dashboard/src/components/add_model/ClassifierVisionConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassifierVisionConfig.tsx new file mode 100644 index 00000000000..34c41fff006 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ClassifierVisionConfig.tsx @@ -0,0 +1,79 @@ +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Switch } from "@/components/ui/switch"; +import React from "react"; + +import type { ClassifierLLMConfigWire } from "./build_complexity_router_config"; + +export const DEFAULT_CLASSIFIER_VISION_ENABLED = false; +export const DEFAULT_CLASSIFIER_VISION_MAX_IMAGES = 1; + +const MAX_IMAGES_ID = "classifier-vision-max-images"; + +interface ClassifierVisionConfigProps { + value: ClassifierLLMConfigWire; + onChange: (value: ClassifierLLMConfigWire) => void; +} + +const ClassifierVisionConfig: React.FC = ({ value, onChange }) => { + const [draftMaxImages, setDraftMaxImages] = React.useState(null); + const enabled = value.vision?.enabled ?? DEFAULT_CLASSIFIER_VISION_ENABLED; + + const handleMaxImagesChange = (raw: string): void => { + setDraftMaxImages(raw); + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onChange({ + ...value, + vision: { ...value.vision, enabled, max_images: Math.max(1, Math.round(parsed)) }, + }); + }; + + return ( +
+
+ { + if (!visionEnabled) { + const { vision: _vision, ...withoutVision } = value; + onChange(withoutVision); + return; + } + onChange({ + ...value, + vision: { + ...value.vision, + enabled: true, + max_images: value.vision?.max_images ?? DEFAULT_CLASSIFIER_VISION_MAX_IMAGES, + }, + }); + }} + aria-label="Use images for classification" + /> + Use images for classification +
+ + Send inline image data to the classifier so it can choose a tier from what the image shows. + + {enabled && ( +
+ + handleMaxImagesChange(event.target.value)} + onBlur={() => setDraftMaxImages(null)} + className="w-full" + /> +
+ )} +
+ ); +}; + +export default ClassifierVisionConfig; diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 0590b524a06..2970e14b335 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -1,5 +1,6 @@ import { fireEvent, renderWithProviders, screen, within } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; +import React from "react"; import { vi } from "vitest"; import ComplexityRouterConfig, { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; vi.mock( @@ -1690,3 +1691,83 @@ describe("ComplexityRouterConfig tier editing", () => { expect(screen.queryByText("Display names rename the built-in tiers", { exact: false })).not.toBeInTheDocument(); }); }); + +describe("classifier vision settings", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + + const VisionFixture = ({ onChange = vi.fn() }: { onChange?: ReturnType }) => { + const [value, setValue] = React.useState(llmValue); + return ( + { + setValue(nextValue); + onChange(nextValue); + }} + /> + ); + }; + + it("starts off and reveals the default cap when enabled", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + const vision = screen.getByRole("switch", { name: "Use images for classification" }); + expect(vision).not.toBeChecked(); + expect(screen.queryByLabelText("Maximum images per request")).not.toBeInTheDocument(); + + fireEvent.click(vision); + + expect(screen.getByLabelText("Maximum images per request")).toHaveValue("1"); + }); + + it("writes the switch and a clamped image cap into the classifier config", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + fireEvent.click(screen.getByRole("switch", { name: "Use images for classification" })); + expect(onChange).toHaveBeenLastCalledWith({ + ...llmValue, + classifier_llm_config: { ...llmValue.classifier_llm_config, vision: { enabled: true, max_images: 1 } }, + }); + + fireEvent.change(screen.getByLabelText("Maximum images per request"), { target: { value: "1.7" } }); + expect(onChange).toHaveBeenLastCalledWith({ + ...llmValue, + classifier_llm_config: { ...llmValue.classifier_llm_config, vision: { enabled: true, max_images: 2 } }, + }); + }); + + it("keeps the image cap draft empty until a valid value is entered", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + fireEvent.click(screen.getByRole("switch", { name: "Use images for classification" })); + onChange.mockClear(); + + const input = screen.getByLabelText("Maximum images per request"); + fireEvent.change(input, { target: { value: "" } }); + + expect(input).toHaveValue(""); + expect(onChange).not.toHaveBeenCalled(); + + fireEvent.change(input, { target: { value: "0" } }); + expect(onChange).toHaveBeenLastCalledWith({ + ...llmValue, + classifier_llm_config: { ...llmValue.classifier_llm_config, vision: { enabled: true, max_images: 1 } }, + }); + }); + + it("is absent when the classifier is heuristic", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + expect(screen.queryByText("Use images for classification")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 1b5bb9e72eb..17c1a83fd67 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -1170,12 +1170,13 @@ describe("buildComplexityRouterConfig stall escalation", () => { }); it("emits the toggle and both knobs when it is on", () => { - const config = buildComplexityRouterConfig({ + const params = { ...baseParams, stallEscalationEnabled: true, stallEscalationWindow: 8, stallEscalationRepeatThreshold: 4, - }); + }; + const config = buildComplexityRouterConfig(params); expect(config.stall_escalation_enabled).toBe(true); expect(config.stall_escalation_window).toBe(8); expect(config.stall_escalation_repeat_threshold).toBe(4); @@ -1207,3 +1208,38 @@ describe("dryRunRejection", () => { expect(dryRunRejection({ valid: true, error: null })).toBeNull(); }); }); + +describe("classifier vision wire payload", () => { + const vision = { enabled: true, max_images: 3 }; + const classifierLlmConfig = { model: "classifier", timeout_ms: 3000, vision }; + + it("keeps vision through the standard-tier payload", () => { + const params = { ...baseParams, classifierType: "llm" as const, classifierLlmConfig }; + const payload = buildComplexityRouterConfig(params); + + expect(payload.classifier_llm_config).toMatchObject({ vision }); + }); + + it("keeps vision through the custom-tier payload", () => { + const customTierSet = { + tiers: [ + { id: "simple", name: "simple", definition: "small talk", models: ["gpt-4o-mini"] }, + { id: "complex", name: "complex", definition: "hard work", models: ["gpt-4o"] }, + ], + fallback_tier_id: "simple", + }; + const payload = buildComplexityRouterConfig({ ...baseParams, customTierSet, classifierLlmConfig }); + + expect(payload.classifier_llm_config).toMatchObject({ vision }); + }); + + it("keeps an untouched classifier config free of vision", () => { + const payload = buildComplexityRouterConfig({ + ...baseParams, + classifierType: "llm", + classifierLlmConfig: { model: "classifier", timeout_ms: 3000 }, + }); + + expect(payload.classifier_llm_config).not.toHaveProperty("vision"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index d7974484970..7769fb832fe 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -1,8 +1,5 @@ -import { KeywordTierRule } from "./KeywordTierRules"; - -type ClassifierLLMConfigWire = ClassifierLLMConfig & { vision?: { enabled?: boolean; max_images?: number } }; - import type { ModelGroup } from "../llm_calls/fetch_models"; +import { KeywordTierRule } from "./KeywordTierRules"; import { type CustomTierSet, type TierRow, @@ -42,6 +39,9 @@ import { usesLlmClassifier, } from "./ComplexityRouterConfig"; +export type ClassifierVisionConfig = { enabled?: boolean; max_images?: number }; +export type ClassifierLLMConfigWire = ClassifierLLMConfig & { vision?: ClassifierVisionConfig }; + /** * Drop an empty system_prompt so the payload carries an override only when there is one. The * backend rejects a blank string rather than reading it as "use the default", and sending `""` @@ -124,7 +124,7 @@ export interface BuildComplexityRouterConfigParams { planModeMinTier: string | undefined; tierLabels: ComplexityTierLabels | undefined; classifierType: ClassifierType; - classifierLlmConfig: ClassifierLLMConfig | undefined; + classifierLlmConfig: ClassifierLLMConfigWire | undefined; classifierContextWindowSize: number | undefined; classifierContextBudgetChars: number | undefined; classifierContextIncludeAssistantTurns: boolean | undefined; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 9a55d3c0703..e45f84fa646 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -1036,3 +1036,60 @@ describe("EditAutoRouterModal with a stored custom tier set", () => { expect(savedConfig().tier_model_configs).toEqual(CUSTOM_STORED.tier_model_configs); }); }); + +describe("EditAutoRouterModal classifier vision", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const STORED_CONFIG = { + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: ["gpt-4o-mini"], COMPLEX: ["gpt-4o-mini"], REASONING: ["gpt-4o-mini"] }, + classifier_type: "llm", + classifier_llm_config: { + model: "gpt-4o-mini", + timeout_ms: 3000, + vision: { enabled: true, max_images: 2 }, + }, + }; + + const renderModal = () => + renderWithProviders( + , + ); + + it("hydrates and keeps a stored vision setting through an untouched save", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.click(await screen.findByText("Advanced: Classification Method")); + expect(screen.getByRole("switch", { name: "Use images for classification" })).toBeChecked(); + expect(screen.getByLabelText("Maximum images per request")).toHaveValue("2"); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classifier_llm_config).toMatchObject({ vision: { enabled: true, max_images: 2 } }); + }); + + it("removes vision when the operator turns it off", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(screen.getByRole("switch", { name: "Use images for classification" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classifier_llm_config).not.toHaveProperty("vision"); + }); +}); From 29ac88ebc6bb93bda02138699c5ebe08bdd4e8cc Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 4 Sep 2026 23:59:51 -0700 Subject: [PATCH 30/38] fix(batches): register ownership for every batch create path (#39810) * fix(batches): register ownership for every batch create path Since the team isolation change, the managed files hook decided whether a response came from a create by looking for the managed input file id on it, which only the unified input path sets. Batches created from a model-encoded input file id, a model param, or a raw provider id with ?provider= never got an ownership row, so they vanished from GET /v1/batches for the key that created them. The create endpoint now stamps a create marker on the response before the hooks run, and the hook keys ownership registration and the batch-created metric on that marker instead of on the input id format. * test(batches): assert ownership registration through the managed files hook The endpoint tests asserted the private create marker, which is wiring, not behaviour. They now run the create through the real managed files hook and assert the ownership row is written for the creating key on every create path, with the unified path driven by a genuine encoded input file id instead of patched decoders. --- .../proxy/hooks/managed_files.py | 8 +-- litellm/proxy/batches_endpoints/endpoints.py | 3 + .../openai_files_endpoints/common_utils.py | 2 + .../proxy/hooks/test_managed_files.py | 16 ++--- .../proxy/test_managed_files_hook.py | 52 +++++++++++++++ .../proxy/batches_endpoints/test_endpoints.py | 64 ++++++++++++++++++- 6 files changed, 129 insertions(+), 16 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index e11c6e70540..bc1eb6cebc2 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -46,6 +46,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.openai_files_endpoints.common_utils import ( + BATCH_CREATE_HIDDEN_PARAM, FILE_LIST_CONTINUATION_CHUNK_SIZE, MAX_FILE_LIST_LIMIT, _is_base64_encoded_unified_file_id, @@ -1321,7 +1322,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ## Check if unified_file_id is in the response unified_file_id = response._hidden_params.get("unified_file_id") # managed file id unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id - is_batch_create: Final = unified_file_id is not None + is_batch_create: Final = response._hidden_params.get(BATCH_CREATE_HIDDEN_PARAM) is True model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) @@ -1410,10 +1411,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) - # Only record batch creation metric on actual create (not retrieve/cancel). - # unified_file_id in _hidden_params is only set by the create_batch endpoint. - original_unified_file_id = response._hidden_params.get("unified_file_id") - if original_unified_file_id: + if is_batch_create: prom_logger = self._get_prometheus_logger() if prom_logger: batch_provider = "" diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index be889a22cae..5c4bacd757c 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_query, ) from litellm.proxy.openai_files_endpoints.common_utils import ( + BATCH_CREATE_HIDDEN_PARAM, _is_base64_encoded_unified_file_id, add_internal_model_credentials, apply_team_provider_credentials, @@ -347,6 +348,8 @@ async def create_batch( **_create_batch_data, ) + response._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True + ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 992ed0d814d..15eeddbc489 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -37,6 +37,8 @@ MAX_FILE_LIST_LIMIT: Final = 10000 FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500 +BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create" + def validate_file_list_limit(limit: int | None) -> None: """Reject a ``limit`` outside the range OpenAI documents for GET /v1/files.""" diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 57394f1cebe..7e96c956664 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -10,6 +10,7 @@ from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFi from litellm.caching import DualCache from litellm.proxy._types import CallTypes from litellm.proxy.openai_files_endpoints.common_utils import ( + BATCH_CREATE_HIDDEN_PARAM, _is_base64_encoded_unified_file_id, encode_file_id_with_model, ) @@ -3185,7 +3186,7 @@ def _batch_response(batch_id, output_file_id=None, is_create=False): output_file_id=output_file_id, ) if is_create: - batch._hidden_params["unified_file_id"] = "unified-input-file-id" + batch._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True return batch @@ -3411,11 +3412,8 @@ async def test_provider_format_file_without_ownership_row_stays_accessible(): @pytest.mark.asyncio -async def test_post_call_batch_create_stores_ownership_row(): - """ - Batch creation (response hidden params carry the unified input file id) - must write an ownership row attributed to the creating key. - """ +@pytest.mark.parametrize("batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID]) +async def test_post_call_batch_create_stores_ownership_row(batch_id): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() @@ -3432,13 +3430,11 @@ async def test_post_call_batch_create_stores_ownership_row(): user_api_key_dict=UserAPIKeyAuth( user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() ), - response=_batch_response(MODEL_ENCODED_BATCH_ID, is_create=True), + response=_batch_response(batch_id, is_create=True), ) upsert_call = prisma_client.db.litellm_managedobjecttable.upsert.await_args - assert upsert_call.kwargs["where"] == { - "unified_object_id": MODEL_ENCODED_BATCH_ID - } + assert upsert_call.kwargs["where"] == {"unified_object_id": batch_id} create_data = upsert_call.kwargs["data"]["create"] assert create_data["created_by"] == "user_a" assert create_data["team_id"] == "team_a" diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index f3ad8a8592e..091b958d7c3 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -15,6 +15,7 @@ from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM from litellm.types.llms.openai import FileListPage, OpenAIFileObject from litellm.types.utils import LiteLLMBatch @@ -1540,6 +1541,11 @@ async def test_batch_create_hook_persists_creating_key_and_tags(): managed_files = _make_managed_files_instance() creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None) create_response = _make_batch_response(status="validating", output_file_id=None) + create_response._hidden_params = { + BATCH_CREATE_HIDDEN_PARAM: True, + "model_id": "model-deploy-xyz", + "model_name": "azure/gpt-4", + } await managed_files.async_post_call_success_hook( data={"litellm_metadata": {"tags": ["env:prod", "team:ml"], "user_api_key": creator.api_key}}, @@ -1554,6 +1560,52 @@ async def test_batch_create_hook_persists_creating_key_and_tags(): assert stored["user_api_key_dict"] is creator +@pytest.mark.asyncio +async def test_batch_create_hook_records_created_metric_once(): + managed_files = _make_managed_files_instance() + prometheus_logger = MagicMock() + managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger) + create_response = _make_batch_response(status="validating", output_file_id=None) + create_response._hidden_params = { + BATCH_CREATE_HIDDEN_PARAM: True, + "model_id": "model-deploy-xyz", + "model_name": "azure/gpt-4", + } + + await managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None), + response=create_response, + ) + + prometheus_logger.record_managed_batch_created.assert_called_once() + recorded = prometheus_logger.record_managed_batch_created.call_args.kwargs + assert recorded["model"] == "azure/gpt-4" + assert recorded["api_provider"] == "azure" + assert recorded["user"] == "alice" + + +@pytest.mark.asyncio +async def test_batch_retrieve_hook_does_not_record_created_metric(): + managed_files = _make_managed_files_instance() + prometheus_logger = MagicMock() + managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger) + retrieve_response = _make_batch_response(status="in_progress", output_file_id=None) + retrieve_response._hidden_params = { + "unified_batch_id": "some-unified-batch-id", + "model_id": "model-deploy-xyz", + "model_name": "azure/gpt-4", + } + + await managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None), + response=retrieve_response, + ) + + prometheus_logger.record_managed_batch_created.assert_not_called() + + @pytest.mark.asyncio async def test_batch_retrieve_hook_does_not_claim_attribution(): """A retrieve carries unified_batch_id but no unified_file_id, so it must not rewrite diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 2f6a5a3b0e0..a37c8ff2bb4 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -29,6 +29,7 @@ added to this layer raises instead of silently passing - the inventory of seams cannot drift without a test failure. """ +import base64 import json from contextlib import ExitStack from dataclasses import dataclass @@ -36,7 +37,7 @@ from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest - +from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles import litellm import litellm.proxy.batches_endpoints.endpoints as endpoints @@ -989,6 +990,67 @@ async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds): assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch" +def install_managed_files_hook(harness: Harness) -> AsyncMock: + prisma_client = AsyncMock() + managed_files = _PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client) + harness.logging.post_call_success_hook = AsyncMock(side_effect=managed_files.async_post_call_success_hook) + harness.router.model_list = [] + return prisma_client + + +TEAM_A_KEY = UserAPIKeyAuth(api_key="sk-team-a", user_id="user_a", team_id="team_a") + + +def assert_ownership_registered_for_team_a(prisma_client: AsyncMock, batch_id: str) -> None: + upsert = prisma_client.db.litellm_managedobjecttable.upsert + upsert.assert_awaited_once() + assert upsert.await_args.kwargs["where"] == {"unified_object_id": batch_id} + created = upsert.await_args.kwargs["data"]["create"] + assert created["created_by"] == "user_a" + assert created["team_id"] == "team_a" + prisma_client.db.litellm_managedobjecttable.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + {"input_file_id": AZURE_FILE_ID}, + {"input_file_id": "file-plain", "model": "vertex-model"}, + {"input_file_id": "file-plain"}, + ], + ids=["model_encoded_file_id", "model_param", "provider_fallback"], +) +async def test_create__registers_ownership_for_creator(harness, openai_env_creds, body): + set_body(harness, {**body, "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + prisma_client = install_managed_files_hook(harness) + + resp = await call_create(harness, user=TEAM_A_KEY) + + assert_ownership_registered_for_team_a(prisma_client, resp.id) + + +@pytest.mark.asyncio +async def test_create__unified_file_id_registers_ownership_for_creator(harness): + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,input-uuid;target_model_names,gpt-4o-mini" + ).decode() + set_body( + harness, + { + "input_file_id": unified_input_file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + prisma_client = install_managed_files_hook(harness) + + resp = await call_create(harness, user=TEAM_A_KEY) + + assert harness.router_acreate.call_count == 1 + assert_ownership_registered_for_team_a(prisma_client, resp.id) + + @pytest.mark.asyncio async def test_create__metadata_sanitized_before_forwarding(harness, openai_env_creds): set_body( From af3ddb477a852f20898aeefd7bc35713188f98da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:22:20 -0700 Subject: [PATCH 31/38] fix(realtime): release the budget reservation on a failed session and scrub relayed close details A refused or failed /v1/realtime session never ran the success cost callback or a failure hook, so its pre-call budget reservation stayed open and kept the key/team/user spend counters pinned above real spend, 429ing later requests on the same key until the counter's TTL expired. The endpoint now reconciles the reservation in a finally, reusing a shared release_or_invalidate_budget_reservation helper that mirrors the success/failure paths (release to zero, else invalidate the reserved counters and finalize). The relayed upstream close message and reason also go through the proxy's client-facing redaction, so a credential, internal hostname, private IP, or server path echoed by the upstream never reaches the client verbatim. --- .../litellm_core_utils/realtime_streaming.py | 6 +- litellm/proxy/proxy_server.py | 12 +++ .../spend_tracking/budget_reservation.py | 25 ++++++ .../test_realtime_streaming.py | 21 +++-- tests/test_litellm/proxy/test_proxy_server.py | 83 +++++++++++++++++++ 5 files changed, 137 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index b448cb7c9ff..c670278d3fb 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cas from typing_extensions import ReadOnly import litellm -from litellm._logging import _redact_string, verbose_logger +from litellm._logging import redact_internal_details_from_client_message, verbose_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.types.llms.openai import ( @@ -1567,8 +1567,8 @@ class RealTimeStreaming: await asyncio.gather(forward_task, client_task, return_exceptions=True) async def _close_client(self, close: BackendClose) -> None: - redacted_message: Final = _redact_string(close.message) - redacted_reason: Final = _redact_string(close.reason) + redacted_message: Final = redact_internal_details_from_client_message(close.message) + redacted_reason: Final = redact_internal_details_from_client_message(close.reason) try: if close.code != 1000: await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error")) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61bcdea94d4..06cfee45918 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11453,6 +11453,16 @@ def _realtime_query_params_template(model: str | None, intent: str | None) -> tu return tuple(params) +async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth) -> None: + from litellm.proxy.spend_tracking.budget_reservation import ( + release_or_invalidate_budget_reservation, + ) + + await release_or_invalidate_budget_reservation( + budget_reservation=user_api_key_dict.budget_reservation, + ) + + @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") @@ -11592,6 +11602,8 @@ async def realtime_websocket_endpoint( ) except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") + finally: + await _release_realtime_budget_reservation(user_api_key_dict) ###################################################################### diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 91d2ece7a51..2ee7320b82c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -373,6 +373,31 @@ async def invalidate_budget_reservation_counters( await _invalidate_spend_counter(counter_key=counter_key) +async def release_or_invalidate_budget_reservation( + budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict +) -> None: + """Reconcile a still-open reservation on a terminal path that settles no cost. + + A failed or upstream-refused request never runs the success cost callback, so + its pre-call reservation stays open and keeps the spend counter pinned above + real spend until the counter's TTL expires, 429ing later requests on the same + key. Release it to zero; if the release itself fails (e.g. the counter store is + unreachable) drop the reserved counters directly and mark the reservation + finalized so nothing reprocesses it. Idempotent: the finalized guard makes a + second call a no-op once success or failure handling already reconciled. + """ + if budget_reservation is None or budget_reservation.get("finalized") is True: + return + try: + await release_budget_reservation(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead + verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters") + try: + await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + finally: + budget_reservation["finalized"] = True + + async def _get_budget_counters( request_body: dict, valid_token: UserAPIKeyAuth, diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index bcfacf16205..5352c894f87 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3229,21 +3229,28 @@ async def test_bidirectional_forward_relays_upstream_policy_close_to_client(): client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) +@pytest.mark.parametrize( + "leaked_detail", + ( + pytest.param("sk-live-abcdef0123456789abcdef0123", id="credential"), + pytest.param("vertex-int.svc.cluster.local", id="internal-hostname"), + pytest.param("/etc/litellm/service-account.json", id="filesystem-path"), + ), +) @pytest.mark.asyncio -async def test_upstream_close_reason_with_a_secret_is_redacted_before_reaching_the_client(): - """LIT-6973: the relayed close mirrors the handshake path and scrubs credential - patterns, so an upstream error echoing a token never reaches the client verbatim.""" - secret: Final = "sk-live-abcdef0123456789abcdef0123" +async def test_upstream_close_details_are_scrubbed_before_reaching_the_client(leaked_detail: str): + """LIT-6973: the relayed close goes through the proxy's client-facing redaction, so an upstream + error echoing a credential, an internal host, or a server path never reaches the client verbatim.""" client_ws: Final = _client_ws_that_never_sends() - upstream_close: Final = ConnectionClosed(Close(1008, f"auth failed for {secret}"), None) + upstream_close: Final = ConnectionClosed(Close(1008, f"upstream rejected: {leaked_detail}"), None) session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) await session.run() (error_event,) = _error_events_sent_to(client_ws) - assert secret not in error_event["error"]["message"] + assert leaked_detail not in error_event["error"]["message"] relayed_reason: Final = client_ws.close.await_args.kwargs["reason"] - assert secret not in relayed_reason + assert leaked_detail not in relayed_reason assert "REDACTED" in relayed_reason diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d91928a203e..57e2cfc3332 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9521,6 +9521,89 @@ def test_realtime_websocket_route_aliases_registered(): ) +def _lit6973_fake_realtime_ws() -> MagicMock: + ws = MagicMock() + ws.headers = {} + ws.scope = {"headers": [], "type": "websocket"} + ws.url = "ws://testserver/v1/realtime" + ws.accept = AsyncMock() + ws.send_text = AsyncMock() + ws.close = AsyncMock() + return ws + + +async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: + """Drive realtime_websocket_endpoint through a session the upstream refused. + + route_request resolves normally because the relay handles the refusal + internally (sends the error event, closes the client), so neither the + success cost callback nor a failure hook runs on _ProxyDBLogger. The + endpoint itself must reconcile the pre-call budget reservation, so the + real release runs (entries is empty, so it touches no counter store) and + the caller asserts on the observable reservation state afterwards.""" + from litellm.proxy import proxy_server as ps + + user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") + user_api_key_dict.budget_reservation = reservation + + completed: Final = asyncio.get_running_loop().create_future() + completed.set_result(None) + + pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, MagicMock())) + can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally under test + pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state + route = patch.object(ps, "route_request", new=AsyncMock(return_value=completed)) # test-quality-ok: fakes the relay that already handled the refusal so the session returns normally + with can_call, pre, route: + await ps.realtime_websocket_endpoint( + websocket=_lit6973_fake_realtime_ws(), + model="vertex_ai/gemini-live-2.5-flash", + intent=None, + guardrails=None, + user_api_key_dict=user_api_key_dict, + ) + + +@pytest.mark.asyncio +async def test_refused_realtime_session_releases_the_budget_reservation(): + """LIT-6973: reclassifying a refused realtime session as a failure removed the + success-path reservation release, so the pre-call reservation stayed open and + pinned the key/team/user spend counters, locking the key after a couple of + refusals. The endpoint must reconcile it: the reservation ends up finalized.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + await _lit6973_drive_refused_realtime_session(reservation) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): + """If releasing the reservation itself fails (e.g. the counter store is down), + the reserved counters must be invalidated directly so the estimate does not + stay pinned, and the reservation is finalized so nothing reprocesses it.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy.spend_tracking import budget_reservation as br + + reservation: Final = { + "reserved_cost": 0.55, + "input_cost": 0.0, + "finalized": False, + "entries": [{"counter_key": "spend:key:hashed-token"}], + } + invalidated: Final[list[str]] = [] + + async def _record(counter_key: str) -> None: + invalidated.append(counter_key) + + failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated + sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable + with failing_release, sink: + await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) + + assert invalidated == ["spend:key:hashed-token"] + assert reservation["finalized"] is True + + class TestTransformRequestBannedParams: """ /utils/transform_request applies the same banned-param check as LLM endpoints. From 1fe87e8e25206c039be5e87a5808a30f47cc3183 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:11:24 -0700 Subject: [PATCH 32/38] fix(realtime): settle the budget reservation only for sessions the success log does not own The blanket finally release from the previous commit also zeroed the reservation of successful sessions. Success settlement is enqueued on the logging worker, not awaited, so the endpoint's finally ran first and released the reservation the cost callback still had to reconcile, dropping the real spend from the key/team/user counters. The relay now stamps a synchronous marker (REALTIME_SESSION_SUCCESS_LOGGED_KEY) on the shared logging object at the single success-dispatch site, and the endpoint releases the reservation only when that marker is absent. Refused or failed sessions, which never log success, still release; successful sessions leave the reservation for the cost callback to settle to actual spend. Exactly one settler touches each reservation, so the idempotent reconcile never double-adjusts. --- .../litellm_core_utils/realtime_streaming.py | 4 ++ litellm/proxy/proxy_server.py | 7 ++- .../test_realtime_streaming.py | 32 +++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 55 +++++++++++++------ 4 files changed, 80 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c670278d3fb..75046f2cf87 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -35,6 +35,9 @@ else: CLIENT_CONNECTION_CLASS = Any +REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" + + @dataclass(frozen=True, slots=True) class BackendClose: code: int @@ -421,6 +424,7 @@ class RealTimeStreaming: self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) ) + self.logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True async def _send_to_backend(self, message: str) -> bool: """Send a message to the backend WebSocket. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 06cfee45918..7d59dfa86c4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11603,7 +11603,12 @@ async def realtime_websocket_endpoint( except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") finally: - await _release_realtime_budget_reservation(user_api_key_dict) + from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + ) + + if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): + await _release_realtime_budget_reservation(user_api_key_dict) ###################################################################### diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 5352c894f87..9c0f6f59463 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -14,6 +14,7 @@ import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, RealTimeStreaming, client_sent_openai_beta_realtime_header, ) @@ -3380,3 +3381,34 @@ async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the assert session.logging.logged_sessions == ((),) assert session.logging.logged_failures == () client_ws.close.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_success_logging_stamps_the_reservation_ownership_marker(): + """LIT-6973: only the success path enqueues the cost callback that settles the + session's budget reservation, so it stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on + the shared logging object. The proxy endpoint reads that stamp to decide whether to + release the reservation itself, so a logged-as-success session must carry it.""" + client_ws: Final = _client_ws_that_never_sends() + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + upstream_close: Final = ConnectionClosed(Close(1000, ""), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close)) + + await session.run() + + assert session.logging.logged_sessions != () + assert session.logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + + +@pytest.mark.asyncio +async def test_refused_session_does_not_stamp_the_reservation_ownership_marker(): + """A refused session logs a failure, not a success, so it must not stamp + REALTIME_SESSION_SUCCESS_LOGGED_KEY. If it did, the proxy endpoint would skip its + own reservation release and the refused session's reservation would stay pinned.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 57e2cfc3332..697d4c182c7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9532,27 +9532,34 @@ def _lit6973_fake_realtime_ws() -> MagicMock: return ws -async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: - """Drive realtime_websocket_endpoint through a session the upstream refused. +async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_success: bool) -> None: + """Drive realtime_websocket_endpoint to just before its budget-reservation finally. - route_request resolves normally because the relay handles the refusal - internally (sends the error event, closes the client), so neither the - success cost callback nor a failure hook runs on _ProxyDBLogger. The - endpoint itself must reconcile the pre-call budget reservation, so the - real release runs (entries is empty, so it touches no counter store) and - the caller asserts on the observable reservation state afterwards.""" + route_request resolves normally in both cases: the relay owns the session + once route_request returns. A successful session enqueues its success cost + callback and stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on the shared logging + object; a refused one does neither. The endpoint keys its reservation cleanup + off that stamp, so backend_logged_success reproduces both branches. The fake + logging object carries a real model_call_details dict so the stamp is + observable, and the reservation has empty entries so the real release touches + no counter store.""" + from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.proxy import proxy_server as ps user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") user_api_key_dict.budget_reservation = reservation - completed: Final = asyncio.get_running_loop().create_future() - completed.set_result(None) + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} - pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, MagicMock())) + async def fake_llm_call() -> None: + if backend_logged_success: + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + + pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)) can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally under test pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state - route = patch.object(ps, "route_request", new=AsyncMock(return_value=completed)) # test-quality-ok: fakes the relay that already handled the refusal so the session returns normally + route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object with can_call, pre, route: await ps.realtime_websocket_endpoint( websocket=_lit6973_fake_realtime_ws(), @@ -9565,17 +9572,31 @@ async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: @pytest.mark.asyncio async def test_refused_realtime_session_releases_the_budget_reservation(): - """LIT-6973: reclassifying a refused realtime session as a failure removed the - success-path reservation release, so the pre-call reservation stayed open and - pinned the key/team/user spend counters, locking the key after a couple of - refusals. The endpoint must reconcile it: the reservation ends up finalized.""" + """LIT-6973: a refused realtime session enqueues no success cost callback, so + the pre-call reservation would stay open and pin the key/team/user spend + counters, locking the key after a couple of refusals. The endpoint sees no + success stamp and reconciles it: the reservation ends up finalized.""" reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} - await _lit6973_drive_refused_realtime_session(reservation) + await _lit6973_drive_realtime_session(reservation, backend_logged_success=False) assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): + """A billable realtime session settles its reservation through the enqueued + success cost callback, not the endpoint. The endpoint must not finalize it in + its finally, or it would reconcile the reservation to zero before the cost + callback applies real spend, so billable sessions stop counting against budget. + With the success stamp present, the endpoint leaves the reservation untouched.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + await _lit6973_drive_realtime_session(reservation, backend_logged_success=True) + + assert reservation["finalized"] is False + + @pytest.mark.asyncio async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): """If releasing the reservation itself fails (e.g. the counter store is down), From aca1c54391ceef578c31603fcd7872b267b451ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:26:36 -0700 Subject: [PATCH 33/38] refactor(proxy): build the OpenAI websocket refusal frame from a TypedDict The two dict literals behind the refusal event counted against the LIT002 ceiling once the base branch used up its headroom, so the frame is now a ReadOnly TypedDict built in one shot. Importing Literal explicitly also makes the UP037 suppression on the Vertex discovery signature unnecessary, so it goes. --- .../llm_passthrough_endpoints.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index e92c949299c..32da2658b99 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -16,12 +16,13 @@ import re from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket from fastapi.responses import StreamingResponse from starlette.websockets import WebSocketState +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider @@ -1775,7 +1776,7 @@ def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str] def get_vertex_pass_through_handler( - call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here + call_type: Literal["discovery", "aiplatform"], ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2352,6 +2353,16 @@ class _OpenAIWebsocketRefusal: message: str +class _OpenAIWebsocketErrorDetail(TypedDict): + type: ReadOnly[Literal["invalid_request_error"]] + message: ReadOnly[str] + + +class _OpenAIWebsocketErrorFrame(TypedDict): + type: ReadOnly[Literal["error"]] + error: ReadOnly[_OpenAIWebsocketErrorDetail] + + _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( close_reason="OpenAI websocket passthrough is disabled", message=( @@ -2451,14 +2462,11 @@ async def openai_websocket_proxy_route( refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) if refusal is not None: await websocket.accept(subprotocol=negotiated_subprotocol) - await websocket.send_text( - json.dumps( - { - "type": "error", - "error": {"type": "invalid_request_error", "message": refusal.message}, - } - ) - ) + error_frame: Final[_OpenAIWebsocketErrorFrame] = { + "type": "error", + "error": {"type": "invalid_request_error", "message": refusal.message}, + } + await websocket.send_text(json.dumps(error_frame)) await websocket.close(code=1008, reason=refusal.close_reason) return From 5a35e6d41f76d2b09a258b3e2b051e3e7e745c79 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:18:33 -0700 Subject: [PATCH 34/38] fix(realtime): release the budget reservation when a session is rejected before the relay starts The three pre-relay exits of realtime_websocket_endpoint (missing model, key/model access denied, pre-call rejection such as a rate limit or a guardrail) returned before the finally that releases the auth-time budget reservation, so a rejected session pinned the key at the reserved amount until the counter TTL expired and its next requests got budget_exceeded while /key/info showed spend 0. A single _reject_realtime_session helper now releases the reservation before sending the error event and closing, and release_or_invalidate_budget_reservation shields the release from a second cancellation and logs, rather than raises, a failing invalidate fallback so it can never mask the session's own outcome. --- litellm/proxy/proxy_server.py | 43 ++++++---- .../spend_tracking/budget_reservation.py | 4 +- tests/test_litellm/proxy/test_proxy_server.py | 79 +++++++++++++++++-- 3 files changed, 103 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d59dfa86c4..65fe3ede822 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11463,6 +11463,25 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth ) +async def _reject_realtime_session( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, + *, + code: int, + reason: str, + error_message: str | None = None, +) -> None: + await _release_realtime_budget_reservation(user_api_key_dict) + if error_message is not None: + try: + await websocket.send_text( + json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) + ) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") + await websocket.close(code=code, reason=reason) + + @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") @@ -11488,7 +11507,9 @@ async def realtime_websocket_endpoint( if intent == "transcription": route_model = "gpt-realtime-whisper" else: - await websocket.close(code=1008, reason="model query parameter is required") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1008, reason="model query parameter is required" + ) return assert route_model is not None try: @@ -11499,7 +11520,7 @@ async def realtime_websocket_endpoint( llm_router=llm_router, ) except ProxyException as e: - await websocket.close(code=1008, reason=e.message[:120]) + await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120]) return await websocket.accept(**accept_kwargs) @@ -11558,21 +11579,9 @@ async def realtime_websocket_endpoint( ) except Exception as e: verbose_proxy_logger.exception("Realtime pre-call error") - try: - await websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_error", - "message": str(e), - }, - } - ) - ) - except Exception: - pass - await websocket.close(code=1011, reason="Pre-call error") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) + ) return # Phase 2: route to upstream LLM. diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 2ee7320b82c..ed2bc87597c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -389,11 +389,13 @@ async def release_or_invalidate_budget_reservation( if budget_reservation is None or budget_reservation.get("finalized") is True: return try: - await release_budget_reservation(budget_reservation=budget_reservation) + await asyncio.shield(release_budget_reservation(budget_reservation=budget_reservation)) except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters") try: await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed + verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after release failed") finally: budget_reservation["finalized"] = True diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 697d4c182c7..8b151d9e1f6 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9532,8 +9532,15 @@ def _lit6973_fake_realtime_ws() -> MagicMock: return ws -async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_success: bool) -> None: - """Drive realtime_websocket_endpoint to just before its budget-reservation finally. +async def _lit6973_drive_realtime_session( + reservation: dict, *, backend_logged_success: bool, phase_one_exit: str | None = None +) -> MagicMock: + """Drive realtime_websocket_endpoint through one of its reservation-settling exits. + + phase_one_exit picks a rejection before the relay: "model_access" makes the + key/model check raise ProxyException, "pre_call" makes pre-call processing + (rate limits, guardrails) raise. Neither reaches route_request, so no success + log can own the reservation and the endpoint has to release it on that exit. route_request resolves normally in both cases: the relay owns the session once route_request returns. A successful session enqueues its success cost @@ -9556,18 +9563,30 @@ async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_s if backend_logged_success: logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True - pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)) - can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally under test + from litellm.proxy._types import ProxyException + + model_access_error: Final = ( + ProxyException(message="key cannot access model", type="auth_error", param="model", code=401) + if phase_one_exit == "model_access" + else None + ) + pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None + pre_call: Final = AsyncMock( + side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) + ) + ws: Final = _lit6973_fake_realtime_ws() + can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object with can_call, pre, route: await ps.realtime_websocket_endpoint( - websocket=_lit6973_fake_realtime_ws(), + websocket=ws, model="vertex_ai/gemini-live-2.5-flash", intent=None, guardrails=None, user_api_key_dict=user_api_key_dict, ) + return ws @pytest.mark.asyncio @@ -9583,6 +9602,39 @@ async def test_refused_realtime_session_releases_the_budget_reservation(): assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_realtime_session_rejected_in_pre_call_releases_the_budget_reservation(): + """A rate-limit or guardrail rejection happens before route_request, so the + relay never runs and no success log can own the reservation. The endpoint + must release it on that exit too, or the key stays pinned at the reserved + amount and its next requests 429 with budget_exceeded while /key/info shows + spend 0 (reproduced live with rpm_limit=1). The client still gets the + pre-call error event and the 1011 close it got before.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call" + ) + + assert reservation["finalized"] is True + assert json.loads(ws.send_text.await_args.args[0])["error"]["message"] == "Rate limit exceeded" + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + + +@pytest.mark.asyncio +async def test_realtime_session_denied_model_access_releases_the_budget_reservation(): + """The key/model access check rejects before the socket is even accepted; + that exit skipped the release as well, pinning the reservation.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access" + ) + + assert reservation["finalized"] is True + ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model") + + @pytest.mark.asyncio async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): """A billable realtime session settles its reservation through the enqueued @@ -9625,6 +9677,23 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback_fails(): + """Both counter-store calls failing must not raise out of the realtime + endpoint's finally (it would mask the session's own outcome) and must still + stamp finalized so nothing retries the same reservation.""" + from litellm.proxy.spend_tracking import budget_reservation as br + + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch + failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail + + with failing_release, failing_invalidate: + await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) + + assert reservation["finalized"] is True + + class TestTransformRequestBannedParams: """ /utils/transform_request applies the same banned-param check as LLM endpoints. From 37722eba68149c5f3e59ed0f3c12798a84aa1bc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:35:26 -0700 Subject: [PATCH 35/38] fix(realtime): close a rejected client before releasing its budget reservation A slow or unreachable counter store made a pre-relay rejection wait behind the reservation release before the client saw the error event and the close. Close first and release in finally, mirroring the relay's own failure path, so a client that already hung up still gets its reservation released. --- litellm/proxy/proxy_server.py | 20 ++++---- tests/test_litellm/proxy/test_proxy_server.py | 47 ++++++++++++++++++- 2 files changed, 56 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 65fe3ede822..a5c60d8e976 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11471,15 +11471,17 @@ async def _reject_realtime_session( reason: str, error_message: str | None = None, ) -> None: - await _release_realtime_budget_reservation(user_api_key_dict) - if error_message is not None: - try: - await websocket.send_text( - json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) - ) - except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below - verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") - await websocket.close(code=code, reason=reason) + try: + if error_message is not None: + try: + await websocket.send_text( + json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) + ) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") + await websocket.close(code=code, reason=reason) + finally: + await _release_realtime_budget_reservation(user_api_key_dict) @app.websocket("/openai/v1/realtime") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8b151d9e1f6..4d70c9d436f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9533,7 +9533,11 @@ def _lit6973_fake_realtime_ws() -> MagicMock: async def _lit6973_drive_realtime_session( - reservation: dict, *, backend_logged_success: bool, phase_one_exit: str | None = None + reservation: dict, + *, + backend_logged_success: bool, + phase_one_exit: str | None = None, + websocket: MagicMock | None = None, ) -> MagicMock: """Drive realtime_websocket_endpoint through one of its reservation-settling exits. @@ -9574,7 +9578,7 @@ async def _lit6973_drive_realtime_session( pre_call: Final = AsyncMock( side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) ) - ws: Final = _lit6973_fake_realtime_ws() + ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws() can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object @@ -9635,6 +9639,45 @@ async def test_realtime_session_denied_model_access_releases_the_budget_reservat ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model") +@pytest.mark.asyncio +async def test_rejected_realtime_session_closes_the_client_before_releasing_the_reservation(): + """The counter release can block on a slow or unreachable store, and a + rejected client must not sit behind it: the relay's own failure path closes + the client first and releases in its finally, so the pre-relay rejection + has to close first as well. The fake close checks the reservation is still + open when the client is closed.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + + async def close_while_reservation_is_still_open(**_: object) -> None: + assert reservation["finalized"] is False, "client was closed only after the reservation release" + + ws.close = AsyncMock(side_effect=close_while_reservation_is_still_open) + + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call", websocket=ws + ) + + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_rejected_realtime_session_releases_the_reservation_when_the_client_is_already_gone(): + """A client that hung up before the rejection makes the close raise; the + reservation must still be released, or the key stays pinned.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + ws.close = AsyncMock(side_effect=RuntimeError("client already disconnected")) + + with pytest.raises(RuntimeError, match="client already disconnected"): + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access", websocket=ws + ) + + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): """A billable realtime session settles its reservation through the enqueued From 03da725ee4de2414056765f1968794e4c0634ce2 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:30:42 -0700 Subject: [PATCH 36/38] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/integrations/test_shadow_eval_logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index d67e957ca16..371f6f75a05 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -159,7 +159,7 @@ def _reasoning_judge_router( if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} return {"choices": [{"message": {"content": "shadow answer"}}]} - budget_for_the_answer = kwargs["max_tokens"] - reasoning_tokens + budget_for_the_answer: Final = kwargs["max_tokens"] - reasoning_tokens return {"choices": [{"message": {"content": verdict[: max(0, budget_for_the_answer)]}}]} router.acompletion = MagicMock(side_effect=acompletion) From 0b3687ec56153225d7b8f2a0c2652bf2f589ce2e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:49:00 -0700 Subject: [PATCH 37/38] fix(shadow_eval): import Final for the test helper's annotation --- tests/test_litellm/integrations/test_shadow_eval_logger.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 371f6f75a05..eecd876219e 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -3,6 +3,7 @@ the detached pipeline's single attempt-row write, and the cache-first job lookup import asyncio from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest From d0d09e53438d51b25cb0e0f8a29a329e8d93a7e9 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 09:51:23 -0700 Subject: [PATCH 38/38] feat(router): meter auto-router tier and prompt customization against the auto_router license feature (#39674) Generalizes the heuristic_v2 ceiling from #39468 into a capability table whose records own their in-process predicate, SQL spelling and refusal wording. The existing heuristic_v2 capability keeps its own one-router ceiling. A single customization capability combines operator-defined tier definitions with every operator-written part of the classifier prompt. The prompt half only applies to classifier types that call an LLM. The shipped default prompt, classification rubric presets, tier-label renames and tier model choices remain ungated. Scope every enforcement point to actual complexity routers. A model-less PATCH or legacy update now decrypts the stored model before accepting strategy-router settings, so a regular model cannot acquire a router config or spend a license slot. Under the existing advisory lock, the cross-pod candidate query returns only model scalars and the count decrypts and classifies them in process; old non-router rows carrying a capability-shaped config no longer block a real complexity router. The signed auto_router license feature makes both ceilings unlimited. --- litellm/constants.py | 2 +- litellm/proxy/auth/litellm_license.py | 11 +- .../model_management_endpoints.py | 157 +++++++--- litellm/proxy/proxy_server.py | 34 +- litellm/router.py | 45 +-- .../router_utils/auto_router_model_naming.py | 134 +++++++- litellm/types/router.py | 4 +- .../proxy/auth/test_litellm_license.py | 18 +- .../test_model_management_endpoints.py | 292 +++++++++++++++--- .../proxy/proxy_server/test_proxy_config.py | 91 +++++- .../router_strategy/test_complexity_router.py | 232 +++++++++++++- .../test_auto_router_model_naming.py | 172 +++++++++-- 12 files changed, 987 insertions(+), 205 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index da731cb5eb2..7d6de612349 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -40,7 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", - "heuristic_v2_router_limit", + "auto_router_capability_limit", } ) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index 55bb1e3925a..067ac7905c5 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router" -HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." +AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." class LicenseCheck: @@ -153,11 +153,12 @@ class LicenseCheck: return False return team_count > _max_teams_in_license - def heuristic_v2_router_limit(self) -> int | None: + def auto_router_capability_limit(self) -> int | None: """ - How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the - signed license lists the auto_router feature, otherwise one. A license verified through - the API carries no feature list, so it does not lift the limit either. + How many auto-routers may claim each licensed capability (heuristic_v2, operator-defined + tier_definitions): unlimited (None) only when the signed license lists the auto_router + feature, otherwise one per capability. A license verified through the API carries no + feature list, so it does not lift the limit either. """ if self.airgapped_license_data is None: return 1 diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index d4e03a05c52..b77108911aa 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -50,7 +50,7 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UserAPIKeyAuth, ) -from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY +from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -98,11 +98,13 @@ from litellm.router_strategy.complexity_router import ( normalize_classification_prompt, ) from litellm.router_utils.auto_router_model_naming import ( + GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, + capability_limit_violation, carries_complexity_router_settings, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, - uses_heuristic_v2_classifier, + count_capability_routers, + gated_capability_of, + is_complexity_router_model, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -237,11 +239,13 @@ def _strategy_router_write_violation( An auto-router deployment's ``litellm_params.model`` (``auto_router/...``) is the discriminator the router loads it by; a write that mangles it makes the router drop the deployment silently under ``ignore_invalid_deployments``. - Only writes that supply ``litellm_params.model`` are judged on the naming - contract, against the merged (stored + incoming) params, so partial patches - and restores of an already-corrupted row stay legal. A config is judged only - when the write carries one, for the same reason: a rename must not be held - hostage by a stored config it does not touch. Returns the violation, or None. + A patch adding auto-router settings is judged against the effective model, + decrypting the stored model when the patch omits it, so a regular deployment + cannot claim a strategy-router configuration. Unrelated partial patches and + restores that do not touch strategy-router settings stay legal. A config is + judged only when the write carries one, for the same reason: a rename must + not be held hostage by a stored config it does not touch. Returns the + violation, or None. """ if incoming_params is None: return None @@ -256,14 +260,18 @@ def _strategy_router_write_violation( for source in (incoming_params, existing_params) if source is not None and getattr(source, field, None) is not None ) - # Scope reads the incoming model because the stored one is encrypted at rest. - if carries_complexity_router_settings(incoming_params.model, present_fields): + effective_params: Final = _effective_complexity_router_params(incoming_params, existing_params) + effective_model: Final = effective_params.get("model") + if carries_complexity_router_settings( + effective_model if isinstance(effective_model, str) else None, present_fields + ): placement_violation: Final = validate_complexity_router_config_placement(incoming_params.model_extra) if placement_violation is not None: return placement_violation - if incoming_params.model is None: - return None - return validate_strategy_router_model_write(model=incoming_params.model, present_fields=present_fields) + return validate_strategy_router_model_write( + model=effective_model if isinstance(effective_model, str) else "", + present_fields=present_fields, + ) def _raise_on_strategy_router_write_violation( @@ -281,14 +289,23 @@ def _raise_on_strategy_router_write_violation( ) -HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301 -_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" -_HEURISTIC_V2_DB_ROWS_SQL: Final = """ -SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable" +AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 +_CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_STORED_LITELLM_PARAMS_SQL: Final = ( + "(CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)" +) +_STORED_COMPLEXITY_CONFIG_SQL: Final = f"{_STORED_LITELLM_PARAMS_SQL} -> 'complexity_router_config'" +_CAPABILITY_DB_ROWS_SQL: Final[Mapping[str, str]] = MappingProxyType( + { + capability.key: f""" +SELECT {_STORED_LITELLM_PARAMS_SQL} ->> 'model' AS model +FROM "LiteLLM_ProxyModelTable" WHERE model_id <> $1 - AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END) - -> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2' + AND ({capability.sql_config_predicate.format(config=_STORED_COMPLEXITY_CONFIG_SQL)}) """ + for capability in GATED_AUTO_ROUTER_CAPABILITIES + } +) def _effective_complexity_router_config( @@ -301,13 +318,44 @@ def _effective_complexity_router_config( return existing_params.complexity_router_config -@asynccontextmanager -async def _heuristic_v2_slot( - prisma_client: PrismaClient, *, effective_config: object, model_id: str | None -) -> AsyncGenerator[_ProxyModelTable, None]: - """Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled. +def _effective_model( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> str | None: + """The model a write leaves on the row, decrypting an existing value only when the patch omits it.""" + incoming: Final = None if incoming_params is None else incoming_params.model + if incoming is not None: + return incoming + existing: Final = None if existing_params is None else existing_params.model + if existing is None: + return None + decrypted: Final = decrypt_value_helper( + value=existing, + key="model", + exception_type="debug", + return_original_value=True, + ) + return decrypted if isinstance(decrypted, str) else None - A write that leaves the row on classifier_type heuristic_v2 under a limited license runs + +def _effective_complexity_router_params( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> Mapping[str, object]: + """The model and complexity config a write leaves, for placement and capability decisions.""" + return MappingProxyType( + { + "model": _effective_model(incoming_params, existing_params), + "complexity_router_config": _effective_complexity_router_config(incoming_params, existing_params), + } + ) + + +@asynccontextmanager +async def _auto_router_capability_slot( + prisma_client: PrismaClient, *, effective_params: Mapping[str, object], model_id: str | None +) -> AsyncGenerator[_ProxyModelTable, None]: + """Hand out the model table to write through while the row's claim on a licensed capability is settled. + + A write that leaves the row claiming a licensed capability under a limited license runs inside one transaction that takes an advisory lock in its own statement before counting (a statement's snapshot predates anything it locks), so pods cannot both pass the count: the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged @@ -321,21 +369,37 @@ async def _heuristic_v2_slot( """ from litellm.proxy.proxy_server import _license_check, llm_router - limit: Final = _license_check.heuristic_v2_router_limit() - if limit is None or not uses_heuristic_v2_classifier(effective_config): + limit: Final = _license_check.auto_router_capability_limit() + capability: Final = gated_capability_of(effective_params) + if limit is None or capability is None: yield _proxy_model_table(prisma_client) return async with prisma_client.db.tx() as tx_ctx: tables: Final[_TxModelTables] = tx_ctx - await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY) - rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "") - db_held: Final = rows[0].get("held") if rows else 0 + await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw( + _CAPABILITY_DB_ROWS_SQL[capability.key], model_id or "" + ) + db_held: Final = sum( + 1 + for row in rows + for stored_model in (row.get("model"),) + if isinstance(stored_model, str) + and is_complexity_router_model( + decrypt_value_helper( + value=stored_model, + key="model", + exception_type="debug", + return_original_value=True, + ) + ) + ) config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) - held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows) - violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit) + held: Final = db_held + count_capability_routers(config_rows, capability=capability) + violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit) if violation is not None: raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}" + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" ) yield tables.litellm_proxymodeltable await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") @@ -791,6 +855,9 @@ async def patch_model( existing_params=db_model.litellm_params, ) + effective_params: Final = _effective_complexity_router_params( + patch_data.litellm_params, db_model.litellm_params + ) requested_model_name: Final = patch_data.model_name stored_model_name: str | None = None @@ -799,11 +866,9 @@ async def patch_model( stored_model_name = update_data.get("model_name") update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name update_data["updated_at"] = cast(str, get_utc_datetime()) - async with _heuristic_v2_slot( + async with _auto_router_capability_slot( prisma_client, - effective_config=_effective_complexity_router_config( - patch_data.litellm_params, db_model.litellm_params - ), + effective_params=effective_params, model_id=model_id, ) as table: return await table.update(where={"model_id": model_id}, data=update_data) @@ -1959,9 +2024,12 @@ async def add_new_model( model_params=priced_model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, - slot=_heuristic_v2_slot( + slot=_auto_router_capability_slot( prisma_client, - effective_config=priced_model_params.litellm_params.complexity_router_config, + effective_params=_effective_complexity_router_params( + priced_model_params.litellm_params, + None, + ), model_id=priced_model_params.model_info.id, ), ) @@ -2110,6 +2178,9 @@ async def update_model( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, ) + effective_params: Final = _effective_complexity_router_params( + model_params.litellm_params, deployment.litellm_params + ) # update DB if store_model_in_db is True: @@ -2147,11 +2218,9 @@ async def update_model( "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, **({} if renamed_to is None else {"model_name": renamed_to}), } - async with _heuristic_v2_slot( + async with _auto_router_capability_slot( prisma_client, - effective_config=_effective_complexity_router_config( - model_params.litellm_params, deployment.litellm_params - ), + effective_params=effective_params, model_id=_model_id, ) as table: model_response: Final = await table.update( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0f0abd0eeae..88e3f79ca52 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -118,10 +118,11 @@ from litellm.router_utils.add_retry_fallback_headers import ( get_hidden_params_dict, ) from litellm.router_utils.auto_router_model_naming import ( + GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, + capability_limit_violation, carries_complexity_router_settings, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, + count_capability_routers, validate_complexity_router_config_placement, ) from litellm.types.utils import ( @@ -303,7 +304,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck +from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, @@ -4340,17 +4341,28 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object]) raise ValueError(f"model {model.get('model_name', '')!r}: {violation}") -def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: +def validate_auto_router_capability_limits(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: """ - Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows. + Refuse to start when config.yaml defines more auto-routers claiming a licensed capability than allowed. Checked here rather than left to router registration for the same reason as the two validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so the router's own refusal would turn the extra router into a silently missing model. """ - violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit) - if violation is not None: - raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}") + violations: Final = tuple( + message + for capability in GATED_AUTO_ROUTER_CAPABILITIES + if ( + message := capability_limit_violation( + capability=capability, + held=count_capability_routers(model_list, capability=capability), + limit=limit, + ) + ) + is not None + ) + if violations: + raise ValueError(f"config.yaml model_list: {' '.join(violations)} {AUTO_ROUTER_LICENSE_REMEDY}") def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place @@ -5758,7 +5770,7 @@ class ProxyConfig: model_list: Final = config.get("model_list", None) if model_list: router_params["model_list"] = model_list - validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit()) + validate_auto_router_capability_limits(model_list, limit=_license_check.auto_router_capability_limit()) print( # noqa: T201 "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" ) @@ -5848,7 +5860,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, - heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, + auto_router_capability_limit=_license_check.auto_router_capability_limit, ) if redis_usage_cache is not None and router.cache.redis_cache is None: @@ -6309,7 +6321,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, - heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, + auto_router_capability_limit=_license_check.auto_router_capability_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) else: diff --git a/litellm/router.py b/litellm/router.py index 6c7611c6236..6943eece90f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -116,10 +116,11 @@ from litellm.router_utils.add_retry_fallback_headers import ( ) from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, + GatedAutoRouterCapability, + capability_limit_violation, + claimed_capability, classify_strategy_router_model, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, - uses_heuristic_v2_classifier, + count_capability_routers, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -208,6 +209,7 @@ from litellm.types.router import ( AlertingConfig, AllowedFailsPolicy, AssistantsTypedDict, + AutoRouterCapabilityLimit, ConsumedRequestTagsStamp, CredentialLiteLLMParams, CustomRoutingStrategyBase, @@ -215,7 +217,6 @@ from litellm.types.router import ( DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, - HeuristicV2RouterLimit, LiteLLM_Params, MockRouterTestingParams, ModelGroupInfo, @@ -692,7 +693,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, - heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None, + auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -769,7 +770,7 @@ class Router: self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments - self.heuristic_v2_router_limit = heuristic_v2_router_limit + self.auto_router_capability_limit = auto_router_capability_limit self.fallback_access_check: Final = fallback_access_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks @@ -8811,20 +8812,21 @@ class Router: if not (isinstance(model_info, Mapping) and model_info.get("db_model")): yield deployment - def heuristic_v2_router_limit_violation(self) -> str | None: + def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None: """ - Why one more heuristic_v2 router cannot join this router, or None when it can. + Why one more router claiming ``capability`` cannot join this router, or None when it can. Judged against every deployment currently on the model_list; an upsert pops the row being - edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is - resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which - is the SDK default, and the proxy injects a resolver backed by its license. + edited first, so an edit of an existing gated router keeps its own slot. The limit is + resolved on every call through ``auto_router_capability_limit``; unset means unlimited, + which is the SDK default, and the proxy injects a resolver backed by its license. """ - limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None - others: Final = count_heuristic_v2_routers( - deployment for deployment in self.model_list if isinstance(deployment, Mapping) + limit: Final = self.auto_router_capability_limit() if self.auto_router_capability_limit is not None else None + others: Final = count_capability_routers( + (deployment for deployment in self.model_list if isinstance(deployment, Mapping)), + capability=capability, ) - return heuristic_v2_limit_violation(held=others + 1, limit=limit) + return capability_limit_violation(capability=capability, held=others + 1, limit=limit) def init_complexity_router_deployment(self, deployment: Deployment): """ @@ -8843,8 +8845,9 @@ class Router: ) complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config - if uses_heuristic_v2_classifier(complexity_router_config): - limit_violation: Final = self.heuristic_v2_router_limit_violation() + capability: Final = claimed_capability(complexity_router_config) + if capability is not None: + limit_violation: Final = self.auto_router_capability_violation(capability) if limit_violation is not None: raise ValueError(limit_violation) @@ -9674,13 +9677,13 @@ class Router: """Put a deployment back the way it was before a failed upsert popped it. A rollback re-admits state that was already serving, so it does not go through the - heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first + capability ceiling a newcomer gets: with the ceiling tightened since the deployment first registered, judging the rollback would drop a serving router over an unrelated failed edit. """ if previous_deployment is None or self.has_model_id(model_id): return - limit_resolver: Final = self.heuristic_v2_router_limit - self.heuristic_v2_router_limit = None + limit_resolver: Final = self.auto_router_capability_limit + self.auto_router_capability_limit = None try: self.add_deployment(deployment=previous_deployment) verbose_router_logger.info( @@ -9696,7 +9699,7 @@ class Router: restore_error, ) finally: - self.heuristic_v2_router_limit = limit_resolver + self.auto_router_capability_limit = limit_resolver @staticmethod def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index 2efbfb5782e..190c4921d5f 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under ``ignore_invalid_deployments``. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -81,6 +81,11 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None: return "semantic" +def is_complexity_router_model(model: str | None) -> bool: + """Whether ``model`` selects the complexity-router implementation.""" + return classify_strategy_router_model(model or "") == "complexity" + + def _named(value: object, role: StrategyRouterDependencyRole) -> tuple[StrategyRouterDependency, ...]: """One dependency from a scalar field, or none when it is absent or not a name.""" return (StrategyRouterDependency(value, role),) if isinstance(value, str) and value else () @@ -168,20 +173,121 @@ def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool: return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2" -def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool: - """Whether this deployment is a complexity router that classifies with heuristic_v2.""" - return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and ( - uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config")) +def defines_custom_tiers(complexity_router_config: object) -> bool: + """Whether this complexity config replaces the built-in tier ladder with operator-defined tier_definitions. + + Mirrors the SQL spelling on the capability record: only an actual array claims the capability, + so an explicit JSON null or a malformed value does not. + """ + return isinstance(_mapping(complexity_router_config).get("tier_definitions"), (list, tuple)) + + +OPERATOR_CLASSIFIER_PROMPT_FIELDS: Final = ("classification_prompt", "classification_examples") + + +def defines_custom_classifier_prompt(complexity_router_config: object) -> bool: + """Whether an operator wrote any part of this router's classifier prompt themselves. + + Three spellings, all metered: a whole replacement prompt (``classifier_llm_config.system_prompt``), + replacement opening instructions (``classification_prompt``), and replacement calibration examples + (``classification_examples``). Choosing a shipped ``classification_rubric`` preset is not authoring. + Scoped to the classifier types that actually call an LLM, which is also where the config validator + accepts these fields: the heuristic scorers never read them. + """ + config: Final = _mapping(complexity_router_config) + if config.get("classifier_type") not in LLM_CLASSIFIER_TYPES: + return False + return _mapping(config.get("classifier_llm_config")).get("system_prompt") is not None or any( + config.get(field) is not None for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS ) -def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int: - """How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers.""" - return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params")))) +def uses_custom_tier_or_classifier_prompt(complexity_router_config: object) -> bool: + """Whether this router replaces shipped tiers or its shipped classifier prompt.""" + return defines_custom_tiers(complexity_router_config) or defines_custom_classifier_prompt(complexity_router_config) -def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: - """Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits. +_LLM_CLASSIFIER_TYPES_SQL: Final = ", ".join(f"'{name}'" for name in sorted(LLM_CLASSIFIER_TYPES)) + + +@dataclass(frozen=True, slots=True) +class GatedAutoRouterCapability: + """A complexity-router capability the license meters, in every spelling an enforcement point needs. + + ``uses`` and ``sql_config_predicate`` answer the same question, in process and in a DB count over + stored ``litellm_params`` (``{config}`` is the caller's expression for the normalized + ``complexity_router_config`` jsonb, substituted as many times as the predicate needs); they live + on one record so they cannot drift apart. ``subject`` and ``remedy`` build the shared refusal + message. A validated config claims at most one capability, and the validator is what makes that + true: tier_definitions rejects every heuristic classifier_type, and it also rejects the + classifier system_prompt, which in turn only applies to the classifier types heuristic_v2 is not. + """ + + key: str + subject: str + remedy: str + uses: Callable[[object], bool] + sql_config_predicate: str + + +HEURISTIC_V2_CAPABILITY: Final = GatedAutoRouterCapability( + key="heuristic_v2", + subject="with classifier_type 'heuristic_v2'", + remedy="Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router.", + uses=uses_heuristic_v2_classifier, + sql_config_predicate="{config} ->> 'classifier_type' = 'heuristic_v2'", +) + +_OPERATOR_PROMPT_FIELDS_SQL: Final = " OR ".join( + f"{{config}} ->> '{field}' IS NOT NULL" for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS +) + +CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability( + key="tier_or_classifier_prompt", + subject="with operator-defined tier_definitions or an operator-written classifier prompt", + remedy=( + "Use the shipped tiers and classifier prompt for this router or remove an existing router " + "with tier_definitions or its own classifier prompt." + ), + uses=uses_custom_tier_or_classifier_prompt, + sql_config_predicate=( + "jsonb_typeof({config} -> 'tier_definitions') = 'array' OR " + f"({{config}} ->> 'classifier_type' IN ({_LLM_CLASSIFIER_TYPES_SQL}) AND (" + "{config} -> 'classifier_llm_config' ->> 'system_prompt' IS NOT NULL OR " + f"{_OPERATOR_PROMPT_FIELDS_SQL}))" + ), +) + +GATED_AUTO_ROUTER_CAPABILITIES: Final = (HEURISTIC_V2_CAPABILITY, CUSTOMIZATION_CAPABILITY) + + +def claimed_capability(complexity_router_config: object) -> GatedAutoRouterCapability | None: + """The licensed capability this complexity config claims, or None.""" + return next( + (capability for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(complexity_router_config)), + None, + ) + + +def gated_capability_of(litellm_params: Mapping[str, object]) -> GatedAutoRouterCapability | None: + """The licensed capability this deployment claims, or None unless it is a complexity router.""" + model: Final = litellm_params.get("model") + if not is_complexity_router_model(model if isinstance(model, str) else None): + return None + return claimed_capability(litellm_params.get("complexity_router_config")) + + +def count_capability_routers( + deployments: Iterable[Mapping[str, object]], *, capability: GatedAutoRouterCapability +) -> int: + """How many of ``deployments`` (router model_list entries or config.yaml rows) claim ``capability``.""" + return sum( + 1 for deployment in deployments if gated_capability_of(_mapping(deployment.get("litellm_params"))) is capability + ) + + +def capability_limit_violation(*, capability: GatedAutoRouterCapability, held: int, limit: int | None) -> str | None: + """Why holding ``held`` routers claiming ``capability`` exceeds ``limit``, or None when it fits. ``limit`` None means unlimited. The message is shared by every enforcement point (config load, model writes, router registration) and stays SDK-neutral: it names the cap and what @@ -190,8 +296,8 @@ def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: if limit is None or held <= limit: return None return ( - f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make " - f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router." + f"At most {limit} auto-router(s) {capability.subject} can be registered but this would make " + f"{held}. {capability.remedy}" ) @@ -237,9 +343,7 @@ def carries_complexity_router_settings(model: str | None, present_fields: frozen ``validate_strategy_router_model_write`` is judged on, so a router named only by its default model is in scope, and a field added to the table above is covered here for free. """ - return classify_strategy_router_model(model or "") == "complexity" or bool( - present_fields & _COMPLEXITY_ROUTER_FIELDS - ) + return is_complexity_router_model(model) or bool(present_fields & _COMPLEXITY_ROUTER_FIELDS) def validate_complexity_router_config_placement(litellm_params: Mapping[str, object] | None) -> str | None: diff --git a/litellm/types/router.py b/litellm/types/router.py index 267e8853db1..728d1037f3d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -887,9 +887,9 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... -class HeuristicV2RouterLimit(Protocol): +class AutoRouterCapabilityLimit(Protocol): """ - Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited. + Resolves how many complexity routers may claim each licensed capability right now; None means unlimited. The Router calls it on every registration and limit query instead of caching the answer, so the proxy can keep the limit on its license object (re-verified on config load) rather than hand diff --git a/tests/test_litellm/proxy/auth/test_litellm_license.py b/tests/test_litellm/proxy/auth/test_litellm_license.py index 1db53638070..d3f80982c7a 100644 --- a/tests/test_litellm/proxy/auth/test_litellm_license.py +++ b/tests/test_litellm/proxy/auth/test_litellm_license.py @@ -34,27 +34,27 @@ def test_is_over_limit(): assert license_check.is_over_limit(99) is False -def test_heuristic_v2_router_limit() -> None: +def test_auto_router_capability_limit() -> None: """Only the signed license's auto_router feature lifts the one-router limit; an API-verified license (no airgapped data) and an airgapped license without the feature keep it.""" license_check = LicenseCheck() license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]} - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None license_check.airgapped_license_data = { "expiration_date": "2999-01-01", "allowed_features": ["sso", "auto_router", "audit_logs"], } - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]} - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 license_check.airgapped_license_data = {"expiration_date": "2999-01-01"} - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 license_check.airgapped_license_data = None - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]: @@ -81,12 +81,12 @@ def test_expired_or_unreadable_license_grants_no_features() -> None: license_check = LicenseCheck() public_key, valid_key = _signed_license("2999-01-01") assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None _, expired_key = _signed_license("2000-01-01") assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True assert license_check.airgapped_license_data is None - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True @@ -98,4 +98,4 @@ def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None: public_key, license_key = _signed_license("2999-01-01") assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 3edeeedbae9..33de2a09626 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -2,6 +2,7 @@ import inspect import asyncio import contextlib import json +from collections.abc import Mapping from typing import Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -4048,6 +4049,72 @@ class TestStrategyRouterWriteValidation: ) assert _strategy_router_write_violation(incoming_params=None, existing_params=None) is None + @pytest.mark.parametrize( + "config", + [ + {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tier_definitions": [ + {"name": "routine", "description": "routine drafting"}, + {"name": "hard", "description": "hard reasoning"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + }, + ], + ) + def test_model_less_patch_cannot_attach_router_config_to_a_regular_model(self, config: dict[str, object]) -> None: + """The license gate applies only to complexity routers, so a partial PATCH cannot poison a regular + model with a capability-shaped config and make it occupy a slot.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + violation = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams(complexity_router_config=config), + existing_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + ) + + assert violation is not None + assert "does not start with 'auto_router/'" in violation + assert "complexity_router_config" in violation + + def test_effective_params_decrypts_a_stored_complexity_router_model(self, monkeypatch) -> None: + """A database row encrypts model, so the model-aware gate must not accidentally rely on plaintext mocks.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _effective_complexity_router_params, + ) + from litellm.types.router import updateLiteLLMParams + + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt") + encrypted_model = encrypt_value_helper("auto_router/complexity_router") + effective_params = _effective_complexity_router_params( + updateLiteLLMParams(complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini"}}), + LiteLLM_Params(model=encrypted_model), + ) + + assert effective_params["model"] == "auto_router/complexity_router" + + def test_model_less_patch_keeps_a_complexity_router_in_scope(self) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} + ), + existing_params=self._stored_complexity_params(), + ) + is None + ) + def test_restore_of_corrupted_row_is_allowed(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( _strategy_router_write_violation, @@ -4354,33 +4421,33 @@ class TestStrategyRouterWriteValidation: ) @staticmethod - def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router: + def _live_router_holding_one_capability(limit: int | None, config: Mapping[str, object]) -> Router: return Router( model_list=[ {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}}, { - "model_name": "held-v2", + "model_name": "held", "litellm_params": { "model": "auto_router/complexity_router", - "complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + "complexity_router_config": config, }, "model_info": {"id": "held-id"}, }, ], - heuristic_v2_router_limit=lambda: limit, + auto_router_capability_limit=lambda: limit, ) class _FakeTx: - """Stands in for a prisma transaction: records the raw statements and exposes the model table.""" + """Stands in for a prisma transaction: records raw statements and returns encrypted-model candidates.""" - def __init__(self, db_held: int) -> None: - self.db_held = db_held + def __init__(self, db_models: list[str]) -> None: + self.db_models = db_models self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: self.raw_calls.append((sql, args)) - return [{"held": self.db_held}] if "count(*)" in sql else [] + return [{"model": model} for model in self.db_models] if "AS model" in sql else [] async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": return self @@ -4391,9 +4458,9 @@ class TestStrategyRouterWriteValidation: class _FakeDb: """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" - def __init__(self, db_held: int, existing_row: object = None) -> None: + def __init__(self, db_models: list[str], existing_row: object = None) -> None: self.db = self - self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held) + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models) self.litellm_proxymodeltable = MagicMock( create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) ) @@ -4403,6 +4470,43 @@ class TestStrategyRouterWriteValidation: _V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} _V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}} + _CUSTOM_TIERS = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tier_definitions": [ + {"name": "routine", "description": "routine drafting"}, + {"name": "hard", "description": "hard reasoning"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + } + _TIER_LABELS_ONLY = { + "classifier_type": "heuristic", + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "tier_labels": {"SIMPLE": "Cheap"}, + } + _CUSTOM_PROMPT = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + } + _OPERATOR_EXAMPLES = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "classification_examples": '- "reset my password" -> SIMPLE', + } + _OPERATOR_OPENING_PROMPT = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "classification_prompt": "Grade by data sensitivity", + } + _SHIPPED_RUBRIC = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "classification_rubric": "agentic"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + } @pytest.mark.parametrize( "incoming,existing,expected", @@ -4431,41 +4535,55 @@ class TestStrategyRouterWriteValidation: @pytest.mark.asyncio @pytest.mark.parametrize( - "limit,effective_config,db_held,config_holds_one,model_id,expected", + "limit,effective_params,db_models,config_config,model_id,expected", [ - (1, _V2, 1, False, None, "refused"), - (1, _V2, 0, True, None, "refused"), - (1, _V2, 0, False, None, "reserved"), - (1, _V2, 0, False, "held-id", "reserved"), - (2, _V2, 1, False, None, "reserved"), - (1, _V1, 5, True, None, "plain"), - (1, None, 5, True, None, "plain"), - (None, _V2, 5, True, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], _V2, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, None, "reserved"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, "held-id", "reserved"), + (2, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "reserved"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, ["openai/gpt-4o"], None, None, "reserved"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, [], _CUSTOM_PROMPT, None, "refused"), + (1, {"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIERS}, ["auto_router/complexity_router"], None, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V1}, ["auto_router/complexity_router"], _V2, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": None}, ["auto_router/complexity_router"], _V2, None, "plain"), + (None, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], _V2, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _TIER_LABELS_ONLY}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_PROMPT}, ["auto_router/complexity_router"], None, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, [], _CUSTOM_TIERS, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_OPENING_PROMPT}, [], _CUSTOM_PROMPT, None, "refused"), + (None, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"), ], ) - async def test_heuristic_v2_slot_matrix( + async def test_auto_router_capability_slot_matrix( self, limit: int | None, - effective_config: object, - db_held: int, - config_holds_one: bool, + effective_params: Mapping[str, object], + db_models: list[str], + config_config: Mapping[str, object] | None, model_id: str | None, expected: str, ) -> None: - """The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows - (other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL - parameter, and every other write runs on the plain client with no lock.""" + """The slot is claimed inside a locked transaction only for a write that claims a licensed capability + under a limit; the DB rows (other pods included) plus config.yaml routers decide, the row being edited + is excluded through the SQL parameter, and every other write runs on the plain client with no lock. + + heuristic_v2 has its own slot, while custom tier definitions and custom prompts count into one shared + customization slot. Renaming built-in tiers through tier_labels claims nothing at all.""" from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import ( - HEURISTIC_V2_SLOT_LOCK_KEY, - _heuristic_v2_slot, + AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY, + _auto_router_capability_slot, ) + from litellm.router_utils.auto_router_model_naming import gated_capability_of - fake = self._FakeDb(db_held) - live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None + capability = gated_capability_of(effective_params) + + fake = self._FakeDb(db_models) + live_router = self._live_router_holding_one_capability(limit, config_config) if config_config is not None else None with ( - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", @@ -4474,13 +4592,15 @@ class TestStrategyRouterWriteValidation: ): if expected == "refused": with pytest.raises(HTTPException) as exc_info: - async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id): + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id): pass assert exc_info.value.status_code == 403 + assert capability is not None assert "At most 1 auto-router" in str(exc_info.value.detail) + assert capability.subject in str(exc_info.value.detail) assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) return - async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables: + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id) as tables: handle = tables if expected == "plain": await handle.create(data={}) @@ -4489,10 +4609,13 @@ class TestStrategyRouterWriteValidation: return assert handle is fake.tx_obj.litellm_proxymodeltable published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable") - (lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls + (lock_sql, lock_params), (count_sql, count_params) = fake.tx_obj.raw_calls assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql - assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,) + assert lock_params == (AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,) assert count_params == (model_id or "",) + assert "AS model" in count_sql + assert capability is not None + assert capability.sql_config_predicate.split("{config}")[-1].strip() in count_sql @pytest.mark.asyncio async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: @@ -4549,14 +4672,14 @@ class TestStrategyRouterWriteValidation: ) admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - fake = self._FakeDb(db_held=1) + fake = self._FakeDb(["auto_router/complexity_router"]) with ( patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -4579,6 +4702,93 @@ class TestStrategyRouterWriteValidation: fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() fake.litellm_proxymodeltable.create.assert_not_awaited() + @pytest.mark.asyncio + async def test_model_less_patch_rejects_router_config_on_a_regular_model(self) -> None: + """PATCH rejects the poison before its row write or the capability slot.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + from litellm.types.router import updateLiteLLMParams + + model_id = "regular-model" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + regular = Deployment( + model_name="regular-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + model_info={"id": model_id}, + ) + fake = self._FakeDb([]) + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag + patch( # test-quality-ok: inject stored regular row without a database + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=regular), + ), + patch( # test-quality-ok: endpoint must reject before database authorization needs a live store + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._CUSTOM_TIERS) + ), + user_api_key_dict=admin, + ) + + assert exc_info.value.code == "400" + assert "does not start with 'auto_router/'" in str(exc_info.value.message) + assert fake.tx_obj.raw_calls == [] + assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0 + assert fake.litellm_proxymodeltable.update.await_count == 0 + + @pytest.mark.asyncio + async def test_model_less_legacy_update_rejects_router_config_on_a_regular_model(self) -> None: + """The legacy update endpoint enforces the same boundary before its row write or slot.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + from litellm.types.router import ModelInfo, updateLiteLLMParams + + model_id = "regular-model" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + regular = Deployment( + model_name="regular-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + model_info={"id": model_id}, + ) + existing_row = MagicMock() + existing_row.model_dump.return_value = regular.model_dump() + existing_row.litellm_params = regular.litellm_params.model_dump() + fake = self._FakeDb([], existing_row=existing_row) + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag + patch( # test-quality-ok: endpoint must reject before database authorization needs a live store + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._CUSTOM_TIERS), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=admin, + ) + + assert exc_info.value.code == "400" + assert "does not start with 'auto_router/'" in str(exc_info.value.message) + assert fake.tx_obj.raw_calls == [] + assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0 + assert fake.litellm_proxymodeltable.update.await_count == 0 + @pytest.mark.asyncio async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" @@ -4591,14 +4801,14 @@ class TestStrategyRouterWriteValidation: model_id = "other-id" admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - fake = self._FakeDb(db_held=1) + fake = self._FakeDb(["auto_router/complexity_router"]) with ( patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: the write must be refused before this DB step runs "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", new=AsyncMock(return_value=self._db_complexity_router(model_id)), @@ -4643,14 +4853,14 @@ class TestStrategyRouterWriteValidation: "model_info": {"id": model_id}, } existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"] - fake = self._FakeDb(db_held=1, existing_row=existing_row) + fake = self._FakeDb(["auto_router/complexity_router"], existing_row=existing_row) with ( patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index dcfad8f6815..2babfe432f3 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -28,7 +28,7 @@ from litellm.proxy.proxy_server import ( resolve_routing_plugins, validate_deployment_complexity_router_placement, validate_deployment_max_agentic_loops, - validate_heuristic_v2_router_limit, + validate_auto_router_capability_limits, ) from .conftest import normalize @@ -204,13 +204,71 @@ def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> } -def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None: +def _custom_tier_row(model_name: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "tier_definitions": [ + {"name": "routine", "description": "routine drafting"}, + {"name": "hard", "description": "hard reasoning"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + }, + }, + } + + +def _operator_examples_row(model_name: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "classification_examples": '- "reset my password" -> SIMPLE', + }, + }, + } + + +def _custom_prompt_row(model_name: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + }, + }, + } + + +@pytest.mark.parametrize( + "over_limit_rows,subject", + [ + ([_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], "heuristic_v2"), + ([_custom_tier_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "tier_definitions"), + ([_custom_prompt_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"), + ([_custom_tier_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"), + ([_operator_examples_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"), + ], +) +def test_validate_auto_router_capability_limits_refuses_to_start_over_the_limit( + over_limit_rows: list[dict[str, object]], subject: str +) -> None: """Same reason as the two validators above: the proxy router swallows registration errors, so an over-limit config.yaml must fail here instead of booting with a silently missing router.""" with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info: - validate_heuristic_v2_router_limit( - [_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1 - ) + validate_auto_router_capability_limits(over_limit_rows, limit=1) + assert subject in str(exc_info.value) assert "'auto_router' feature lifts the limit" in str(exc_info.value) @@ -220,12 +278,15 @@ def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> ([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None), ([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1), ([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1), + ([_custom_tier_row("a"), _custom_tier_row("b")], None), + ([_custom_tier_row("a"), _heuristic_v2_row("b")], 1), ], ) -def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone( +def test_validate_auto_router_capability_limits_leaves_configs_within_the_limit_alone( model_list: list[dict[str, object]], limit: int | None ) -> None: - assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None + """The last case is the separate-ceiling invariant: one router of each capability fits under a limit of one.""" + assert validate_auto_router_capability_limits(model_list, limit=limit) is None _TWO_HEURISTIC_V2_ROUTERS_YAML = ( @@ -247,7 +308,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = ( " classifier_type: heuristic_v2\n" " tiers: {SIMPLE: gpt-4o-mini}\n" "router_settings:\n" - " heuristic_v2_router_limit: 99\n" + " auto_router_capability_limit: 99\n" ) @@ -256,7 +317,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = ( async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only( tmp_path, monkeypatch, license_limit: int | None ) -> None: - """`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator + """`router_settings.auto_router_capability_limit` is managed outside config.yaml: an operator cannot grant the entitlement by editing the config, and a licensed proxy boots both routers.""" f = tmp_path / "c.yaml" f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML) @@ -264,15 +325,15 @@ async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_lic monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) monkeypatch.setattr( - "litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: license_limit ) if license_limit is None: router, _model_list, _general_settings = await ProxyConfig().load_config( router=None, config_file_path=str(f) ) - assert router.heuristic_v2_router_limit is not None - assert router.heuristic_v2_router_limit() is None + assert router.auto_router_capability_limit is not None + assert router.auto_router_capability_limit() is None assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] return @@ -296,12 +357,12 @@ async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_b monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) - monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1) router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f)) - assert router.heuristic_v2_router_limit is not None - assert router.heuristic_v2_router_limit() == 1 + assert router.auto_router_capability_limit is not None + assert router.auto_router_capability_limit() == 1 assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"}) assert router.upsert_deployment(db_row) is None diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 52e58304476..918ec7bc100 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -15,6 +15,12 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.router_utils.auto_router_model_naming import ( + CUSTOMIZATION_CAPABILITY, + GATED_AUTO_ROUTER_CAPABILITIES, + HEURISTIC_V2_CAPABILITY, + count_capability_routers, +) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY @@ -46,7 +52,6 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, TrainedTierArtifact, ) -from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers from litellm.types.router import ( Deployment, LiteLLM_Params, @@ -1124,7 +1129,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-b", "id-b", "heuristic_v2"), self._router_row("v1-c", "id-c", "heuristic"), ], - heuristic_v2_router_limit=lambda: 1, + auto_router_capability_limit=lambda: 1, ignore_invalid_deployments=True, ) @@ -1139,7 +1144,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: 1, + auto_router_capability_limit=lambda: 1, ) def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None: @@ -1152,14 +1157,14 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: limits["value"], + auto_router_capability_limit=lambda: limits["value"], ignore_invalid_deployments=True, ) assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] - assert router.heuristic_v2_router_limit_violation() is None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None limits["value"] = 1 - assert router.heuristic_v2_router_limit_violation() is not None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] @@ -1174,7 +1179,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: limits["value"], + auto_router_capability_limit=lambda: limits["value"], ignore_invalid_deployments=True, ) limits["value"] = 1 @@ -1195,7 +1200,7 @@ class TestRouterComplexityDeploymentMethods: assert router.upsert_deployment(Deployment(**db_row)) is not None assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"] - assert count_heuristic_v2_routers(router.config_deployments()) == 1 + assert count_capability_routers(router.config_deployments(), capability=HEURISTIC_V2_CAPABILITY) == 1 def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None: """A rollback after a failed upsert re-admits state that was already serving, so it must not be @@ -1208,7 +1213,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: limits["value"], + auto_router_capability_limit=lambda: limits["value"], ignore_invalid_deployments=True, ) limits["value"] = 1 @@ -1220,7 +1225,7 @@ class TestRouterComplexityDeploymentMethods: assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] live = router.get_deployment(model_id="id-a") assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2" - assert router.heuristic_v2_router_limit_violation() is not None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None: router = Router( @@ -1232,18 +1237,18 @@ class TestRouterComplexityDeploymentMethods: ) assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] - assert router.heuristic_v2_router_limit_violation() is None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None - def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None: + def test_auto_router_capability_violation_frees_the_slot_of_the_router_being_edited(self) -> None: """A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot while a different deployment switching to heuristic_v2 is refused.""" router = Router( model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")], - heuristic_v2_router_limit=lambda: 1, + auto_router_capability_limit=lambda: 1, ignore_invalid_deployments=True, ) - assert router.heuristic_v2_router_limit_violation() is not None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2") assert router.upsert_deployment(Deployment(**edited)) is not None @@ -1254,6 +1259,205 @@ class TestRouterComplexityDeploymentMethods: assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"] + @staticmethod + def _custom_tier_row(model_name: str, model_id: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tier_definitions": [ + {"name": "routine", "description": "routine drafting and lookups"}, + {"name": "hard", "description": "multi-step reasoning under tradeoffs"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + }, + }, + "model_info": {"id": model_id}, + } + + @staticmethod + def _custom_prompt_row(model_name: str, model_id: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + } | {"model_info": {"id": model_id}} + + def test_a_second_custom_prompt_router_is_refused_under_the_ceiling(self) -> None: + """An operator-written classifier system_prompt is metered like the other licensed capabilities.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._custom_prompt_row("prompt-a", "id-a"), + self._custom_prompt_row("prompt-b", "id-b"), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None: + """Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no + prompt at all, leaves a router unmetered, so several of them register under a ceiling of one.""" + def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]: + llm_config: dict[str, object] = {"model": "gpt-4o-mini"} + if preset is not None: + llm_config["classification_rubric"] = preset + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": llm_config, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + "model_info": {"id": model_id}, + } + + router = Router( + model_list=[ + self._POOL, + rubric("default-a", "id-a", None), + rubric("preset-b", "id-b", "agentic"), + rubric("preset-c", "id-c", "chat"), + ], + auto_router_capability_limit=lambda: 1, + ) + + assert sorted(router.complexity_routers) == ["default-a", "preset-b", "preset-c"] + + def test_a_second_custom_tier_router_is_refused_under_the_ceiling(self) -> None: + """Operator-defined tier sets are metered like heuristic_v2: one per proxy without the license.""" + with pytest.raises(ValueError, match="tier_definitions"): + Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._custom_tier_row("tiers-b", "id-b"), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_custom_tier_routers_are_unlimited_with_the_license_feature(self) -> None: + router = Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._custom_tier_row("tiers-b", "id-b"), + ], + auto_router_capability_limit=lambda: None, + ) + + assert sorted(router.complexity_routers) == ["tiers-a", "tiers-b"] + assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is None + + def test_each_capability_holds_its_own_slot(self) -> None: + """heuristic_v2 has its own slot, while custom tiers and custom prompts share one customization + slot: one v2 plus EITHER customization fits, but a second customization of any form is refused.""" + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._custom_tier_row("tiers-a", "id-t"), + ], + auto_router_capability_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"] + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None + assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is not None + + assert router.upsert_deployment(Deployment(**self._custom_tier_row("tiers-b", "id-t2"))) is None + assert router.upsert_deployment(Deployment(**self._custom_prompt_row("prompt-b", "id-p2"))) is None + assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"] + + @staticmethod + def _operator_prompt_row(model_name: str, model_id: str, field: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + field: '- "reset my password" -> SIMPLE', + }, + }, + "model_info": {"id": model_id}, + } + + @pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"]) + def test_operator_written_prompt_sections_claim_the_customization_slot(self, field: str) -> None: + """The dashboard prompt editor writes opening instructions and calibration examples as their own + fields on a BUILT-IN tier router, so each must claim the slot on its own.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._operator_prompt_row("prompt-a", "id-a", field), + self._operator_prompt_row("prompt-b", "id-b", field), + ], + auto_router_capability_limit=lambda: 1, + ) + + @pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"]) + def test_an_operator_prompt_section_claims_the_slot_held_by_custom_tiers(self, field: str) -> None: + """Switching the FORM of customization cannot buy a second unlicensed router.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._operator_prompt_row("prompt-b", "id-b", field), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_a_custom_prompt_claims_the_slot_held_by_custom_tiers(self) -> None: + """The customization ceiling is shared: changing its form cannot get a second unlicensed router.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._custom_prompt_row("prompt-b", "id-b"), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None: + """tier_labels renames the built-in ladder without defining one, so it stays ungated: two such + routers register under a ceiling of one.""" + def labeled(model_name: str, model_id: str) -> dict[str, object]: + row = self._router_row(model_name, model_id, "heuristic") + row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"} + return row + + router = Router( + model_list=[self._POOL, labeled("labels-a", "id-a"), labeled("labels-b", "id-b")], + auto_router_capability_limit=lambda: 1, + ) + + assert sorted(router.complexity_routers) == ["labels-a", "labels-b"] + def test_hybrid_initialization_waits_for_later_pool_deployments(self): router = Router( model_list=[ diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 238d0546518..8dede941a14 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -5,9 +5,11 @@ import pytest from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, - is_heuristic_v2_router, + GATED_AUTO_ROUTER_CAPABILITIES, + capability_limit_violation, + claimed_capability, + count_capability_routers, + gated_capability_of, strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, @@ -376,38 +378,122 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie assert carries_complexity_router_settings(model, present_fields) is scoped +_HV2_CONFIG: Mapping[str, object] = {"classifier_type": "heuristic_v2"} +_CUSTOM_TIER_CONFIG: Mapping[str, object] = { + "classifier_type": "llm", + "tier_definitions": [{"name": "routine", "description": "easy"}, {"name": "hard", "description": "hard"}], +} +_CUSTOM_PROMPT_CONFIG: Mapping[str, object] = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, +} + + @pytest.mark.parametrize( - "litellm_params,expected", + "config,expected_key", [ - ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), - ({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), - ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False), - ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False), - ({"model": "auto_router/complexity_router"}, False), - ({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), - ({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), - ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False), - ({}, False), + (_CUSTOM_PROMPT_CONFIG, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": "grade it"}, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_examples": '- "x" -> SIMPLE'}, "tier_or_classifier_prompt"), + ({"classifier_type": "hybrid", "classification_examples": "- y -> MEDIUM"}, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": None, "classification_examples": None}, None), + ({"classifier_type": "heuristic", "classification_examples": "- x -> SIMPLE"}, None), + ({"classifier_type": "hybrid", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"), + ({"classifier_type": "heuristic_first", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "classification_rubric": "chat"}}, None), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}}, None), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": None}}, None), + ({"classifier_type": "heuristic", "classifier_llm_config": {"system_prompt": "p"}}, None), + ({"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}}, "heuristic_v2"), + ({"classifier_type": "llm", "classifier_llm_config": "not a mapping"}, None), ], ) -def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None: - """Only a complexity router whose config selects heuristic_v2 counts toward the license limit.""" - assert is_heuristic_v2_router(litellm_params) is expected +def test_custom_classifier_prompt_capability(config: Mapping[str, object], expected_key: str | None) -> None: + """Every operator-written part of the classifier prompt claims the customization slot: a whole + replacement system_prompt, replacement opening instructions (classification_prompt), or replacement + calibration examples (classification_examples). + + A shipped rubric preset stays free, and the heuristic scorers never read system_prompt, so a + value sitting on one is inert and claims nothing (heuristic_v2 still claims its own capability). + """ + claimed = claimed_capability(config) + assert (None if claimed is None else claimed.key) == expected_key -def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None: - v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}} +@pytest.mark.parametrize( + "model,expected", + [ + ("auto_router/complexity_router", True), + ("auto_router/complexity_router-eu", True), + ("auto_router/semantic_router", False), + ("auto_router/adaptive_router", False), + ("auto_router/quality_router", False), + ("openai/gpt-4o", False), + (None, False), + ], +) +def test_is_complexity_router_model(model: str | None, expected: bool) -> None: + from litellm.router_utils.auto_router_model_naming import is_complexity_router_model + + assert is_complexity_router_model(model) is expected + + +@pytest.mark.parametrize( + "litellm_params,expected_key", + [ + ({"model": "auto_router/complexity_router", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"), + ({"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_definitions": None}}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_labels": {"SIMPLE": "Cheap"}}}, None), + ({"model": "auto_router/complexity_router"}, None), + ({"model": "auto_router/quality_router", "complexity_router_config": _HV2_CONFIG}, None), + ({"model": "auto_router/quality_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None), + ({"model": "openai/gpt-4o", "complexity_router_config": _HV2_CONFIG}, None), + ({"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, None), + ({}, None), + ], +) +def test_gated_capability_of(litellm_params: Mapping[str, object], expected_key: str | None) -> None: + """Only a complexity router claiming a licensed capability counts toward that capability's limit. + + Renaming the built-in tiers through tier_labels is not a custom tier set, so it stays ungated. + """ + capability = gated_capability_of(litellm_params) + assert (None if capability is None else capability.key) == expected_key + + +@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key) +def test_count_capability_routers_counts_only_its_own_capability(capability) -> None: + """Each capability has its own ceiling, so a router claiming the sibling capability never counts, + while a custom tier set and a custom classifier prompt count into the SAME customization slot.""" + def row(name: str, config: Mapping[str, object] | None) -> Mapping[str, object]: + params = {"model": "auto_router/complexity_router"} | ({} if config is None else {"complexity_router_config": config}) + return {"model_name": name, "litellm_params": params} + + by_key = { + "heuristic_v2": (_HV2_CONFIG, _HV2_CONFIG), + "tier_or_classifier_prompt": (_CUSTOM_TIER_CONFIG, _CUSTOM_PROMPT_CONFIG), + } + mine_first, mine_second = by_key[capability.key] + theirs = next(configs[0] for key, configs in by_key.items() if key != capability.key) rows: list[Mapping[str, object]] = [ - {"model_name": "a", "litellm_params": v2}, - {"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}}, - {"model_name": "c", "litellm_params": v2}, - {"model_name": "d"}, - {"model_name": "e", "litellm_params": "not a mapping"}, + row("a", mine_first), + row("b", theirs), + {"model_name": "c", "litellm_params": {"model": "openai/gpt-4o"}}, + row("d", mine_second), + {"model_name": "e"}, + {"model_name": "f", "litellm_params": "not a mapping"}, ] - assert count_heuristic_v2_routers(rows) == 2 - assert count_heuristic_v2_routers(()) == 0 + assert count_capability_routers(rows, capability=capability) == 2 + assert count_capability_routers((), capability=capability) == 0 +@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key) @pytest.mark.parametrize( "held,limit,violates", [ @@ -419,10 +505,42 @@ def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ (4, 3, True), ], ) -def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None: - violation = heuristic_v2_limit_violation(held=held, limit=limit) +def test_capability_limit_violation(held: int, limit: int | None, violates: bool, capability) -> None: + violation = capability_limit_violation(capability=capability, held=held, limit=limit) assert (violation is not None) is violates if violation is not None: assert f"At most {limit} auto-router" in violation assert f"would make {held}" in violation + assert capability.subject in violation + assert capability.remedy in violation assert "license" not in violation + + +def test_every_gated_capability_has_a_distinct_predicate_and_sql_spelling() -> None: + """The in-process and SQL halves of a capability must stay paired, and no two capabilities may collide.""" + keys = tuple(capability.key for capability in GATED_AUTO_ROUTER_CAPABILITIES) + assert len(set(keys)) == len(keys) + for capability in GATED_AUTO_ROUTER_CAPABILITIES: + assert "{config}" in capability.sql_config_predicate + assert capability.uses is not None + + +@pytest.mark.parametrize( + "config", + [ + _HV2_CONFIG, + _CUSTOM_TIER_CONFIG, + _CUSTOM_PROMPT_CONFIG, + {"classifier_type": "heuristic"}, + {"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}}, + {"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": "p"}, "tier_labels": {"SIMPLE": "Cheap"}}, + ], +) +def test_capabilities_are_mutually_exclusive_on_one_config(config: Mapping[str, object]) -> None: + """No config claims two capabilities, which is what lets one lock and one count serve them all. + + The config validator is what makes this true and is pinned separately in test_complexity_router: + tier_definitions rejects every heuristic classifier_type and rejects the classifier system_prompt, + and system_prompt only counts for the classifier types heuristic_v2 is not one of. + """ + assert sum(1 for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(config)) <= 1