mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
merge: origin/litellm_internal_staging into litellm_lite_login_prefill_code
This commit is contained in:
mateo-berri
commit
62badd58fb
55 changed files with 1990 additions and 211 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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": <its type, or "ephemeral">}``
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"}]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ const dedupe = (models: string[]): string[] => Array.from(new Set(models));
|
|||
const COMPLEXITY_TYPE_LABELS: Record<string, string> = {
|
||||
llm: "LLM Classifier",
|
||||
heuristic_first: "Heuristic first",
|
||||
hybrid: "Hybrid",
|
||||
custom: "Custom classifier",
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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<{
|
|||
</span>
|
||||
</Label>
|
||||
</SimpleTooltip>
|
||||
<SimpleTooltip content={scorerLockedReason}>
|
||||
<Label className="items-start font-normal leading-normal has-data-disabled:cursor-not-allowed has-data-disabled:opacity-50">
|
||||
<RadioGroupItem value="hybrid" className="mt-0.5" disabled={scorerLocked} />
|
||||
<span>
|
||||
<strong className="font-semibold">Hybrid</strong>{" "}
|
||||
<span className="text-muted-foreground">
|
||||
keeps the local score at any tier, and only pays for the classifier when that score lands near a tier
|
||||
boundary
|
||||
</span>
|
||||
</span>
|
||||
</Label>
|
||||
</SimpleTooltip>
|
||||
</div>
|
||||
</RadioGroup>
|
||||
);
|
||||
|
|
@ -263,6 +277,8 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
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<ClassificationMethodConfigProps> = ({
|
|||
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<ClassificationMethodConfigProps> = ({
|
|||
</div>
|
||||
)}
|
||||
|
||||
{classifierType === "hybrid" && (
|
||||
<div className="mt-4 space-y-2">
|
||||
<strong className="block font-semibold">Boundary margin</strong>
|
||||
<Input
|
||||
id={HYBRID_BOUNDARY_MARGIN_ID}
|
||||
type="text"
|
||||
inputMode="decimal"
|
||||
value={
|
||||
draft?.id === HYBRID_BOUNDARY_MARGIN_ID
|
||||
? draft.raw
|
||||
: String(value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN)
|
||||
}
|
||||
onChange={(event) => handleHybridBoundaryMarginChange(event.target.value)}
|
||||
onBlur={() => setDraft(null)}
|
||||
className="w-full"
|
||||
/>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
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
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="mt-4 space-y-2">
|
||||
<strong className="block font-semibold">How often to classify</strong>
|
||||
<RadioGroup
|
||||
|
|
|
|||
|
|
@ -128,7 +128,7 @@ export interface ClassifierLLMConfig {
|
|||
system_prompt?: string;
|
||||
}
|
||||
|
||||
export type ClassifierType = "heuristic" | "heuristic_v2" | "llm" | "heuristic_first";
|
||||
export type ClassifierType = "heuristic" | "heuristic_v2" | "llm" | "heuristic_first" | "hybrid";
|
||||
|
||||
/**
|
||||
* Whether this router can call classifier_llm_config.model. Mirrors the backend's
|
||||
|
|
@ -136,7 +136,7 @@ export type ClassifierType = "heuristic" | "heuristic_v2" | "llm" | "heuristic_f
|
|||
* control and payload key, so a new chaining type cannot strip knobs the operator set.
|
||||
*/
|
||||
export const usesLlmClassifier = (classifierType: ClassifierType): boolean =>
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -342,6 +342,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
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,
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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: [
|
||||
|
|
|
|||
|
|
@ -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<string, unknown> = {
|
||||
...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");
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ const CONSTANT_CAUSE_LABELS: Record<string, string> = {
|
|||
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",
|
||||
|
|
|
|||
13
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
13
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue