mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_responses_guardrail_namespace_tools
# Conflicts: # type-discipline-budget.json
This commit is contained in:
mateo-berri
commit
2842e90836
66 changed files with 6472 additions and 263 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:
|
||||
|
|
|
|||
|
|
@ -1565,6 +1565,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
optional_params.pop("thinking", None)
|
||||
else:
|
||||
optional_params["thinking"] = value
|
||||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider
|
||||
)
|
||||
elif param == "reasoning_effort":
|
||||
# Accept both string ("low") and dict ({"effort": "low",
|
||||
# "summary": "concise"}). The Responses->Chat parser keeps the
|
||||
|
|
|
|||
|
|
@ -13,7 +13,12 @@ import httpx
|
|||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME
|
||||
from litellm.constants import (
|
||||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_file_ids_from_messages,
|
||||
)
|
||||
|
|
@ -534,6 +539,51 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
)
|
||||
optional_params.pop("thinking", None)
|
||||
|
||||
@staticmethod
|
||||
def translate_legacy_thinking_for_adaptive_model(
|
||||
model: str,
|
||||
optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers
|
||||
custom_llm_provider: str,
|
||||
) -> None:
|
||||
"""Translate legacy ``thinking.type=enabled`` to adaptive for the
|
||||
adaptive-thinking models that reject it (4.7+ and the 5 families).
|
||||
Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the
|
||||
legacy shape natively, so it is forwarded verbatim and the caller's
|
||||
``budget_tokens`` cap keeps applying. Caller-provided
|
||||
``output_config.effort`` is never overridden.
|
||||
"""
|
||||
if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
|
||||
return
|
||||
if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider):
|
||||
return
|
||||
thinking: Final = optional_params.get("thinking")
|
||||
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
|
||||
return
|
||||
|
||||
effort: Final = AnthropicModelInfo._legacy_budget_to_effort(
|
||||
model=model,
|
||||
budget_tokens=int(thinking.get("budget_tokens") or 0),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
existing_output_config: Final = optional_params.get("output_config")
|
||||
optional_params["thinking"] = {"type": "adaptive"}
|
||||
optional_params["output_config"] = {
|
||||
"effort": effort,
|
||||
**(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _legacy_budget_to_effort(model: str, budget_tokens: int, custom_llm_provider: str) -> str:
|
||||
if budget_tokens >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and (
|
||||
AnthropicModelInfo._supports_model_capability(model, "supports_xhigh_reasoning_effort", custom_llm_provider)
|
||||
):
|
||||
return "xhigh"
|
||||
if budget_tokens >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET:
|
||||
return "high"
|
||||
if budget_tokens >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET:
|
||||
return "medium"
|
||||
return "low"
|
||||
|
||||
def is_effort_used(
|
||||
self,
|
||||
optional_params: dict | None,
|
||||
|
|
@ -1361,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):
|
||||
|
|
|
|||
|
|
@ -3,11 +3,6 @@ from typing import Any, Final
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import (
|
||||
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
|
||||
)
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import verbose_logger
|
||||
|
|
@ -400,46 +395,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
existing_output_config.setdefault("effort", mapped_effort)
|
||||
optional_params["output_config"] = existing_output_config
|
||||
|
||||
@staticmethod
|
||||
def _translate_legacy_thinking_for_adaptive_model(
|
||||
model: str, optional_params: dict, custom_llm_provider: str
|
||||
) -> None:
|
||||
"""Translate legacy ``thinking.type=enabled`` to adaptive for the
|
||||
adaptive-thinking models that reject it (4.7+ and the 5 families).
|
||||
Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the
|
||||
legacy shape natively, so it is forwarded verbatim and the caller's
|
||||
``budget_tokens`` cap keeps applying. Caller-provided
|
||||
``output_config.effort`` is never overridden.
|
||||
"""
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
|
||||
return
|
||||
if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider):
|
||||
return
|
||||
thinking: Final = optional_params.get("thinking")
|
||||
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
|
||||
return
|
||||
|
||||
budget: Final = int(thinking.get("budget_tokens") or 0)
|
||||
if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and (
|
||||
AnthropicConfig._supports_effort_level(model, "xhigh", custom_llm_provider)
|
||||
):
|
||||
effort = "xhigh"
|
||||
elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET:
|
||||
effort = "high"
|
||||
elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET:
|
||||
effort = "medium"
|
||||
else:
|
||||
effort = "low"
|
||||
|
||||
optional_params["thinking"] = {"type": "adaptive"}
|
||||
existing_output_config = optional_params.get("output_config")
|
||||
if not isinstance(existing_output_config, dict):
|
||||
existing_output_config = {}
|
||||
existing_output_config.setdefault("effort", effort)
|
||||
optional_params["output_config"] = existing_output_config
|
||||
|
||||
@staticmethod
|
||||
def _translate_adaptive_effort_for_non_adaptive_model(
|
||||
model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str
|
||||
|
|
@ -606,7 +561,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
custom_llm_provider=self._resolved_provider,
|
||||
)
|
||||
|
||||
self._translate_legacy_thinking_for_adaptive_model(
|
||||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=model,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
custom_llm_provider=self._resolved_provider,
|
||||
|
|
|
|||
|
|
@ -17,13 +17,14 @@ def _promote_extra_body_to_optional_params(optional_params: dict) -> None:
|
|||
``output_config`` get auto-routed into ``extra_body`` by
|
||||
``add_provider_specific_params_to_optional_params``. For the Azure→Anthropic
|
||||
route those keys must reach the request body and be validated, so promote
|
||||
them. ``setdefault`` keeps explicit top-level values authoritative.
|
||||
them. The caller's values overwrite mapped top-level duplicates, matching
|
||||
the native ``anthropic`` provider, where the same passthrough lands on
|
||||
top-level ``optional_params`` after mapping.
|
||||
"""
|
||||
extra_body: Final = optional_params.get("extra_body")
|
||||
if not isinstance(extra_body, dict) or not extra_body:
|
||||
return
|
||||
for k, v in extra_body.items():
|
||||
optional_params.setdefault(k, v)
|
||||
optional_params.update(extra_body)
|
||||
optional_params.pop("extra_body", None)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -943,6 +943,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model)
|
||||
else:
|
||||
optional_params["thinking"] = value
|
||||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=model, optional_params=optional_params, custom_llm_provider="bedrock"
|
||||
)
|
||||
elif param == "reasoning_effort" and isinstance(value, str):
|
||||
self._handle_reasoning_effort_parameter(
|
||||
model=model, reasoning_effort=value, optional_params=optional_params
|
||||
|
|
@ -1334,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))
|
||||
|
|
|
|||
|
|
@ -107,6 +107,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
# Restore original model name
|
||||
model = original_model
|
||||
|
||||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=original_model, optional_params=optional_params, custom_llm_provider="bedrock"
|
||||
)
|
||||
|
||||
# The stub model hides the original model from the parent's forced-tool-use backstop
|
||||
response_format_tool_choice: Final = optional_params.get("tool_choice")
|
||||
if (
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -330,6 +330,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
) -> dict:
|
||||
is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params)
|
||||
mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
|
||||
if "claude" in model:
|
||||
AnthropicConfig.translate_legacy_thinking_for_adaptive_model(
|
||||
model=model, optional_params=mapped_params, custom_llm_provider="databricks"
|
||||
)
|
||||
if "tools" in mapped_params:
|
||||
mapped_params["tools"] = self._map_openai_to_dbrx_tool(model=model, tools=mapped_params["tools"])
|
||||
if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -177,6 +177,10 @@ class VertexAIAnthropicConfig(AnthropicConfig):
|
|||
# Restore original model name for any other processing
|
||||
model = original_model
|
||||
|
||||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=original_model, optional_params=optional_params, custom_llm_provider="vertex_ai"
|
||||
)
|
||||
|
||||
return optional_params
|
||||
|
||||
def transform_response(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -68,6 +68,36 @@ still resolve to a deployment in `model_list`; this configuration does not creat
|
|||
- abc
|
||||
```
|
||||
|
||||
### Heuristic v2
|
||||
|
||||
Set `classifier_type: heuristic_v2` to classify with the bundled calibrated
|
||||
success-probability model instead of the hand-written weighted scorer
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: smart-router
|
||||
litellm_params:
|
||||
model: auto_router/complexity_router
|
||||
complexity_router_config:
|
||||
classifier_type: heuristic_v2
|
||||
tiers:
|
||||
SIMPLE: luna
|
||||
MEDIUM: terra
|
||||
COMPLEX: sol
|
||||
REASONING: sol-ultra
|
||||
```
|
||||
|
||||
No classifier model call or per-model training data is required. The classifier
|
||||
uses global tier quality, request-type quality, and similar-request cohorts from
|
||||
the bundled UltraFeedback artifact. It estimates success at every tier, enforces
|
||||
monotonic probabilities, and returns the first tier meeting the trained 0.75
|
||||
threshold. The existing complexity-router tier pool then selects and dispatches
|
||||
a model from that tier
|
||||
|
||||
Spend logs record `routing_decision.cause: heuristic_v2`, the detected request
|
||||
type, and all four predicted probabilities. Existing `classifier_type: heuristic`
|
||||
configurations keep the original weighted scorer unchanged
|
||||
|
||||
### Renaming the tiers
|
||||
|
||||
`tier_labels` puts your own vocabulary on the four tiers:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -33,6 +33,11 @@ from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
|
||||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
|
||||
from litellm.router_strategy.complexity_router.tier_predictor import (
|
||||
TierSuccessPredictor,
|
||||
resolve_tier_artifact,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
|
||||
ModelResponse,
|
||||
|
|
@ -790,6 +795,7 @@ class ClassificationOutcome(NamedTuple):
|
|||
signals: tuple[str, ...]
|
||||
cause: Literal[
|
||||
"heuristic_scorer",
|
||||
"heuristic_v2",
|
||||
"reasoning_override",
|
||||
"llm_classifier",
|
||||
"heuristic_first_short_circuit",
|
||||
|
|
@ -978,6 +984,11 @@ class ComplexityRouter(CustomLogger):
|
|||
if llm_classifier_configured
|
||||
else None
|
||||
)
|
||||
self._tier_success_predictor: TierSuccessPredictor | None = (
|
||||
TierSuccessPredictor(resolve_tier_artifact(self.config.heuristic_v2_artifact))
|
||||
if self.config.classifier_type == "heuristic_v2"
|
||||
else None
|
||||
)
|
||||
|
||||
verbose_router_logger.debug("ComplexityRouter initialized for %s with tiers: %s", model_name, self.config.tiers)
|
||||
|
||||
|
|
@ -1350,6 +1361,8 @@ class ComplexityRouter(CustomLogger):
|
|||
custom tier set, and classifier_fallback otherwise decides between the heuristic scorer and
|
||||
default_model. The outcome's `cause` reports which path actually ran.
|
||||
"""
|
||||
if self.config.classifier_type == "heuristic_v2":
|
||||
return self._classify_with_heuristic_v2(prompt)
|
||||
if self.config.classifier_type == "custom":
|
||||
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:
|
||||
|
|
@ -1359,6 +1372,24 @@ class ComplexityRouter(CustomLogger):
|
|||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages)
|
||||
|
||||
def _classify_with_heuristic_v2(self, prompt: str) -> ClassificationOutcome:
|
||||
predictor: Final = self._tier_success_predictor
|
||||
if predictor is None:
|
||||
raise ValueError("heuristic v2 predictor is not configured")
|
||||
request_type: Final = classify_prompt(prompt)
|
||||
prediction: Final = predictor.predict(prompt, request_type)
|
||||
tier: Final = TIER_SEVERITY_ORDER[prediction.required_tier - 1]
|
||||
probability_signals: Final = tuple(
|
||||
f"tier-probability:{candidate.value.lower()}={prediction.probabilities[index]:.6f}"
|
||||
for index, candidate in enumerate(TIER_SEVERITY_ORDER, start=1)
|
||||
)
|
||||
return ClassificationOutcome(
|
||||
tier=tier,
|
||||
score=None,
|
||||
signals=(f"request-type:{request_type.value}", *probability_signals),
|
||||
cause="heuristic_v2",
|
||||
)
|
||||
|
||||
async def _classify_heuristic_first(
|
||||
self,
|
||||
prompt: str,
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_seriali
|
|||
|
||||
from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
|
||||
|
||||
from .tier_predictor import TrainedTierArtifact
|
||||
|
||||
|
||||
class ComplexityTier(str, Enum):
|
||||
"""Complexity tiers for routing decisions."""
|
||||
|
|
@ -625,12 +627,19 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
|
||||
# Classifier strategy
|
||||
classifier_type: Literal["heuristic", "llm", "custom", "heuristic_first"] = Field(
|
||||
classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first"] = Field(
|
||||
default="heuristic",
|
||||
description=(
|
||||
"Classification strategy: local regex/keyword scoring, 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"
|
||||
"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"
|
||||
),
|
||||
)
|
||||
heuristic_v2_artifact: TrainedTierArtifact | Literal["ultrafeedback"] = Field(
|
||||
default="ultrafeedback",
|
||||
description=(
|
||||
"Success-probability artifact used by classifier_type 'heuristic_v2'. The bundled "
|
||||
"UltraFeedback artifact is selected by default; an inline trained artifact may replace it"
|
||||
),
|
||||
)
|
||||
classifier_llm_config: ClassifierLLMConfig | None = Field(
|
||||
|
|
@ -1248,10 +1257,10 @@ 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_first"):
|
||||
if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first"):
|
||||
raise ValueError(
|
||||
"tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only "
|
||||
"produces the built-in tiers"
|
||||
"produces the four built-in tiers, as does heuristic_v2"
|
||||
)
|
||||
conflicts: Final = self._tier_definition_conflicts()
|
||||
if conflicts:
|
||||
|
|
|
|||
156
litellm/router_strategy/complexity_router/tier_predictor.py
Normal file
156
litellm/router_strategy/complexity_router/tier_predictor.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from litellm.types.router import RequestType
|
||||
|
||||
|
||||
class TierGlobalStatistic(BaseModel):
|
||||
tier: int = Field(ge=1, le=4)
|
||||
successes: float = Field(ge=0.0)
|
||||
observations: float = Field(gt=0.0)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _successes_do_not_exceed_observations(self) -> TierGlobalStatistic:
|
||||
if self.successes > self.observations:
|
||||
raise ValueError("successes cannot exceed observations")
|
||||
return self
|
||||
|
||||
|
||||
class TierDomainStatistic(TierGlobalStatistic):
|
||||
request_type: RequestType
|
||||
|
||||
|
||||
class TierCohortStatistic(TierGlobalStatistic):
|
||||
cohort: str = Field(min_length=1)
|
||||
|
||||
|
||||
class TierDataset(BaseModel):
|
||||
name: str = Field(min_length=1)
|
||||
url: str = Field(min_length=1)
|
||||
license: str = Field(min_length=1)
|
||||
rows: int = Field(gt=0)
|
||||
success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1)
|
||||
|
||||
|
||||
class TrainedTierArtifact(BaseModel):
|
||||
schema_version: Literal[1] = 1
|
||||
global_statistics: tuple[TierGlobalStatistic, ...]
|
||||
domain_statistics: tuple[TierDomainStatistic, ...] = ()
|
||||
cohort_statistics: tuple[TierCohortStatistic, ...] = ()
|
||||
domain_prior_mass: float = Field(default=200.0, gt=0.0)
|
||||
cohort_prior_mass: float = Field(default=20.0, gt=0.0)
|
||||
routing_threshold: float = Field(default=0.75, ge=0.0, le=1.0)
|
||||
datasets: tuple[TierDataset, ...] = ()
|
||||
success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1)
|
||||
split_method: str = Field(default="sha256(prompt): 70% train, 15% validation, 15% test", min_length=1)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _statistics_are_unique(self) -> TrainedTierArtifact:
|
||||
global_tiers: Final = tuple(stat.tier for stat in self.global_statistics)
|
||||
if frozenset(global_tiers) != frozenset((1, 2, 3, 4)) or len(global_tiers) != 4:
|
||||
raise ValueError("global statistics must contain each tier exactly once")
|
||||
domain_keys: Final = tuple((stat.request_type, stat.tier) for stat in self.domain_statistics)
|
||||
if len(domain_keys) != len(frozenset(domain_keys)):
|
||||
raise ValueError("domain statistics must contain unique request_type and tier pairs")
|
||||
cohort_keys: Final = tuple((stat.cohort, stat.tier) for stat in self.cohort_statistics)
|
||||
if len(cohort_keys) != len(frozenset(cohort_keys)):
|
||||
raise ValueError("cohort statistics must contain unique cohort and tier pairs")
|
||||
return self
|
||||
|
||||
|
||||
_CODE_PATTERN: Final = re.compile(
|
||||
r"```|\b(def|class|function|python|javascript|typescript|sql|code)\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_MATH_PATTERN: Final = re.compile(
|
||||
r"\b(solve|calculate|equation|probability|theorem|proof|integral)\b|[$=]",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_MULTIPLE_CHOICE_PATTERN: Final = re.compile(r"(?:^|\s)[A-D][.)]\s")
|
||||
_TIERS: Final = (1, 2, 3, 4)
|
||||
_BUILTIN_ARTIFACTS: Final = MappingProxyType({"ultrafeedback": "ultrafeedback_tiers.json"})
|
||||
|
||||
|
||||
def resolve_tier_artifact(artifact: TrainedTierArtifact | str) -> TrainedTierArtifact:
|
||||
if isinstance(artifact, TrainedTierArtifact):
|
||||
return artifact
|
||||
filename: Final = _BUILTIN_ARTIFACTS.get(artifact)
|
||||
if filename is None:
|
||||
raise ValueError(f"unknown complexity router tier artifact: {artifact}")
|
||||
path: Final = Path(__file__).with_name("artifacts") / filename
|
||||
return TrainedTierArtifact.model_validate_json(path.read_text())
|
||||
|
||||
|
||||
def similarity_cohort(prompt: str, request_type: RequestType) -> str:
|
||||
length: Final = len(prompt)
|
||||
length_bucket: Final = (
|
||||
"short" if length < 200 else "medium" if length < 800 else "long" if length < 2000 else "very_long"
|
||||
)
|
||||
code: Final = int(bool(_CODE_PATTERN.search(prompt)))
|
||||
math: Final = int(bool(_MATH_PATTERN.search(prompt)))
|
||||
multiple_choice: Final = int(bool(_MULTIPLE_CHOICE_PATTERN.search(prompt)))
|
||||
non_ascii: Final = int(sum(ord(character) > 127 for character in prompt) / max(1, length) > 0.1)
|
||||
return f"{request_type.value}|{length_bucket}|code={code}|math={math}|mc={multiple_choice}|intl={non_ascii}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TierPrediction:
|
||||
probabilities: Mapping[int, float]
|
||||
required_tier: int
|
||||
|
||||
|
||||
class TierSuccessPredictor:
|
||||
def __init__(self, artifact: TrainedTierArtifact) -> None:
|
||||
self._artifact = artifact
|
||||
self._global: Mapping[int, TierGlobalStatistic] = MappingProxyType(
|
||||
{stat.tier: stat for stat in artifact.global_statistics}
|
||||
)
|
||||
self._domain: Mapping[tuple[RequestType, int], TierDomainStatistic] = MappingProxyType(
|
||||
{(stat.request_type, stat.tier): stat for stat in artifact.domain_statistics}
|
||||
)
|
||||
self._cohort: Mapping[tuple[str, int], TierCohortStatistic] = MappingProxyType(
|
||||
{(stat.cohort, stat.tier): stat for stat in artifact.cohort_statistics}
|
||||
)
|
||||
|
||||
@property
|
||||
def routing_threshold(self) -> float:
|
||||
return self._artifact.routing_threshold
|
||||
|
||||
def predict(self, prompt: str, request_type: RequestType) -> TierPrediction:
|
||||
cohort: Final = similarity_cohort(prompt, request_type)
|
||||
raw: Final = tuple(self._probability(tier, request_type, cohort) for tier in _TIERS)
|
||||
monotonic: Final = tuple(max(raw[:index]) for index in range(1, len(raw) + 1))
|
||||
probabilities: Final[Mapping[int, float]] = MappingProxyType(
|
||||
{int(tier): probability for tier, probability in zip(_TIERS, monotonic)}
|
||||
)
|
||||
required_tier: Final = next(
|
||||
(tier for tier in _TIERS if probabilities[tier] >= self._artifact.routing_threshold),
|
||||
4,
|
||||
)
|
||||
return TierPrediction(probabilities=probabilities, required_tier=required_tier)
|
||||
|
||||
def _probability(self, tier: int, request_type: RequestType, cohort: str) -> float:
|
||||
global_stat: Final = self._global[tier]
|
||||
global_mean: Final = (global_stat.successes + 1.0) / (global_stat.observations + 2.0)
|
||||
domain_stat: Final = self._domain.get((request_type, tier))
|
||||
domain_mean: Final = self._posterior_mean(domain_stat, self._artifact.domain_prior_mass, global_mean)
|
||||
cohort_stat: Final = self._cohort.get((cohort, tier))
|
||||
return self._posterior_mean(cohort_stat, self._artifact.cohort_prior_mass, domain_mean)
|
||||
|
||||
@staticmethod
|
||||
def _posterior_mean(
|
||||
statistic: TierGlobalStatistic | None,
|
||||
prior_mass: float,
|
||||
prior_mean: float,
|
||||
) -> float:
|
||||
if statistic is None:
|
||||
return prior_mean
|
||||
return (statistic.successes + prior_mass * prior_mean) / (statistic.observations + prior_mass)
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -2839,6 +2839,7 @@ class StandardLoggingRoutingDecisionTierBoundaries(TypedDict):
|
|||
|
||||
RoutingDecisionCause = Literal[
|
||||
"heuristic_scorer",
|
||||
"heuristic_v2",
|
||||
# The scorer found 2+ reasoning markers and forced REASONING regardless of score.
|
||||
# A distinct cause rather than a marker inside `signals`, because it is the fact
|
||||
# that tells a reader the score did NOT choose the tier; encoding it as free text
|
||||
|
|
|
|||
|
|
@ -278,7 +278,10 @@ bindings = "pyo3"
|
|||
features = ["extension-module"]
|
||||
profile = "release"
|
||||
editable-profile = "dev"
|
||||
include = ["litellm/proxy/_experimental/out/**"]
|
||||
include = [
|
||||
"litellm/proxy/_experimental/out/**",
|
||||
"litellm/router_strategy/complexity_router/artifacts/*.json",
|
||||
]
|
||||
exclude = [
|
||||
"litellm/proxy/enterprise",
|
||||
"litellm/proxy/enterprise/**",
|
||||
|
|
|
|||
|
|
@ -3127,6 +3127,31 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,budget_tokens,expected",
|
||||
[
|
||||
("claude-opus-4-8", 4096, ({"type": "adaptive"}, {"effort": "high"})),
|
||||
("claude-opus-4-7", 24000, ({"type": "adaptive"}, {"effort": "xhigh"})),
|
||||
("claude-opus-4-6", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)),
|
||||
("claude-sonnet-4-5-20250929", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)),
|
||||
],
|
||||
)
|
||||
def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_models(model, budget_tokens, expected):
|
||||
"""Adaptive-only models reject thinking={type: enabled} with a 400, so the
|
||||
legacy shape must be upgraded to adaptive + output_config.effort on
|
||||
/chat/completions too, while models that accept it keep the caller's budget."""
|
||||
config = AnthropicConfig()
|
||||
|
||||
result = config.map_openai_params(
|
||||
non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000},
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert (result["thinking"], result.get("output_config")) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad_value",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -362,8 +362,8 @@ class TestAzureAnthropicConfig:
|
|||
)
|
||||
assert "xhigh" in str(exc_info.value)
|
||||
|
||||
def test_extra_body_promotion_does_not_clobber_top_level(self):
|
||||
"""Top-level ``optional_params`` wins over duplicates in ``extra_body``."""
|
||||
def test_extra_body_promotion_overrides_mapped_top_level(self):
|
||||
"""The caller's ``extra_body`` wins over a mapped top-level duplicate, like the native ``anthropic`` passthrough."""
|
||||
config = AzureAnthropicConfig()
|
||||
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
|
|
@ -383,7 +383,31 @@ class TestAzureAnthropicConfig:
|
|||
headers=headers,
|
||||
)
|
||||
|
||||
assert result["output_config"] == {"effort": "low"}
|
||||
assert result["output_config"] == {"effort": "high"}
|
||||
|
||||
def test_legacy_thinking_upgrade_keeps_caller_effort_from_extra_body(self, local_model_cost_map):
|
||||
config = AzureAnthropicConfig()
|
||||
|
||||
mapped = config.map_openai_params(
|
||||
non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}, "max_tokens": 100},
|
||||
optional_params={},
|
||||
model="claude-opus-4-8",
|
||||
drop_params=False,
|
||||
)
|
||||
assert mapped["thinking"] == {"type": "adaptive"}
|
||||
assert mapped["output_config"] == {"effort": "low"}
|
||||
|
||||
result = config.transform_request(
|
||||
model="claude-opus-4-8",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
optional_params={**mapped, "extra_body": {"output_config": {"effort": "high"}}},
|
||||
litellm_params={"api_key": "test-key"},
|
||||
headers={"api-key": "test-key", "anthropic-version": "2023-06-01"},
|
||||
)
|
||||
|
||||
assert result["thinking"] == {"type": "adaptive"}
|
||||
assert result["output_config"] == {"effort": "high"}
|
||||
assert "extra_body" not in result
|
||||
|
||||
def test_context_management_mixed_edits_beta_headers(self):
|
||||
"""Test that context_management with both compact and other edits adds both beta headers"""
|
||||
|
|
|
|||
|
|
@ -671,3 +671,46 @@ def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice
|
|||
assert "output_format" not in result
|
||||
assert "tools" in result
|
||||
assert "tool_choice" not in result
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["us.anthropic.claude-sonnet-5", "us.anthropic.claude-fable-5-1"])
|
||||
def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map, model):
|
||||
result = AmazonAnthropicClaudeConfig().map_openai_params(
|
||||
non_default_params={
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "test_schema",
|
||||
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
|
||||
},
|
||||
},
|
||||
"thinking": {"type": "enabled", "budget_tokens": 4096},
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert "tools" in result
|
||||
assert result["thinking"] == {"type": "adaptive"}
|
||||
assert result["output_config"] == {"effort": "high"}
|
||||
|
||||
|
||||
def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map):
|
||||
"""Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub
|
||||
model before the shared Anthropic mapping, which hid the adaptive-only model
|
||||
from the legacy ``thinking`` upgrade and left ``type=enabled`` on the wire."""
|
||||
result = AmazonAnthropicClaudeConfig().map_openai_params(
|
||||
non_default_params={
|
||||
"response_format": {"type": "json_object"},
|
||||
"thinking": {"type": "enabled", "budget_tokens": 4096},
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
optional_params={},
|
||||
model="us.anthropic.claude-fable-5-1",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert result["thinking"] == {"type": "adaptive"}
|
||||
assert result["output_config"] == {"effort": "high"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -6347,6 +6375,82 @@ def test_adaptive_thinking_passes_through_on_46_plus_converse(model):
|
|||
assert optional_params.get("thinking") == {"type": "adaptive"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,budget_tokens,expected_effort",
|
||||
[
|
||||
("anthropic.claude-opus-4-8", 4096, "high"),
|
||||
("us.anthropic.claude-opus-4-8", 2000, "low"),
|
||||
("global.anthropic.claude-opus-4-8", 12000, "xhigh"),
|
||||
("us.anthropic.claude-opus-4-7", 3000, "medium"),
|
||||
("anthropic.claude-fable-5", 4096, "high"),
|
||||
],
|
||||
)
|
||||
def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_converse(model, budget_tokens, expected_effort):
|
||||
"""Adaptive-only models (4.7+, 5 families) reject thinking={type: enabled}
|
||||
with a 400 on Bedrock Converse, so the legacy shape from callers like Claude
|
||||
Code must be upgraded to thinking={type: adaptive} + output_config.effort
|
||||
derived from budget_tokens, matching the /v1/messages passthrough."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
optional_params = config.map_openai_params(
|
||||
non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000},
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
)
|
||||
request = config.transform_request(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["additionalModelRequestFields"]["thinking"] == {"type": "adaptive"}
|
||||
assert request["additionalModelRequestFields"]["output_config"] == {"effort": expected_effort}
|
||||
|
||||
|
||||
def test_legacy_thinking_translation_keeps_caller_output_config_effort_converse():
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
optional_params = config.map_openai_params(
|
||||
non_default_params={
|
||||
"output_config": {"effort": "low"},
|
||||
"thinking": {"type": "enabled", "budget_tokens": 12000},
|
||||
"max_tokens": 64000,
|
||||
},
|
||||
optional_params={},
|
||||
model="anthropic.claude-opus-4-8",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert optional_params["thinking"] == {"type": "adaptive"}
|
||||
assert optional_params["output_config"] == {"effort": "low"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"us.anthropic.claude-opus-4-6",
|
||||
"anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
],
|
||||
)
|
||||
def test_legacy_thinking_forwarded_verbatim_when_model_accepts_it_converse(model):
|
||||
"""The 4.6 family and pre-adaptive models accept thinking={type: enabled}
|
||||
natively, so the caller's budget_tokens cap must keep applying."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
optional_params = config.map_openai_params(
|
||||
non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}, "max_tokens": 8192},
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert optional_params["thinking"] == {"type": "enabled", "budget_tokens": 4096}
|
||||
assert "output_config" not in optional_params
|
||||
|
||||
|
||||
def test_adaptive_thinking_dropped_when_max_tokens_too_small_converse():
|
||||
"""When max_tokens can't fit even the minimum thinking budget, the raw
|
||||
adaptive block must be dropped entirely rather than translated, so the
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -422,6 +422,27 @@ def test_databricks_config_probes_capabilities_under_databricks_namespace():
|
|||
assert DatabricksConfig().custom_llm_provider == "databricks"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, expected_thinking, expected_output_config",
|
||||
[
|
||||
("databricks-claude-opus-4-8", {"type": "adaptive"}, {"effort": "high"}),
|
||||
("databricks-claude-opus-4-6", {"type": "enabled", "budget_tokens": 4096}, None),
|
||||
],
|
||||
ids=["adaptive_only_upgrades_to_adaptive", "legacy_capable_forwards_verbatim"],
|
||||
)
|
||||
def test_map_openai_params_upgrades_legacy_thinking_on_adaptive_only_claude(
|
||||
model, expected_thinking, expected_output_config
|
||||
):
|
||||
mapped = DatabricksConfig().map_openai_params(
|
||||
non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}},
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
)
|
||||
assert mapped["thinking"] == expected_thinking
|
||||
assert mapped.get("output_config") == expected_output_config
|
||||
|
||||
|
||||
def _streaming_chunk(usage=None, choices=None):
|
||||
base = {
|
||||
"id": "chatcmpl-test",
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -752,3 +752,26 @@ def test_vertex_ai_fable_5_1_response_format_uses_native_output_format(local_mod
|
|||
assert "output_format" in result_params
|
||||
assert "tool_choice" not in result_params
|
||||
assert "tools" not in result_params
|
||||
|
||||
|
||||
def test_vertex_ai_anthropic_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map):
|
||||
result_params = VertexAIAnthropicConfig().map_openai_params(
|
||||
non_default_params={
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "test_schema",
|
||||
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
|
||||
},
|
||||
},
|
||||
"thinking": {"type": "enabled", "budget_tokens": 4096},
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
optional_params={},
|
||||
model="claude-opus-4-8",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert "tools" in result_params
|
||||
assert result_params["thinking"] == {"type": "adaptive"}
|
||||
assert result_params["output_config"] == {"effort": "high"}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -12,7 +12,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm._logging import verbose_router_logger
|
||||
|
|
@ -34,10 +33,14 @@ from litellm.router_strategy.complexity_router.config import (
|
|||
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
|
||||
DEFAULT_COMPLEXITY_CONFIG,
|
||||
DEFAULT_TECHNICAL_KEYWORDS,
|
||||
ClassificationRubric,
|
||||
ClassifierLLMConfig,
|
||||
ComplexityRouterConfig,
|
||||
ComplexityTier,
|
||||
ClassificationRubric,
|
||||
)
|
||||
from litellm.router_strategy.complexity_router.tier_predictor import (
|
||||
TierGlobalStatistic,
|
||||
TrainedTierArtifact,
|
||||
)
|
||||
from litellm.types.router import (
|
||||
Deployment,
|
||||
|
|
@ -46,6 +49,16 @@ from litellm.types.router import (
|
|||
)
|
||||
|
||||
|
||||
def _heuristic_v2_artifact() -> TrainedTierArtifact:
|
||||
return TrainedTierArtifact(
|
||||
global_statistics=tuple(
|
||||
TierGlobalStatistic(tier=tier, successes=successes, observations=100)
|
||||
for tier, successes in enumerate((10, 20, 90, 99), start=1)
|
||||
),
|
||||
routing_threshold=0.8,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_router_instance():
|
||||
"""Create a mock LiteLLM Router instance."""
|
||||
|
|
@ -1696,6 +1709,59 @@ class TestLLMClassifier:
|
|||
assert outcome.cause == "heuristic_scorer"
|
||||
assert outcome.score is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier(self, mock_router_instance):
|
||||
router = ComplexityRouter(
|
||||
model_name="tier-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={
|
||||
"classifier_type": "heuristic_v2",
|
||||
"heuristic_v2_artifact": _heuristic_v2_artifact(),
|
||||
"tiers": {
|
||||
"SIMPLE": "simple-model",
|
||||
"MEDIUM": "medium-model",
|
||||
"COMPLEX": "complex-model",
|
||||
"REASONING": "reasoning-model",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="tier-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "Handle this new request"}],
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert response.model == "complex-model"
|
||||
assert response.routing_decision["tier"] == "COMPLEX"
|
||||
assert response.routing_decision["cause"] == "heuristic_v2"
|
||||
assert response.routing_decision["signals"] == [
|
||||
"request-type:general",
|
||||
"tier-probability:simple=0.107843",
|
||||
"tier-probability:medium=0.205882",
|
||||
"tier-probability:complex=0.892157",
|
||||
"tier-probability:reasoning=0.980392",
|
||||
]
|
||||
|
||||
def test_heuristic_v2_needs_no_classifier_model(self):
|
||||
config = ComplexityRouterConfig(classifier_type="heuristic_v2")
|
||||
|
||||
assert config.classifier_llm_config is None
|
||||
assert config.heuristic_v2_artifact == "ultrafeedback"
|
||||
|
||||
def test_heuristic_v2_rejects_custom_tier_definitions(self):
|
||||
with pytest.raises(ValidationError, match="as does heuristic_v2"):
|
||||
ComplexityRouterConfig(
|
||||
classifier_type="heuristic_v2",
|
||||
tier_definitions=(
|
||||
{"name": "low", "description": "easy work"},
|
||||
{"name": "high", "description": "hard work"},
|
||||
),
|
||||
tiers={"low": "cheap", "high": "expensive"},
|
||||
fallback_tier="high",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclassify_llm_success_routes_by_llm_verdict(self, llm_complexity_router, mock_router_instance):
|
||||
"""A well-formed structured LLM response should decide the tier directly.
|
||||
|
|
@ -10268,7 +10334,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
|
||||
|
|
@ -10291,7 +10358,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"}]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,91 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.complexity_router.tier_predictor import (
|
||||
TierCohortStatistic,
|
||||
TierDomainStatistic,
|
||||
TierGlobalStatistic,
|
||||
TierSuccessPredictor,
|
||||
TrainedTierArtifact,
|
||||
resolve_tier_artifact,
|
||||
similarity_cohort,
|
||||
)
|
||||
from litellm.types.router import RequestType
|
||||
|
||||
|
||||
def _artifact(
|
||||
global_successes: tuple[float, float, float, float] = (4.0, 5.0, 6.0, 7.0),
|
||||
threshold: float = 0.75,
|
||||
domain_statistics: tuple[TierDomainStatistic, ...] = (),
|
||||
cohort_statistics: tuple[TierCohortStatistic, ...] = (),
|
||||
) -> TrainedTierArtifact:
|
||||
return TrainedTierArtifact(
|
||||
global_statistics=tuple(
|
||||
TierGlobalStatistic(tier=tier, successes=successes, observations=10.0)
|
||||
for tier, successes in enumerate(global_successes, start=1)
|
||||
),
|
||||
domain_statistics=domain_statistics,
|
||||
cohort_statistics=cohort_statistics,
|
||||
domain_prior_mass=10.0,
|
||||
cohort_prior_mass=10.0,
|
||||
routing_threshold=threshold,
|
||||
)
|
||||
|
||||
|
||||
def test_predictions_are_monotonic_across_tiers() -> None:
|
||||
predictor: Final = TierSuccessPredictor(_artifact(global_successes=(9.0, 2.0, 7.0, 6.0)))
|
||||
|
||||
prediction: Final = predictor.predict("hello", RequestType.GENERAL)
|
||||
|
||||
probabilities: Final = tuple(prediction.probabilities.values())
|
||||
assert probabilities == tuple(sorted(probabilities))
|
||||
|
||||
|
||||
def test_domain_and_cohort_statistics_back_off_hierarchically() -> None:
|
||||
matching_cohort: Final = similarity_cohort("hello", RequestType.GENERAL)
|
||||
artifact: Final = _artifact(
|
||||
global_successes=(1.0, 5.0, 6.0, 7.0),
|
||||
domain_statistics=(
|
||||
TierDomainStatistic(
|
||||
tier=1,
|
||||
request_type=RequestType.GENERAL,
|
||||
successes=10.0,
|
||||
observations=10.0,
|
||||
),
|
||||
),
|
||||
cohort_statistics=(
|
||||
TierCohortStatistic(
|
||||
tier=1,
|
||||
cohort=matching_cohort,
|
||||
successes=0.0,
|
||||
observations=10.0,
|
||||
),
|
||||
),
|
||||
)
|
||||
predictor: Final = TierSuccessPredictor(artifact)
|
||||
|
||||
cohort_probability: Final = predictor.predict("hello", RequestType.GENERAL).probabilities[1]
|
||||
domain_probability: Final = predictor.predict("hello " * 100, RequestType.GENERAL).probabilities[1]
|
||||
global_probability: Final = predictor.predict("hello", RequestType.WRITING).probabilities[1]
|
||||
|
||||
assert cohort_probability == pytest.approx(7.0 / 24.0)
|
||||
assert domain_probability == pytest.approx(7.0 / 12.0)
|
||||
assert global_probability == pytest.approx(1.0 / 6.0)
|
||||
|
||||
|
||||
def test_selects_first_tier_above_probability_threshold() -> None:
|
||||
predictor: Final = TierSuccessPredictor(_artifact(global_successes=(4.0, 6.0, 8.0, 9.0), threshold=0.7))
|
||||
|
||||
prediction: Final = predictor.predict("hello", RequestType.GENERAL)
|
||||
|
||||
assert prediction.required_tier == 3
|
||||
|
||||
|
||||
def test_builtin_ultrafeedback_artifact_is_loadable() -> None:
|
||||
artifact: Final = resolve_tier_artifact("ultrafeedback")
|
||||
|
||||
assert artifact.routing_threshold == 0.75
|
||||
assert artifact.domain_prior_mass == 200.0
|
||||
assert artifact.cohort_prior_mass == 20.0
|
||||
assert artifact.datasets[0].license == "MIT"
|
||||
|
|
@ -6,7 +6,7 @@
|
|||
"limit": 26765
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 269
|
||||
"limit": 261
|
||||
},
|
||||
"LIT004": {
|
||||
"limit": 40
|
||||
|
|
@ -27,10 +27,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16490
|
||||
"limit": 16482
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5534
|
||||
"limit": 5520
|
||||
},
|
||||
"LIT012": {
|
||||
"limit": 4495
|
||||
|
|
|
|||
|
|
@ -43,6 +43,10 @@ const DEFAULT_SCORING_EXPLANATION =
|
|||
"The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " +
|
||||
"terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:";
|
||||
|
||||
const HEURISTIC_V2_EXPLANATION =
|
||||
"The router estimates success probability for all four tiers with the bundled calibrated model, then selects " +
|
||||
"the first tier that meets its trained threshold. It runs locally with no classifier API call.";
|
||||
|
||||
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";
|
||||
|
|
@ -62,6 +66,7 @@ const CUSTOM_PROMPT_WITH_DEFAULT_MODEL_FALLBACK =
|
|||
* at all, so the panel must not keep implying a score is involved on either router.
|
||||
*/
|
||||
const scoringExplanation = (value: ComplexityRouterConfigValue): string => {
|
||||
if (value.classifier_type === "heuristic_v2") return HEURISTIC_V2_EXPLANATION;
|
||||
const usesCustomPrompt =
|
||||
usesLlmClassifier(value.classifier_type) && Boolean(value.classifier_llm_config?.system_prompt?.trim());
|
||||
if (!usesCustomPrompt) return DEFAULT_SCORING_EXPLANATION;
|
||||
|
|
@ -179,6 +184,17 @@ 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="heuristic_v2" className="mt-0.5" disabled={scorerLocked} />
|
||||
<span>
|
||||
<strong className="font-semibold">Heuristic v2</strong>{" "}
|
||||
<span className="text-muted-foreground">
|
||||
uses bundled calibrated four-tier probabilities with no API call
|
||||
</span>
|
||||
</span>
|
||||
</Label>
|
||||
</SimpleTooltip>
|
||||
<Label className="items-start font-normal leading-normal">
|
||||
<RadioGroupItem value="llm" className="mt-0.5" />
|
||||
<span>
|
||||
|
|
|
|||
|
|
@ -127,6 +127,31 @@ describe("ComplexityRouterConfig", () => {
|
|||
expect(onChange).toHaveBeenCalledWith(expectedValue);
|
||||
});
|
||||
|
||||
it("selects heuristic v2 without requiring a classifier model or showing weighted scoring", () => {
|
||||
const onChange = vi.fn();
|
||||
const { rerender } = renderWithProviders(
|
||||
<ComplexityRouterConfig modelInfo={mockModelInfo} value={defaultValue} onChange={onChange} />,
|
||||
);
|
||||
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
fireEvent.click(screen.getByText("Heuristic v2"));
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
classifier_type: "heuristic_v2",
|
||||
classifier_llm_config: undefined,
|
||||
}),
|
||||
);
|
||||
|
||||
const heuristicV2Value: ComplexityRouterConfigValue = { ...defaultValue, classifier_type: "heuristic_v2" };
|
||||
rerender(<ComplexityRouterConfig modelInfo={mockModelInfo} value={heuristicV2Value} onChange={onChange} />);
|
||||
|
||||
expect(screen.queryByText("Classifier Model")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Advanced scoring")).not.toBeInTheDocument();
|
||||
expect(screen.getByText(/estimates success probability for all four tiers/)).toBeInTheDocument();
|
||||
expect(screen.queryByText(/Score < 0.15/)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show classifier fields and use the configured values when classifier_type is llm", () => {
|
||||
const llmValue: ComplexityRouterConfigValue = {
|
||||
...defaultValue,
|
||||
|
|
|
|||
|
|
@ -128,7 +128,7 @@ export interface ClassifierLLMConfig {
|
|||
system_prompt?: string;
|
||||
}
|
||||
|
||||
export type ClassifierType = "heuristic" | "llm" | "heuristic_first";
|
||||
export type ClassifierType = "heuristic" | "heuristic_v2" | "llm" | "heuristic_first";
|
||||
|
||||
/**
|
||||
* Whether this router can call classifier_llm_config.model. Mirrors the backend's
|
||||
|
|
@ -161,6 +161,7 @@ export const heuristicScoringRoleFor = (
|
|||
classifierType: ClassifierType,
|
||||
classifierFallback: ClassifierFallback | undefined,
|
||||
): HeuristicScoringRole => {
|
||||
if (classifierType === "heuristic_v2") return "never";
|
||||
if (classifierType === "heuristic" || classifierType === "heuristic_first") return "decides";
|
||||
return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never";
|
||||
};
|
||||
|
|
@ -188,13 +189,19 @@ const builtInTierInfo = (rowId: string): { label: string; description: string; e
|
|||
return builtIn ? TIER_DESCRIPTIONS[builtIn] : undefined;
|
||||
};
|
||||
|
||||
const tierConfigIntroText = (value: ComplexityRouterConfigValue): string => {
|
||||
if (value.classifier_type === "heuristic_v2") {
|
||||
return "The complexity router classifies each request with a calibrated local four-tier model (no API calls). Configure which model(s) handle each tier.";
|
||||
}
|
||||
if (heuristicScoringRole(value) === "never") {
|
||||
return "The complexity router classifies each request with your classifier model and routes it to that tier. Configure which model(s) handle each tier.";
|
||||
}
|
||||
return "The complexity router automatically classifies requests by complexity using rule-based scoring (no API calls, <1ms latency). Configure which model(s) handle each tier.";
|
||||
};
|
||||
|
||||
const TierConfigIntro: React.FC<{ value: ComplexityRouterConfigValue }> = ({ value }) => (
|
||||
<>
|
||||
<span className="block mb-6 text-muted-foreground">
|
||||
{heuristicScoringRole(value) === "never"
|
||||
? "The complexity router classifies each request with your classifier model and routes it to that tier. Configure which model(s) handle each tier."
|
||||
: "The complexity router automatically classifies requests by complexity using rule-based scoring (no API calls, <1ms latency). Configure which model(s) handle each tier."}
|
||||
</span>
|
||||
<span className="block mb-6 text-muted-foreground">{tierConfigIntroText(value)}</span>
|
||||
|
||||
<span className="block mb-4 text-xs text-muted-foreground">
|
||||
{restrictedBy(value, "displayNames")?.reason ??
|
||||
|
|
|
|||
|
|
@ -165,6 +165,7 @@ describe("ClassificationMethodConfig scorer gating", () => {
|
|||
|
||||
it.each([
|
||||
["heuristic decides the tier", "heuristic" as ClassifierType, undefined, true],
|
||||
["heuristic v2 decides without the weighted scorer", "heuristic_v2" as ClassifierType, undefined, false],
|
||||
["an LLM classifier falls back to the heuristic", "llm" as ClassifierType, "heuristic" as ClassifierFallback, true],
|
||||
[
|
||||
"an LLM classifier falls back to the default model",
|
||||
|
|
|
|||
|
|
@ -111,6 +111,21 @@ describe("buildComplexityRouterConfig", () => {
|
|||
expect(config.classifier_llm_config).toBeUndefined();
|
||||
});
|
||||
|
||||
it("emits heuristic_v2 without classifier-only fields", () => {
|
||||
const trainedParams: BuildComplexityRouterConfigParams = {
|
||||
...baseParams,
|
||||
classifierType: "heuristic_v2",
|
||||
classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
|
||||
classifierContextWindowSize: 5,
|
||||
classifierFallback: "heuristic",
|
||||
};
|
||||
const config = buildComplexityRouterConfig(trainedParams);
|
||||
expect(config.classifier_type).toBe("heuristic_v2");
|
||||
expect(config.classifier_llm_config).toBeUndefined();
|
||||
expect(config.classifier_context_window_size).toBeUndefined();
|
||||
expect(config.classifier_fallback).toBeUndefined();
|
||||
});
|
||||
|
||||
it("includes classifier_context_window_size and classifier_context_budget_chars only when classifier_type is llm", () => {
|
||||
const params: BuildComplexityRouterConfigParams = {
|
||||
...baseParams,
|
||||
|
|
@ -713,6 +728,10 @@ describe("getClassifierModelError", () => {
|
|||
expect(getClassifierModelError({ classifier_type: "heuristic" })).toBeNull();
|
||||
});
|
||||
|
||||
it("stays quiet for a heuristic v2 router, which runs locally", () => {
|
||||
expect(getClassifierModelError({ classifier_type: "heuristic_v2" })).toBeNull();
|
||||
});
|
||||
|
||||
it("blocks an LLM classifier with no model, which the router cannot start without", () => {
|
||||
expect(getClassifierModelError({ classifier_type: "llm" })).toBe(
|
||||
"Please select a classifier model, or switch back to Heuristic",
|
||||
|
|
@ -791,7 +810,7 @@ describe("heuristic_first", () => {
|
|||
});
|
||||
|
||||
it("omits heuristic_first_max_tier on every other classifier type, which the backend rejects it on", () => {
|
||||
for (const classifierType of ["heuristic", "llm"] as const) {
|
||||
for (const classifierType of ["heuristic", "heuristic_v2", "llm"] as const) {
|
||||
const config = buildComplexityRouterConfig({ ...heuristicFirstParams, classifierType });
|
||||
expect(config.heuristic_first_max_tier).toBeUndefined();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -84,6 +84,7 @@ function describeReasoningOverride(tierLabel: string | undefined, floor: number
|
|||
|
||||
const CONSTANT_CAUSE_LABELS: Record<string, string> = {
|
||||
heuristic_scorer: "Heuristic scorer",
|
||||
heuristic_v2: "Heuristic v2",
|
||||
heuristic_first_short_circuit: "Heuristic scorer, classifier skipped",
|
||||
classifier_plugin: "Custom classifier plugin",
|
||||
semantic_keyword_match: "Semantic keyword match",
|
||||
|
|
|
|||
115
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
115
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -34575,11 +34575,11 @@ export interface components {
|
|||
classifier_plugin_timeout_ms: number;
|
||||
/**
|
||||
* Classifier Type
|
||||
* @description Classification strategy: local regex/keyword scoring, 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, or 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier
|
||||
* @default heuristic
|
||||
* @enum {string}
|
||||
*/
|
||||
classifier_type: "heuristic" | "llm" | "custom" | "heuristic_first";
|
||||
classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first";
|
||||
/**
|
||||
* Code Keywords
|
||||
* @description Keywords indicating code-related content
|
||||
|
|
@ -34640,6 +34640,12 @@ export interface components {
|
|||
* @description The highest tier the local scorer may decide on its own; required when classifier_type is 'heuristic_first' and rejected otherwise. A request whose heuristic tier is at or below this one skips the LLM classifier and routes straight to that heuristic tier, so the classifier call is only paid for on traffic the scorer could not place cheaply. The scorer must also have produced at least one signal: a prompt where no dimension fired scores 0.0 and would otherwise land SIMPLE by default rather than by evidence, which is how a chained router would silently send unclassified traffic to the cheapest model. Names a built-in tier, and may not name the highest one, since that would make the LLM classifier unreachable.
|
||||
*/
|
||||
heuristic_first_max_tier?: string | null;
|
||||
/**
|
||||
* Heuristic V2 Artifact
|
||||
* @description Success-probability artifact used by classifier_type 'heuristic_v2'. The bundled UltraFeedback artifact is selected by default; an inline trained artifact may replace it
|
||||
* @default ultrafeedback
|
||||
*/
|
||||
heuristic_v2_artifact: components["schemas"]["TrainedTierArtifact"] | "ultrafeedback";
|
||||
/**
|
||||
* Housekeeping Patterns
|
||||
* @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.
|
||||
|
|
@ -34778,6 +34784,12 @@ export interface components {
|
|||
} & {
|
||||
[key: string]: unknown;
|
||||
};
|
||||
/**
|
||||
* RequestType
|
||||
* @description Fixed v0 taxonomy. User-extensible types come in v1.
|
||||
* @enum {string}
|
||||
*/
|
||||
RequestType: "code_generation" | "code_understanding" | "technical_design" | "analytical_reasoning" | "writing" | "factual_lookup" | "general";
|
||||
/** ResetSpendRequest */
|
||||
ResetSpendRequest: {
|
||||
/** Reset To */
|
||||
|
|
@ -35855,7 +35867,7 @@ export interface components {
|
|||
* Cause
|
||||
* @enum {string}
|
||||
*/
|
||||
cause?: "heuristic_scorer" | "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" | "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 */
|
||||
|
|
@ -36738,6 +36750,33 @@ export interface components {
|
|||
[key: string]: unknown;
|
||||
};
|
||||
};
|
||||
/** TierCohortStatistic */
|
||||
TierCohortStatistic: {
|
||||
/** Cohort */
|
||||
cohort: string;
|
||||
/** Observations */
|
||||
observations: number;
|
||||
/** Successes */
|
||||
successes: number;
|
||||
/** Tier */
|
||||
tier: number;
|
||||
};
|
||||
/** TierDataset */
|
||||
TierDataset: {
|
||||
/** License */
|
||||
license: string;
|
||||
/** Name */
|
||||
name: string;
|
||||
/** Rows */
|
||||
rows: number;
|
||||
/**
|
||||
* Success Definition
|
||||
* @default quality score meets the dataset success threshold
|
||||
*/
|
||||
success_definition: string;
|
||||
/** Url */
|
||||
url: string;
|
||||
};
|
||||
/**
|
||||
* TierDefinition
|
||||
* @description An operator-defined tier: the name the LLM classifier must return and its rubric description.
|
||||
|
|
@ -36754,6 +36793,25 @@ export interface components {
|
|||
*/
|
||||
name: string;
|
||||
};
|
||||
/** TierDomainStatistic */
|
||||
TierDomainStatistic: {
|
||||
/** Observations */
|
||||
observations: number;
|
||||
request_type: components["schemas"]["RequestType"];
|
||||
/** Successes */
|
||||
successes: number;
|
||||
/** Tier */
|
||||
tier: number;
|
||||
};
|
||||
/** TierGlobalStatistic */
|
||||
TierGlobalStatistic: {
|
||||
/** Observations */
|
||||
observations: number;
|
||||
/** Successes */
|
||||
successes: number;
|
||||
/** Tier */
|
||||
tier: number;
|
||||
};
|
||||
/**
|
||||
* TokenCountDetailsResponse
|
||||
* @description Response structure for token count details with modality breakdown.
|
||||
|
|
@ -37023,6 +37081,57 @@ export interface components {
|
|||
} & {
|
||||
[key: string]: unknown;
|
||||
};
|
||||
/** TrainedTierArtifact */
|
||||
TrainedTierArtifact: {
|
||||
/**
|
||||
* Cohort Prior Mass
|
||||
* @default 20
|
||||
*/
|
||||
cohort_prior_mass: number;
|
||||
/**
|
||||
* Cohort Statistics
|
||||
* @default []
|
||||
*/
|
||||
cohort_statistics: components["schemas"]["TierCohortStatistic"][];
|
||||
/**
|
||||
* Datasets
|
||||
* @default []
|
||||
*/
|
||||
datasets: components["schemas"]["TierDataset"][];
|
||||
/**
|
||||
* Domain Prior Mass
|
||||
* @default 200
|
||||
*/
|
||||
domain_prior_mass: number;
|
||||
/**
|
||||
* Domain Statistics
|
||||
* @default []
|
||||
*/
|
||||
domain_statistics: components["schemas"]["TierDomainStatistic"][];
|
||||
/** Global Statistics */
|
||||
global_statistics: components["schemas"]["TierGlobalStatistic"][];
|
||||
/**
|
||||
* Routing Threshold
|
||||
* @default 0.75
|
||||
*/
|
||||
routing_threshold: number;
|
||||
/**
|
||||
* Schema Version
|
||||
* @default 1
|
||||
* @constant
|
||||
*/
|
||||
schema_version: 1;
|
||||
/**
|
||||
* Split Method
|
||||
* @default sha256(prompt): 70% train, 15% validation, 15% test
|
||||
*/
|
||||
split_method: string;
|
||||
/**
|
||||
* Success Definition
|
||||
* @default quality score meets the dataset success threshold
|
||||
*/
|
||||
success_definition: string;
|
||||
};
|
||||
/** TransformRequestBody */
|
||||
TransformRequestBody: {
|
||||
call_type: components["schemas"]["CallTypes"];
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue