diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 60e691e0ef5..aa2cb96a9f7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15306 + "limit": 15290 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38350 + "limit": 38332 }, "reportUnknownParameterType": { "limit": 19625 }, "reportUnknownVariableType": { - "limit": 29877 + "limit": 29861 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 748ef938cea..dc41c7dadc8 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1955,11 +1955,10 @@ Model Info: if not thresholds_enabled and not anomalies_enabled: return - if prisma_client is None: - from litellm.proxy.proxy_server import prisma_client as global_prisma_client + from litellm.proxy.proxy_server import prisma_client as global_prisma_client - prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client - if prisma_client is None: + client: Final = prisma_client if prisma_client is not None else global_prisma_client + if client is None: return from litellm.integrations.SlackAlerting.user_spend_alerts import ( @@ -1970,7 +1969,7 @@ Model Info: try: today: Final = datetime.datetime.now(datetime.timezone.utc).date() rows: Final = await fetch_user_spend_rows( - prisma_client=prisma_client, + prisma_client=client, today=today, baseline_days=self.alerting_args.spend_anomaly_baseline_days, ) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index dc61ee38a8c..2d737bc34e7 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -419,7 +419,6 @@ class WebSearchInterceptionLogger(CustomLogger): if call_type in (CallTypes.responses, CallTypes.aresponses): return self._convert_responses_tools(kwargs=kwargs, tools=tools) - # Check if any tool is a web search tool (native or already LiteLLM standard) has_websearch: Final = any(is_web_search_tool(t) for t in tools) if not has_websearch: diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 11a4f68ece0..6079b709bcc 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1411,6 +1411,97 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format +def _normalized_cache_control(cache_control: object) -> dict[str, str] | None: # mutable-ok: JSON wire format + if not isinstance(cache_control, Mapping): + return None + cache_type: Final = cache_control.get("type") + return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"} # mutable-ok: JSON wire format + + +def _with_portable_cache_control(block: Mapping[str, object]) -> dict[str, object]: # mutable-ok: JSON wire format + if "cache_control" not in block: + return dict(block) # mutable-ok: JSON wire format + normalized: Final = _normalized_cache_control(block["cache_control"]) + rest: Final = {key: value for key, value in block.items() if key != "cache_control"} # mutable-ok: JSON wire format + return rest if normalized is None else {**rest, "cache_control": normalized} # mutable-ok: JSON wire format + + +def _with_portable_cache_control_in_blocks(blocks: object) -> object: + if isinstance(blocks, str) or not isinstance(blocks, Sequence): + return blocks + return [ # mutable-ok: JSON wire format + _with_portable_cache_control(block) if isinstance(block, Mapping) else block for block in blocks + ] + + +def _with_portable_cache_control_in_content_block(block: object) -> object: + if not isinstance(block, Mapping): + return block + portable: Final = _with_portable_cache_control(block) + if portable.get("type") != "tool_result" or "content" not in portable: + return portable + return { # mutable-ok: JSON wire format + **portable, + "content": _with_portable_cache_control_in_blocks(portable["content"]), + } + + +def _with_portable_cache_control_in_message(message: object) -> object: + if not isinstance(message, Mapping) or "content" not in message: + return message + content: Final = message["content"] + if isinstance(content, str) or not isinstance(content, Sequence): + return message + return { # mutable-ok: JSON wire format + **message, + "content": [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_content_block(block) for block in content + ], + } + + +def _with_portable_cache_control_in_messages(messages: object) -> object: + if isinstance(messages, str) or not isinstance(messages, Sequence): + return messages + return [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_message(message) for message in messages + ] + + +def _with_portable_cache_control_in_scoped_value(key: str, value: object) -> object: + match key: + case "system" | "tools": + return _with_portable_cache_control_in_blocks(value) + case "messages": + return _with_portable_cache_control_in_messages(value) + case _: + return value + + +def normalize_cache_control_in_anthropic_payload( + payload: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: JSON wire format + """ + Return a copy of an Anthropic /v1/messages payload with every + ``cache_control`` entry reduced to ``{"type": }`` + at the places the Messages API defines it: the request itself, system + blocks, tools, message content blocks, and ``tool_result`` content blocks. + Application data such as ``tool_use.input`` and tool ``input_schema`` is + never touched, even when it happens to contain a ``cache_control`` key. + + Anthropic itself accepts prompt-caching extensions such as ``ttl``, but + strict non-Anthropic implementations of the Messages API validate the field + literally and reject the whole request (``cache_control.ttl: 1h is not + supported``, ``cache_control.type is required``), which 400s clients like + Claude Code that send cache hints. Non-dict ``cache_control`` values are + dropped entirely. The caller's payload is never mutated. + """ + portable: Final = _with_portable_cache_control(payload) + return { # mutable-ok: JSON wire format + key: _with_portable_cache_control_in_scoped_value(key, value) for key, value in portable.items() + } + + def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "anthropic-ratelimit-requests-limit" in headers: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 69985bcdaa3..b82903d6f87 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -99,6 +99,10 @@ def _deployment_passes_through_anthropic_messages(model_info: object) -> bool: return isinstance(supported_endpoints, (list, tuple)) and "/v1/messages" in supported_endpoints +def _deployment_supports_cache_control_ttl(model_info: object) -> bool: + return isinstance(model_info, dict) and model_info.get("cache_control_ttl") is True + + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() @@ -568,7 +572,9 @@ def anthropic_messages_handler( OpenAILikeAnthropicMessagesConfig, ) - anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig() + anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig( + cache_control_ttl=_deployment_supports_cache_control_ttl(kwargs.get("model_info")), + ) if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. if _should_route_to_responses_api(custom_llm_provider, original_model, model): diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index a8c5e97e7bf..e097805f54a 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1337,6 +1337,7 @@ class AmazonConverseConfig(BaseConfig): ) additional_request_params.pop("parallel_tool_calls", None) + additional_request_params.pop("client_metadata", None) # Only set the topK value in for models that support it additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 66ee5f10679..1e5329c90dd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -748,6 +748,15 @@ def strip_bedrock_throughput_suffix(model: str) -> str: MANTLE_MESSAGES_PATH: Final = "/anthropic/v1/messages" +_MANTLE_OPENAI_BASE_SUFFIXES: Final = ("/openai/v1", "/v1") + + +def _mantle_api_base_from_env() -> str | None: + env_base: Final = get_secret_str("BEDROCK_MANTLE_API_BASE") + if env_base is None: + return None + base: Final = env_base.rstrip("/") + return next((base[: -len(suffix)] for suffix in _MANTLE_OPENAI_BASE_SUFFIXES if base.endswith(suffix)), base) def build_mantle_messages_url( @@ -758,12 +767,15 @@ def build_mantle_messages_url( """Build the bedrock-mantle Anthropic /messages URL. Honors an explicit endpoint override (``api_base``, then - ``aws_bedrock_runtime_endpoint``) so private VPC / VPCE / GovCloud Mantle - endpoints are reachable; otherwise falls back to the public regional host. + ``aws_bedrock_runtime_endpoint``, then ``BEDROCK_MANTLE_API_BASE``) so + private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise + falls back to the public regional host. The mantle messages path is appended unless the override already carries it, - so callers can pass either the host or the full messages URL. + so callers can pass either the host or the full messages URL. The env var is + shared with the OpenAI-surface ``bedrock_mantle/*`` routes, which need it to + carry their ``/v1`` or ``/openai/v1`` base, so that suffix is dropped first. """ - override: Final = api_base or aws_bedrock_runtime_endpoint + override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env() if override: base: Final = override.rstrip("/") if base.endswith(MANTLE_MESSAGES_PATH): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3cd6ee54069..26c085a95e6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, ) +from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import realtime_error_event, websocket_close_reason from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming @@ -274,6 +275,16 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False +def _aws_signing_overrides(optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any]) -> Mapping[str, Any]: + return MappingProxyType( + { + key: litellm_params[key] + for key in AWS_CREDENTIAL_KWARGS_KEYS + if optional_params.get(key) is None and litellm_params.get(key) is not None + } + ) + + def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]: """Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM enforcement, so the Responses WebSocket loop can charge every @@ -538,7 +549,10 @@ class BaseLLMHTTPHandler: headers, signed_json_body = provider_config.sign_request( headers=headers, - optional_params=optional_params, + optional_params={ + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, request_data=data, api_base=api_base, api_key=api_key, diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index d6b217d5746..73086ba395b 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -8,6 +8,7 @@ Based on official GigaChat SDK authentication flow. import time import uuid from collections.abc import Mapping +from types import MappingProxyType from typing import Final import httpx @@ -32,8 +33,8 @@ GIGACHAT_SCOPE: Final = "GIGACHAT_API_PERS" # Token expiry buffer in milliseconds (refresh token 60s before expiry) TOKEN_EXPIRY_BUFFER_MS: Final = 60000 -# Cache for access tokens _token_cache: Final = InMemoryCache() +_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) class GigaChatAuthError(BaseLLMException): @@ -80,10 +81,9 @@ def get_access_token( Raises: GigaChatAuthError: If authentication fails """ - if not litellm_params: - litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default + params: Final = litellm_params or _NO_LITELLM_PARAMS - access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token @@ -94,24 +94,20 @@ def get_access_token( message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() - effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url() - # Check cache cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: _token, _expires_at = cached - # Check if token is still valid (with buffer) if time.time() * 1000 < _expires_at - TOKEN_EXPIRY_BUFFER_MS: verbose_logger.debug("Using cached GigaChat access token") return _token - # Request new token new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: - # Cache token ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) @@ -126,10 +122,9 @@ async def get_access_token_async( litellm_params: Mapping[str, object] | None = None, ) -> str: """Async version of get_access_token.""" - if not litellm_params: - litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default + params: Final = litellm_params or _NO_LITELLM_PARAMS - access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token @@ -140,10 +135,9 @@ async def get_access_token_async( message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() - effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url() - # Check cache cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: @@ -152,11 +146,9 @@ async def get_access_token_async( verbose_logger.debug("Using cached GigaChat access token") return _token - # Request new token new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: - # Cache token ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 2875b30232e..0a4cbd8e520 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -52,7 +52,6 @@ class GigaChatModelResponseIterator: tool_use: ChatCompletionToolCallChunk | None = None # rebind-ok: conditionally assigned on function_call finish_reason: str | None = chunk_finish_reason - # Handle function_call in stream raw_function_call: Final = delta.get("function_call") if chunk_finish_reason == "function_call" and isinstance(raw_function_call, Mapping) and raw_function_call: func_call: Final[Mapping[str, object]] = raw_function_call diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 8f23c5175ec..89920ebd27b 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -10,6 +10,7 @@ import json import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -34,6 +35,9 @@ else: LiteLLMLoggingObj = Any +_EMPTY_FUNCTION: Final[Mapping[str, object]] = MappingProxyType({}) + + def is_valid_json(value: str) -> bool: """Checks whether the value passed is a valid serialized JSON string""" try: @@ -111,11 +115,9 @@ class GigaChatConfig(BaseConfig): """ Set up headers with OAuth token. """ - # Get access token credentials: Final = api_key or get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY") access_token: Final = get_access_token(credentials=credentials, litellm_params=litellm_params) - # Store credentials for image uploads self._current_credentials = credentials self._current_api_base = api_base @@ -208,18 +210,18 @@ class GigaChatConfig(BaseConfig): def _convert_tools_to_functions(self, tools: Sequence) -> Sequence[dict]: """Convert OpenAI tools format to GigaChat functions format.""" - functions: Final[list[dict]] = [] # mutable-ok: accumulator for building functions list - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "function": - func = tool.get("function", {}) - functions.append( - { - "name": func.get("name", ""), - "description": func.get("description", ""), - "parameters": func.get("parameters", {}), - } - ) - return functions + return [ + { + "name": function.get("name", ""), + "description": function.get("description", ""), + "parameters": function.get("parameters", {}), + } + for function in ( + tool.get("function", _EMPTY_FUNCTION) + for tool in tools + if isinstance(tool, dict) and tool.get("type") == "function" + ) + ] def _map_tool_choice(self, tool_choice: str | Mapping[str, object]) -> str | Mapping[str, object] | None: """ @@ -299,7 +301,6 @@ class GigaChatConfig(BaseConfig): if part.get("type") == "text": texts.append(part.get("text", "")) elif part.get("type") == "image_url": - # Extract image URL and upload to GigaChat image_url: object = part.get("image_url", {}) upload_url: str if isinstance(image_url, str): @@ -322,16 +323,13 @@ class GigaChatConfig(BaseConfig): headers: Mapping[str, object], ) -> dict: # mutable-ok: request payload sent to httpx """Transform OpenAI request to GigaChat format.""" - # Transform messages giga_messages: Final = self._transform_messages(messages) - # Build request request_data: Final[dict[str, object]] = { "model": model.replace("gigachat/", ""), "messages": giga_messages, } - # Add optional params for key in [ "temperature", "top_p", @@ -343,7 +341,6 @@ class GigaChatConfig(BaseConfig): if key in optional_params: request_data[key] = optional_params[key] - # Add functions if present if "functions" in optional_params: request_data["functions"] = optional_params["functions"] if "function_call" in optional_params: @@ -358,10 +355,8 @@ class GigaChatConfig(BaseConfig): for i, msg in enumerate(messages): message = dict(msg) - # Remove unsupported fields message.pop("name", None) - # Transform roles role = message.get("role", "user") if role == "developer": message["role"] = "system" @@ -374,18 +369,15 @@ class GigaChatConfig(BaseConfig): if not isinstance(content, str) or not is_valid_json(content): message["content"] = json.dumps(content, ensure_ascii=False) - # Handle None content if message.get("content") is None: message["content"] = "" - # Handle list content (multimodal) - extract text and images content = message.get("content") if isinstance(content, list): message["content"], attachments = self._transform_list_content(content) if attachments: message["attachments"] = attachments - # Transform tool_calls to function_call tool_calls = message.get("tool_calls") if tool_calls and isinstance(tool_calls, list) and len(tool_calls) > 0: tool_call = tool_calls[0] @@ -436,13 +428,11 @@ class GigaChatConfig(BaseConfig): message_data = choice.get("message", {}) finish_reason = choice.get("finish_reason", "stop") - # Transform function_call to tool_calls or content if finish_reason == "function_call" and message_data.get("function_call"): func_call = message_data["function_call"] args = func_call.get("arguments", {}) if is_structured_output: - # Convert to content for structured output if isinstance(args, dict): content = json.dumps(args, ensure_ascii=False) else: @@ -452,7 +442,6 @@ class GigaChatConfig(BaseConfig): message_data.pop("functions_state_id", None) finish_reason = "stop" else: - # Convert to tool_calls format if isinstance(args, dict): args = json.dumps(args, ensure_ascii=False) message_data["tool_calls"] = [ @@ -468,7 +457,6 @@ class GigaChatConfig(BaseConfig): message_data.pop("function_call", None) finish_reason = "tool_calls" - # Clean up GigaChat-specific fields message_data.pop("functions_state_id", None) choices.append( diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 2ec8324e33c..0db4475be8f 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -112,18 +112,10 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): "input": ["text1", "text2", ...] } """ - # Normalize input to list - if isinstance(input, str): - input_list: list = [input] # rebind-ok: locally scoped conversion - else: - input_list = input - - # Remove gigachat/ prefix from model if present - model = model.removeprefix("gigachat/") # rebind-ok: parameter reassignment for normalization - + normalized_input: Final = [input] if isinstance(input, str) else input # mutable-ok: preserve list API return { - "model": model, - "input": input_list, + "model": model.removeprefix("gigachat/"), + "input": normalized_input, } def transform_embedding_response( diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index a0edc6f5682..e1f73d04275 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -60,7 +60,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): """ Set up headers with OAuth token. """ - # Get access token access_token: Final = get_access_token(credentials=api_key, litellm_params=litellm_params) headers["Authorization"] = f"Bearer {access_token}" # rebind-ok: mutating for OAuth setup @@ -82,7 +81,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager - # cost tracking only for completions and embeddings if "completions" in endpoint: provider_chat_config: Final = ProviderConfigManager.get_provider_chat_config( provider=LlmProviders(custom_llm_provider), diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index cbb35cd1b57..ce7e848ed7f 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -4,7 +4,6 @@ from typing import Final from litellm.secret_managers.main import get_secret_str from litellm.types.utils import PromptTokensDetailsWrapper, Usage -# GigaChat API endpoint GIGACHAT_BASE_URL: Final = "https://gigachat.devices.sberbank.ru/api/v1" diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0e8fa294f5d..764d80c6f82 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -3,16 +3,20 @@ Transformation logic for Hosted VLLM rerank """ from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final import httpx +from pydantic import ValidationError from litellm._uuid import uuid +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( + HostedVLLMRerankTruncationParams, OptionalRerankParams, RerankBilledUnits, RerankRequest, @@ -34,6 +38,13 @@ class HostedVLLMRerankError(BaseLLMException): super().__init__(status_code=status_code, message=message, headers=headers) +def validated_truncation_params(non_default_params: Mapping[str, object] | None) -> HostedVLLMRerankTruncationParams: + try: + return HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + except ValidationError as error: + raise UnsupportedParamsError(status_code=400, message=f"hosted_vllm rerank: {error}") from error + + class HostedVLLMRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass @@ -62,7 +73,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig): "top_n", "rank_fields", "return_documents", + "max_tokens_per_doc", "instruction", + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", ] def map_cohere_rerank_params( @@ -100,7 +115,15 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if instruction is not None: mapped_params["instruction"] = instruction - return dict(mapped_params) + truncation: Final = validated_truncation_params(non_default_params) + forwarded: Final[OptionalRerankParams] = { + **mapped_params, + "max_tokens_per_doc": max_tokens_per_doc, + "truncate_prompt_tokens": truncation.truncate_prompt_tokens, + "truncation_side": truncation.truncation_side, + "max_tokens_per_query": truncation.max_tokens_per_query, + } + return dict(forwarded) def validate_environment( self, @@ -138,6 +161,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if "documents" not in optional_rerank_params: raise ValueError("documents is required for Hosted VLLM rerank") + truncation: Final = HostedVLLMRerankTruncationParams.model_validate(optional_rerank_params) rerank_request: Final = RerankRequest( model=model, query=optional_rerank_params["query"], @@ -146,6 +170,10 @@ class HostedVLLMRerankConfig(BaseRerankConfig): rank_fields=optional_rerank_params.get("rank_fields", None), return_documents=optional_rerank_params.get("return_documents", None), instruction=optional_rerank_params.get("instruction", None), + max_tokens_per_doc=truncation.max_tokens_per_doc, + truncate_prompt_tokens=truncation.truncate_prompt_tokens, + truncation_side=truncation.truncation_side, + max_tokens_per_query=truncation.max_tokens_per_query, ) return rerank_request.model_dump(exclude_none=True) diff --git a/litellm/llms/openai_like/README.md b/litellm/llms/openai_like/README.md index e9aaafe48a1..e1409b81c35 100644 --- a/litellm/llms/openai_like/README.md +++ b/litellm/llms/openai_like/README.md @@ -54,7 +54,10 @@ That's it! The provider will be automatically loaded and available. "constraints": { "temperature_max": 1.0, "temperature_min": 0.0, - "temperature_min_with_n_gt_1": 0.3 + "temperature_min_with_n_gt_1": 0.3, + // /v1/messages providers only: keep Anthropic cache_control extensions + // such as ttl instead of stripping them down to {"type": ...} + "cache_control_ttl": true }, // Optional: Special handling flags diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 11dc236064d..ac99617521c 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,11 +1,13 @@ from typing import Any, Final import litellm +from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -19,10 +21,17 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``"/v1/messages"``. The inbound Anthropic payload (system, cache_control, thinking, tools, ...) is forwarded essentially unchanged to ``{api_base}/v1/messages``, so Anthropic-only features that the - Anthropic->OpenAI translation would otherwise drop are preserved. Response - parsing and streaming are inherited from the native Anthropic config. + Anthropic->OpenAI translation would otherwise drop are preserved. The one + exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) + are stripped unless the deployment opts in with + ``model_info.cache_control_ttl: true``. Response parsing and streaming are + inherited from the native Anthropic config. """ + def __init__(self, cache_control_ttl: bool = False) -> None: + super().__init__() + self._cache_control_ttl: Final = cache_control_ttl + def validate_anthropic_messages_environment( self, headers: dict[str, str], @@ -53,6 +62,35 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return False + def supports_cache_control_ttl(self) -> bool: + return self._cache_control_ttl + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict], # mutable-ok: matches dict-typed base signature + anthropic_messages_optional_request_params: dict, # mutable-ok: matches dict-typed base signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: matches dict-typed base signature + ) -> dict: # mutable-ok: matches dict-typed base signature + """ + Anthropic ignores prompt-caching hints it cannot honor, but strict + non-Anthropic implementations of the Messages API 400 the whole request + on Anthropic-only ``cache_control`` extensions (``cache_control.ttl: 1h + is not supported``), so unless the provider declares ttl support the + hints are reduced to their portable ``{"type": ...}`` core. + """ + request: Final = super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + if self.supports_cache_control_ttl(): + return request + return normalize_cache_control_in_anthropic_payload(request) + def get_complete_url( self, api_base: str | None, @@ -81,7 +119,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): """ def __init__(self, provider: SimpleProviderConfig): - super().__init__() + super().__init__(cache_control_ttl=bool(provider.constraints.get("cache_control_ttl"))) self._provider = provider @property diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index c4bd03fb1c3..7076683f294 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -113,10 +113,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): ) ) - # Compliant: Save a strong reference to prevent GC self._background_tasks.add(task) - # Remove the task from the set when it finishes to avoid memory leaks task.add_done_callback(self._background_tasks.discard) except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( @@ -578,7 +576,6 @@ def llm_passthrough_route( else: return response except Exception as e: - # provider_config is guaranteed non-None here due to the earlier guard assert provider_config is not None raise base_llm_http_handler._handle_error( e=e, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d1ef73a15cd..37474f85fe7 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -74,7 +74,13 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( ) -def _connection_error_message(exc: BaseException) -> str: +def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + if isinstance(exc, TimeoutError): + return ( + f"Failed to connect to MCP server: no response from {url or 'the server'} " + f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " + "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." + ) if isinstance(exc, httpx.LocalProtocolError): return ( "Failed to connect to MCP server: a request header is malformed. " @@ -92,6 +98,9 @@ def _connection_error_message(exc: BaseException) -> str: if MCP_AVAILABLE: + from mcp.types import Tool as MCPTool + + from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -876,7 +885,6 @@ if MCP_AVAILABLE: return (), classify_list_exception(e) return tools_result, ServerListOk(tool_count=len(tools_result)) - # Query all servers the user has access to queried_servers: Final = tuple( server for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids) @@ -1141,12 +1149,18 @@ if MCP_AVAILABLE: scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes + async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: + with anyio.move_on_after(deadline): + return await client.list_tools(raise_on_error=True) + return None + async def _execute_with_mcp_client( request: NewMCPServerRequest, operation: Callable[..., Awaitable[Mapping[str, object]]], mcp_auth_header: str | dict[str, str] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + timeout_seconds: float = MCP_TOOL_LISTING_TIMEOUT, ) -> Mapping[str, object]: """ Create a temporary MCP client from *request*, run *operation*, and return the result. @@ -1162,6 +1176,10 @@ if MCP_AVAILABLE: oauth2_headers: Headers extracted from the incoming request (may contain the litellm API key — must NOT be forwarded for M2M servers). raw_headers: Raw request headers forwarded for stdio env construction. + timeout_seconds: Cap on OAuth discovery, connect, handshake, and *operation* + combined. Defaults to ``MCP_TOOL_LISTING_TIMEOUT`` (30s, below common LB + timeouts) so an unreachable upstream yields this endpoint's JSON error + instead of an opaque load-balancer 504 with an empty body. Returns: The dict returned by *operation*, or an error dict on failure. @@ -1252,15 +1270,16 @@ if MCP_AVAILABLE: static_headers=request.static_headers, ) - client: Final = await global_mcp_server_manager._create_mcp_client( - server=server_model, - mcp_auth_header=mcp_auth_header, - extra_headers=merged_headers, - stdio_env=stdio_env, - cred_provider=preview_cred_provider, - ) + with anyio.fail_after(timeout_seconds): + client: Final = await global_mcp_server_manager._create_mcp_client( + server=server_model, + mcp_auth_header=mcp_auth_header, + extra_headers=merged_headers, + stdio_env=stdio_env, + cred_provider=preview_cred_provider, + ) - return await operation(client) + return await operation(client) except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise @@ -1269,7 +1288,7 @@ if MCP_AVAILABLE: return { "status": "error", "error": True, - "message": _connection_error_message(e), + "message": _connection_error_message(e, request.url, timeout_seconds), } async def _preview_openapi_tools(spec_path: str) -> dict: @@ -1422,9 +1441,7 @@ if MCP_AVAILABLE: getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT, ) - list_tools_result = None # rebind-ok: set inside the timeout scope below - with anyio.move_on_after(listing_deadline): - list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above + list_tools_result: Final = await _list_tools_within(client, listing_deadline) if list_tools_result is None: verbose_logger.warning( "MCP tools/list preview timed out after %s seconds while paginating upstream tools", diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 558ea54495f..c2614b85016 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -6,8 +6,11 @@ External callers (public IPs) only see servers with available_on_public_internet """ import ipaddress +import os +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Final +from urllib.parse import urlparse from fastapi import Request from pydantic import TypeAdapter, ValidationError @@ -137,7 +140,7 @@ class IPAddressUtils: @staticmethod def is_request_from_trusted_proxy( request: Request, - general_settings: dict[str, Any] | None = None, + general_settings: Mapping[str, Any] | None = None, ) -> bool: """ Return True if X-Forwarded-* headers on this request should be trusted. @@ -190,6 +193,36 @@ class IPAddressUtils: trusted_networks: Final = IPAddressUtils.parse_trusted_proxy_networks(trusted_ranges) return IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks) + @staticmethod + def is_request_https( + request: Request, + general_settings: Mapping[str, Any] | None = None, + ) -> bool: + """ + Whether this request's PUBLIC-facing origin is HTTPS, for deciding + whether a cookie set on the response should be marked ``Secure``. + + litellm only sees a plain-HTTP hop whenever TLS terminates at a + reverse proxy, so ``request.url.scheme`` alone cannot answer this in + that deployment shape. Resolved from the first trusted signal: + 1. ``PROXY_BASE_URL`` (operator-declared public origin). + 2. ``X-Forwarded-Proto``, only when the request's direct peer is a + configured trusted proxy -- see ``is_request_from_trusted_proxy``. + An untrusted caller cannot spoof this header to strip Secure. + 3. The request's own literal scheme (direct TLS termination, or no + reverse proxy in front of litellm). + """ + configured_base_url: Final = os.environ.get("PROXY_BASE_URL", "").strip() + if configured_base_url: + return urlparse(configured_base_url).scheme == "https" + + if IPAddressUtils.is_request_from_trusted_proxy(request, general_settings=general_settings): + forwarded_proto: Final = request.headers.get("X-Forwarded-Proto") + if forwarded_proto: + return forwarded_proto.split(",")[0].strip().lower() == "https" + + return request.url.scheme == "https" + @staticmethod def extract_client_ip_from_xff_hops( xff_header: str, diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py index 27018769909..9cabac2d0fa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py +++ b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py @@ -8,6 +8,7 @@ import json import os from collections.abc import Mapping +from itertools import islice from typing import ( TYPE_CHECKING, Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml @@ -341,19 +342,18 @@ def _json_safe( if depth >= _MAX_DEPTH or id(value) in seen: return None - nested: Final = seen | {id(value)} # mutable-ok: one-shot set literal, unioned into a frozenset immediately + nested: Final = seen | frozenset((id(value),)) if isinstance(value, dict): - out: dict[str, object] = {} # mutable-ok: bounded accumulator local to this call, never escapes as-is - for key, item in list(value.items())[:_MAX_ITEMS]: # mutable-ok: list() only to slice an unordered view - if isinstance(key, str) and key not in strip_keys: - out[key] = _json_safe(item, depth + 1, nested, strip_keys) - return out + return { + key: _json_safe(item, depth + 1, nested, strip_keys) + for key, item in islice(value.items(), _MAX_ITEMS) + if isinstance(key, str) and key not in strip_keys + } if isinstance(value, (list, tuple, set, frozenset)): return [ # mutable-ok: return value is a one-shot list, discarded by the caller after use - _json_safe(item, depth + 1, nested, strip_keys) - for item in list(value)[:_MAX_ITEMS] # mutable-ok: list() only to slice an unordered view + _json_safe(item, depth + 1, nested, strip_keys) for item in islice(value, _MAX_ITEMS) ] dump: Final = getattr(value, "model_dump", None) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 3e67b211f62..466b100ea1f 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -36,6 +36,7 @@ from pydantic import ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role from litellm.proxy.utils import get_custom_url @@ -131,7 +132,7 @@ class SAMLAuthHandler: @staticmethod def _is_https(request: Request) -> bool: - return SAMLAuthHandler._base_url(request).startswith("https") + return IPAddressUtils.is_request_https(request) @staticmethod def _acs_url(request: Request) -> str: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 6b98d9f9a26..1feefa5725d 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -92,6 +92,7 @@ from litellm.proxy.auth.auth_utils import ( has_user_setup_sso, ) from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import ( admin_ui_disabled, @@ -1118,7 +1119,7 @@ async def google_login( request=request, ) if sso_redirect is not None: - _persist_return_to_cookie(sso_redirect, return_to) + _persist_return_to_cookie(sso_redirect, return_to, request) return sso_redirect from fastapi.responses import HTMLResponse @@ -1138,7 +1139,7 @@ async def google_login( # helper the SSO branch uses, so /login can resume the connect flow instead of dead-ending at the # dashboard. One implementation → the two sign-in branches cannot diverge (and the login form always # renders, since the helper never raises on a bad return_to). - _persist_return_to_cookie(form_response, return_to) + _persist_return_to_cookie(form_response, return_to, request) return form_response @@ -2741,6 +2742,7 @@ async def _sso_return_to_redirect( jwt_token: str, redis_usage_cache, user_api_key_cache, + request: Request, ) -> RedirectResponse | None: """Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard. @@ -2759,7 +2761,7 @@ async def _sso_return_to_redirect( if _is_same_origin_return_path(return_to): redirect_response = RedirectResponse(url=return_to, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(redirect_response, request, jwt_token) redirect_response.delete_cookie("litellm_cp_return_to") return redirect_response @@ -2782,7 +2784,25 @@ async def _sso_return_to_redirect( return None -def _persist_return_to_cookie(response: Response, return_to: str | None) -> None: +def set_session_token_cookie(response: Response, request: Request, jwt_token: str) -> None: + """Set the ``token`` session cookie shared by every sign-in path. + + Not HttpOnly: the dashboard reads this cookie via ``document.cookie`` to + populate its own Authorization headers (see + ``ui/litellm-dashboard/src/utils/cookieUtils.ts``), so marking it + HttpOnly would break login. Secure is still required whenever the public + origin is HTTPS, resolved the same trust-aware way as every other + litellm cookie.""" + response.set_cookie( + key="token", + value=jwt_token, + secure=IPAddressUtils.is_request_https(request), + httponly=False, + samesite="lax", + ) + + +def _persist_return_to_cookie(response: Response, return_to: str | None, request: Request) -> None: """Best-effort: persist a SAFE ``return_to`` on ``response`` as the one-shot ``litellm_cp_return_to`` cookie so ANY sign-in path — SSO / Okta / generic OR the username/password form — can resume there afterwards. THIS is the single source of truth, called by every sign-in branch so they cannot @@ -2803,6 +2823,7 @@ def _persist_return_to_cookie(response: Response, return_to: str | None) -> None max_age=600, httponly=True, samesite="lax", + secure=IPAddressUtils.is_request_https(request), ) @@ -3079,8 +3100,11 @@ class SSOAuthenticationHandler: # incoming request is HTTP (local dev). Without # ``Secure`` the cookie is sent over plain HTTP, # letting a network observer read and replay the - # state value and bypass this protection. - secure_flag: Final = request is None or request.url.scheme == "https" + # state value and bypass this protection. Trust-aware: + # honors PROXY_BASE_URL / a trusted reverse proxy's + # X-Forwarded-Proto instead of only the literal scheme + # litellm sees on the wire. + secure_flag: Final = request is None or IPAddressUtils.is_request_https(request) redirect_response.set_cookie( key="litellm_oauth_state", value=state_value, @@ -3628,6 +3652,7 @@ class SSOAuthenticationHandler: jwt_token=jwt_token, redis_usage_cache=redis_usage_cache, user_api_key_cache=user_api_key_cache, + request=request, ) if return_to_redirect is not None: return return_to_redirect @@ -3636,7 +3661,7 @@ class SSOAuthenticationHandler: litellm_dashboard_ui += "?login=success" verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui) redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(redirect_response, request, jwt_token) return redirect_response @staticmethod diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 78d8ce296b8..b48b8d81494 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1731,7 +1731,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: def get_vertex_pass_through_handler( - call_type: Literal["discovery", "aiplatform"], # noqa: UP037 + call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2961,7 +2961,6 @@ async def handle_gigachat_passthrough_router_model( """ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - # Detect streaming based on request body is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown] data: dict[str, Any] = await _read_request_body( @@ -2997,7 +2996,6 @@ async def handle_gigachat_passthrough_router_model( data["json"] = request_body data["custom_llm_provider"] = "gigachat" - # Remove sensitive keys from data keys: Final = [ # mutable-ok: list of keys to remove from data "gigachat_auth_url", "gigachat_access_token", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1e0792cb8a5..27132c90e05 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15329,7 +15329,10 @@ async def login(request: Request): # authorize round-trip), mirroring the SSO callback; otherwise land on the dashboard. Gated by # _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the # one-shot cookie is cleared after use. - from litellm.proxy.management_endpoints.ui_sso import _sso_return_to_redirect + from litellm.proxy.management_endpoints.ui_sso import ( + _sso_return_to_redirect, + set_session_token_cookie, + ) # Resume through the SAME resumer the SSO callback uses, rather than a second, narrower arm. # _persist_return_to_cookie stores both shapes it accepts (a relative same-origin path AND a @@ -15346,6 +15349,7 @@ async def login(request: Request): jwt_token=jwt_token, redis_usage_cache=redis_usage_cache, user_api_key_cache=user_api_key_cache, + request=request, ) except Exception: # noqa: BLE001 # resuming must NEVER block a completed sign-in # The symmetric half of _persist_return_to_cookie's "never raises" contract. The resumer @@ -15360,7 +15364,7 @@ async def login(request: Request): # Create redirect response with cookie redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(redirect_response, request, jwt_token) if cp_return_to: redirect_response.delete_cookie(key="litellm_cp_return_to") return redirect_response @@ -15370,6 +15374,7 @@ async def login(request: Request): async def login_v2(request: Request): global premium_user, general_settings, master_key from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie from litellm.proxy.utils import get_custom_url try: @@ -15404,7 +15409,7 @@ async def login_v2(request: Request): content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, status_code=status.HTTP_200_OK, ) - json_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(json_response, request, jwt_token) return json_response except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.login_v2(): Exception occurred - %s", e) @@ -15504,6 +15509,8 @@ async def login_v3(request: Request): @router.post("/v3/login/exchange", include_in_schema=False) # exchange single-use opaque code for JWT async def login_v3_exchange(request: Request): + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + try: if not general_settings.get("control_plane_url"): raise ProxyException( @@ -15550,7 +15557,7 @@ async def login_v3_exchange(request: Request): }, status_code=status.HTTP_200_OK, ) - json_response.set_cookie(key="token", value=cached_data["token"]) + set_session_token_cookie(json_response, request, cached_data["token"]) return json_response except ProxyException: raise diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 8e0cad39561..ad8b67d5e8f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -274,6 +274,49 @@ except that the heuristic outcome is the one already computed rather than a seco Spend logs record `routing_decision.cause` as `heuristic_first_short_circuit` when the classifier was skipped, and `llm_classifier` when it ran, so the two are told apart per request. +### Hybrid + +`classifier_type: hybrid` also scores locally first, but it asks a different question than +`heuristic_first`. Where heuristic-first asks how CHEAP the scorer's tier is and pays for the +classifier on everything above a ceiling, hybrid asks how DECIDED the score is and pays for the +classifier only where the score lands near a tier boundary. A confident score keeps its tier at +every tier, the most expensive one included: + +```yaml +model_list: + - model_name: smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + classifier_type: hybrid + hybrid_boundary_margin: 0.03 + classifier_llm_config: + model: gpt-4o-mini + tiers: + SIMPLE: gpt-4o-mini + MEDIUM: gpt-4o + COMPLEX: claude-sonnet-4 + REASONING: o1-preview +``` + +A request routes on the scorer's own tier when its score is further than `hybrid_boundary_margin` +from every active boundary. Everything else goes to the classifier: a score inside the band, where a +hair's difference would have named the adjacent tier and its model pool, and a prompt where no +dimension fired at all, which has no opinion to be confident about. `hybrid_boundary_margin` is +required for this type and rejected on the others, the same way `heuristic_first_max_tier` is +required for heuristic-first, so the two modes are told apart by the knob each one takes rather than +by a shared field that means something different per type. + +Pick the margin against the score distribution rather than by intuition. The scorer combines a small +set of discretely weighted dimensions, so achievable scores cluster on a lumpy grid instead of +spreading smoothly, and widening the margin admits whole clusters at once rather than a few more +requests. Spend logs record `routing_decision.cause` as `hybrid_short_circuit` when the classifier +was skipped and `llm_classifier` when it ran. + +Operator-defined tier sets (`tier_definitions`) are not supported here, for the same reason they are +not supported under heuristic-first: the scorer only produces the built-in tiers. Classifier failure +behaves exactly as it does under `classifier_type: llm`. + ### Reasoning Override If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone. diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a123a75dd81..d205db90607 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -799,6 +799,7 @@ class ClassificationOutcome(NamedTuple): "reasoning_override", "llm_classifier", "heuristic_first_short_circuit", + "hybrid_short_circuit", "housekeeping", "classifier_plugin", "classifier_fallback", @@ -1241,6 +1242,15 @@ class ComplexityRouter(CustomLogger): return tier, weighted_score, tuple(signals), "heuristic_scorer" + def _is_near_tier_boundary(self, score: float, margin: float) -> bool: + boundaries: Final = self._effective_tier_boundaries() + active_boundaries: Final = ( + boundaries["simple_medium"], + boundaries["medium_complex"], + boundaries["complex_reasoning"], + ) + return any(abs(score - boundary) <= margin for boundary in active_boundaries) + def _effective_reasoning_override_min_score(self) -> float: """The score a request must reach before the reasoning-marker override may promote it. @@ -1367,6 +1377,8 @@ class ComplexityRouter(CustomLogger): return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages) if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None: return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages) + if self.config.classifier_type == "hybrid" and self.config.classifier_llm_config is not None: + return await self._classify_hybrid(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type != "llm" or self.config.classifier_llm_config is None: tier, score, signals, cause = self._score_and_classify(prompt, system_prompt) return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) @@ -1418,6 +1430,29 @@ class ComplexityRouter(CustomLogger): return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="heuristic_first_short_circuit") return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored) + async def _classify_hybrid( + self, + prompt: str, + system_prompt: str | None, + request_kwargs: dict[str, Any] | None, # mutable-ok: handed to _classify_with_llm as-is + messages: Sequence[Mapping[str, object]] | None, + ) -> ClassificationOutcome: + """Score locally, and only pay for the classifier when the score sits near a tier boundary. + + Where heuristic_first asks how CHEAP the scorer's tier is, this asks how DECIDED it is, so a + confident score keeps its tier at every tier including the most expensive one. Two things make + a score undecided: landing within hybrid_boundary_margin of an active boundary, where a + hair's difference in score would have named the adjacent tier and its model pool, and firing + no dimension at all, which scores 0.0 and lands SIMPLE by default rather than by evidence. + """ + tier, score, signals, cause = self._score_and_classify(prompt, system_prompt) + scored: Final = ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) + margin: Final = self.config.hybrid_boundary_margin + decided: Final = margin is not None and bool(signals) and not self._is_near_tier_boundary(score, margin) + if decided: + return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="hybrid_short_circuit") + return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored) + async def _llm_classifier_outcome( self, prompt: str, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index fdf2a3a0b39..0ae0db63fad 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -43,7 +43,7 @@ DEFAULT_CLASSIFICATION_RUBRIC: Final[ClassificationRubric] = ClassificationRubri # The classifier_type values that can call classifier_llm_config.model. Every consumer asking # "is the classifier model a real dependency of this router" resolves it here, including the ones # that only hold the raw config mapping and cannot reach ComplexityRouterConfig.uses_llm_classifier. -LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first"}) +LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first", "hybrid"}) TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = ( @@ -627,12 +627,13 @@ class ComplexityRouterConfig(BaseModel): ) # Classifier strategy - classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first"] = Field( + classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first", "hybrid"] = Field( default="heuristic", description=( "Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, " - "an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays " - "for the LLM classifier when the local scorer does not confidently land a cheap tier" + "an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays " + "for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', " + "which trusts the local scorer everywhere except when its score lands near a tier boundary" ), ) heuristic_v2_artifact: TrainedTierArtifact | Literal["ultrafeedback"] = Field( @@ -644,7 +645,10 @@ class ComplexityRouterConfig(BaseModel): ) classifier_llm_config: ClassifierLLMConfig | None = Field( default=None, - description="Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first'", + description=( + "Configuration for the LLM classifier; required when classifier_type is 'llm', " + "'heuristic_first' or 'hybrid'" + ), ) heuristic_first_max_tier: str | None = Field( default=None, @@ -659,6 +663,19 @@ class ComplexityRouterConfig(BaseModel): "may not name the highest one, since that would make the LLM classifier unreachable." ), ) + hybrid_boundary_margin: float | None = Field( + default=None, + ge=0, + le=1, + description=( + "How close to a tier boundary a heuristic score has to land before the LLM classifier breaks the " + "tie; required when classifier_type is 'hybrid' and rejected otherwise. Everything further than " + "this from every active boundary routes on the scorer's own tier with no classifier call, at any " + "tier, which is what separates 'hybrid' from 'heuristic_first' and its cheap-tier ceiling. A " + "prompt where no dimension fired still goes to the classifier, since the scorer has no opinion " + "to be near a boundary with. 0 escalates only scores sitting exactly on a boundary." + ), + ) classifier_plugin: ClassifierPlugin | None = Field( default=None, description=( @@ -1135,6 +1152,23 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_hybrid_boundary_margin(self) -> "ComplexityRouterConfig": + if self.classifier_type != "hybrid": + if self.hybrid_boundary_margin is not None: + raise ValueError( + f"hybrid_boundary_margin is set but classifier_type is {self.classifier_type!r}; " + "the scorer would never consult the classifier on a near-boundary score. Set " + "classifier_type 'hybrid' or remove hybrid_boundary_margin" + ) + return self + if self.hybrid_boundary_margin is None: + raise ValueError( + "hybrid_boundary_margin is required when classifier_type is 'hybrid': without a margin no " + "score is ever near enough to a boundary to escalate, which is classifier_type 'heuristic'" + ) + return self + @field_validator("fallback_tier") @classmethod def _reject_blank_optional_text(cls, value: str | None) -> str | None: @@ -1257,7 +1291,7 @@ class ComplexityRouterConfig(BaseModel): ) if duplicated: raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}") - if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first"): + if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first", "hybrid"): raise ValueError( "tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only " "produces the four built-in tiers, as does heuristic_v2" diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 903781b2ccd..a76e6cf1187 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,8 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from pydantic import BaseModel, PrivateAttr -from typing_extensions import Required, TypedDict +from typing import Literal + +from pydantic import BaseModel, ConfigDict, PrivateAttr +from typing_extensions import ReadOnly, Required, TypedDict class RerankRequest(BaseModel): @@ -21,6 +23,18 @@ class RerankRequest(BaseModel): # (e.g. hosted vLLM / Qwen3-Reranker, DeepInfra). Omitted from the outgoing # request when None, so this is fully backward-compatible. instruction: str | None = None + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + + +class HostedVLLMRerankTruncationParams(BaseModel): + model_config = ConfigDict(frozen=True) + + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + max_tokens_per_doc: int | None = None class OptionalRerankParams(TypedDict, total=False): @@ -32,6 +46,9 @@ class OptionalRerankParams(TypedDict, total=False): max_chunks_per_doc: int | None max_tokens_per_doc: int | None instruction: str | None + truncate_prompt_tokens: ReadOnly[int | None] + truncation_side: ReadOnly[Literal["left", "right"] | None] + max_tokens_per_query: ReadOnly[int | None] class RerankBilledUnits(TypedDict, total=False): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 09c01873a9d..ee6f09e05dc 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2852,6 +2852,7 @@ RoutingDecisionCause = Literal[ # scorer, and from "classifier_fallback", which is the scorer running because a call failed: # only this cause means an LLM classifier was configured, reachable, and deliberately skipped. "heuristic_first_short_circuit", + "hybrid_short_circuit", # The operator's classifier plugin (classifier_type 'custom') decided the tier. "classifier_plugin", # The LLM classifier or classifier plugin failed on a router with an operator-defined diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index ad4c3d6bfbb..e819433c269 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -296,21 +296,15 @@ async def test_bedrock_converse_budget_tokens_preserved(): mock_acompletion.assert_called_once() call_kwargs = mock_acompletion.call_args.kwargs - print( - "acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str) - ) + print("acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str)) # Verify thinking parameter is passed through with budget_tokens preserved thinking_param = call_kwargs.get("thinking") - assert ( - thinking_param is not None - ), "thinking parameter should be passed to acompletion" - assert ( - thinking_param.get("type") == "enabled" - ), "thinking.type should be 'enabled'" - assert ( - thinking_param.get("budget_tokens") == 1024 - ), f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + assert thinking_param is not None, "thinking parameter should be passed to acompletion" + assert thinking_param.get("type") == "enabled", "thinking.type should be 'enabled'" + assert thinking_param.get("budget_tokens") == 1024, ( + f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + ) def test_openai_model_with_thinking_converts_to_reasoning(): @@ -342,23 +336,18 @@ def test_openai_model_with_thinking_converts_to_reasoning(): call_kwargs = mock_responses.call_args.kwargs # Verify reasoning is set (converted from thinking) - assert ( - "reasoning" in call_kwargs - ), "reasoning should be passed to litellm.responses" + assert "reasoning" in call_kwargs, "reasoning should be passed to litellm.responses" # budget_tokens=1024 -> effort="low" (at the LOW budget threshold) # reasoning_auto_summary is False by default, so no summary key expected_reasoning = {"effort": "low"} assert call_kwargs["reasoning"] == expected_reasoning, ( - f"reasoning should be {expected_reasoning} for budget_tokens=1024, " - f"got {call_kwargs.get('reasoning')}" + f"reasoning should be {expected_reasoning} for budget_tokens=1024, got {call_kwargs.get('reasoning')}" ) assert "summary" not in call_kwargs["reasoning"] # Verify thinking is NOT passed directly to the Responses API - assert ( - "thinking" not in call_kwargs - ), "thinking should NOT be passed directly to litellm.responses" + assert "thinking" not in call_kwargs, "thinking should NOT be passed directly to litellm.responses" class TestThinkingParameterTransformation: @@ -411,9 +400,7 @@ class TestThinkingParameterTransformation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "detailed"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "detailed"}} finally: litellm.reasoning_auto_summary = original @@ -611,9 +598,9 @@ class TestThinkingSummaryPreservation: mock_responses.assert_called_once() call_kwargs = mock_responses.call_args.kwargs reasoning = call_kwargs["reasoning"] - assert ( - reasoning["summary"] == "concise" - ), f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + assert reasoning["summary"] == "concise", ( + f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + ) def test_responses_adapter_preserves_summary(self): """translate_thinking_to_reasoning should include summary when user provides it.""" @@ -622,9 +609,7 @@ class TestThinkingSummaryPreservation: ) thinking = {"type": "enabled", "budget_tokens": 5000, "summary": "concise"} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high", "summary": "concise"} def test_responses_adapter_no_summary_by_default(self): @@ -638,11 +623,7 @@ class TestThinkingSummaryPreservation: try: litellm.reasoning_auto_summary = False thinking = {"type": "enabled", "budget_tokens": 5000} - result = ( - LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high"} assert result is not None and "summary" not in result finally: @@ -659,9 +640,7 @@ class TestThinkingSummaryPreservation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "concise"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "concise"}} def test_translate_thinking_for_model_disabled_stays_plain_string_when_auto_summary_enabled(self): """Disabled thinking must stay a plain string even when reasoning_auto_summary is on.""" @@ -807,9 +786,7 @@ def test_presanitized_flag_not_leaked_to_provider_params(): def fake_base_handler(*args, **kwargs): captured.update(kwargs) - captured["optional"] = kwargs.get( - "anthropic_messages_optional_request_params", {} - ) + captured["optional"] = kwargs.get("anthropic_messages_optional_request_params", {}) return "stub" with patch.object( @@ -974,6 +951,38 @@ def test_gate_passthrough_skipped_when_only_chat_completions_supported(monkeypat assert "config" not in captured +@pytest.mark.parametrize( + "model_info, expected_ttl_support", + [ + ({"supported_endpoints": ["/v1/messages"]}, False), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": True}, True), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": "yes"}, False), + ], +) +def test_gate_passthrough_forwards_cache_control_ttl_only_when_deployment_opts_in( + monkeypatch, model_info, expected_ttl_support +): + """The passthrough config strips cache_control.ttl unless the deployment sets + model_info.cache_control_ttl to exactly true.""" + from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( + anthropic_messages_handler, + ) + + captured, _ = _gate_stubs(monkeypatch) + + result = anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "Hello"}], + model="openai/some-model", + api_key="sk-test", + api_base="https://host/v1", + model_info=model_info, + ) + + assert result == "native-passthrough" + assert captured["config"].supports_cache_control_ttl() is expected_ttl_support + + def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): """Regional and provider-prefixed Claude 4.8+/5 entries carry ``supports_mid_conversation_system``, but the bare first-party keys @@ -987,9 +996,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys import litellm - cost_map_path = os.path.join( - os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json" - ) + cost_map_path = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") with open(cost_map_path) as f: cost_map = json.load(f) rules = cost_map["fallback_generalizations"]["rules"] @@ -1028,9 +1035,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys ("perplexity/sonar", "sonar", "https://api.perplexity.ai/chat/completions"), ], ) -async def test_messages_strips_provider_prefix_exactly_once( - requested_model, expected_wire_model, expected_url -): +async def test_messages_strips_provider_prefix_exactly_once(requested_model, expected_wire_model, expected_url): """ BerriAI/litellm#37716: only the leading provider segment may be stripped on the way upstream. diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index ca913e7fe29..cb05cdb9451 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,6 +979,34 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} +@pytest.mark.parametrize( + "model", + [ + "anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-4-8", + "amazon.nova-pro-v1:0", + "us.meta.llama4-maverick-17b-instruct-v1:0", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", + "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", + ], +) +def test_client_metadata_stripped_from_converse_request(model): + data = AmazonConverseConfig()._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "anthropic_beta": ["computer-use-2025-01-24"], + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + fields = data["additionalModelRequestFields"] + assert "client_metadata" not in fields + assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index d34517f61f6..09be2118001 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -128,6 +128,12 @@ def test_mantle_messages_url_construction(): _VPC_ENDPOINT = "https://vpce-0a1b2c3d.bedrock-mantle.us-gov-west-1.vpce.amazonaws.com" +@pytest.fixture(autouse=True) +def no_ambient_mantle_api_base(monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + + def test_mantle_chat_url_honors_api_base_host(): config = AmazonMantleConfig() url = config.get_complete_url( @@ -193,6 +199,48 @@ def test_mantle_messages_url_honors_aws_bedrock_runtime_endpoint(): assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" +_ENV_ENDPOINT = "https://bedrock-mantle.us-east-1.api.aws.internal.example.com" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +@pytest.mark.parametrize( + "env_value", + [_ENV_ENDPOINT, f"{_ENV_ENDPOINT}/", f"{_ENV_ENDPOINT}/v1", f"{_ENV_ENDPOINT}/openai/v1"], +) +def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls, env_value): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", env_value) + url = config_cls().get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + ) + assert url == f"{_ENV_ENDPOINT}/anthropic/v1/messages" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +@pytest.mark.parametrize( + ("api_base", "optional_params"), + [ + (_VPC_ENDPOINT, {"aws_region_name": "us-gov-west-1"}), + (None, {"aws_region_name": "us-gov-west-1", "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT}), + ], +) +def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env( + monkeypatch, config_cls, api_base, optional_params +): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) + url = config_cls().get_complete_url( + api_base=api_base, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params=optional_params, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + def test_mantle_transform_request_strips_prefix_and_adds_model(): config = AmazonMantleConfig() request = config.transform_request( diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index cd775abf136..56b111f294e 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -486,6 +486,71 @@ class TestBedrockMantleChatAuth: assert "/us-east-2/bedrock/aws4_request" in authorization assert requests[0]["url"].startswith("https://bedrock-mantle.us-east-2.api.aws") + def test_completion_per_request_role_reaches_signer_and_not_the_body(self, monkeypatch): + from unittest.mock import MagicMock, Mock + + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.utils import ModelResponse + + for var in ("BEDROCK_MANTLE_API_KEY", "AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE"): + monkeypatch.delenv(var, raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + return_value=Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + ) + url = "https://bedrock-mantle.us-east-1.api.aws/openai/v1/chat/completions" + client = HTTPHandler(client=httpx.Client()) + client.post = Mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "google.gemma-4-31b", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + request=httpx.Request("POST", url), + ) + ) + + BaseLLMHTTPHandler().completion( + model="google.gemma-4-31b", + messages=[{"role": "user", "content": "hello"}], + api_base=None, + custom_llm_provider="bedrock_mantle", + model_response=ModelResponse(), + encoding=None, + logging_obj=Mock(), + optional_params={}, + timeout=10, + litellm_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/attributed-role", + "aws_session_name": "user-123", + "aws_region_name": "us-east-1", + }, + acompletion=False, + client=client, + provider_config=BedrockMantleChatConfig(aws_signer=signer), + ) + + credential_kwargs = signer.get_credentials.call_args.kwargs + assert credential_kwargs["aws_role_name"] == "arn:aws:iam::000000000000:role/attributed-role" + assert credential_kwargs["aws_session_name"] == "user-123" + sent = client.post.call_args.kwargs + assert sent["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert not [key for key in json.loads(sent["data"]) if key.startswith("aws_")] + class TestBedrockMantleProjectHeader: def test_validate_environment_sets_openai_project_header(self): diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 26f841c1146..1d583c16ad7 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2295,6 +2295,25 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques assert retry_authorization != first_attempt_headers["Authorization"] +def test_aws_signing_overrides_only_fills_missing_credentials(): + from litellm.llms.custom_httpx.llm_http_handler import _aws_signing_overrides + + overrides = _aws_signing_overrides( + {"temperature": 0.2, "aws_region_name": "us-west-2"}, + { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + "aws_region_name": "us-east-1", + "api_key": "not-an-aws-param", + }, + ) + + assert dict(overrides) == { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + } + + class TestServerFulfilledToolsInRequest: """_server_fulfilled_tools_in_request gates the buffered (non-leaking) streaming mode for server-fulfilled tools like headroom_retrieve.""" diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index e6e6aa946d5..9a62fcf6f0f 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,8 +1,14 @@ +import json import os import sys +from typing import Final +from unittest.mock import MagicMock, patch +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.types.rerank import ( @@ -87,9 +93,7 @@ class TestHostedVLLMRerankTransform: assert "instruction" not in body def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self): - with pytest.raises( - ValueError, match="Hosted VLLM does not support max_chunks_per_doc" - ): + with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"): self.config.map_cohere_rerank_params( non_default_params=None, model=self.model, @@ -104,12 +108,10 @@ class TestHostedVLLMRerankTransform: url = self.config.get_complete_url(base, self.model) assert url == "https://api.example.com/rerank" # Already ends with /rerank - url2 = self.config.get_complete_url( - "https://api.example.com/rerank", self.model - ) + url2 = self.config.get_complete_url("https://api.example.com/rerank", self.model) assert url2 == "https://api.example.com/rerank" # Raises if api_base is None - with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'): + with pytest.raises(ValueError, match="api_base must be provided for Hosted VLLM rerank"): self.config.get_complete_url(None, self.model) def test_transform_response(self): @@ -173,3 +175,121 @@ class TestGetOptionalRerankParamsInstruction: documents=["doc1", "doc2"], ) assert "instruction" not in params + + +class TestHostedVLLMRerankTruncationParams: + def setup_method(self): + self.config = HostedVLLMRerankConfig() + self.model = "hosted-vllm-model" + + def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): + params: Final = self.config.map_cohere_rerank_params( + non_default_params={ + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "metadata": {"user_api_key": "sk-test"}, + }, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + max_tokens_per_doc=128, + ) + assert params["truncate_prompt_tokens"] == 512 + assert params["truncation_side"] == "left" + assert params["max_tokens_per_query"] == 64 + assert params["max_tokens_per_doc"] == 128 + assert "metadata" not in params + + @pytest.mark.parametrize( + "bad_params", + [{"truncation_side": "middle"}, {"truncate_prompt_tokens": "lots"}, {"max_tokens_per_query": -1.5}], + ) + def test_map_cohere_rerank_params_rejects_invalid_truncation_params_as_400(self, bad_params: dict[str, object]): + with pytest.raises(litellm.UnsupportedParamsError) as raised: + self.config.map_cohere_rerank_params( + non_default_params=dict(bad_params), + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + assert raised.value.status_code == 400 + assert next(iter(bad_params)) in str(raised.value) + + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): + params: Final = self.config.map_cohere_rerank_params( + non_default_params={"metadata": {"user_api_key": "sk-test"}}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + body: Final = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys: Final = { + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", + "max_tokens_per_doc", + } + assert not truncation_keys & body.keys() + assert body == { + "model": self.model, + "query": "test query", + "documents": ["doc1", "doc2"], + "return_documents": True, + } + + def test_transform_request_forwards_truncation_params(self): + body: Final = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={ + "query": "test query", + "documents": ["doc1", "doc2"], + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "max_tokens_per_doc": 128, + }, + headers={}, + ) + assert body["truncate_prompt_tokens"] == 512 + assert body["truncation_side"] == "left" + assert body["max_tokens_per_query"] == 64 + assert body["max_tokens_per_doc"] == 128 + + def test_transform_request_omits_truncation_params_when_absent(self): + body: Final = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, + headers={}, + ) + assert "truncate_prompt_tokens" not in body + assert "truncation_side" not in body + assert "max_tokens_per_query" not in body + assert "max_tokens_per_doc" not in body + + def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): + client: Final = HTTPHandler() + mock_response: Final = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "score-1", + "results": [{"index": 0, "relevance_score": 0.5}], + "usage": {"total_tokens": 512}, + } + with patch.object(client, "post", return_value=mock_response) as mock_post: + litellm.rerank( + model="hosted_vllm/BAAI/bge-reranker-base", + api_base="http://vllm.local:8000", + query="List all the unique case ids", + documents=["a document longer than the reranker context window"], + truncate_prompt_tokens=512, + truncation_side="left", + client=client, + ) + sent_body: Final = json.loads(mock_post.call_args.kwargs["data"]) + assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" + assert sent_body["truncate_prompt_tokens"] == 512 + assert sent_body["truncation_side"] == "left" diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 9a6a039a470..67a56fdcd79 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -318,3 +318,203 @@ def test_json_provider_messages_config_probes_capabilities_under_provider_slug() ) assert JSONProviderAnthropicMessagesConfig(provider).custom_llm_provider == "exampleprovider" assert OpenAILikeAnthropicMessagesConfig().custom_llm_provider == "anthropic" + + +def _cache_control_request_params() -> tuple[list, dict]: + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "write a regex for a US phone number", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + optional_params = { + "max_tokens": 256, + "system": [ + { + "type": "text", + "text": "You are Claude Code.", + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + } + ], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object"}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + return messages, optional_params + + +def test_request_strips_cache_control_ttl_everywhere(config): + """Regression: Claude Code always sends ``cache_control: {type: ephemeral, + ttl: 1h}``, and strict non-Anthropic /v1/messages validators 400 the whole + request on the ttl extension (``cache_control.ttl: 1h is not supported``).""" + messages, optional_params = _cache_control_request_params() + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_defaults_missing_cache_control_type_and_drops_non_dict(config): + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "a", "cache_control": {"ttl": "1h"}}, + {"type": "text", "text": "b", "cache_control": None}, + ], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + blocks = payload["messages"][0]["content"] + assert blocks[0]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in blocks[1] + + +def test_native_anthropic_config_keeps_cache_control_ttl(): + """Anthropic itself accepts ttl, so the normalization must stay scoped to + the OpenAI-like passthrough and never reach the native Anthropic path.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + messages, optional_params = _cache_control_request_params() + payload = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-20250514", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} + + +def test_deployment_opt_in_keeps_cache_control_ttl(): + config = OpenAILikeAnthropicMessagesConfig(cache_control_ttl=True) + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 16}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_json_provider_constraint_opts_into_cache_control_ttl(): + from litellm.llms.openai_like.json_loader import SimpleProviderConfig + from litellm.llms.openai_like.messages.transformation import ( + JSONProviderAnthropicMessagesConfig, + ) + + base_data = {"base_url": "https://api.example.com/v1", "api_key_env": "EXAMPLE_API_KEY"} + strict = JSONProviderAnthropicMessagesConfig(SimpleProviderConfig(slug="strictprov", data=base_data)) + lenient = JSONProviderAnthropicMessagesConfig( + SimpleProviderConfig(slug="lenientprov", data={**base_data, "constraints": {"cache_control_ttl": True}}) + ) + + def transform(provider_config): + messages, optional_params = _cache_control_request_params() + return provider_config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(config): + """Regression: the sanitizer must only touch ``cache_control`` where the + Messages API defines it (request, system, tools, content blocks, tool_result + content), never application data such as ``tool_use.input`` or a tool's + ``input_schema`` that happens to contain a ``cache_control`` key.""" + tool_input = {"cache_control": {"type": "ephemeral", "ttl": "1h"}, "query": "x"} + input_schema = { + "type": "object", + "properties": {"cache_control": {"type": "string", "ttl": "1h"}}, + } + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": tool_input}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "content": [ + {"type": "text", "text": "result", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + }, + {"type": "text", "text": "plain string content stays", "cache_control": {"ttl": "1h"}}, + ], + }, + {"role": "user", "content": "a plain string message"}, + ] + optional_params = { + "max_tokens": 64, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "tools": [ + { + "name": "lookup", + "input_schema": input_schema, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["input_schema"] == input_schema + assert payload["messages"][0]["content"][0]["input"] == tool_input + tool_result = payload["messages"][1]["content"][0] + assert tool_result["cache_control"] == {"type": "ephemeral"} + assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 0480bbc40a7..d441c05090b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,4 +1,5 @@ import asyncio +import inspect import json import sys from datetime import datetime @@ -13,6 +14,7 @@ import pytest from fastapi import HTTPException from starlette.requests import Request +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, @@ -109,6 +111,71 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert "stack_trace" not in result + @pytest.mark.asyncio + async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def hanging_operation(client): + await asyncio.Event().wait() + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, hanging_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + @pytest.mark.asyncio + async def test_timeout_covers_client_creation(self, monkeypatch): + async def hanging_create_client(*args, **kwargs): + await asyncio.Event().wait() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + hanging_create_client, + ) + + async def unreached_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, unreached_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + def test_timeout_defaults_to_tool_listing_timeout(self): + default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default + assert default == MCP_TOOL_LISTING_TIMEOUT + + def test_connection_error_message_timeout_names_url_and_budget(self): + message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) + assert "https://api.example.com/mcp/" in message + assert "30s" in message + @pytest.mark.asyncio async def test_forwards_static_headers(self, monkeypatch): """Ensure static_headers are forwarded to the MCP client during test calls. @@ -3168,17 +3235,21 @@ class TestConnectionErrorMessage: secret = "Bearer sk-super-secret-token" exc = httpx.LocalProtocolError(f"Illegal header value b' {secret}'") - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "header" in message.lower() assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed")) + message = rest_endpoints._connection_error_message( + httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out")) + message = rest_endpoints._connection_error_message( + httpx.ConnectTimeout("timed out"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): @@ -3188,11 +3259,11 @@ class TestConnectionErrorMessage: request=httpx.Request("POST", "http://x/"), response=response, ) - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message def test_unknown_error_falls_back_to_generic(self): - message = rest_endpoints._connection_error_message(RuntimeError("weird")) + message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message assert "proxy logs" in message.lower() diff --git a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py index e0585ab04f1..4c219760762 100644 --- a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py +++ b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py @@ -591,3 +591,100 @@ class TestFilterServerIdsByIpWithInfo: ) assert allowed == [] assert blocked == 2 + + +def _make_scheme_request( + scheme: str, client_host: str = "203.0.113.5", headers: dict[str, str] | None = None +) -> Request: + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = client_host + request.headers = headers or {} + request.url = MagicMock() + request.url.scheme = scheme + return request + + +class TestIsRequestHttps: + """Regression tests for the cookie Secure trust-boundary resolution. + + litellm only sees a plain-HTTP hop when TLS terminates at a reverse + proxy, so a cookie's Secure attribute must not be derived from the + literal request scheme alone. It must also not blindly trust a + client-spoofable X-Forwarded-Proto header with no trust boundary. + """ + + def test_direct_https_is_secure(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request("https") + assert IPAddressUtils.is_request_https(request, general_settings={}) is True + + def test_direct_http_is_not_secure(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request("http") + assert IPAddressUtils.is_request_https(request, general_settings={}) is False + + def test_spoofed_forwarded_proto_without_trusted_proxy_config_is_ignored( + self, monkeypatch + ): + # Regression: an internal HTTP hop with an attacker-supplied + # X-Forwarded-Proto: https must NOT flip Secure on, because no + # trust boundary (use_x_forwarded_for + mcp_trusted_proxy_ranges) + # is configured. Blindly trusting this header is itself a + # vulnerability. + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "http", headers={"X-Forwarded-Proto": "https"} + ) + assert IPAddressUtils.is_request_https(request, general_settings={}) is False + + def test_forwarded_proto_honored_only_from_trusted_proxy(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "http", + client_host="10.0.0.5", + headers={"X-Forwarded-Proto": "https"}, + ) + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + assert IPAddressUtils.is_request_https(request, general_settings=general_settings) is True + + def test_forwarded_proto_http_from_trusted_proxy_is_not_secure(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "https", + client_host="10.0.0.5", + headers={"X-Forwarded-Proto": "http"}, + ) + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + assert IPAddressUtils.is_request_https(request, general_settings=general_settings) is False + + def test_untrusted_direct_peer_falls_back_to_literal_scheme(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "http", + client_host="203.0.113.5", + headers={"X-Forwarded-Proto": "https"}, + ) + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + assert IPAddressUtils.is_request_https(request, general_settings=general_settings) is False + + def test_proxy_base_url_https_overrides_literal_http_scheme(self, monkeypatch): + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + request = _make_scheme_request("http") + assert IPAddressUtils.is_request_https(request, general_settings={}) is True + + def test_proxy_base_url_http_overrides_literal_https_scheme(self, monkeypatch): + # An explicit operator-configured plain-http public origin wins over + # the literal connection scheme, same as the https direction above. + monkeypatch.setenv("PROXY_BASE_URL", "http://litellm.internal") + request = _make_scheme_request("https") + assert IPAddressUtils.is_request_https(request, general_settings={}) is False diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py index 57decc7d458..36d6414ba9f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -642,3 +642,77 @@ async def test_read_acs_post_data_rejects_oversized_stream_without_content_lengt with pytest.raises(HTTPException) as exc: await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) assert exc.value.status_code == 413 + + +def _fake_request_with_scheme(scheme, headers=None, client_host="203.0.113.5"): + """A fuller fake Request than ``_fake_request``: adds ``url``, ``headers`` and + ``client``, which ``IPAddressUtils.is_request_https`` reads directly instead of + going through ``PROXY_BASE_URL``.""" + return type( + "Req", + (), + { + "base_url": URL(f"{scheme}://proxy.example.com/"), + "url": URL(f"{scheme}://proxy.example.com/sso/saml/login"), + "query_params": {}, + "cookies": {}, + "headers": headers or {}, + "client": type("Client", (), {"host": client_host})(), + }, + )() + + +class TestSAMLAuthnCookieSecureFlag: + """Regression tests for the litellm_saml_authn cookie's Secure attribute. + litellm only sees a plain-HTTP hop whenever TLS terminates at a reverse + proxy, so Secure must not be derived from the literal request scheme alone.""" + + @pytest.mark.asyncio + async def test_secure_over_direct_https(self, saml_env, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + cache = DualCache() + request = _fake_request_with_scheme("https") + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" in cookie + assert "SameSite=none" in cookie + + @pytest.mark.asyncio + async def test_not_secure_over_direct_http(self, saml_env, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + cache = DualCache() + request = _fake_request_with_scheme("http") + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" not in cookie + assert "SameSite=lax" in cookie + + @pytest.mark.asyncio + async def test_secure_behind_trusted_tls_terminating_proxy(self, saml_env, monkeypatch): + """THE regression: TLS terminates at a reverse proxy, litellm only sees a + plain-HTTP hop, but the cookie must still be marked Secure when the operator + has configured a trusted proxy reporting X-Forwarded-Proto: https.""" + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + cache = DualCache() + request = _fake_request_with_scheme( + "http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5" + ) + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" in cookie + + @pytest.mark.asyncio + async def test_untrusted_spoofed_forwarded_proto_is_ignored(self, saml_env, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + cache = DualCache() + request = _fake_request_with_scheme( + "http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5" + ) + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" not in cookie diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index e648bd09734..dd8c752a868 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -7604,6 +7604,112 @@ class TestPKCEStateCookieBinding: assert cookie_str is not None assert "Secure" not in cookie_str + @pytest.mark.asyncio + async def test_redirect_response_sets_secure_flag_behind_trusted_tls_terminating_proxy( + self, monkeypatch + ): + """Regression: litellm sees a plain-HTTP hop when TLS terminates at a reverse + proxy. The Secure flag must still be set when the direct peer is a configured + trusted proxy and it reports X-Forwarded-Proto: https -- but NOT from an + unconfigured/untrusted caller spoofing the same header (see the sibling test + below).""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="http://idp.internal/authorize?state=behind-proxy-state" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + proxied_request = MagicMock(spec=Request) + proxied_request.url.scheme = "http" + proxied_request.headers = {"X-Forwarded-Proto": "https"} + proxied_request.client = MagicMock() + proxied_request.client.host = "10.0.0.5" + + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "behind-proxy-state", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="http://idp.internal/authorize", + request=proxied_request, + ) + + cookie_headers = response.headers.getlist("set-cookie") + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert cookie_str is not None + assert "Secure" in cookie_str + + @pytest.mark.asyncio + async def test_redirect_response_ignores_spoofed_forwarded_proto_without_trust_config( + self, monkeypatch + ): + """The same X-Forwarded-Proto: https header must NOT flip Secure on when no + trusted-proxy config is present -- honoring it unconditionally would let any + client spoof the header and would not itself be the vulnerability the ticket + warns against.""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="http://idp.internal/authorize?state=spoofed-state" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + spoofed_request = MagicMock(spec=Request) + spoofed_request.url.scheme = "http" + spoofed_request.headers = {"X-Forwarded-Proto": "https"} + spoofed_request.client = MagicMock() + spoofed_request.client.host = "203.0.113.5" + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "spoofed-state", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="http://idp.internal/authorize", + request=spoofed_request, + ) + + cookie_headers = response.headers.getlist("set-cookie") + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert cookie_str is not None + assert "Secure" not in cookie_str + @pytest.mark.asyncio async def test_pkce_callback_rejects_missing_cookie(self): """When PKCE is enabled and a code_verifier is in the cache, the @@ -8586,6 +8692,24 @@ class TestSameOriginReturnPath: assert _is_same_origin_return_path("") is False +def _make_https_request() -> Request: + request = MagicMock(spec=Request) + request.url.scheme = "https" + request.headers = {} + request.client = MagicMock() + request.client.host = "203.0.113.5" + return request + + +def _make_http_request() -> Request: + request = MagicMock(spec=Request) + request.url.scheme = "http" + request.headers = {} + request.client = MagicMock() + request.client.host = "203.0.113.5" + return request + + class TestPersistReturnToCookieSharedHelper: """The single shared return_to helper used by EVERY sign-in branch (SSO / Okta / generic AND the username/password form). It must be best-effort and NEVER raise — a bad return_to can never block @@ -8603,7 +8727,7 @@ class TestPersistReturnToCookieSharedHelper: monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) resp = Response() - _persist_return_to_cookie(resp, "/mcp/authorize?client_id=llm_dcrc_abc") + _persist_return_to_cookie(resp, "/mcp/authorize?client_id=llm_dcrc_abc", _make_https_request()) assert "litellm_cp_return_to=" in self._cookie(resp) def test_bad_absolute_with_control_plane_configured_does_not_raise_and_is_not_stored(self, monkeypatch): @@ -8617,7 +8741,7 @@ class TestPersistReturnToCookieSharedHelper: "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"} ) resp = Response() - _persist_return_to_cookie(resp, "https://evil.example.com/steal") # must not raise + _persist_return_to_cookie(resp, "https://evil.example.com/steal", _make_https_request()) # must not raise assert "litellm_cp_return_to=" not in self._cookie(resp) def test_none_return_to_is_a_noop(self): @@ -8626,7 +8750,7 @@ class TestPersistReturnToCookieSharedHelper: from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie resp = Response() - _persist_return_to_cookie(resp, None) + _persist_return_to_cookie(resp, None, _make_https_request()) assert "litellm_cp_return_to=" not in self._cookie(resp) def test_control_plane_matching_absolute_is_stored(self, monkeypatch): @@ -8638,5 +8762,126 @@ class TestPersistReturnToCookieSharedHelper: "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"} ) resp = Response() - _persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models") + _persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models", _make_https_request()) assert "litellm_cp_return_to=" in self._cookie(resp) + + def test_cookie_is_secure_and_httponly_over_https(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + resp = Response() + _persist_return_to_cookie(resp, "/mcp/authorize", _make_https_request()) + cookie = self._cookie(resp) + assert "Secure" in cookie + assert "HttpOnly" in cookie + assert "SameSite=lax" in cookie + + def test_cookie_is_not_secure_over_plain_http_direct(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + resp = Response() + _persist_return_to_cookie(resp, "/mcp/authorize", _make_http_request()) + assert "Secure" not in self._cookie(resp) + + def test_cookie_is_secure_behind_trusted_tls_terminating_proxy(self, monkeypatch): + """Regression for the reported bug: TLS terminates at a reverse proxy, litellm only + sees a plain-HTTP hop, but a trusted X-Forwarded-Proto: https must still mark the + cookie Secure.""" + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + resp = Response() + request = _make_http_request() + request.client.host = "10.0.0.5" + request.headers = {"X-Forwarded-Proto": "https"} + _persist_return_to_cookie(resp, "/mcp/authorize", request) + assert "Secure" in self._cookie(resp) + + +class TestSessionTokenCookie: + """Regression tests for the ``token`` session cookie set by every sign-in path + (username/password login, SSO callback, the CLI /v2, /v3 login exchange helpers). + It was previously set with no Secure/HttpOnly/SameSite attributes at all -- always + sent over plain HTTP and readable by any script on the page. HttpOnly must stay off + deliberately: the dashboard reads this cookie via document.cookie.""" + + @staticmethod + def _cookie(resp) -> str: + return resp.headers.get("set-cookie", "") + + def test_secure_over_direct_https(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + resp = Response() + set_session_token_cookie(resp, _make_https_request(), "jwt-token-value") + cookie = self._cookie(resp) + assert "token=jwt-token-value" in cookie + assert "Secure" in cookie + assert "SameSite=lax" in cookie + assert "HttpOnly" not in cookie + + def test_not_secure_over_direct_http(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + resp = Response() + set_session_token_cookie(resp, _make_http_request(), "jwt-token-value") + assert "Secure" not in self._cookie(resp) + + def test_secure_behind_trusted_tls_terminating_proxy(self, monkeypatch): + """THE regression: TLS terminates at a reverse proxy, litellm only sees a + plain-HTTP hop, but the session cookie must still be marked Secure when the + operator has configured a trusted proxy that reports X-Forwarded-Proto: https.""" + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + request = _make_http_request() + request.client.host = "10.0.0.5" + request.headers = {"X-Forwarded-Proto": "https"} + resp = Response() + set_session_token_cookie(resp, request, "jwt-token-value") + assert "Secure" in self._cookie(resp) + + def test_untrusted_spoofed_forwarded_proto_is_ignored(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + request = _make_http_request() + request.headers = {"X-Forwarded-Proto": "https"} + resp = Response() + set_session_token_cookie(resp, request, "jwt-token-value") + assert "Secure" not in self._cookie(resp) + + def test_proxy_base_url_https_overrides_literal_http_scheme(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + resp = Response() + set_session_token_cookie(resp, _make_http_request(), "jwt-token-value") + assert "Secure" in self._cookie(resp) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5f692f8f109..4ed6a468371 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -148,6 +148,72 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): assert mock_jwt_encode.call_args.kwargs == {"algorithm": "HS256"} +def _mock_login_v2_deps(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + + +def test_login_v2_sets_secure_cookie_over_direct_https(monkeypatch): + """Regression: the token cookie previously carried no Secure/HttpOnly/SameSite + attributes at all, so it was always sent over plain HTTP.""" + _mock_login_v2_deps(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + client = TestClient(app, base_url="https://testserver") + response = client.post("/v2/login", json={"username": "alice", "password": "secret"}) + + assert response.status_code == 200 + cookie = response.headers.get("set-cookie") + assert "Secure" in cookie + assert "HttpOnly" not in cookie # deliberate: the dashboard reads this cookie via JS + assert "samesite=lax" in cookie.lower() + + +def test_login_v2_does_not_set_secure_cookie_over_direct_http(monkeypatch): + _mock_login_v2_deps(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + client = TestClient(app, base_url="http://testserver") + response = client.post("/v2/login", json={"username": "alice", "password": "secret"}) + + assert response.status_code == 200 + assert "Secure" not in response.headers.get("set-cookie") + + +def test_login_v2_sets_secure_cookie_behind_trusted_tls_terminating_proxy(monkeypatch): + """THE regression: litellm only sees a plain-HTTP hop when TLS terminates at a + reverse proxy, but the token cookie must still be Secure when the direct peer is + a configured trusted proxy reporting X-Forwarded-Proto: https.""" + _mock_login_v2_deps(monkeypatch) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + + client = TestClient(app, base_url="http://testserver", client=("10.0.0.5", 50000)) + response = client.post( + "/v2/login", + json={"username": "alice", "password": "secret"}, + headers={"X-Forwarded-Proto": "https"}, + ) + + assert response.status_code == 200 + assert "Secure" in response.headers.get("set-cookie") + + def test_login_v2_returns_json_on_proxy_exception(monkeypatch): """Test that /v2/login returns JSON error when ProxyException is raised""" from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -356,6 +422,51 @@ def test_login_v3_exchange_happy_path(monkeypatch): assert exchange_response.cookies.get("token") == "signed-token" +def test_login_v3_exchange_sets_secure_cookie_behind_trusted_tls_terminating_proxy(monkeypatch): + """Regression: /v3/login/exchange's token cookie must be Secure behind a trusted + TLS-terminating reverse proxy even though litellm only sees a plain-HTTP hop.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + { + "control_plane_url": "https://cp.example.com", + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_config = MagicMock() + mock_config.worker_registry = [] + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + + client = TestClient(app, base_url="http://testserver", client=("10.0.0.5", 50000)) + + login_response = client.post("/v3/login", json={"username": "alice", "password": "secret"}) + code = login_response.json()["code"] + + exchange_response = client.post( + "/v3/login/exchange", + json={"code": code}, + headers={"X-Forwarded-Proto": "https"}, + ) + assert exchange_response.status_code == 200 + assert "Secure" in exchange_response.headers.get("set-cookie") + + def test_login_v3_exchange_single_use(monkeypatch): """Code can only be redeemed once.""" mock_prisma_client = MagicMock() diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 82659571de4..941d78085e4 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10089,6 +10089,207 @@ class TestHeuristicFirst: assert outcome.cause == "default_model_fallback" +# Scores 0.175 with one signal, so it sits 0.025 from simple_medium: the pair of tiers either side of +# that boundary are different model pools, and a hair's difference in score picks the other one. +NEAR_BOUNDARY_PROMPT = ( + "design a distributed cache with consistent hashing, then explain the failure modes step by step" +) + +# Scores 0.075 with signals, the far side of any margin under 0.075: the scorer is decided here. +CLEAR_OF_BOUNDARY_PROMPT = "explain step by step how consistent hashing rebalances keys" + + +def _hybrid_router(mock_router_instance, **config_overrides): + config = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "tier_boundaries": dict(HEURISTIC_FIRST_BOUNDARIES), + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.03, + "classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400}, + **config_overrides, + } + return ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + + +class TestHybridConfig: + """Config validation for classifier_type='hybrid'.""" + + @pytest.mark.parametrize( + "overrides, expected", + [ + ({"classifier_llm_config": None}, "classifier_llm_config is required"), + ({"hybrid_boundary_margin": None}, "hybrid_boundary_margin is required"), + ({"hybrid_boundary_margin": -0.01}, "greater than or equal to 0"), + ({"hybrid_boundary_margin": 1.01}, "less than or equal to 1"), + ], + ) + def test_rejects_incoherent_config(self, overrides, expected): + config = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.03, + "classifier_llm_config": {"model": "haiku-classifier"}, + **overrides, + } + with pytest.raises(ValidationError, match=expected): + ComplexityRouterConfig(**config) + + @pytest.mark.parametrize("classifier_type", ["heuristic", "llm", "custom", "heuristic_first"]) + def test_margin_rejected_on_every_other_classifier_type(self, classifier_type): + """A margin on a router that never compares a score to a boundary is a silent no-op, so it is + refused rather than accepted and ignored. heuristic_first is in this list on purpose: its + ceiling is a different question from proximity, and accepting both on one router would make + two modes out of one classifier_type.""" + config: dict[str, object] = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "classifier_type": classifier_type, + "hybrid_boundary_margin": 0.03, + } + if classifier_type in ("llm", "heuristic_first"): + config["classifier_llm_config"] = {"model": "haiku-classifier"} + if classifier_type == "heuristic_first": + config["heuristic_first_max_tier"] = "SIMPLE" + if classifier_type == "custom": + config["classifier_plugin"] = _FixedTierClassifier("SIMPLE") + with pytest.raises(ValidationError, match="hybrid_boundary_margin is set but classifier_type"): + ComplexityRouterConfig(**config) + + def test_the_cheap_tier_ceiling_is_rejected_here(self): + """The two modes are told apart by which knob they take, so the ceiling is refused on hybrid + exactly as the margin is refused on heuristic_first.""" + with pytest.raises(ValidationError, match="heuristic_first_max_tier is set but classifier_type"): + ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + heuristic_first_max_tier="SIMPLE", + classifier_llm_config={"model": "haiku-classifier"}, + ) + + def test_custom_tier_set_is_rejected(self): + """The scorer only emits the four built-in tiers, so it cannot judge proximity on a replaced set.""" + with pytest.raises(ValidationError, match="tier_definitions requires classifier_type"): + ComplexityRouterConfig( + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + classifier_llm_config={"model": "haiku-classifier"}, + tier_definitions=[{"name": "lo", "description": "x"}, {"name": "hi", "description": "y"}], + tiers={"lo": "gpt-4o-mini", "hi": "gpt-4o"}, + ) + + def test_classifier_model_is_a_dependency(self): + config = ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + classifier_llm_config={"model": "haiku-classifier"}, + ) + assert config.uses_llm_classifier is True + + +class TestHybrid: + """Behavior of the hybrid chain: the scorer keeps its tier unless the score is near a boundary.""" + + @pytest.mark.asyncio + async def test_near_boundary_prompt_escalates(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + router = _hybrid_router(mock_router_instance) + + _tier, score, signals, _cause = router._score_and_classify(NEAR_BOUNDARY_PROMPT) + assert signals and abs(score - HEURISTIC_FIRST_BOUNDARIES["simple_medium"]) < 0.03 + + outcome = await router.aclassify(NEAR_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.tier == ComplexityTier.COMPLEX + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_score_clear_of_every_boundary_keeps_the_heuristic_tier(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock() + router = _hybrid_router(mock_router_instance) + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_not_called() + assert outcome.tier == ComplexityTier.SIMPLE + assert outcome.cause == "hybrid_short_circuit" + + @pytest.mark.asyncio + async def test_an_expensive_tier_short_circuits_too(self, mock_router_instance): + """This is the whole difference from heuristic_first, which would have escalated this by tier + alone. Hybrid asks whether the score is DECIDED, not whether the tier is cheap.""" + mock_router_instance.acompletion = AsyncMock() + router = _hybrid_router( + mock_router_instance, + tier_boundaries={"simple_medium": -0.9, "medium_complex": -0.8, "complex_reasoning": -0.7}, + ) + + tier, _score, signals, _cause = router._score_and_classify(CLEAR_OF_BOUNDARY_PROMPT) + assert (tier, bool(signals)) == (ComplexityTier.REASONING, True) + + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_not_called() + assert outcome.tier == ComplexityTier.REASONING + assert outcome.cause == "hybrid_short_circuit" + + @pytest.mark.asyncio + async def test_widening_the_margin_escalates_what_a_narrow_one_kept(self, mock_router_instance): + """The margin is the knob: the same prompt short-circuits at 0.03 and escalates at 0.08.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}')) + router = _hybrid_router(mock_router_instance, hybrid_boundary_margin=0.08) + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_a_zero_margin_escalates_only_an_exact_boundary_score(self, mock_router_instance): + """0 is a real margin, not an off switch: a score sitting exactly on the line still escalates. + + The boundary is spelled as the scorer's own accumulated float rather than the 0.075 it prints + as, because the comparison is on raw floats: a boundary written 0.075 sits 1.4e-17 away from + this score and a zero margin correctly declines to call that exact.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}')) + on_the_line = 0.07499999999999998 + router = _hybrid_router( + mock_router_instance, + tier_boundaries={"simple_medium": on_the_line, "medium_complex": 0.35, "complex_reasoning": 0.60}, + hybrid_boundary_margin=0, + ) + + _tier, score, _signals, _cause = router._score_and_classify(CLEAR_OF_BOUNDARY_PROMPT) + assert score == on_the_line + + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_no_signal_prompt_escalates_however_far_from_a_boundary(self, mock_router_instance): + """The scorer with no opinion has no tier to be confident about, so proximity cannot save it.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + router = _hybrid_router(mock_router_instance) + + tier, score, signals, _cause = router._score_and_classify(NO_SIGNAL_PROMPT) + assert (tier, score, signals) == (ComplexityTier.SIMPLE, 0.0, ()) + + outcome = await router.aclassify(NO_SIGNAL_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_classifier_failure_falls_back_to_the_scorer(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock(side_effect=RuntimeError("classifier exploded")) + router = _hybrid_router(mock_router_instance) + expected_tier, expected_score, expected_signals, _cause = router._score_and_classify(NEAR_BOUNDARY_PROMPT) + + outcome = await router.aclassify(NEAR_BOUNDARY_PROMPT) + + assert (outcome.tier, outcome.score, outcome.signals) == (expected_tier, expected_score, expected_signals) + assert outcome.cause == "heuristic_scorer" + + def _windowed_router(*deployments: tuple) -> Router: """Real Router; each deployment is (group, provider_model, declared window or None). None means no declared override on a model the cost map does not know: unresolvable.""" @@ -10334,7 +10535,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} first = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS @@ -10357,7 +10559,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} pinned = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 610feeb1c26..041f19558dd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -6,7 +6,7 @@ "limit": 26771 }, "LIT003": { - "limit": 269 + "limit": 261 }, "LIT004": { "limit": 40 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16493 + "limit": 16485 }, "LIT011": { - "limit": 5535 + "limit": 5521 }, "LIT012": { "limit": 4492 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts index bbdf4697315..c4d7f45b7cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts @@ -58,6 +58,7 @@ const dedupe = (models: string[]): string[] => Array.from(new Set(models)); const COMPLEXITY_TYPE_LABELS: Record = { llm: "LLM Classifier", heuristic_first: "Heuristic first", + hybrid: "Hybrid", custom: "Custom classifier", }; diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index c2ed4b1f55a..a29527a20fa 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -35,6 +35,7 @@ import { heuristicScoringRole, usesLlmClassifier, DEFAULT_HEURISTIC_FIRST_MAX_TIER, + DEFAULT_HYBRID_BOUNDARY_MARGIN, HEURISTIC_FIRST_MAX_TIER_KEYS, effectiveClassifierType, } from "./ComplexityRouterConfig"; @@ -50,6 +51,7 @@ const HEURISTIC_V2_EXPLANATION = const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms"; const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; +const HYBRID_BOUNDARY_MARGIN_ID = "hybrid-boundary-margin"; const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + @@ -213,6 +215,18 @@ const ClassifierTypeRadios: React.FC<{ + + + ); @@ -263,6 +277,8 @@ const ClassificationMethodConfig: React.FC = ({ classifierType === "heuristic_first" ? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER : undefined, + hybrid_boundary_margin: + classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined, }; onChange(nextValue); }; @@ -271,6 +287,13 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, heuristic_first_max_tier: tier }); }; + const handleHybridBoundaryMarginChange = (raw: string) => { + setDraft({ id: HYBRID_BOUNDARY_MARGIN_ID, raw }); + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onChange({ ...value, hybrid_boundary_margin: Math.min(1, Math.max(0, parsed)) }); + }; + const handleClassificationPromptChange = (classificationPrompt: string | undefined) => { onChange({ ...value, classification_prompt: classificationPrompt }); }; @@ -391,6 +414,30 @@ const ClassificationMethodConfig: React.FC = ({ )} + {classifierType === "hybrid" && ( +
+ Boundary margin + handleHybridBoundaryMarginChange(event.target.value)} + onBlur={() => setDraft(null)} + className="w-full" + /> +

+ A score further than this from every tier boundary routes on the scorer's own tier, however expensive + that tier is. A score closer than this, and anything the scorer found no signal for at all, goes to the + classifier to break the tie +

+
+ )} +
How often to classify - classifierType === "llm" || classifierType === "heuristic_first"; + classifierType === "llm" || classifierType === "heuristic_first" || classifierType === "hybrid"; export type ClassifierFallback = "heuristic" | "default_model"; @@ -162,7 +162,8 @@ export const heuristicScoringRoleFor = ( classifierFallback: ClassifierFallback | undefined, ): HeuristicScoringRole => { if (classifierType === "heuristic_v2") return "never"; - if (classifierType === "heuristic" || classifierType === "heuristic_first") return "decides"; + if (classifierType === "heuristic" || classifierType === "heuristic_first" || classifierType === "hybrid") + return "decides"; return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never"; }; @@ -404,6 +405,8 @@ export interface ComplexityRouterConfigValue { classification_prompt?: string; /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */ heuristic_first_max_tier?: string; + /** How near a tier boundary a score may land before hybrid defers to the classifier. Required by that type, rejected by the others. */ + hybrid_boundary_margin?: number; classification_mode?: ClassificationMode; session_affinity?: boolean; modality_routing?: boolean; @@ -516,6 +519,9 @@ export const effectiveTierLabel = (tier: keyof ComplexityTiers, tierLabels: Comp export const DEFAULT_HEURISTIC_FIRST_MAX_TIER = "SIMPLE"; +/** What the Hybrid radio starts at. Required by that type, so the form always has a value to send. */ +export const DEFAULT_HYBRID_BOUNDARY_MARGIN = 0.03; + /** * Tiers the heuristic_first threshold may name. The top tier is excluded because it would short * circuit every request and leave the classifier unreachable, which the backend rejects. diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 03b7432e4ec..81acb7b7e50 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -342,6 +342,7 @@ const AddAutoRouterTab: React.FC = ({ planModeMinTier: complexityRouterConfig.plan_mode_min_tier, classificationPrompt: complexityRouterConfig.classification_prompt, heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, + hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin, classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, classifierType: complexityRouterConfig.classifier_type, 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 2d8f55b2b64..c6370541a45 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 @@ -817,6 +817,39 @@ describe("heuristic_first", () => { }); }); +describe("hybrid", () => { + const hybridParams: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "hybrid", + hybridBoundaryMargin: 0.03, + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifierFallback: "default_model", + }; + + it("emits hybrid_boundary_margin, zero included since exactly-on-a-boundary is a real setting", () => { + expect(buildComplexityRouterConfig(hybridParams).hybrid_boundary_margin).toBe(0.03); + expect(buildComplexityRouterConfig({ ...hybridParams, hybridBoundaryMargin: 0 }).hybrid_boundary_margin).toBe(0); + }); + + it("keeps every classifier key the operator set, since hybrid still calls the classifier", () => { + const config = buildComplexityRouterConfig(hybridParams); + expect(config.classifier_type).toBe("hybrid"); + expect(config.classifier_llm_config).toEqual({ model: "gpt-4o-mini", timeout_ms: 3000 }); + expect(config.classifier_fallback).toBe("default_model"); + }); + + it("omits hybrid_boundary_margin on every other classifier type, which the backend rejects it on", () => { + for (const classifierType of ["heuristic", "llm", "heuristic_first"] as const) { + const config = buildComplexityRouterConfig({ + ...hybridParams, + classifierType, + ...(classifierType === "heuristic_first" && { heuristicFirstMaxTier: "SIMPLE" }), + }); + expect(config.hybrid_boundary_margin).toBeUndefined(); + } + }); +}); + describe("classification_mode", () => { it("emits user_turn", () => { const config = buildComplexityRouterConfig({ ...baseParams, classificationMode: "user_turn" }); @@ -924,10 +957,12 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { dimensionWeights: { length: 1 }, reasoningOverrideMinScore: 0.5, heuristicFirstMaxTier: "SIMPLE", + hybridBoundaryMargin: 0.03, customTechnicalKeywords: ["kubernetes"], }; const emittingType = key === "heuristic_first_max_tier" ? "heuristic_first" : "llm"; - expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: emittingType })).toHaveProperty(key); + const typeForKey = key === "hybrid_boundary_margin" ? "hybrid" : emittingType; + expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: typeForKey })).toHaveProperty(key); expect(build(loaded)).not.toHaveProperty(key); }); 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 af34dc92c0f..94c5badf6d1 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 @@ -107,6 +107,7 @@ export interface BuildComplexityRouterConfigParams { classifierFallback: ClassifierFallback | undefined; classificationPrompt: string | undefined; heuristicFirstMaxTier: string | undefined; + hybridBoundaryMargin?: number; classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; modalityRouting?: boolean; @@ -163,6 +164,7 @@ export interface ComplexityRouterConfigPayload { classifier_fallback?: ClassifierFallback; classification_prompt?: string; heuristic_first_max_tier?: string; + hybrid_boundary_margin?: number; classification_mode: ClassificationMode; session_affinity: boolean; deployment_affinity: boolean; @@ -352,6 +354,7 @@ const classifierWireFields = ( classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, @@ -360,6 +363,7 @@ const classifierWireFields = ( | "classifierLlmConfig" | "classifierFallback" | "heuristicFirstMaxTier" + | "hybridBoundaryMargin" | "classifierContextWindowSize" | "classifierContextBudgetChars" | "classifierContextIncludeAssistantTurns" @@ -371,6 +375,8 @@ const classifierWireFields = ( classifierFallback !== undefined && { classifier_fallback: classifierFallback }), ...(effectiveType === "heuristic_first" && heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }), + ...(effectiveType === "hybrid" && + hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }), ...(usesLlmClassifier(effectiveType) && classifierContextWindowSize !== undefined && { classifier_context_window_size: classifierContextWindowSize, @@ -399,6 +405,7 @@ export const buildComplexityRouterConfig = ({ classifierFallback, classificationPrompt, heuristicFirstMaxTier, + hybridBoundaryMargin, classificationMode, sessionAffinity, modalityRouting, @@ -444,6 +451,7 @@ export const buildComplexityRouterConfig = ({ classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, diff --git a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts index d2706a82037..3c5b149f4da 100644 --- a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts +++ b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts @@ -122,10 +122,10 @@ export const CUSTOM_TIER_RESTRICTIONS = { reason: "Session pinning escalates along the built-in tier ladder, which your tier set replaces", }, heuristicClassifier: { - omit: ["heuristic_first_max_tier"], + omit: ["heuristic_first_max_tier", "hybrid_boundary_margin"], reason: "The heuristic scorer only produces the built-in tiers, so an edited set needs the LLM classifier. " + - "Heuristic first is out for the same reason: its local scorer decides the cheap traffic", + "Heuristic first and hybrid are out for the same reason: their local scorer decides the traffic it is sure of", }, heuristicScoring: { omit: [ diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 3e63cbd3b32..fae6744d3ac 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -520,16 +520,21 @@ describe("managed keys survive an untouched open-and-save", () => { }; // tier_definitions, fallback_tier and classification_prompt cannot sit beside heuristic_first, which - // this fixture uses, so no single stored config can hold every managed key. They get their own round - // trip below. - const CUSTOM_TIER_ONLY_KEYS = new Set(["tier_definitions", "fallback_tier", "classification_prompt"]); + // this fixture uses, and hybrid_boundary_margin belongs to the sibling hybrid type, so no single + // stored config can hold every managed key. Each gets its own round trip below. + const KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS = new Set([ + "tier_definitions", + "fallback_tier", + "classification_prompt", + "hybrid_boundary_margin", + ]); it("carries every managed key a built-in router can hold through hydrate then save", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); const saved = buildUpdatedComplexityRouterConfig(STORED_ALL_MANAGED, hydrated); const dropped = [...MANAGED_COMPLEXITY_ROUTER_KEYS] - .filter((key) => !CUSTOM_TIER_ONLY_KEYS.has(key)) + .filter((key) => !KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS.has(key)) .filter((key) => saved[key] === undefined); expect(dropped).toEqual([]); }); @@ -583,6 +588,18 @@ describe("managed keys survive an untouched open-and-save", () => { ); }); + it("round-trips a hybrid router's margin, which save requires and the backend rejects without", () => { + const storedHybrid: Record = { + ...STORED_ALL_MANAGED, + classifier_type: "hybrid", + hybrid_boundary_margin: 0.05, + }; + delete storedHybrid.heuristic_first_max_tier; + const hydrated = hydrateComplexityRouterConfig(storedHybrid, undefined); + expect(hydrated.hybrid_boundary_margin).toBe(0.05); + expect(buildUpdatedComplexityRouterConfig(storedHybrid, hydrated).hybrid_boundary_margin).toBe(0.05); + }); + it("round-trips the heuristic_first threshold, which save requires and the backend rejects without", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); expect(hydrated.heuristic_first_max_tier).toBe("SIMPLE"); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index e18582f77a0..7aa5f4beb1f 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -89,6 +89,7 @@ export interface StoredComplexityRouterConfig { plan_mode_min_tier?: unknown; classification_prompt?: unknown; heuristic_first_max_tier?: unknown; + hybrid_boundary_margin?: unknown; tier_labels?: unknown; classifier_type?: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; @@ -167,6 +168,8 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" ? parsedConfig.heuristic_first_max_tier : undefined, + hybrid_boundary_margin: + typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, classification_mode: parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" ? parsedConfig.classification_mode @@ -214,6 +217,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_fallback", "classification_prompt", "heuristic_first_max_tier", + "hybrid_boundary_margin", "classification_mode", "session_affinity", "modality_routing", @@ -303,6 +307,7 @@ export const buildUpdatedComplexityRouterConfig = ( planModeMinTier: value.plan_mode_min_tier, classificationPrompt: value.classification_prompt, heuristicFirstMaxTier: value.heuristic_first_max_tier, + hybridBoundaryMargin: value.hybrid_boundary_margin, classificationMode: value.classification_mode, tierLabels: value.tier_labels, classifierType: value.classifier_type, diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index 66e80dcc655..ffef01010b7 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -86,6 +86,7 @@ const CONSTANT_CAUSE_LABELS: Record = { heuristic_scorer: "Heuristic scorer", heuristic_v2: "Heuristic v2", heuristic_first_short_circuit: "Heuristic scorer, classifier skipped", + hybrid_short_circuit: "Heuristic scorer, score clear of every boundary", classifier_plugin: "Custom classifier plugin", semantic_keyword_match: "Semantic keyword match", session_affinity_pin: "Pinned to session", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 79e44de4f1f..be81450b941 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34560,7 +34560,7 @@ export interface components { * @enum {string} */ classifier_fallback: "heuristic" | "default_model"; - /** @description Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first' */ + /** @description Configuration for the LLM classifier; required when classifier_type is 'llm', 'heuristic_first' or 'hybrid' */ classifier_llm_config?: components["schemas"]["ClassifierLLMConfig"] | null; /** * Classifier Plugin @@ -34575,11 +34575,11 @@ export interface components { classifier_plugin_timeout_ms: number; /** * Classifier Type - * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier + * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary * @default heuristic * @enum {string} */ - classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first"; + classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first" | "hybrid"; /** * Code Keywords * @description Keywords indicating code-related content @@ -34651,6 +34651,11 @@ export interface components { * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. */ housekeeping_patterns?: string[] | null; + /** + * Hybrid Boundary Margin + * @description How close to a tier boundary a heuristic score has to land before the LLM classifier breaks the tie; required when classifier_type is 'hybrid' and rejected otherwise. Everything further than this from every active boundary routes on the scorer's own tier with no classifier call, at any tier, which is what separates 'hybrid' from 'heuristic_first' and its cheap-tier ceiling. A prompt where no dimension fired still goes to the classifier, since the scorer has no opinion to be near a boundary with. 0 escalates only scores sitting exactly on a boundary. + */ + hybrid_boundary_margin?: number | null; /** * Keyword Tier Rules * @description Rules that force a specific tier when their keywords match the prompt @@ -35867,7 +35872,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; /** Classifier Cost */ classifier_cost?: number; /** Classifier Model */