mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
Merge remote-tracking branch 'upstream/litellm_internal_staging' into resolve/pr-39182-greptile-p1
This commit is contained in:
commit
68f653b705
87 changed files with 4854 additions and 294 deletions
|
|
@ -101,6 +101,12 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
# The base image only configures Chainguard's authenticated apk repo, which
|
||||
# requires an enterprise subscription. Add the public Wolfi repo so `apk add`
|
||||
# also works for anyone installing extra packages into a running container.
|
||||
# https://github.com/BerriAI/litellm/issues/33518
|
||||
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
|
||||
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
|
||||
|
||||
|
|
|
|||
|
|
@ -136,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
|
|||
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
|
||||
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
|
||||
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
|
||||
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
|
||||
|
||||
# Allowlist of commands permitted for MCP stdio transport.
|
||||
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
|
||||
|
|
@ -1589,6 +1590,7 @@ KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
|
|||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
|
||||
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
|
||||
MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job"
|
||||
USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job"
|
||||
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
|
||||
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
|
||||
SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning"
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import base64
|
|||
import os
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from datetime import timedelta
|
||||
from functools import partial
|
||||
from importlib import metadata
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
|
|
@ -47,7 +48,8 @@ from mcp.types import Tool as MCPTool
|
|||
from pydantic import AnyUrl
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
|
||||
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
||||
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
from litellm.types.llms.custom_http import VerifyTypes
|
||||
from litellm.types.mcp import (
|
||||
|
|
@ -603,17 +605,19 @@ class MCPClient:
|
|||
"""
|
||||
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
|
||||
|
||||
async def _list_tools_operation(session: ClientSession):
|
||||
return await session.list_tools()
|
||||
|
||||
try:
|
||||
result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
|
||||
tool_count: Final = len(result.tools)
|
||||
tool_names: Final = [tool.name for tool in result.tools]
|
||||
# A per-server timeout above the global default extends the whole-walk deadline
|
||||
listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)
|
||||
tools: Final = await self.run_with_session(
|
||||
partial(list_tools_with_pagination, listing_deadline=listing_deadline),
|
||||
quiet_on_error=raise_on_error,
|
||||
)
|
||||
tool_count: Final = len(tools)
|
||||
tool_names: Final = tuple(tool.name for tool in tools)
|
||||
verbose_logger.info(
|
||||
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
|
||||
)
|
||||
return result.tools
|
||||
return tools
|
||||
except asyncio.CancelledError:
|
||||
verbose_logger.warning("MCP client list_tools was cancelled")
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -1,14 +1,22 @@
|
|||
import json
|
||||
from typing import Final, Literal
|
||||
|
||||
import anyio
|
||||
from mcp import ClientSession
|
||||
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
|
||||
from mcp.types import CallToolResult as MCPCallToolResult
|
||||
from mcp.types import PaginatedRequestParams
|
||||
from mcp.types import Tool as MCPTool
|
||||
from openai.types.chat import ChatCompletionToolParam
|
||||
from openai.types.responses.function_tool_param import FunctionToolParam
|
||||
from openai.types.shared_params.function_definition import FunctionDefinition
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
MCP_CLIENT_TIMEOUT,
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesTool
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
|
||||
|
|
@ -90,6 +98,64 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
|
|||
)
|
||||
|
||||
|
||||
async def list_tools_with_pagination(
|
||||
session: ClientSession, listing_deadline: float | None = None
|
||||
) -> list[MCPTool]: # mutable-ok: list return contract
|
||||
"""Collect tools from every tools/list page by following nextCursor.
|
||||
|
||||
Stops and returns the tools collected so far when the upstream repeats a
|
||||
cursor, the page cap is reached, or the whole-walk deadline expires, so a
|
||||
buggy or slow upstream yields a partial catalog instead of an error.
|
||||
listing_deadline overrides the default whole-walk deadline; callers with a
|
||||
per-server timeout above the global default pass it through here.
|
||||
"""
|
||||
tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools
|
||||
seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops
|
||||
cursor: str | None = None # rebind-ok: advances to each page's nextCursor
|
||||
# The per-request session read timeout restarts on every page, so a multi-page
|
||||
# walk needs its own overall deadline. max() keeps the pre-pagination guarantee
|
||||
# that a single page slower than the listing timeout but within the client
|
||||
# timeout still succeeds.
|
||||
effective_deadline: Final = (
|
||||
listing_deadline if listing_deadline is not None else max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT)
|
||||
)
|
||||
|
||||
with anyio.move_on_after(effective_deadline):
|
||||
for _ in range(MCP_TOOL_LISTING_MAX_PAGES):
|
||||
result = (
|
||||
await session.list_tools()
|
||||
if cursor is None
|
||||
else await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
|
||||
)
|
||||
tools.extend(result.tools)
|
||||
|
||||
next_cursor = getattr(result, "nextCursor", None)
|
||||
if not isinstance(next_cursor, str) or not next_cursor:
|
||||
return tools
|
||||
if next_cursor in seen_cursors:
|
||||
verbose_logger.warning(
|
||||
"MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far",
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
seen_cursors.add(next_cursor)
|
||||
cursor = next_cursor
|
||||
|
||||
verbose_logger.warning(
|
||||
"MCP server tools/list pagination exceeded the maximum of %s pages; returning %s tools collected so far",
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
|
||||
verbose_logger.warning(
|
||||
"MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far",
|
||||
effective_deadline,
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
|
||||
|
||||
async def load_mcp_tools(
|
||||
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
|
||||
) -> list[MCPTool] | list[ChatCompletionToolParam]:
|
||||
|
|
@ -103,10 +169,12 @@ async def load_mcp_tools(
|
|||
|
||||
If format is set to "openai", the tools are converted to OpenAI API compatible tools.
|
||||
"""
|
||||
tools: Final = await session.list_tools()
|
||||
tools: Final = await list_tools_with_pagination(session)
|
||||
if format == "openai":
|
||||
return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools]
|
||||
return tools.tools
|
||||
return [ # mutable-ok: public API returns a list
|
||||
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools
|
||||
]
|
||||
return tools
|
||||
|
||||
|
||||
########################################################
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ from .utils import process_slack_alerting_variables
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.router import Router as _Router
|
||||
|
||||
Router = _Router
|
||||
|
|
@ -1944,6 +1945,69 @@ Model Info:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error sending weekly spend report %s", e)
|
||||
|
||||
async def send_user_spend_alerts(self, prisma_client: "PrismaClient | None" = None) -> None:
|
||||
"""Check per-user daily/monthly spend thresholds and spend anomalies, alerting once per user per period."""
|
||||
if self.alerting is None or "slack" not in self.alerting:
|
||||
return
|
||||
|
||||
thresholds_enabled: Final = AlertType.user_spend_thresholds in self.alert_types
|
||||
anomalies_enabled: Final = AlertType.user_spend_anomalies in self.alert_types
|
||||
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
|
||||
|
||||
prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client
|
||||
if prisma_client is None:
|
||||
return
|
||||
|
||||
from litellm.integrations.SlackAlerting.user_spend_alerts import (
|
||||
evaluate_user_spend,
|
||||
fetch_user_spend_rows,
|
||||
)
|
||||
|
||||
try:
|
||||
today: Final = datetime.datetime.now(datetime.timezone.utc).date()
|
||||
rows: Final = await fetch_user_spend_rows(
|
||||
prisma_client=prisma_client,
|
||||
today=today,
|
||||
baseline_days=self.alerting_args.spend_anomaly_baseline_days,
|
||||
)
|
||||
all_events: Final = tuple(
|
||||
event
|
||||
for row in rows
|
||||
for event in evaluate_user_spend(
|
||||
row=row,
|
||||
args=self.alerting_args,
|
||||
today=today,
|
||||
thresholds_enabled=thresholds_enabled,
|
||||
anomalies_enabled=anomalies_enabled,
|
||||
)
|
||||
)
|
||||
cached_flags: Final = await asyncio.gather(
|
||||
*(self.internal_usage_cache.async_get_cache(key=event.cache_key) for event in all_events)
|
||||
)
|
||||
new_events: Final = tuple(event for event, cached in zip(all_events, cached_flags) if not cached)
|
||||
for alert_type in (AlertType.user_spend_thresholds, AlertType.user_spend_anomalies):
|
||||
typed_events = tuple(event for event in new_events if event.alert_type == alert_type)
|
||||
if not typed_events:
|
||||
continue
|
||||
await self.send_alert(
|
||||
message="\n\n".join(event.message for event in typed_events),
|
||||
level="High",
|
||||
alert_type=alert_type,
|
||||
alerting_metadata={}, # mutable-ok: send_alert takes a dict payload
|
||||
)
|
||||
for event in typed_events:
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=event.cache_key,
|
||||
value="SENT",
|
||||
ttl=event.cache_ttl,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # background job must not crash the scheduler
|
||||
verbose_proxy_logger.exception("Error sending user spend alerts: %s", e)
|
||||
|
||||
async def send_fallback_stats_from_prometheus(self):
|
||||
"""
|
||||
Helper to send fallback statistics from prometheus server -> to slack
|
||||
|
|
|
|||
139
litellm/integrations/SlackAlerting/user_spend_alerts.py
Normal file
139
litellm/integrations/SlackAlerting/user_spend_alerts.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
"""Per-user daily/monthly spend threshold alerts and spend anomaly detection."""
|
||||
|
||||
import datetime
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.constants import HOURS_IN_A_DAY
|
||||
from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
DAY_SECONDS: Final = HOURS_IN_A_DAY * 60 * 60
|
||||
MONTHLY_ALERT_TTL_SECONDS: Final = 32 * DAY_SECONDS
|
||||
|
||||
USER_SPEND_QUERY: Final = """
|
||||
SELECT
|
||||
user_id,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date = $1), 0)::float AS daily_spend,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0)::float AS monthly_spend,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date >= $3 AND date < $1), 0)::float AS baseline_spend
|
||||
FROM "LiteLLM_DailyUserSpend"
|
||||
WHERE date >= LEAST($2, $3) AND user_id IS NOT NULL
|
||||
GROUP BY user_id
|
||||
HAVING COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0) > 0
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserSpendRow:
|
||||
user_id: str
|
||||
daily_spend: float
|
||||
monthly_spend: float
|
||||
baseline_spend: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserSpendAlertEvent:
|
||||
kind: Literal["daily_threshold", "monthly_threshold", "anomaly"]
|
||||
alert_type: AlertType
|
||||
message: str
|
||||
cache_key: str
|
||||
cache_ttl: int
|
||||
|
||||
|
||||
USER_SPEND_ROWS_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...])
|
||||
|
||||
|
||||
async def fetch_user_spend_rows(
|
||||
prisma_client: "PrismaClient",
|
||||
today: datetime.date,
|
||||
baseline_days: int,
|
||||
) -> tuple[UserSpendRow, ...]:
|
||||
today_str: Final = today.strftime("%Y-%m-%d")
|
||||
month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d")
|
||||
baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d")
|
||||
raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
|
||||
return USER_SPEND_ROWS_ADAPTER.validate_python(raw)
|
||||
|
||||
|
||||
def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
|
||||
threshold: Final = args.daily_spend_per_user_threshold
|
||||
if threshold is None or row.daily_spend < threshold:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="daily_threshold",
|
||||
alert_type=AlertType.user_spend_thresholds,
|
||||
message=(
|
||||
f"User Daily Spend Threshold Crossed:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend Today: `${row.daily_spend:.2f}`\n"
|
||||
f"Daily Threshold: `${threshold:.2f}`"
|
||||
),
|
||||
cache_key=f"user_spend_alert_daily_{row.user_id}_{today_str}",
|
||||
cache_ttl=DAY_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _monthly_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, month_str: str) -> UserSpendAlertEvent | None:
|
||||
threshold: Final = args.monthly_spend_per_user_threshold
|
||||
if threshold is None or row.monthly_spend < threshold:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="monthly_threshold",
|
||||
alert_type=AlertType.user_spend_thresholds,
|
||||
message=(
|
||||
f"User Monthly Spend Threshold Crossed:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend This Month: `${row.monthly_spend:.2f}`\n"
|
||||
f"Monthly Threshold: `${threshold:.2f}`"
|
||||
),
|
||||
cache_key=f"user_spend_alert_monthly_{row.user_id}_{month_str}",
|
||||
cache_ttl=MONTHLY_ALERT_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _anomaly_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
|
||||
if row.daily_spend < args.spend_anomaly_min_spend:
|
||||
return None
|
||||
baseline_daily_avg: Final = row.baseline_spend / args.spend_anomaly_baseline_days
|
||||
if row.baseline_spend > 0 and row.daily_spend <= args.spend_anomaly_multiplier * baseline_daily_avg:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="anomaly",
|
||||
alert_type=AlertType.user_spend_anomalies,
|
||||
message=(
|
||||
f"User Spend Anomaly Detected:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend Today: `${row.daily_spend:.2f}`\n"
|
||||
f"Daily Average (last {args.spend_anomaly_baseline_days} days): `${baseline_daily_avg:.2f}`\n"
|
||||
f"Trigger: spend above `{args.spend_anomaly_multiplier}x` the daily average "
|
||||
f"(minimum `${args.spend_anomaly_min_spend:.2f}`)"
|
||||
),
|
||||
cache_key=f"user_spend_alert_anomaly_{row.user_id}_{today_str}",
|
||||
cache_ttl=DAY_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def evaluate_user_spend(
|
||||
row: UserSpendRow,
|
||||
args: SlackAlertingArgs,
|
||||
today: datetime.date,
|
||||
thresholds_enabled: bool,
|
||||
anomalies_enabled: bool,
|
||||
) -> tuple[UserSpendAlertEvent, ...]:
|
||||
today_str: Final = today.strftime("%Y-%m-%d")
|
||||
month_str: Final = today.strftime("%Y-%m")
|
||||
threshold_events: Final = (
|
||||
(
|
||||
_daily_threshold_event(row=row, args=args, today_str=today_str),
|
||||
_monthly_threshold_event(row=row, args=args, month_str=month_str),
|
||||
)
|
||||
if thresholds_enabled
|
||||
else ()
|
||||
)
|
||||
anomaly_events: Final = (_anomaly_event(row=row, args=args, today_str=today_str),) if anomalies_enabled else ()
|
||||
return tuple(event for event in (*threshold_events, *anomaly_events) if event is not None)
|
||||
|
|
@ -127,7 +127,7 @@ def handle_anthropic_text_model_custom_llm_provider(
|
|||
return model, custom_llm_provider
|
||||
|
||||
|
||||
def declared_authenticating_provider(model: str, custom_llm_provider: str | None = None) -> str | None:
|
||||
def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
|
||||
"""The authenticating provider this pair already names, or None.
|
||||
|
||||
get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, because their
|
||||
|
|
@ -135,7 +135,7 @@ def declared_authenticating_provider(model: str, custom_llm_provider: str | None
|
|||
and for a declared pair the resolver's answer is the declaration itself, so metadata callers
|
||||
adopt the declaration instead of resolving.
|
||||
"""
|
||||
declared: Final = custom_llm_provider or model.split("/", 1)[0]
|
||||
declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
|
||||
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2959,13 +2959,25 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"Model=%s not found in completion cost map. Setting 'response_cost' to None", self.model
|
||||
)
|
||||
self.model_call_details["response_cost"] = None
|
||||
except Exception: # noqa: BLE001 # cost calculation must never block later callbacks (slot release)
|
||||
verbose_logger.exception(
|
||||
"Error calculating streaming response cost for model=%s. Setting 'response_cost' to None",
|
||||
self.model,
|
||||
)
|
||||
self.model_call_details["response_cost"] = None
|
||||
|
||||
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
|
||||
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
|
||||
complete_streaming_response, start_time, end_time
|
||||
)
|
||||
try:
|
||||
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
|
||||
complete_streaming_response, start_time, end_time
|
||||
)
|
||||
except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release)
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception building the standard logging payload "
|
||||
"for a streaming response; callbacks still run without it"
|
||||
)
|
||||
|
||||
# print standard logging payload
|
||||
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
|
||||
|
|
@ -3005,32 +3017,39 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
## LOGGING HOOK ##
|
||||
|
||||
for callback in callbacks:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
try:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=self.model_call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=self.model_call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
continue
|
||||
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
continue
|
||||
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
elif isinstance(callback, CustomLogger):
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
result=result, litellm_logging_obj=self, custom_logger=callback
|
||||
)
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
elif isinstance(callback, CustomLogger):
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
result=result, litellm_logging_obj=self, custom_logger=callback
|
||||
)
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release)
|
||||
verbose_logger.error(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s",
|
||||
traceback.format_exc(),
|
||||
)
|
||||
self._handle_callback_failure(callback=callback)
|
||||
|
||||
self.has_run_logging(event_type="async_success")
|
||||
|
||||
|
|
|
|||
|
|
@ -1281,6 +1281,19 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin
|
|||
return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo
|
||||
|
||||
|
||||
def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]:
|
||||
function: Final = tool.get("function")
|
||||
if not isinstance(function, dict):
|
||||
return tool
|
||||
parameters: Final = function.get("parameters")
|
||||
if not isinstance(parameters, dict):
|
||||
return tool
|
||||
flattened: Final = flatten_top_level_schema_combinators(parameters)
|
||||
if flattened is parameters:
|
||||
return tool
|
||||
return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts
|
||||
|
||||
|
||||
def _get_image_mime_type_from_url(url: str) -> str | None:
|
||||
"""
|
||||
Get mime type for common image URLs
|
||||
|
|
|
|||
|
|
@ -86,22 +86,41 @@ def _decoded_sse_data_line(line: bytes) -> object | None:
|
|||
return None
|
||||
|
||||
|
||||
def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
|
||||
def _anthropic_event_payload(chunk: object, event_type: str) -> Mapping[str, object] | None:
|
||||
if isinstance(chunk, dict):
|
||||
return chunk if chunk.get("type") == "error" else None
|
||||
return chunk if chunk.get("type") == event_type else None
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines())
|
||||
return next(
|
||||
(
|
||||
candidate
|
||||
for candidate in decoded_lines
|
||||
if isinstance(candidate, dict) and candidate.get("type") == "error"
|
||||
if isinstance(candidate, dict) and candidate.get("type") == event_type
|
||||
),
|
||||
None,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
|
||||
return _anthropic_event_payload(chunk, "error")
|
||||
|
||||
|
||||
def parse_anthropic_refusal_stop_details(chunk: object) -> Mapping[str, object] | None:
|
||||
"""
|
||||
Return the ``stop_details`` object of an Anthropic SSE ``message_delta``
|
||||
chunk whose delta carries ``stop_reason: "refusal"`` (a safeguard refusal:
|
||||
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback),
|
||||
or None for any other chunk, a plain refusal without ``stop_details`` included.
|
||||
"""
|
||||
payload: Final = _anthropic_event_payload(chunk, "message_delta")
|
||||
delta: Final = payload.get("delta") if payload is not None else None
|
||||
if not isinstance(delta, dict) or delta.get("stop_reason") != "refusal":
|
||||
return None
|
||||
stop_details: Final = delta.get("stop_details")
|
||||
return stop_details if isinstance(stop_details, dict) else None
|
||||
|
||||
|
||||
def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None:
|
||||
"""Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None."""
|
||||
payload: Final = _anthropic_error_event_payload(chunk)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,40 @@
|
|||
from collections.abc import Mapping
|
||||
from functools import lru_cache
|
||||
from typing import Any, Final, cast, get_type_hints
|
||||
from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints
|
||||
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.exceptions import ContentPolicyViolationError
|
||||
|
||||
|
||||
def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None:
|
||||
"""
|
||||
Return the ``stop_details`` of an Anthropic Messages response refused by a
|
||||
safeguard (``stop_reason: "refusal"`` carrying ``stop_details``:
|
||||
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback),
|
||||
or None for any other response, a plain refusal without ``stop_details`` included.
|
||||
"""
|
||||
if not isinstance(response, dict) or response.get("stop_reason") != "refusal":
|
||||
return None
|
||||
stop_details: Final = response.get("stop_details")
|
||||
return stop_details if isinstance(stop_details, dict) else None
|
||||
|
||||
|
||||
def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError":
|
||||
"""The exception a safeguard-refused Anthropic response converts into so the
|
||||
content-policy fallback chain can re-dispatch it."""
|
||||
from litellm.exceptions import ContentPolicyViolationError
|
||||
|
||||
return ContentPolicyViolationError(
|
||||
message=f"Anthropic safeguard refusal (category: {stop_details.get('category')}).",
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _anthropic_messages_optional_param_keys() -> frozenset[str]:
|
||||
|
|
@ -100,14 +129,12 @@ def mock_response(
|
|||
model=model,
|
||||
)
|
||||
return AnthropicMessagesResponse(
|
||||
**{
|
||||
"content": [{"text": mock_response, "type": "text"}],
|
||||
"id": "msg_013Zva2CMHLNnXjNJJKqJ2EF",
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"role": "assistant",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"type": "message",
|
||||
"usage": {"input_tokens": 2095, "output_tokens": 503},
|
||||
}
|
||||
content=[{"text": mock_response, "type": "text"}],
|
||||
id="msg_013Zva2CMHLNnXjNJJKqJ2EF",
|
||||
model="claude-sonnet-4-20250514",
|
||||
role="assistant",
|
||||
stop_reason="end_turn",
|
||||
stop_sequence=None,
|
||||
type="message",
|
||||
usage={"input_tokens": 2095, "output_tokens": 503},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from httpx._models import Headers, Response
|
||||
|
|
@ -6,6 +8,7 @@ import litellm
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
drop_tool_reference_parts_from_tool_messages,
|
||||
hoist_images_from_tool_messages,
|
||||
tool_with_flattened_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_azure_openai_messages,
|
||||
|
|
@ -32,6 +35,19 @@ else:
|
|||
LoggingClass = Any
|
||||
|
||||
|
||||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
tools: Final = optional_params.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return _NO_TOOLS_UPDATE
|
||||
flattened: Final = [ # mutable-ok: request tools are a JSON list
|
||||
tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools
|
||||
]
|
||||
return MappingProxyType({"tools": flattened})
|
||||
|
||||
|
||||
class AzureOpenAIConfig(BaseConfig):
|
||||
"""
|
||||
Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions
|
||||
|
|
@ -261,6 +277,7 @@ class AzureOpenAIConfig(BaseConfig):
|
|||
"model": model,
|
||||
"messages": azure_messages,
|
||||
**optional_params,
|
||||
**flattened_tools_update(optional_params),
|
||||
}
|
||||
|
||||
def transform_response(
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.utils import get_model_info, supports_reasoning
|
||||
|
||||
from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig
|
||||
from .gpt_transformation import flattened_tools_update
|
||||
|
||||
|
||||
class AzureOpenAIO1Config(OpenAIOSeriesConfig):
|
||||
|
|
@ -108,4 +109,8 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig):
|
|||
headers: dict,
|
||||
) -> dict:
|
||||
model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name
|
||||
return super().transform_request(model, messages, optional_params, litellm_params, headers)
|
||||
flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict
|
||||
**optional_params,
|
||||
**flattened_tools_update(optional_params),
|
||||
}
|
||||
return super().transform_request(model, messages, flattened_params, litellm_params, headers)
|
||||
|
|
|
|||
|
|
@ -20,9 +20,9 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
|
|||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
drop_tool_reference_parts_from_tool_messages,
|
||||
flatten_top_level_schema_combinators,
|
||||
get_tool_call_names,
|
||||
hoist_images_from_tool_messages,
|
||||
tool_with_flattened_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
async_convert_url_to_base64,
|
||||
|
|
@ -70,19 +70,6 @@ else:
|
|||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]:
|
||||
function: Final = tool.get("function")
|
||||
if not isinstance(function, dict):
|
||||
return tool
|
||||
parameters: Final = function.get("parameters")
|
||||
if not isinstance(parameters, dict):
|
||||
return tool
|
||||
flattened: Final = flatten_top_level_schema_combinators(parameters)
|
||||
if flattened is parameters:
|
||||
return tool
|
||||
return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts
|
||||
|
||||
|
||||
class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
||||
"""
|
||||
Reference: https://platform.openai.com/docs/api-reference/chat/create
|
||||
|
|
@ -466,7 +453,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
):
|
||||
return _NO_TOOLS_UPDATE
|
||||
flattened: Final = [ # mutable-ok: request tools are a JSON list
|
||||
_tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools
|
||||
tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools
|
||||
]
|
||||
return MappingProxyType({"tools": flattened})
|
||||
|
||||
|
|
|
|||
|
|
@ -82,6 +82,10 @@ class ResponsesStreamChunk(TypedDict, total=False):
|
|||
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
delta: ReadOnly[str]
|
||||
item_id: ReadOnly[str]
|
||||
output_index: ReadOnly[int]
|
||||
content_index: ReadOnly[int]
|
||||
|
||||
|
||||
def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
|
||||
|
|
@ -658,8 +662,32 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
|
||||
"""
|
||||
Get the string so far from the responses so far.
|
||||
|
||||
``response.output_text.done`` events carry the whole part in ``text``, while
|
||||
``response.output_text.delta`` events carry fragments in ``delta``. A stream
|
||||
that dies before its done event (``response.failed`` / ``response.incomplete``)
|
||||
has text only in deltas, so per content part the done text wins when present
|
||||
and the joined deltas fill in otherwise, never both.
|
||||
"""
|
||||
return "".join([response.get("text", "") for response in responses_so_far])
|
||||
keyed_events: Final = tuple(
|
||||
(
|
||||
(event.get("item_id"), event.get("output_index"), event.get("content_index")),
|
||||
event.get("text"),
|
||||
event.get("delta"),
|
||||
)
|
||||
for event in responses_so_far
|
||||
if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str)
|
||||
)
|
||||
|
||||
def part_text(part_key: tuple[object, object, object]) -> str:
|
||||
done_texts: Final = tuple(
|
||||
text for key, text, _ in keyed_events if key == part_key and isinstance(text, str)
|
||||
)
|
||||
if done_texts:
|
||||
return done_texts[-1]
|
||||
return "".join(delta for key, _, delta in keyed_events if key == part_key and isinstance(delta, str))
|
||||
|
||||
return "".join(part_text(key) for key in dict.fromkeys(key for key, _, _ in keyed_events))
|
||||
|
||||
def _has_text_content(self, response: "ResponsesAPIResponse") -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Mapping
|
|||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT
|
||||
from litellm.exceptions import (
|
||||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
|
|
@ -86,8 +88,6 @@ def _connection_error_message(exc: BaseException) -> str:
|
|||
|
||||
|
||||
if MCP_AVAILABLE:
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
|
||||
global_mcp_server_manager,
|
||||
|
|
@ -1173,6 +1173,7 @@ if MCP_AVAILABLE:
|
|||
transport=request.transport,
|
||||
auth_type=request.auth_type,
|
||||
mcp_info=request.mcp_info,
|
||||
timeout=request.timeout,
|
||||
command=request.command,
|
||||
args=request.args,
|
||||
env=request.env,
|
||||
|
|
@ -1402,11 +1403,28 @@ if MCP_AVAILABLE:
|
|||
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
|
||||
|
||||
async def _list_tools_operation(client):
|
||||
async def _list_tools_session_operation(session):
|
||||
return await session.list_tools()
|
||||
|
||||
list_tools_response: Final = await client.run_with_session(_list_tools_session_operation)
|
||||
list_tools_result: Final[list[MCPTool]] = list_tools_response.tools
|
||||
# Bound the whole pagination walk: without this the preview is limited only by the
|
||||
# per-request timeout times the page cap. max() keeps the pre-pagination guarantee
|
||||
# that a single slow page within the client timeout still succeeds, and a
|
||||
# per-server timeout above the global default extends the deadline with it.
|
||||
listing_deadline: Final = max(
|
||||
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
|
||||
if list_tools_result is None:
|
||||
verbose_logger.warning(
|
||||
"MCP tools/list preview timed out after %s seconds while paginating upstream tools",
|
||||
listing_deadline,
|
||||
)
|
||||
return { # mutable-ok: error response payload
|
||||
"status": "error",
|
||||
"error": True,
|
||||
"message": f"Timed out listing tools after {listing_deadline} seconds. "
|
||||
"The MCP server may be responding slowly or paginating excessively.",
|
||||
}
|
||||
model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result]
|
||||
return {
|
||||
"tools": model_dumped_tools,
|
||||
|
|
|
|||
|
|
@ -1216,6 +1216,13 @@ class GenerateKeyRequest(KeyRequestBase):
|
|||
organization_id: str | None = None
|
||||
project_id: str | None = None
|
||||
|
||||
@field_validator("team_id", mode="before")
|
||||
@classmethod
|
||||
def treat_cleared_team_id_as_unset(cls, v: object) -> object:
|
||||
if v == "":
|
||||
return None
|
||||
return v
|
||||
|
||||
|
||||
class GenerateKeyResponse(KeyRequestBase):
|
||||
key: str
|
||||
|
|
|
|||
174
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
Normal file
174
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
"""
|
||||
Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only
|
||||
stream event carrying a model, so streamed responses report the requested model like
|
||||
non-streaming ones do.
|
||||
|
||||
Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the
|
||||
provider passthrough path) or as event dicts (fake-stream and agentic paths).
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
_MESSAGE_START_EVENT: Final = "message_start"
|
||||
_MESSAGE_START_MARKER: Final = b"message_start"
|
||||
_SSE_DATA_FIELD: Final = "data:"
|
||||
_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n")
|
||||
_MAX_HELD_BYTES: Final = 65536
|
||||
_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"')
|
||||
|
||||
_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _restamped_event(event: Mapping[str, object], requested_model: str) -> Mapping[str, object] | None:
|
||||
message: Final = event.get("message")
|
||||
if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict):
|
||||
return None
|
||||
if message.get("model") == requested_model:
|
||||
return None
|
||||
return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is
|
||||
|
||||
|
||||
def _restamped_data_line(line: str, requested_model: str) -> str | None:
|
||||
stripped: Final = line.strip()
|
||||
if not stripped.startswith(_SSE_DATA_FIELD):
|
||||
return None
|
||||
payload: Final = stripped[len(_SSE_DATA_FIELD) :].strip()
|
||||
if not payload or payload == "[DONE]":
|
||||
return None
|
||||
try:
|
||||
event: Final = _EVENT_ADAPTER.validate_json(payload)
|
||||
except ValidationError:
|
||||
return None
|
||||
restamped: Final = _restamped_event(event, requested_model)
|
||||
if restamped is None:
|
||||
return None
|
||||
terminator: Final = line[len(line.rstrip("\r\n")) :]
|
||||
return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}"
|
||||
|
||||
|
||||
def _restamped_frame(frame: str, requested_model: str) -> str | None:
|
||||
lines: Final = frame.splitlines(keepends=True)
|
||||
restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines)
|
||||
if all(line is None for line in restamped):
|
||||
return None
|
||||
return "".join(new if new is not None else old for new, old in zip(restamped, lines))
|
||||
|
||||
|
||||
def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object:
|
||||
"""
|
||||
Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``.
|
||||
|
||||
Chunks that carry no model are returned unchanged.
|
||||
"""
|
||||
if isinstance(chunk, dict):
|
||||
try:
|
||||
event: Final = _EVENT_ADAPTER.validate_python(chunk)
|
||||
except ValidationError:
|
||||
return chunk
|
||||
return _restamped_event(event, requested_model) or chunk
|
||||
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
if _MESSAGE_START_EVENT.encode() not in chunk:
|
||||
return chunk
|
||||
restamped_bytes: Final = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model)
|
||||
return chunk if restamped_bytes is None else restamped_bytes.encode("utf-8")
|
||||
|
||||
if isinstance(chunk, str):
|
||||
if _MESSAGE_START_EVENT not in chunk:
|
||||
return chunk
|
||||
restamped_text: Final = _restamped_frame(chunk, requested_model)
|
||||
return chunk if restamped_text is None else restamped_text
|
||||
|
||||
return chunk
|
||||
|
||||
|
||||
def _is_ping_frame(frame: bytes) -> bool:
|
||||
return any(marker in frame for marker in _PING_MARKERS)
|
||||
|
||||
|
||||
class AnthropicStreamModelRestamper:
|
||||
"""
|
||||
Per-stream restamper for the encoded passthrough path, where chunks are raw
|
||||
transport reads: the ``message_start`` SSE frame can arrive split across
|
||||
chunks or coalesced with later frames. Complete frames (``\\n\\n``,
|
||||
``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator
|
||||
closes them and an incomplete tail is held until it completes, so the
|
||||
restamp never misses a torn frame; ``flush`` returns whatever is still held
|
||||
when the stream ends so no bytes are swallowed. Once ``message_start`` has
|
||||
been handled, or the first real event proves the stream carries none, every
|
||||
later chunk passes through untouched.
|
||||
"""
|
||||
|
||||
def __init__(self, requested_model: str) -> None:
|
||||
self._requested_model: Final = requested_model
|
||||
self._held = b""
|
||||
self._armed = True
|
||||
|
||||
def process(self, chunk: object) -> object:
|
||||
if not self._armed:
|
||||
return chunk
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
return self._process_encoded(bytes(chunk))
|
||||
if isinstance(chunk, str):
|
||||
return self._process_encoded(chunk.encode("utf-8"))
|
||||
restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model)
|
||||
if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"):
|
||||
self._armed = False
|
||||
return restamped
|
||||
|
||||
def flush(self) -> bytes:
|
||||
held: Final = self._held
|
||||
self._held = b""
|
||||
self._armed = False
|
||||
if not held:
|
||||
return b""
|
||||
restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model)
|
||||
return restamped if isinstance(restamped, bytes) else held
|
||||
|
||||
def _process_encoded(self, data: bytes) -> bytes:
|
||||
combined: Final = self._held + data
|
||||
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined))
|
||||
if not boundaries:
|
||||
if len(combined) > _MAX_HELD_BYTES:
|
||||
self._held = b""
|
||||
self._armed = False
|
||||
return combined
|
||||
self._held = combined
|
||||
return b""
|
||||
emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]])
|
||||
tail: Final = combined[boundaries[-1] :]
|
||||
if not self._armed:
|
||||
self._held = b""
|
||||
return emitted + tail
|
||||
self._held = tail
|
||||
return emitted
|
||||
|
||||
def _restamped_closed_block(self, closed: bytes) -> bytes:
|
||||
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed))
|
||||
frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries))
|
||||
decider: Final = next(
|
||||
(
|
||||
index
|
||||
for index, frame in enumerate(frames)
|
||||
if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame))
|
||||
),
|
||||
None,
|
||||
)
|
||||
if decider is None:
|
||||
return closed
|
||||
self._armed = False
|
||||
if _MESSAGE_START_MARKER not in frames[decider]:
|
||||
return closed
|
||||
restamped_text: Final = _restamped_frame(
|
||||
frames[decider].decode("utf-8", errors="ignore"), self._requested_model
|
||||
)
|
||||
if restamped_text is None:
|
||||
return closed
|
||||
return b"".join(
|
||||
restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames)
|
||||
)
|
||||
|
|
@ -9,6 +9,7 @@ import click
|
|||
import requests
|
||||
|
||||
from .auth import context_secret_vault, get_stored_api_key, login
|
||||
from .cmd_quoting import quote_for_cmd
|
||||
|
||||
ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN"
|
||||
|
|
@ -151,31 +152,9 @@ def verify_proxy_key(
|
|||
|
||||
|
||||
_WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"})
|
||||
_CMD_PERCENT_GUARD: Final = "%%cd:~,%"
|
||||
_CMD_LINE_BREAKS: Final = ("\r", "\n")
|
||||
|
||||
|
||||
def _double_trailing_backslashes(segment: str) -> str:
|
||||
bare: Final = segment.rstrip("\\")
|
||||
return bare + "\\" * 2 * (len(segment) - len(bare))
|
||||
|
||||
|
||||
def _quote_for_cmd(token: str) -> str:
|
||||
"""Quote one token so both parsers that read it see the original text.
|
||||
|
||||
Follows the algorithm the Rust standard library settled on for batch files
|
||||
after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a
|
||||
quoted string on a lone `"` and so wants an embedded one doubled, and the
|
||||
shim's own interpreter, which re-splits `%*` under C runtime rules where a
|
||||
backslash escapes the quote that follows it, so every backslash run standing
|
||||
before a quote is doubled. Quoting cannot stop cmd expanding `%VAR%`, so each
|
||||
`%` is prefixed with `%%cd:~,`: the zero-length substring of the always
|
||||
defined `cd` expands to nothing and leaves no `%` pair for cmd to match.
|
||||
"""
|
||||
escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"'))
|
||||
return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"'
|
||||
|
||||
|
||||
def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]:
|
||||
"""Build what CreateProcess runs, routing batch shims through cmd.exe.
|
||||
|
||||
|
|
@ -202,7 +181,7 @@ def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]:
|
|||
f"Cannot pass an argument containing a line break to `{os.path.basename(path)}` on "
|
||||
"Windows: cmd.exe ends the command line there, so the agent would silently lose it."
|
||||
)
|
||||
inner: Final = " ".join(_quote_for_cmd(token) for token in (path, *rest))
|
||||
inner: Final = " ".join(quote_for_cmd(token) for token in (path, *rest))
|
||||
return f'cmd.exe /d /e:on /v:off /s /c "{inner}"'
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ live here rather than in either command module.
|
|||
|
||||
import shlex
|
||||
import shutil
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
|
@ -17,6 +18,8 @@ from pydantic import JsonValue, TypeAdapter, ValidationError
|
|||
|
||||
from litellm.litellm_core_utils.private_json import write_private_json
|
||||
|
||||
from .cmd_quoting import quote_for_cmd
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
|
|
@ -87,9 +90,12 @@ def merge_claude_settings(
|
|||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Claude Code hands the string to the system shell, `sh` on POSIX and cmd.exe
|
||||
on Windows, so every token is quoted for the shell that will read it.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
|
|
@ -106,7 +112,8 @@ def resolve_api_key_helper(base_url: str) -> str:
|
|||
raise ClaudeSettingsError(
|
||||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} --base-url {shlex.quote(base_url)} auth print-token"
|
||||
quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote
|
||||
return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token"))
|
||||
|
||||
|
||||
def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
|
||||
|
|
|
|||
26
litellm/proxy/client/cli/commands/cmd_quoting.py
Normal file
26
litellm/proxy/client/cli/commands/cmd_quoting.py
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
"""Quoting for command lines that cmd.exe reads before handing them to a program."""
|
||||
|
||||
from typing import Final
|
||||
|
||||
_CMD_PERCENT_GUARD: Final = "%%cd:~,%"
|
||||
|
||||
|
||||
def _double_trailing_backslashes(segment: str) -> str:
|
||||
bare: Final = segment.rstrip("\\")
|
||||
return bare + "\\" * 2 * (len(segment) - len(bare))
|
||||
|
||||
|
||||
def quote_for_cmd(token: str) -> str:
|
||||
"""Quote one token so both parsers that read it see the original text.
|
||||
|
||||
Follows the algorithm the Rust standard library settled on for batch files
|
||||
after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a
|
||||
quoted string on a lone `"` and so wants an embedded one doubled, and the
|
||||
program's own C runtime argv split, where a backslash escapes the quote that
|
||||
follows it, so every backslash run standing before a quote is doubled.
|
||||
Quoting cannot stop cmd expanding `%VAR%`, so each `%` is prefixed with
|
||||
`%%cd:~,`: the zero-length substring of the always defined `cd` expands to
|
||||
nothing and leaves no `%` pair for cmd to match.
|
||||
"""
|
||||
escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"'))
|
||||
return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"'
|
||||
|
|
@ -176,6 +176,9 @@ if TYPE_CHECKING:
|
|||
ProxyConfig = _ProxyConfig
|
||||
else:
|
||||
ProxyConfig = Any
|
||||
from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
|
||||
AnthropicStreamModelRestamper,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
add_litellm_data_to_request,
|
||||
refresh_proxy_server_request_body_snapshot,
|
||||
|
|
@ -2490,6 +2493,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data=self.data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
request=request,
|
||||
restamp_model=(
|
||||
None if _should_return_raw_model_name(self.data) else requested_model_from_client
|
||||
),
|
||||
)
|
||||
return await create_response(
|
||||
generator=wrap_sse_stream_with_keepalive_pings(
|
||||
|
|
@ -3442,6 +3448,16 @@ class ProxyBaseLLMRequestProcessing:
|
|||
else:
|
||||
return chunk
|
||||
|
||||
@staticmethod
|
||||
def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer:
|
||||
if restamper is None:
|
||||
return ProxyBaseLLMRequestProcessing.return_sse_chunk
|
||||
|
||||
def serialize(chunk: object) -> str:
|
||||
return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk))
|
||||
|
||||
return serialize
|
||||
|
||||
@staticmethod
|
||||
async def _finalize_streaming_generator_cleanup(
|
||||
request: Request | None,
|
||||
|
|
@ -3502,11 +3518,16 @@ class ProxyBaseLLMRequestProcessing:
|
|||
serialize_chunk: StreamChunkSerializer,
|
||||
serialize_error: StreamErrorSerializer,
|
||||
request: Request | None = None,
|
||||
flush_tail: Callable[[], bytes] | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
|
||||
cost injection, then yields chunks via serialize_chunk; on exception runs
|
||||
failure hook and yields via serialize_error. Use for SSE or NDJSON.
|
||||
|
||||
``flush_tail`` runs once after the upstream iterator completes cleanly and
|
||||
its non-empty result is yielded, so a serializer that buffers bytes across
|
||||
chunks can emit anything still held at end of stream.
|
||||
"""
|
||||
verbose_proxy_logger.debug("inside generator")
|
||||
# Resolve per-stream (not per-chunk) whether the heavy per-chunk path
|
||||
|
|
@ -3569,6 +3590,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# so it must not suppress that refund.
|
||||
delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
yield serialize_chunk(chunk)
|
||||
held_tail: Final = flush_tail() if flush_tail is not None else b""
|
||||
if held_tail:
|
||||
yield serialize_chunk(held_tail)
|
||||
stream_completed = True
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
# Client disconnected mid-stream. CancelledError / GeneratorExit
|
||||
|
|
@ -3579,8 +3603,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# billing and release exactly once. This is the outermost generator
|
||||
# Starlette closes on disconnect, so the nested iterator hook (which
|
||||
# only sees GeneratorExit on GC) cannot own the refund.
|
||||
if not stream_completed:
|
||||
client_disconnected = True
|
||||
client_disconnected = not stream_completed
|
||||
if not delivered_chunk and not _withheld_provider_output(response):
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
release_budget_reservation_on_cancel,
|
||||
|
|
@ -3634,6 +3657,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data: dict,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
request: Request | None = None,
|
||||
restamp_model: str | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Anthropic /messages and Google /generateContent streaming data generator require SSE events.
|
||||
|
|
@ -3642,17 +3666,23 @@ class ProxyBaseLLMRequestProcessing:
|
|||
SSE serializers directly (rather than re-wrapping it in another
|
||||
``async for: yield`` trampoline), so a streamed chunk traverses one
|
||||
fewer async-generator layer / coroutine resume on the hot path.
|
||||
|
||||
``restamp_model`` publishes that name on the Anthropic ``message_start``
|
||||
event in place of the provider's model, matching what the non-streaming
|
||||
response reports.
|
||||
"""
|
||||
restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None
|
||||
return ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper),
|
||||
serialize_error=lambda proxy_exc: (
|
||||
f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n"
|
||||
),
|
||||
request=request,
|
||||
flush_tail=None if restamper is None else restamper.flush,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.caching import DualCache
|
|||
from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name
|
||||
|
|
@ -215,6 +216,16 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]:
|
|||
return redacted if isinstance(redacted, list) else assessments
|
||||
|
||||
|
||||
_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
|
||||
|
||||
|
||||
def _is_responses_api_route(request_route: str | None) -> bool:
|
||||
if request_route is None:
|
||||
return False
|
||||
call_types: Final = get_call_types_for_route(request_route)
|
||||
return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types)
|
||||
|
||||
|
||||
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
||||
# During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
|
||||
# OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
|
||||
|
|
@ -2709,6 +2720,24 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
yield streamed_chunk
|
||||
return
|
||||
|
||||
# Responses-API events are neither chat-completions chunks nor raw
|
||||
# Anthropic SSE, so the assembly below cannot scan them; the unified
|
||||
# guardrail's translation layer can, with buffering semantics kept.
|
||||
if _is_responses_api_route(user_api_key_dict.request_route):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
async for translated_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
guardrail_to_apply=self,
|
||||
buffer_until_moderated_default=True,
|
||||
):
|
||||
yield translated_chunk
|
||||
return
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
|
||||
from litellm.main import stream_chunk_builder
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ from typing import (
|
|||
import anyio
|
||||
import websockets
|
||||
import websockets.exceptions
|
||||
from pydantic import BaseModel, Json, JsonValue
|
||||
from pydantic import BaseModel, Json, JsonValue, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, assert_never
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -253,6 +253,7 @@ from litellm.constants import (
|
|||
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
|
||||
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
|
||||
USER_SPEND_ALERTS_JOB_ID,
|
||||
WEEKLY_SPEND_REPORT_JOB_ID,
|
||||
)
|
||||
from litellm.exceptions import RejectedRequestError
|
||||
|
|
@ -9866,6 +9867,35 @@ class ProxyStartupEvent:
|
|||
replace_existing=True,
|
||||
)
|
||||
|
||||
slack_alerting_args: Final = proxy_logging_obj.slack_alerting_instance.alerting_args
|
||||
user_spend_check_interval: Final = (
|
||||
slack_alerting_args.user_spend_check_interval
|
||||
if isinstance(slack_alerting_args, SlackAlertingArgs) # pyright: ignore[reportUnnecessaryIsInstance] # tests inject a mock slack_alerting_instance
|
||||
else SlackAlertingArgs().user_spend_check_interval
|
||||
)
|
||||
|
||||
async def _scheduled_user_spend_alerts() -> None:
|
||||
if (
|
||||
await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=USER_SPEND_ALERTS_JOB_ID,
|
||||
ttl=max(user_spend_check_interval - 60, 60),
|
||||
allow_reentrant=False,
|
||||
)
|
||||
is False
|
||||
):
|
||||
return
|
||||
await proxy_logging_obj.slack_alerting_instance.send_user_spend_alerts()
|
||||
|
||||
scheduler.add_job(
|
||||
_scheduled_user_spend_alerts,
|
||||
"interval",
|
||||
seconds=user_spend_check_interval,
|
||||
next_run_time=datetime.now(timezone.utc) + timedelta(seconds=10 + random.randint(0, 60)),
|
||||
id=USER_SPEND_ALERTS_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
if os.getenv("PROMETHEUS_URL"):
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
|
|
@ -12481,11 +12511,21 @@ async def supported_openai_params(model: str):
|
|||
--header 'Authorization: Bearer sk-1234'
|
||||
```
|
||||
"""
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
|
||||
global llm_router
|
||||
try:
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
|
||||
resolved_models: Final = llm_router.resolved_litellm_models(model) if llm_router is not None else ()
|
||||
target_model: Final = resolved_models[0] if resolved_models else model
|
||||
declared_provider: Final = declared_authenticating_provider(target_model)
|
||||
litellm_model, custom_llm_provider = (
|
||||
(target_model.removeprefix(f"{declared_provider}/"), declared_provider)
|
||||
if declared_provider is not None
|
||||
else litellm.get_llm_provider(model=target_model)[:2]
|
||||
)
|
||||
return {
|
||||
"supported_openai_params": litellm.get_supported_openai_params(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
model=litellm_model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
}
|
||||
except Exception:
|
||||
|
|
@ -14962,17 +15002,25 @@ async def alerting_settings(
|
|||
alerting_args_dict = {}
|
||||
alerting_values = None
|
||||
|
||||
allowed_args: Final = {
|
||||
"slack_alerting": {"type": "Boolean"},
|
||||
"daily_report_frequency": {"type": "Integer"},
|
||||
"report_check_interval": {"type": "Integer"},
|
||||
"budget_alert_ttl": {"type": "Integer"},
|
||||
"outage_alert_ttl": {"type": "Integer"},
|
||||
"region_outage_alert_ttl": {"type": "Integer"},
|
||||
"minor_outage_alert_threshold": {"type": "Integer"},
|
||||
"major_outage_alert_threshold": {"type": "Integer"},
|
||||
"max_outage_alert_list_size": {"type": "Integer"},
|
||||
}
|
||||
allowed_args: Final = MappingProxyType(
|
||||
{
|
||||
"slack_alerting": "Boolean",
|
||||
"daily_report_frequency": "Integer",
|
||||
"report_check_interval": "Integer",
|
||||
"budget_alert_ttl": "Integer",
|
||||
"outage_alert_ttl": "Integer",
|
||||
"region_outage_alert_ttl": "Integer",
|
||||
"minor_outage_alert_threshold": "Integer",
|
||||
"major_outage_alert_threshold": "Integer",
|
||||
"max_outage_alert_list_size": "Integer",
|
||||
"daily_spend_per_user_threshold": "Float",
|
||||
"monthly_spend_per_user_threshold": "Float",
|
||||
"spend_anomaly_multiplier": "Float",
|
||||
"spend_anomaly_baseline_days": "Integer",
|
||||
"spend_anomaly_min_spend": "Float",
|
||||
"user_spend_check_interval": "Integer",
|
||||
}
|
||||
)
|
||||
|
||||
_slack_alerting: Final[SlackAlerting] = proxy_logging_obj.slack_alerting_instance
|
||||
_slack_alerting_args_dict: Final = _slack_alerting.alerting_args.model_dump()
|
||||
|
|
@ -14987,7 +15035,7 @@ async def alerting_settings(
|
|||
|
||||
_response_obj = ConfigList(
|
||||
field_name="slack_alerting",
|
||||
field_type=allowed_args["slack_alerting"]["type"],
|
||||
field_type=allowed_args["slack_alerting"],
|
||||
field_description="Enable slack alerting for monitoring proxy in production: llm outages, budgets, spend tracking failures.",
|
||||
field_value=is_slack_enabled,
|
||||
stored_in_db=True if alerting_values is not None else False,
|
||||
|
|
@ -15006,7 +15054,7 @@ async def alerting_settings(
|
|||
|
||||
_response_obj = ConfigList(
|
||||
field_name=field_name,
|
||||
field_type=allowed_args[field_name]["type"],
|
||||
field_type=allowed_args[field_name],
|
||||
field_description=field_info.description or "",
|
||||
field_value=_slack_alerting_args_dict.get(field_name, None),
|
||||
stored_in_db=_stored_in_db,
|
||||
|
|
@ -16434,6 +16482,16 @@ async def update_config_general_settings(
|
|||
detail={"error": f"Invalid type of field value={type(data.field_value)} passed in."},
|
||||
)
|
||||
|
||||
if data.field_name == "alerting_args":
|
||||
try:
|
||||
SlackAlertingArgs.model_validate(data.field_value)
|
||||
except ValidationError as e:
|
||||
errors: Final = "; ".join(f"{'.'.join(str(loc) for loc in err['loc'])}: {err['msg']}" for err in e.errors())
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Invalid alerting_args: {errors}"},
|
||||
)
|
||||
|
||||
## get general settings from db
|
||||
db_general_settings: Final = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
|
|
|
|||
|
|
@ -120,13 +120,15 @@ async def _apply_over_budget_reservation_policy(
|
|||
applied_entries: list[dict[str, float | str]],
|
||||
reservation_cost: float,
|
||||
current_spend: float,
|
||||
fail_closed_budget_enforcement: bool = False,
|
||||
) -> float:
|
||||
"""
|
||||
Decide what to do when a counter is over budget, and return the reservation
|
||||
cost to carry into the next counter. Three outcomes: an over-budget key that
|
||||
opted into throttling releases its own reservation (the rate limiter slows
|
||||
it) and keeps the cost; a partially-remaining budget resizes the reservation
|
||||
down to what is left; anything else hard-blocks by raising.
|
||||
down to what is left, unless strict enforcement is on, because the known
|
||||
estimate already does not fit; anything else hard-blocks by raising.
|
||||
"""
|
||||
if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token):
|
||||
await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost)
|
||||
|
|
@ -134,21 +136,36 @@ async def _apply_over_budget_reservation_policy(
|
|||
return reservation_cost
|
||||
|
||||
remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost)
|
||||
if remaining_before_reservation > 1e-12:
|
||||
await _resize_applied_reservation(
|
||||
entries=applied_entries,
|
||||
current_reserved_cost=reservation_cost,
|
||||
new_reserved_cost=remaining_before_reservation,
|
||||
if remaining_before_reservation <= 1e-12:
|
||||
_raise_counter_budget_exceeded(counter=counter, current_cost=current_spend)
|
||||
if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12:
|
||||
_raise_counter_budget_exceeded(
|
||||
counter=counter,
|
||||
current_cost=current_spend - reservation_cost,
|
||||
estimated_cost=reservation_cost,
|
||||
)
|
||||
return remaining_before_reservation
|
||||
await _resize_applied_reservation(
|
||||
entries=applied_entries,
|
||||
current_reserved_cost=reservation_cost,
|
||||
new_reserved_cost=remaining_before_reservation,
|
||||
)
|
||||
return remaining_before_reservation
|
||||
|
||||
|
||||
def _raise_counter_budget_exceeded(
|
||||
counter: _BudgetCounter,
|
||||
current_cost: float,
|
||||
estimated_cost: float | None = None,
|
||||
) -> NoReturn:
|
||||
estimate_detail: Final = "" if estimated_cost is None else f"Estimated request cost: {estimated_cost}, "
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=current_spend,
|
||||
current_cost=current_cost,
|
||||
max_budget=counter.max_budget,
|
||||
message=(
|
||||
"Budget has been exceeded! "
|
||||
f"{counter.entity_type}={counter.entity_id} "
|
||||
f"Current cost: {current_spend}, "
|
||||
f"Current cost: {current_cost}, "
|
||||
f"{estimate_detail}"
|
||||
f"Max budget: {counter.max_budget}"
|
||||
),
|
||||
entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type),
|
||||
|
|
@ -258,6 +275,7 @@ async def reserve_budget_for_request(
|
|||
applied_entries=applied_entries,
|
||||
reservation_cost=reservation_cost,
|
||||
current_spend=current_spend,
|
||||
fail_closed_budget_enforcement=fail_closed_budget_enforcement,
|
||||
)
|
||||
continue
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -143,7 +143,11 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
from litellm.router_utils.fallback_event_handlers import (
|
||||
AttemptedFallbackTargets,
|
||||
_check_non_standard_fallback_format,
|
||||
get_fallback_model_group,
|
||||
clear_pre_routing_selection,
|
||||
fallback_lookup_groups,
|
||||
get_fallback_model_group_for_lookup_groups,
|
||||
get_pre_routing_selection,
|
||||
record_pre_routing_selection,
|
||||
run_async_fallback,
|
||||
)
|
||||
from litellm.router_utils.get_retry_from_policy import (
|
||||
|
|
@ -4918,6 +4922,19 @@ class Router:
|
|||
)
|
||||
response = await response
|
||||
|
||||
if self._should_raise_anthropic_refusal_error(
|
||||
model=model,
|
||||
original_generic_function=original_generic_function,
|
||||
response=response,
|
||||
kwargs=kwargs,
|
||||
):
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
|
||||
safeguard_refusal_error,
|
||||
)
|
||||
|
||||
refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape
|
||||
raise safeguard_refusal_error(model=model, stop_details=refusal_details)
|
||||
|
||||
self.success_calls[model_name] += 1
|
||||
verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
|
||||
|
||||
|
|
@ -4964,6 +4981,11 @@ class Router:
|
|||
# fallback to the original reference for any non-picklable value.
|
||||
# The original_generic_function is preserved so the per-attempt
|
||||
# helper knows which underlying API to call on fallback.
|
||||
# The pre-routing hook stamps its tier selection into this bucket during the primary
|
||||
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
|
||||
# bucket, so the post-call carry-over below always has somewhere to read and write.
|
||||
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
|
||||
|
||||
fallback_kwargs: Final[dict[str, object]] = kwargs.copy()
|
||||
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
|
||||
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
|
||||
|
|
@ -4973,6 +4995,14 @@ class Router:
|
|||
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
|
||||
|
||||
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
|
||||
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
|
||||
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
|
||||
clear_pre_routing_selection(fallback_kwargs)
|
||||
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
|
||||
if live_pre_routing_selection is not None:
|
||||
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
|
||||
|
||||
if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator):
|
||||
return await self._aresponses_streaming_iterator(
|
||||
response=response,
|
||||
|
|
@ -5030,6 +5060,10 @@ class Router:
|
|||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
aclose_if_supported,
|
||||
parse_anthropic_error_event,
|
||||
parse_anthropic_refusal_stop_details,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
|
||||
safeguard_refusal_error,
|
||||
)
|
||||
|
||||
source_iterator: Final = response
|
||||
|
|
@ -5068,13 +5102,35 @@ class Router:
|
|||
continue
|
||||
if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
|
||||
has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit
|
||||
error_event = parse_anthropic_error_event(chunk)
|
||||
# A transport can split one SSE data line across byte chunks, so pre-content
|
||||
# detection parses the accumulated buffer plus the current chunk, never the
|
||||
# chunk alone; the buffer is already capped, which bounds this window too.
|
||||
parse_window = ( # rebind-ok: freshly computed each iteration, never carried over
|
||||
b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime
|
||||
if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime
|
||||
else chunk
|
||||
)
|
||||
error_event = parse_anthropic_error_event(parse_window)
|
||||
retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over
|
||||
not has_generated_content
|
||||
and error_event is not None
|
||||
and _is_retriable_anthropic_status(error_event[2])
|
||||
and not _anthropic_stream_error_is_gateway_verdict(chunk)
|
||||
)
|
||||
refusal_stop_details = ( # rebind-ok: freshly computed each iteration, never carried over
|
||||
parse_anthropic_refusal_stop_details(parse_window)
|
||||
if not has_generated_content and error_event is None
|
||||
else None
|
||||
)
|
||||
if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs):
|
||||
refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details)
|
||||
raise MidStreamFallbackError(
|
||||
message=refusal_error.message,
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
original_exception=refusal_error,
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
if not has_generated_content and not retriable_pending_error and error_event is None:
|
||||
buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk)
|
||||
continue
|
||||
|
|
@ -5186,8 +5242,13 @@ class Router:
|
|||
kwargs=initial_kwargs,
|
||||
metadata_variable_name="litellm_metadata",
|
||||
)
|
||||
# The content-policy dispatch branch matches on the trigger's own type, so a refusal's
|
||||
# MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted.
|
||||
fallback_trigger: Final[Exception] = (
|
||||
e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e
|
||||
)
|
||||
fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success
|
||||
e=e,
|
||||
e=fallback_trigger,
|
||||
disable_fallbacks=False,
|
||||
fallbacks=fallbacks,
|
||||
context_window_fallbacks=context_window_fallbacks,
|
||||
|
|
@ -5243,6 +5304,11 @@ class Router:
|
|||
# share, leaking primary-deployment metadata into the mid-stream
|
||||
# fallback request. safe_deep_copy avoids deep-copying the full
|
||||
# kwargs (which can hold non-deepcopyable logging handles/clients).
|
||||
# The pre-routing hook stamps its tier selection into this bucket during the primary
|
||||
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
|
||||
# bucket, so the post-call carry-over below always has somewhere to read and write.
|
||||
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
|
||||
|
||||
fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry
|
||||
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
|
||||
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
|
||||
|
|
@ -5252,6 +5318,14 @@ class Router:
|
|||
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
|
||||
|
||||
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
|
||||
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
|
||||
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
|
||||
clear_pre_routing_selection(fallback_kwargs)
|
||||
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
|
||||
if live_pre_routing_selection is not None:
|
||||
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
|
||||
|
||||
if kwargs.get("stream") and hasattr(response, "__aiter__"):
|
||||
return await self._aanthropic_messages_streaming_iterator(
|
||||
response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator
|
||||
|
|
@ -6807,6 +6881,9 @@ class Router:
|
|||
original_exception: Final = e
|
||||
fallback_model_group = None
|
||||
original_model_group: Final[str | None] = kwargs.get("model")
|
||||
# A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier
|
||||
# behind the router name, and fallbacks are configured per tier, not per router.
|
||||
lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group)
|
||||
fallback_failure_exception_str = ""
|
||||
|
||||
if disable_fallbacks is True or original_model_group is None:
|
||||
|
|
@ -6851,15 +6928,15 @@ class Router:
|
|||
]
|
||||
# Get external fallbacks — handle both standard and non-standard formats
|
||||
external_fallback_group: list | None = None
|
||||
if fallbacks is not None and model_group is not None:
|
||||
if fallbacks is not None and lookup_groups:
|
||||
if _check_non_standard_fallback_format(fallbacks=fallbacks):
|
||||
# Non-standard formats (e.g. ["claude-3-haiku"] or
|
||||
# [{"model": "...", "messages": [...]}]) are passed through directly
|
||||
external_fallback_group = fallbacks
|
||||
else:
|
||||
external_fallback_group, generic_idx = get_fallback_model_group(
|
||||
external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=fallbacks,
|
||||
model_group=cast(str, model_group),
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
if external_fallback_group is None and generic_idx is not None:
|
||||
external_fallback_group = fallbacks[generic_idx]["*"]
|
||||
|
|
@ -6917,9 +6994,9 @@ class Router:
|
|||
if isinstance(e, litellm.ContextWindowExceededError):
|
||||
if context_window_fallbacks is not None:
|
||||
context_window_fallback_model_group: Final[list[str] | None] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
self._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=context_window_fallbacks,
|
||||
model_group=model_group,
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
)
|
||||
if context_window_fallback_model_group is None:
|
||||
|
|
@ -6950,9 +7027,9 @@ class Router:
|
|||
elif isinstance(e, litellm.ContentPolicyViolationError):
|
||||
if content_policy_fallbacks is not None:
|
||||
content_policy_fallback_model_group: Final[list[str] | None] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
self._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
model_group=model_group,
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
)
|
||||
if content_policy_fallback_model_group is None:
|
||||
|
|
@ -6979,14 +7056,14 @@ class Router:
|
|||
|
||||
if litellm.expose_router_debug_in_errors:
|
||||
e.message += f"\n{error_message}"
|
||||
if fallbacks is not None and model_group is not None:
|
||||
if fallbacks is not None and lookup_groups:
|
||||
verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks))
|
||||
(
|
||||
fallback_model_group,
|
||||
generic_fallback_idx,
|
||||
) = get_fallback_model_group(
|
||||
) = get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}]
|
||||
model_group=cast(str, model_group),
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
## if none, check for generic fallback
|
||||
if fallback_model_group is None and generic_fallback_idx is not None:
|
||||
|
|
@ -6995,12 +7072,12 @@ class Router:
|
|||
if fallback_model_group is None:
|
||||
masked_fallbacks: Final = mask_sensitive_structure(fallbacks)
|
||||
verbose_router_logger.info(
|
||||
"No fallback model group found for original model_group=%s. Fallbacks=%s",
|
||||
model_group,
|
||||
"No fallback model group found for lookup_groups=%s. Fallbacks=%s",
|
||||
" -> ".join(lookup_groups),
|
||||
masked_fallbacks,
|
||||
)
|
||||
if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors:
|
||||
original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={masked_fallbacks}"
|
||||
original_exception.message += f"No fallback model group found for lookup_groups={' -> '.join(lookup_groups)}. Fallbacks={masked_fallbacks}"
|
||||
raise original_exception
|
||||
|
||||
input_kwargs.update(
|
||||
|
|
@ -7046,6 +7123,7 @@ class Router:
|
|||
If it fails after num_retries, fall back to another model group
|
||||
"""
|
||||
model_group: Final[str | None] = kwargs.get("model")
|
||||
clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary
|
||||
if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets):
|
||||
_fallback_metadata_key: Final = _get_router_metadata_variable_name(
|
||||
function_name=getattr(kwargs.get("original_function"), "__name__", None)
|
||||
|
|
@ -7471,6 +7549,24 @@ class Router:
|
|||
break
|
||||
return fallback_model_group
|
||||
|
||||
def _get_fallback_model_group_for_lookup_groups(
|
||||
self,
|
||||
fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the sibling resolver's contract
|
||||
lookup_groups: tuple[str, ...],
|
||||
) -> list[str] | None: # mutable-ok: mirrors the sibling resolver's contract
|
||||
"""First lookup group whose exact-key chain resolves (tier first, then requested group)."""
|
||||
return next(
|
||||
(
|
||||
resolved
|
||||
for resolved in (
|
||||
self._get_fallback_model_group_from_fallbacks(fallbacks=fallbacks, model_group=group)
|
||||
for group in lookup_groups
|
||||
)
|
||||
if resolved is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
def _get_first_default_fallback(self) -> str | None:
|
||||
"""
|
||||
Returns the first model from the default_fallbacks list, if it exists.
|
||||
|
|
@ -7886,6 +7982,31 @@ class Router:
|
|||
return True
|
||||
return False
|
||||
|
||||
def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
|
||||
"""
|
||||
Whether a content-policy fallback would resolve for this request, keyed the same way
|
||||
async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook
|
||||
selected wins over the requested group. Raising without this returning True would turn
|
||||
a deliverable response into an error the fallback chain cannot recover from.
|
||||
"""
|
||||
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
|
||||
if content_policy_fallbacks is not None:
|
||||
return (
|
||||
self._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
lookup_groups=fallback_lookup_groups(kwargs, model_group),
|
||||
)
|
||||
is not None
|
||||
)
|
||||
if self._has_default_fallbacks():
|
||||
return True
|
||||
verbose_router_logger.debug(
|
||||
"No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s",
|
||||
model_group,
|
||||
content_policy_fallbacks,
|
||||
)
|
||||
return False
|
||||
|
||||
def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool:
|
||||
"""
|
||||
Determines if a content policy error should be raised.
|
||||
|
|
@ -7898,27 +8019,26 @@ class Router:
|
|||
if response.choices[0].finish_reason != "content_filter":
|
||||
return False
|
||||
|
||||
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
|
||||
return self._has_content_policy_fallback(model, kwargs)
|
||||
|
||||
### ONLY RAISE ERROR IF CP FALLBACK AVAILABLE ###
|
||||
if content_policy_fallbacks is not None:
|
||||
fallback_model_group = None
|
||||
for item in content_policy_fallbacks: # [{"gpt-3.5-turbo": ["gpt-4"]}]
|
||||
if list(item.keys())[0] == model:
|
||||
fallback_model_group = item[model]
|
||||
break
|
||||
|
||||
if fallback_model_group is not None:
|
||||
return True
|
||||
elif self._has_default_fallbacks(): # default fallbacks set
|
||||
return True
|
||||
|
||||
verbose_router_logger.debug(
|
||||
"Content Policy Error occurred. No available fallbacks. Returning original response. model=%s, content_policy_fallbacks=%s",
|
||||
model,
|
||||
content_policy_fallbacks,
|
||||
def _should_raise_anthropic_refusal_error(
|
||||
self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, Any]
|
||||
) -> bool:
|
||||
"""
|
||||
The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard
|
||||
refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only
|
||||
when a content-policy fallback is configured; a plain refusal without stop_details, or
|
||||
any response with nothing configured, is returned to the client unchanged.
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
|
||||
get_safeguard_refusal_stop_details,
|
||||
)
|
||||
return False
|
||||
|
||||
if getattr(original_generic_function, "__name__", "") != "anthropic_messages":
|
||||
return False
|
||||
if get_safeguard_refusal_stop_details(response) is None:
|
||||
return False
|
||||
return self._has_content_policy_fallback(model, kwargs)
|
||||
|
||||
def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None):
|
||||
_all_deployments: list = []
|
||||
|
|
@ -12087,6 +12207,7 @@ class Router:
|
|||
if pre_routing_hook_response is not None:
|
||||
model = pre_routing_hook_response.model
|
||||
messages = pre_routing_hook_response.messages
|
||||
record_pre_routing_selection(request_kwargs, model)
|
||||
if pre_routing_hook_response.litellm_params:
|
||||
accepted_tier_params: Final = self._tier_params_the_target_accepts(
|
||||
model, pre_routing_hook_response.litellm_params, request_kwargs
|
||||
|
|
@ -12202,6 +12323,7 @@ class Router:
|
|||
if pre_routing_hook_response is not None:
|
||||
model = pre_routing_hook_response.model
|
||||
messages = pre_routing_hook_response.messages
|
||||
record_pre_routing_selection(request_kwargs, model)
|
||||
if pre_routing_hook_response.litellm_params:
|
||||
accepted_tier_params: Final = self._tier_params_the_target_accepts(
|
||||
model, pre_routing_hook_response.litellm_params, request_kwargs
|
||||
|
|
|
|||
|
|
@ -214,6 +214,91 @@ def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
PRE_ROUTING_SELECTED_MODEL_KEY: Final = "pre_routing_selected_model"
|
||||
_ROUTER_METADATA_BUCKETS: Final = ("metadata", "litellm_metadata")
|
||||
|
||||
|
||||
def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selected_model: str) -> None:
|
||||
"""
|
||||
Remember which model a pre-routing hook picked, so fallback lookup can key off it.
|
||||
|
||||
Fallback resolution runs on an outer kwargs dict that ``**kwargs`` already copied, so
|
||||
writing the model there is invisible by the time routing picks a tier. The metadata
|
||||
buckets are nested dicts shared by reference across those copies, which is how the
|
||||
router already carries values back up.
|
||||
|
||||
The write goes through the proxy-internal bucket resolver, never into both buckets:
|
||||
on /v1/messages the top-level ``metadata`` dict is the provider's own request field,
|
||||
so a blanket write would forward the tier stamp upstream.
|
||||
"""
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
|
||||
if request_kwargs is None:
|
||||
return
|
||||
bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs))
|
||||
if isinstance(bucket, dict):
|
||||
bucket[PRE_ROUTING_SELECTED_MODEL_KEY] = selected_model
|
||||
|
||||
|
||||
def clear_pre_routing_selection(request_kwargs: Mapping[str, object] | None) -> None:
|
||||
"""
|
||||
Drop any selection the router did not make itself on this hop.
|
||||
|
||||
The buckets carry whatever the caller sent, so an inbound value is the caller
|
||||
choosing a fallback chain rather than the router choosing a tier. A fallback hop
|
||||
also inherits the previous hop's selection, which would key its own failure off
|
||||
the tier that already failed. Clearing at the start of every hop leaves only a
|
||||
value the pre-routing hook wrote while routing that hop.
|
||||
"""
|
||||
if request_kwargs is None:
|
||||
return
|
||||
for bucket in (request_kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS):
|
||||
if isinstance(bucket, dict) and PRE_ROUTING_SELECTED_MODEL_KEY in bucket:
|
||||
del bucket[PRE_ROUTING_SELECTED_MODEL_KEY]
|
||||
|
||||
|
||||
def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None:
|
||||
"""The model a pre-routing hook selected for this request, if one did."""
|
||||
buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS)
|
||||
selections: Final = (bucket.get(PRE_ROUTING_SELECTED_MODEL_KEY) for bucket in buckets if isinstance(bucket, dict))
|
||||
return next((selected for selected in selections if isinstance(selected, str) and selected), None)
|
||||
|
||||
|
||||
def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]:
|
||||
"""
|
||||
Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins,
|
||||
and the requested group still resolves when no tier-keyed chain exists, so configs keyed
|
||||
on the router name (the documented contract) keep working behind auto-routers.
|
||||
"""
|
||||
ordered: Final = (get_pre_routing_selection(kwargs), model_group)
|
||||
return tuple(dict.fromkeys(group for group in ordered if group))
|
||||
|
||||
|
||||
def _resolved_a_specific_chain(
|
||||
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
) -> bool:
|
||||
resolved, generic_idx = result
|
||||
if resolved is None:
|
||||
return False
|
||||
return generic_idx is None or resolved is not fallbacks[generic_idx]["*"]
|
||||
|
||||
|
||||
def get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
lookup_groups: tuple[str, ...],
|
||||
) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
"""
|
||||
First lookup group with a specifically-keyed chain wins; the generic "*" chain applies
|
||||
only after every group missed, so a catch-all cannot shadow a later group's own chain.
|
||||
"""
|
||||
results: Final = tuple(get_fallback_model_group(fallbacks=fallbacks, model_group=group) for group in lookup_groups)
|
||||
specific: Final = next((result for result in results if _resolved_a_specific_chain(fallbacks, result)), None)
|
||||
if specific is not None:
|
||||
return specific
|
||||
return next((result for result in results if result[0] is not None), (None, None))
|
||||
|
||||
|
||||
def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]:
|
||||
"""
|
||||
Returns:
|
||||
|
|
@ -412,6 +497,7 @@ async def run_async_fallback(
|
|||
# LOGGING
|
||||
kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception)
|
||||
verbose_router_logger.info("Falling back to model_group = %s", mask_sensitive_structure(mg))
|
||||
kwargs.pop("_target_order", None) # rebind-ok: next hop must not inherit the previous order target
|
||||
if isinstance(mg, str):
|
||||
kwargs["model"] = mg
|
||||
elif isinstance(mg, dict):
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from re import Match
|
|||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
|
||||
|
||||
|
||||
class PatternUtils:
|
||||
|
|
@ -204,7 +204,7 @@ class PatternMatchRouter:
|
|||
|
||||
return litellm_deployment_litellm_model
|
||||
|
||||
def get_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict] | None:
|
||||
def get_pattern(self, model: str | None, custom_llm_provider: str | None = None) -> list[dict] | None:
|
||||
"""
|
||||
Check if a pattern exists for the given model and custom llm provider
|
||||
|
||||
|
|
@ -215,18 +215,17 @@ class PatternMatchRouter:
|
|||
Returns:
|
||||
bool: True if pattern exists, False otherwise
|
||||
"""
|
||||
if custom_llm_provider is None:
|
||||
try:
|
||||
(
|
||||
_,
|
||||
custom_llm_provider,
|
||||
_,
|
||||
_,
|
||||
) = get_llm_provider(model=model)
|
||||
except Exception:
|
||||
# get_llm_provider raises exception when provider is unknown
|
||||
pass
|
||||
return self.route(model) or self.route(f"{custom_llm_provider}/{model}")
|
||||
provider: Final = (
|
||||
custom_llm_provider or declared_authenticating_provider(model) or self._resolved_provider(model)
|
||||
)
|
||||
return self.route(model) or self.route(f"{provider}/{model}")
|
||||
|
||||
@staticmethod
|
||||
def _resolved_provider(model: str | None) -> str | None:
|
||||
try:
|
||||
return get_llm_provider(model=model)[1] if model else None
|
||||
except Exception: # noqa: BLE001 # get_llm_provider raises when the provider is unknown; the name then routes as-is
|
||||
return None
|
||||
|
||||
def get_deployments_by_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -427,6 +427,8 @@ class DeploymentAffinityCheck(CustomLogger):
|
|||
"""
|
||||
request_kwargs = request_kwargs or {}
|
||||
typed_healthy_deployments: Final = cast(list[dict], healthy_deployments)
|
||||
if request_kwargs.get("_target_order") is not None:
|
||||
return typed_healthy_deployments
|
||||
|
||||
(
|
||||
enable_user_key,
|
||||
|
|
|
|||
|
|
@ -58,6 +58,9 @@ class PromptCachingDeploymentCheck(CustomLogger):
|
|||
request_kwargs: dict | None = None,
|
||||
parent_otel_span: Span | None = None,
|
||||
) -> list[dict]:
|
||||
if request_kwargs is not None and request_kwargs.get("_target_order") is not None:
|
||||
return healthy_deployments
|
||||
|
||||
if messages is not None and is_prompt_caching_valid_prompt(
|
||||
messages=messages,
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -91,6 +91,40 @@ class SlackAlertingArgs(LiteLLMPydanticObjectBase):
|
|||
default=False,
|
||||
description="If true, the alerting payload will be printed to the console.",
|
||||
)
|
||||
daily_spend_per_user_threshold: float | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Alert when a user's spend for the current day (UTC) crosses this USD amount. Off by default.",
|
||||
)
|
||||
monthly_spend_per_user_threshold: float | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Alert when a user's spend for the current calendar month (UTC) crosses this USD amount. Off by default.",
|
||||
)
|
||||
spend_anomaly_multiplier: float = Field(
|
||||
default=3.0,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Flag a user's spend as anomalous when today's spend exceeds this multiple of their trailing daily average.",
|
||||
)
|
||||
spend_anomaly_baseline_days: int = Field(
|
||||
default=7,
|
||||
ge=1,
|
||||
description="Number of trailing days used to compute a user's daily average spend for anomaly detection.",
|
||||
)
|
||||
spend_anomaly_min_spend: float = Field(
|
||||
default=10.0,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Minimum spend (USD) a user must reach today before an anomaly alert can fire. Reduces false positives.",
|
||||
)
|
||||
user_spend_check_interval: int = Field(
|
||||
default=3600,
|
||||
ge=60,
|
||||
description="How often (in seconds) to check per-user spend thresholds and anomalies. Default is hourly.",
|
||||
)
|
||||
|
||||
|
||||
class DeploymentMetrics(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -138,6 +172,8 @@ class AlertType(str, Enum):
|
|||
budget_alerts = "budget_alerts"
|
||||
spend_reports = "spend_reports"
|
||||
failed_tracking_spend = "failed_tracking_spend"
|
||||
user_spend_thresholds = "user_spend_thresholds"
|
||||
user_spend_anomalies = "user_spend_anomalies"
|
||||
|
||||
# Database alerts
|
||||
db_exceptions = "db_exceptions"
|
||||
|
|
@ -182,6 +218,7 @@ DEFAULT_ALERT_TYPES: Final[list[AlertType]] = [
|
|||
AlertType.budget_alerts,
|
||||
AlertType.spend_reports,
|
||||
AlertType.failed_tracking_spend,
|
||||
AlertType.user_spend_thresholds,
|
||||
# Database alerts
|
||||
AlertType.db_exceptions,
|
||||
# Report alerts
|
||||
|
|
|
|||
|
|
@ -78,6 +78,16 @@ class AnthropicUsage(TypedDict, total=False):
|
|||
server_tool_use: NotRequired[ReadOnly[ServerToolUsage]]
|
||||
|
||||
|
||||
class AnthropicStopDetails(TypedDict, total=False):
|
||||
"""
|
||||
Safeguard verdict accompanying a `stop_reason: "refusal"` response:
|
||||
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback
|
||||
"""
|
||||
|
||||
category: ReadOnly[str | None]
|
||||
explanation: ReadOnly[str | None]
|
||||
|
||||
|
||||
class AnthropicMessagesResponse(TypedDict, total=False):
|
||||
"""
|
||||
Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages
|
||||
|
|
@ -90,7 +100,8 @@ class AnthropicMessagesResponse(TypedDict, total=False):
|
|||
id: str
|
||||
model: str | None # This represents the Model type from Anthropic
|
||||
role: Literal["assistant"] | None
|
||||
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None
|
||||
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] | None
|
||||
stop_details: NotRequired[ReadOnly[AnthropicStopDetails | None]]
|
||||
stop_sequence: str | None
|
||||
type: Literal["message"] | None
|
||||
usage: AnthropicUsage | None
|
||||
|
|
|
|||
|
|
@ -4876,11 +4876,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None:
|
|||
|
||||
def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list:
|
||||
if target_order is not None:
|
||||
filtered: Final = [d for d in healthy_deployments if _get_deployment_order(d) == target_order]
|
||||
if filtered:
|
||||
return filtered
|
||||
# target_order doesn't match any deployment (e.g., external fallback model) — return all
|
||||
return healthy_deployments
|
||||
return [d for d in healthy_deployments if _get_deployment_order(d) == target_order]
|
||||
|
||||
# Default: pick min order group
|
||||
_valid_orders: Final[list[int]] = [
|
||||
|
|
|
|||
|
|
@ -81,6 +81,7 @@ ignored_function_names = [
|
|||
"_merge_tools_from_deployment", # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name)
|
||||
"_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name)
|
||||
"has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call
|
||||
"_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name)
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,16 @@
|
|||
|
||||
The e2e suite ships results to Loki/Grafana from a standard pytest JUnit report
|
||||
(`--junitxml=e2e-report.xml`), not a bespoke log line. JUnit already records
|
||||
outcome, duration, and node id for every `<testcase>`; the only signals it cannot
|
||||
derive on its own are the normalized suite package and the coverage-registry cell
|
||||
ids a test covers. Those ride along as JUnit `<property>` entries via each item's
|
||||
`user_properties`, attached in `conftest.py::pytest_collection_modifyitems`.
|
||||
outcome, duration, and node id for every `<testcase>`; the signals it cannot
|
||||
derive on its own are the normalized suite package, the coverage-registry cell
|
||||
ids a test covers, and where the test's source lives. Those ride along as JUnit
|
||||
`<property>` entries via each item's `user_properties`, attached in
|
||||
`conftest.py::pytest_collection_modifyitems`.
|
||||
|
||||
`source` is a property rather than the `file=` / `line=` attributes pytest used
|
||||
to write, because the `xunit2` family this suite runs on drops those, and
|
||||
switching families would change the XML for every consumer of it -- the
|
||||
Buildkite Test Engine upload and the Loki pipeline included.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -14,22 +20,61 @@ from collections.abc import Iterable
|
|||
|
||||
import pytest
|
||||
|
||||
# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing
|
||||
# at runtime names this suite's place in the repo. test_junit_properties.py
|
||||
# fails from a checkout if it moves.
|
||||
SUITE_ROOT = "tests/e2e"
|
||||
|
||||
|
||||
def suite_parts(path_part: str) -> tuple[str, ...]:
|
||||
"""Path components of a suite file relative to tests/e2e, however it ran.
|
||||
|
||||
Pytest paths are rootdir-relative, and rootdir moves with the invocation: a
|
||||
repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd run (the
|
||||
runner image) gives `logging/test_x.py`. Both collapse to the same tuple.
|
||||
"""
|
||||
raw = tuple(p for p in path_part.replace("\\", "/").split("/") if p and p != ".")
|
||||
return raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw
|
||||
|
||||
|
||||
def package_from_nodeid(nodeid: str) -> str:
|
||||
"""Top-level suite package under tests/e2e/, or 'root' for top-level files.
|
||||
|
||||
Pytest nodeids are relative to the invocation cwd. Repo-root runs look like
|
||||
`tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the
|
||||
`tests/e2e` prefix so package is the suite dir either way.
|
||||
"""
|
||||
path_part = nodeid.split("::", 1)[0].replace("\\", "/")
|
||||
raw = tuple(p for p in path_part.split("/") if p and p != ".")
|
||||
parts = raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw
|
||||
"""Top-level suite package under tests/e2e/, or 'root' for top-level files."""
|
||||
parts = suite_parts(nodeid.split("::", 1)[0])
|
||||
if len(parts) <= 1:
|
||||
return "root"
|
||||
return parts[0]
|
||||
|
||||
|
||||
def source_from_location(path: str, lineno: int | None) -> str:
|
||||
"""Repo-relative `path:line` for a test, or '' when nothing is linkable.
|
||||
|
||||
`pytest.Item.location` gives a rootdir-relative path and a ZERO-based line.
|
||||
The path is re-rooted at SUITE_ROOT so consumers need not know how pytest was
|
||||
started, and the line is emitted ONE-based to match editors, tracebacks and
|
||||
code hosts. A decorated test anchors at its first decorator, which is where
|
||||
pytest reports it.
|
||||
|
||||
Empty rather than a guess for anything unlinkable: no line, a path reaching
|
||||
upward, or a path carrying a colon, which is both how an absolute Windows
|
||||
path arrives and a character `path:line` has no way to represent.
|
||||
"""
|
||||
if lineno is None:
|
||||
return ""
|
||||
normalized = path.replace("\\", "/")
|
||||
if normalized.startswith("/") or ":" in normalized or ".." in normalized.split("/"):
|
||||
return ""
|
||||
parts = suite_parts(normalized)
|
||||
if not parts:
|
||||
return ""
|
||||
return f"{'/'.join((SUITE_ROOT, *parts))}:{lineno + 1}"
|
||||
|
||||
|
||||
def source_from_item(item: pytest.Item) -> str:
|
||||
"""Read the repo-relative `path:line` off a pytest Item's reported location."""
|
||||
path, lineno, _ = item.location
|
||||
return source_from_location(path, lineno)
|
||||
|
||||
|
||||
def dedupe_covers(marker_args: Iterable[tuple[object, ...]]) -> tuple[str, ...]:
|
||||
"""Flatten @pytest.mark.covers arg lists into unique, order-preserving cell
|
||||
ids, dropping anything that is not a non-empty string."""
|
||||
|
|
@ -43,10 +88,12 @@ def covers_from_item(item: pytest.Item) -> tuple[str, ...]:
|
|||
|
||||
def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]:
|
||||
"""The custom signals a standard reporter cannot derive: the normalized suite
|
||||
package and the comma-joined coverage-registry cell ids this test covers."""
|
||||
package, the comma-joined coverage-registry cell ids this test covers, and the
|
||||
repo-relative `path:line` its source sits at."""
|
||||
return (
|
||||
("package", package_from_nodeid(item.nodeid)),
|
||||
("covers", ",".join(covers_from_item(item))),
|
||||
("source", source_from_item(item)),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -806,6 +806,7 @@ class LiteLLMParamsBody(BaseModel):
|
|||
mock_response: str | None = None
|
||||
timeout: float | None = None
|
||||
tpm: int | None = None
|
||||
weight: int | None = None
|
||||
|
||||
|
||||
ModelMode = Literal["batch", "realtime", "image_generation"]
|
||||
|
|
@ -820,6 +821,7 @@ class ModelInfoBody(BaseModel):
|
|||
mode: ModelMode | None = None
|
||||
access_groups: list[str] | None = None
|
||||
team_id: str | None = None
|
||||
allowed_fails_policy: dict[str, int] | None = None
|
||||
|
||||
|
||||
class ModelNewBody(BaseModel):
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ from models import (
|
|||
ChatMessage,
|
||||
ChatResponse,
|
||||
LiteLLMParamsBody,
|
||||
ModelInfoBody,
|
||||
ModelNewBody,
|
||||
ReliabilityChatBody,
|
||||
RouterSettingsOverride,
|
||||
)
|
||||
|
|
@ -26,6 +28,18 @@ from models import (
|
|||
REAL_MODEL = "openai/gpt-5.5"
|
||||
REAL_KEY = "os.environ/OPENAI_API_KEY"
|
||||
|
||||
# The smallest-context chat model OpenAI still serves (16385 tokens). A prompt
|
||||
# past that limit comes back as a real `context_length_exceeded` 400, which is
|
||||
# what litellm maps to ContextWindowExceededError.
|
||||
SMALL_CONTEXT_MODEL = "openai/gpt-3.5-turbo"
|
||||
SMALL_CONTEXT_LIMIT_TOKENS = 16385
|
||||
|
||||
|
||||
def oversized_prompt(marker: str) -> str:
|
||||
"""A prompt comfortably past SMALL_CONTEXT_MODEL's context limit, so the
|
||||
provider refuses it on length rather than answering a truncated version."""
|
||||
return f"{marker} " + ("token " * (SMALL_CONTEXT_LIMIT_TOKENS + 4000))
|
||||
|
||||
|
||||
def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""Register a deployment pointing at an unreachable base, so every call to it
|
||||
|
|
@ -40,6 +54,38 @@ def create_timeout_deployment(proxy: ProxyClient, name: str) -> str:
|
|||
return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001))
|
||||
|
||||
|
||||
def create_small_context_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""Register a deployment on the smallest-context model OpenAI still serves, so an
|
||||
oversized prompt earns a real context-window refusal from the provider."""
|
||||
return proxy.create_model(name, LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY))
|
||||
|
||||
|
||||
def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""The always-picked half of a retry pair: a 1ms deadline the backend always
|
||||
exceeds, all of the model group's shuffle weight, and a cooldown policy that
|
||||
benches it on its first Timeout so the retry cannot land on it again."""
|
||||
return proxy.register_model(
|
||||
ModelNewBody(
|
||||
model_name=name,
|
||||
litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001, weight=1),
|
||||
model_info=ModelInfoBody(allowed_fails_policy={"TimeoutErrorAllowedFails": 0}),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""The other half of a retry pair: healthy, but weight 0, so the weighted shuffle
|
||||
never opens on it. It is reachable only once its sibling is benched and the
|
||||
weighted pick falls through to a uniform one over what is left."""
|
||||
return proxy.register_model(
|
||||
ModelNewBody(
|
||||
model_name=name,
|
||||
litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, weight=0),
|
||||
model_info=ModelInfoBody(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def chat_override(
|
||||
proxy: ProxyClient,
|
||||
key: str,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,10 @@ in the x-litellm-attempted-fallbacks header. Empty content is accepted only when
|
|||
`finish_reason == "length"` and the response billed completion tokens, since
|
||||
gpt-5.5 counts reasoning against max_tokens and can consume the whole budget
|
||||
before emitting any text; a fallback that produced nothing at all still fails.
|
||||
|
||||
The context-window case is a different reroute from a plain failure: the provider
|
||||
refuses the prompt on length, and `context_window_fallbacks` is the setting that
|
||||
reroutes it, not `fallbacks`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -25,8 +29,10 @@ from reliability_support import (
|
|||
completion_tokens_of,
|
||||
content_of,
|
||||
create_bad_base_deployment,
|
||||
create_small_context_deployment,
|
||||
create_timeout_deployment,
|
||||
finish_reason_of,
|
||||
oversized_prompt,
|
||||
reasoning_tokens_of,
|
||||
)
|
||||
|
||||
|
|
@ -82,3 +88,17 @@ class TestReliabilityFallbacks:
|
|||
override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]),
|
||||
)
|
||||
_assert_served_by_fallback(resp)
|
||||
|
||||
@pytest.mark.covers("reliability.fallback.context_window.routes_to_fallback")
|
||||
def test_context_window_routes_to_fallback(
|
||||
self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
primary = f"reliability-ctxfail-{unique_marker()}"
|
||||
model_id = create_small_context_deployment(client.proxy, primary)
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
|
||||
resp = chat_override(
|
||||
client.proxy, scoped_key, primary, oversized_prompt(unique_marker()),
|
||||
override=RouterSettingsOverride(context_window_fallbacks=[{primary: ["gpt-5.5"]}]),
|
||||
)
|
||||
_assert_served_by_fallback(resp)
|
||||
|
|
|
|||
73
tests/e2e/router/test_reliability_retries_e2e.py
Normal file
73
tests/e2e/router/test_reliability_retries_e2e.py
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
"""Live e2e: a request that fails on its first deployment is retried inside its own
|
||||
model group and still comes back a completion.
|
||||
|
||||
The model group is a pair: an always-timing-out deployment that holds all of the
|
||||
group's shuffle weight, and a healthy backup at weight 0. The weighted pick always
|
||||
opens on the timing-out one, its first Timeout benches it (an
|
||||
`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls
|
||||
through to the only deployment left. So the customer sees a completion and the
|
||||
proxy reports that it took a retry to get there, with no random first pick in the
|
||||
middle of it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from complexity_router_client import ComplexityRouterClient
|
||||
from e2e_config import unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
from models import RouterSettingsOverride
|
||||
from reliability_support import (
|
||||
chat_override,
|
||||
completion_tokens_of,
|
||||
content_of,
|
||||
create_always_timing_out_deployment,
|
||||
create_zero_weight_backup_deployment,
|
||||
finish_reason_of,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
||||
class TestReliabilityRetries:
|
||||
@pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries")
|
||||
def test_timeout_on_first_deployment_succeeds_on_retry(
|
||||
self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
group = f"reliability-retry-{unique_marker()}"
|
||||
timing_out = create_always_timing_out_deployment(client.proxy, group)
|
||||
resources.defer(lambda: client.proxy.delete_model(timing_out))
|
||||
backup = create_zero_weight_backup_deployment(client.proxy, group)
|
||||
resources.defer(lambda: client.proxy.delete_model(backup))
|
||||
|
||||
resp = chat_override(
|
||||
client.proxy,
|
||||
scoped_key,
|
||||
group,
|
||||
f"say hi {unique_marker()}",
|
||||
override=RouterSettingsOverride(num_retries=2),
|
||||
)
|
||||
|
||||
assert resp.status_code == 200, (
|
||||
f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}"
|
||||
)
|
||||
|
||||
attempted = resp.headers.get("x-litellm-attempted-retries")
|
||||
assert attempted is not None, "response is missing the x-litellm-attempted-retries header"
|
||||
assert int(attempted) >= 1, (
|
||||
f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never "
|
||||
"opened on the timing-out deployment, so this proves nothing about retries"
|
||||
)
|
||||
|
||||
content = content_of(resp)
|
||||
finish_reason = finish_reason_of(resp)
|
||||
completion_tokens = completion_tokens_of(resp) or 0
|
||||
assert isinstance(content, str), (
|
||||
f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})"
|
||||
)
|
||||
assert content or (finish_reason == "length" and completion_tokens > 0), (
|
||||
f"the retry returned empty content with finish_reason={finish_reason!r}, "
|
||||
f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget "
|
||||
f"was spent on non-visible reasoning (body={resp.body[:300]})"
|
||||
)
|
||||
146
tests/e2e/test_junit_properties.py
Normal file
146
tests/e2e/test_junit_properties.py
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
"""Harness coverage for the custom JUnit properties.
|
||||
|
||||
No proxy and no ``e2e`` marker. Pins the two normalizations that have to agree
|
||||
about where a suite file lives -- ``package_from_nodeid`` (strip the suite root)
|
||||
and ``source_from_location`` (re-root at it) -- across both ways the suite is
|
||||
launched, plus the one-based line offset and the refusal to emit a path that
|
||||
escapes the suite. The consumers of these properties are the Loki/Grafana
|
||||
rollups and, for ``source``, the status page's per-test links to GitHub.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from junit_properties import (
|
||||
SUITE_ROOT,
|
||||
attach_result_properties,
|
||||
dedupe_covers,
|
||||
package_from_nodeid,
|
||||
result_properties,
|
||||
source_from_location,
|
||||
suite_parts,
|
||||
)
|
||||
|
||||
|
||||
class FakeMarker:
|
||||
def __init__(self, name: str, *args: object) -> None:
|
||||
self.name = name
|
||||
self.args = args
|
||||
|
||||
|
||||
class FakeItem:
|
||||
"""The three attributes junit_properties reads off a pytest Item."""
|
||||
|
||||
def __init__(
|
||||
self, nodeid: str, location: tuple[str, int | None, str], markers: tuple[FakeMarker, ...] = ()
|
||||
) -> None:
|
||||
self.nodeid = nodeid
|
||||
self.location = location
|
||||
self.user_properties: list[tuple[str, str]] = []
|
||||
self._markers = markers
|
||||
|
||||
def iter_markers(self, name: str):
|
||||
return (marker for marker in self._markers if marker.name == name)
|
||||
|
||||
|
||||
def repo_root() -> Path | None:
|
||||
"""The litellm checkout above this file, or None when there isn't one."""
|
||||
return next((p for p in Path(__file__).resolve().parents if (p / ".git").exists()), None)
|
||||
|
||||
|
||||
class TestSuiteParts:
|
||||
@pytest.mark.parametrize(
|
||||
"path",
|
||||
["logging/test_x.py", "tests/e2e/logging/test_x.py", "./logging/test_x.py", "tests\\e2e\\logging\\test_x.py"],
|
||||
)
|
||||
def test_both_invocation_shapes_collapse_to_the_same_components(self, path: str) -> None:
|
||||
"""A repo-root run and a suite-cwd run report the same file differently;
|
||||
every downstream signal has to see one spelling."""
|
||||
assert suite_parts(path) == ("logging", "test_x.py")
|
||||
|
||||
def test_top_level_suite_file_keeps_its_single_component(self) -> None:
|
||||
assert suite_parts("tests/e2e/test_fixture_mode.py") == ("test_fixture_mode.py",)
|
||||
|
||||
|
||||
class TestPackageFromNodeid:
|
||||
@pytest.mark.parametrize(
|
||||
("nodeid", "expected"),
|
||||
[
|
||||
("logging/test_x.py::TestFoo::test_bar", "logging"),
|
||||
("tests/e2e/logging/test_x.py::TestFoo::test_bar", "logging"),
|
||||
("quota_management/spend_tracking/test_x.py::test_bar", "quota_management"),
|
||||
("test_fixture_mode.py::TestParseFixtureMode::test_known_values_normalize", "root"),
|
||||
("tests/e2e/test_fixture_mode.py::test_bar", "root"),
|
||||
],
|
||||
)
|
||||
def test_package_is_the_first_dir_under_the_suite_root(self, nodeid: str, expected: str) -> None:
|
||||
assert package_from_nodeid(nodeid) == expected
|
||||
|
||||
|
||||
class TestSourceFromLocation:
|
||||
@pytest.mark.parametrize("path", ["a2a/test_a2a_agent_e2e.py", "tests/e2e/a2a/test_a2a_agent_e2e.py"])
|
||||
def test_path_is_repo_relative_however_pytest_was_started(self, path: str) -> None:
|
||||
assert source_from_location(path, 40) == "tests/e2e/a2a/test_a2a_agent_e2e.py:41"
|
||||
|
||||
def test_line_is_emitted_one_based(self) -> None:
|
||||
"""pytest.Item.location counts from 0; editors, tracebacks and GitHub's
|
||||
#L anchor all count from 1, and an off-by-one lands on the decorator."""
|
||||
assert source_from_location("a2a/test_x.py", 0) == "tests/e2e/a2a/test_x.py:1"
|
||||
|
||||
def test_top_level_suite_file_sits_directly_under_the_suite_root(self) -> None:
|
||||
assert source_from_location("test_fixture_mode.py", 39) == "tests/e2e/test_fixture_mode.py:40"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "lineno"),
|
||||
[
|
||||
("a2a/test_x.py", None),
|
||||
("/app/e2e/a2a/test_x.py", 40),
|
||||
("C:\\app\\e2e\\a2a\\test_x.py", 40),
|
||||
("../conftest.py", 40),
|
||||
("", 40),
|
||||
],
|
||||
)
|
||||
def test_nothing_linkable_yields_empty_rather_than_a_guess(self, path: str, lineno: int | None) -> None:
|
||||
"""A colon is rejected on two counts: it is how a Windows absolute path
|
||||
arrives, and `path:line` cannot represent one in the path half."""
|
||||
assert source_from_location(path, lineno) == ""
|
||||
|
||||
|
||||
class TestResultProperties:
|
||||
def test_every_test_carries_package_covers_and_source(self) -> None:
|
||||
item = FakeItem(
|
||||
"logging/test_x.py::TestFoo::test_bar",
|
||||
("logging/test_x.py", 40, "TestFoo.test_bar"),
|
||||
(FakeMarker("covers", "LOG-1", "LOG-2"),),
|
||||
)
|
||||
assert result_properties(item) == (
|
||||
("package", "logging"),
|
||||
("covers", "LOG-1,LOG-2"),
|
||||
("source", "tests/e2e/logging/test_x.py:41"),
|
||||
)
|
||||
|
||||
def test_attach_is_idempotent(self) -> None:
|
||||
"""Collection can run the hook more than once; a second pass must not
|
||||
double the <property> entries in the report."""
|
||||
item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar"))
|
||||
attach_result_properties(item)
|
||||
attach_result_properties(item)
|
||||
assert [name for name, _ in item.user_properties] == ["package", "covers", "source"]
|
||||
|
||||
|
||||
class TestSuiteRoot:
|
||||
def test_suite_root_names_this_file_s_real_home(self) -> None:
|
||||
"""SUITE_ROOT is hardcoded because the runner image has no repo to read it
|
||||
from. Where there IS a checkout, prove the constant still points at us --
|
||||
otherwise a moved tests/e2e/ ships links that 404."""
|
||||
root = repo_root()
|
||||
if root is None:
|
||||
pytest.skip("no checkout above this file (the runner image copies tests/e2e/ to /app/e2e)")
|
||||
assert (root / SUITE_ROOT / Path(__file__).name).resolve() == Path(__file__).resolve()
|
||||
|
||||
|
||||
class TestDedupeCovers:
|
||||
def test_ids_are_unique_order_preserving_and_non_empty_strings(self) -> None:
|
||||
assert dedupe_covers([("A", "B"), ("B", ""), ("C", 7)]) == ("A", "B", "C")
|
||||
|
|
@ -11,7 +11,9 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY
|
|||
import litellm.experimental_mcp_client.client as mcp_client_module
|
||||
from litellm.experimental_mcp_client.client import MCPClient
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult
|
||||
from mcp.types import CallToolResult as MCPCallToolResult
|
||||
from mcp.types import ListToolsResult, PaginatedRequestParams
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
|
||||
def test_mcp_client_uses_configurable_default_timeout():
|
||||
|
|
@ -185,6 +187,80 @@ class TestMCPClientUnitTests:
|
|||
mock_session_instance.initialize.assert_called_once()
|
||||
mock_session_instance.list_tools.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
|
||||
@patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
|
||||
async def test_list_tools_follows_next_cursor_until_exhausted(
|
||||
self,
|
||||
mock_session_class,
|
||||
mock_transport,
|
||||
):
|
||||
"""Test listing tools follows MCP pagination cursors until exhausted."""
|
||||
mock_transport_ctx = AsyncMock()
|
||||
mock_transport.return_value = mock_transport_ctx
|
||||
mock_transport_instance = MagicMock()
|
||||
mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance)
|
||||
|
||||
mock_session_ctx = AsyncMock()
|
||||
mock_session_class.return_value = mock_session_ctx
|
||||
mock_session_instance = AsyncMock()
|
||||
mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance)
|
||||
|
||||
first_page_tools = [
|
||||
MCPTool(name=f"tool_{idx}", description=f"Tool {idx}", inputSchema={}) for idx in range(100)
|
||||
]
|
||||
second_page_tool = MCPTool(
|
||||
name="tool_100",
|
||||
description="Tool 100",
|
||||
inputSchema={},
|
||||
)
|
||||
mock_session_instance.list_tools.side_effect = [
|
||||
ListToolsResult(tools=first_page_tools, nextCursor="page-2"),
|
||||
ListToolsResult(tools=[second_page_tool]),
|
||||
]
|
||||
|
||||
client = MCPClient("http://example.com")
|
||||
result = await client.list_tools()
|
||||
|
||||
assert result == [*first_page_tools, second_page_tool]
|
||||
assert mock_session_instance.list_tools.call_count == 2
|
||||
second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"]
|
||||
assert isinstance(second_call_params, PaginatedRequestParams)
|
||||
assert second_call_params.cursor == "page-2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
|
||||
@patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py
|
||||
async def test_list_tools_swallows_mid_walk_error_without_raise_on_error(
|
||||
self,
|
||||
mock_session_class,
|
||||
mock_transport,
|
||||
):
|
||||
"""Test a mid-walk failure returns [] when raise_on_error is False."""
|
||||
mock_transport_ctx = AsyncMock()
|
||||
mock_transport.return_value = mock_transport_ctx
|
||||
mock_transport_instance = MagicMock()
|
||||
mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance)
|
||||
|
||||
mock_session_ctx = AsyncMock()
|
||||
mock_session_class.return_value = mock_session_ctx
|
||||
mock_session_instance = AsyncMock()
|
||||
mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance)
|
||||
|
||||
mock_session_instance.list_tools.side_effect = [
|
||||
ListToolsResult(
|
||||
tools=[MCPTool(name="tool_0", description="Tool 0", inputSchema={})],
|
||||
nextCursor="page-2",
|
||||
),
|
||||
RuntimeError("transient upstream failure"),
|
||||
]
|
||||
|
||||
client = MCPClient("http://example.com")
|
||||
result = await client.list_tools()
|
||||
|
||||
assert result == []
|
||||
assert mock_session_instance.list_tools.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch.object(mcp_client_module, "streamable_http_client")
|
||||
@patch.object(mcp_client_module, "ClientSession")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,402 @@
|
|||
"""
|
||||
Unit tests for safeguard-refusal fallback on the /v1/messages router surface.
|
||||
|
||||
An Anthropic safeguard refusal is an HTTP 200 whose body carries
|
||||
stop_reason "refusal" plus a stop_details object; the router converts it
|
||||
into a ContentPolicyViolationError so the content-policy fallback chain
|
||||
runs, but only when a matching fallback is configured. A plain refusal
|
||||
without stop_details, or any refusal with nothing configured, must reach
|
||||
the client byte-identical.
|
||||
|
||||
The upstream is faked at the HTTP boundary by intercepting the third-party
|
||||
transport (httpx.AsyncClient.send), so requests run litellm's real
|
||||
transformation, allowlist, and streaming pipeline end to end.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, AsyncIterator
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm import Router
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
PRE_ROUTING_SELECTED_MODEL_KEY,
|
||||
record_pre_routing_selection,
|
||||
)
|
||||
|
||||
REFUSAL_RESPONSE: dict[str, Any] = {
|
||||
"id": "msg_refusal",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-fable-5",
|
||||
"content": [],
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": {"category": "cyber", "explanation": "flagged"},
|
||||
"usage": {"input_tokens": 25, "output_tokens": 1},
|
||||
}
|
||||
|
||||
PLAIN_REFUSAL_RESPONSE: dict[str, Any] = {k: v for k, v in REFUSAL_RESPONSE.items() if k != "stop_details"}
|
||||
|
||||
OK_RESPONSE: dict[str, Any] = {
|
||||
"id": "msg_ok",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-opus-5",
|
||||
"content": [{"type": "text", "text": "hello"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 25, "output_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
def _sse(event: str, data: dict[str, Any]) -> bytes:
|
||||
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
|
||||
|
||||
|
||||
REFUSAL_STREAM_FRAMES: tuple[bytes, ...] = (
|
||||
_sse("message_start", {"type": "message_start", "message": {**REFUSAL_RESPONSE, "stop_reason": None}}),
|
||||
_sse(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "refusal", "stop_details": {"category": "cyber"}},
|
||||
"usage": {"output_tokens": 1},
|
||||
},
|
||||
),
|
||||
_sse("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
OK_STREAM_FRAMES: tuple[bytes, ...] = (
|
||||
_sse("message_start", {"type": "message_start", "message": {**OK_RESPONSE, "stop_reason": None}}),
|
||||
_sse(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}},
|
||||
),
|
||||
_sse("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
|
||||
def _split_frames_mid_data_line(frames: tuple[bytes, ...]) -> tuple[bytes, ...]:
|
||||
"""Split each frame's data line in half, modeling a transport chunk boundary."""
|
||||
return tuple(part for frame in frames for part in (frame[: len(frame) // 2], frame[len(frame) // 2 :]))
|
||||
|
||||
|
||||
class _FrameStream(httpx.AsyncByteStream):
|
||||
def __init__(self, frames: tuple[bytes, ...]) -> None:
|
||||
self._frames = frames
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
for frame in self._frames:
|
||||
yield frame
|
||||
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class FakeAnthropicUpstream:
|
||||
"""Intercepts the third-party transport (httpx.AsyncClient.send): refuses on fable
|
||||
models, answers on others. The router deliberately does not forward caller-injected
|
||||
clients, so the transport is the seam that exercises the real litellm pipeline."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
refusal_body: dict[str, Any] = REFUSAL_RESPONSE,
|
||||
refusal_frames: tuple[bytes, ...] = REFUSAL_STREAM_FRAMES,
|
||||
) -> None:
|
||||
self.refusal_body = refusal_body
|
||||
self.refusal_frames = refusal_frames
|
||||
self.calls: list[str] = []
|
||||
self.bodies: list[dict[str, Any]] = []
|
||||
|
||||
async def send(self, request: httpx.Request, **kwargs: Any) -> httpx.Response:
|
||||
body = json.loads(request.content or b"{}")
|
||||
model = body.get("model", "")
|
||||
self.calls.append(model)
|
||||
self.bodies.append(body)
|
||||
refuses = "fable" in model
|
||||
if body.get("stream"):
|
||||
frames = self.refusal_frames if refuses else OK_STREAM_FRAMES
|
||||
return httpx.Response(
|
||||
200,
|
||||
stream=_FrameStream(frames),
|
||||
headers={"content-type": "text/event-stream"},
|
||||
request=request,
|
||||
)
|
||||
return httpx.Response(200, json=self.refusal_body if refuses else OK_RESPONSE, request=request)
|
||||
|
||||
def install(self):
|
||||
async def _send(_client: httpx.AsyncClient, request: httpx.Request, **kwargs: Any) -> httpx.Response:
|
||||
return await self.send(request, **kwargs)
|
||||
|
||||
return patch("httpx.AsyncClient.send", new=_send)
|
||||
|
||||
|
||||
FABLE_TIER = {
|
||||
"model_name": "fable-tier",
|
||||
"litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test"},
|
||||
}
|
||||
OPUS_TARGET = {
|
||||
"model_name": "opus-target",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test"},
|
||||
}
|
||||
|
||||
|
||||
def _router(content_policy_fallbacks: list | None) -> Router:
|
||||
return Router(model_list=[FABLE_TIER, OPUS_TARGET], content_policy_fallbacks=content_policy_fallbacks)
|
||||
|
||||
|
||||
async def _collect(stream: AsyncIterator[bytes]) -> bytes:
|
||||
return b"".join([chunk async for chunk in stream])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response():
|
||||
fake = FakeAnthropicUpstream()
|
||||
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
|
||||
|
||||
with fake.install():
|
||||
response = await router.aanthropic_messages(
|
||||
model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
assert response["stop_reason"] == "end_turn"
|
||||
assert response["id"] == "msg_ok"
|
||||
assert len(fake.calls) == 2
|
||||
assert "claude-opus-5" in fake.calls[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"content_policy_fallbacks, upstream_body",
|
||||
[
|
||||
(None, REFUSAL_RESPONSE),
|
||||
([{"unrelated-group": ["opus-target"]}], REFUSAL_RESPONSE),
|
||||
([{"fable-tier": ["opus-target"]}], PLAIN_REFUSAL_RESPONSE),
|
||||
],
|
||||
ids=["nothing-configured", "row-for-other-group", "refusal-without-stop-details"],
|
||||
)
|
||||
async def test_non_streaming_refusal_passes_through_untouched(content_policy_fallbacks, upstream_body):
|
||||
fake = FakeAnthropicUpstream(refusal_body=upstream_body)
|
||||
router = _router(content_policy_fallbacks=content_policy_fallbacks)
|
||||
|
||||
with fake.install():
|
||||
response = await router.aanthropic_messages(
|
||||
model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
assert response["stop_reason"] == "refusal"
|
||||
assert response.get("stop_details") == upstream_body.get("stop_details")
|
||||
assert len(fake.calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_refusal_with_fallback_row_streams_fallback_frames():
|
||||
fake = FakeAnthropicUpstream()
|
||||
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
|
||||
|
||||
with fake.install():
|
||||
stream = await router.aanthropic_messages(
|
||||
model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
body = await _collect(stream)
|
||||
|
||||
assert b'"refusal"' not in body
|
||||
assert b"text_delta" in body
|
||||
assert len(fake.calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_refusal_split_across_chunks_still_falls_back():
|
||||
fake = FakeAnthropicUpstream(refusal_frames=_split_frames_mid_data_line(REFUSAL_STREAM_FRAMES))
|
||||
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
|
||||
|
||||
with fake.install():
|
||||
stream = await router.aanthropic_messages(
|
||||
model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
body = await _collect(stream)
|
||||
|
||||
assert b'"refusal"' not in body
|
||||
assert b"text_delta" in body
|
||||
assert len(fake.calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_refusal_without_fallback_row_passes_frames_through():
|
||||
fake = FakeAnthropicUpstream()
|
||||
router = _router(content_policy_fallbacks=None)
|
||||
|
||||
with fake.install():
|
||||
stream = await router.aanthropic_messages(
|
||||
model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
body = await _collect(stream)
|
||||
|
||||
assert b'"stop_reason": "refusal"' in body
|
||||
assert b"stop_details" in body
|
||||
assert len(fake.calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata():
|
||||
"""The pre-routing hook's tier stamp must reach the mid-stream fallback lookup even when the
|
||||
request carries no metadata bucket at all (the snapshot is taken before the request runs)."""
|
||||
fake = FakeAnthropicUpstream()
|
||||
smart_router = {
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"}
|
||||
},
|
||||
"complexity_router_default_model": "fable-tier",
|
||||
},
|
||||
"model_info": {"id": "router-1", "db_model": True},
|
||||
}
|
||||
router = Router(
|
||||
model_list=[FABLE_TIER, OPUS_TARGET, smart_router],
|
||||
content_policy_fallbacks=[{"fable-tier": ["opus-target"]}],
|
||||
ignore_invalid_deployments=True,
|
||||
)
|
||||
|
||||
with fake.install():
|
||||
stream = await router.aanthropic_messages(
|
||||
model="smart-router", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
body = await _collect(stream)
|
||||
|
||||
assert b'"refusal"' not in body
|
||||
assert b"text_delta" in body
|
||||
assert len(fake.calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain():
|
||||
fake = FakeAnthropicUpstream()
|
||||
router = _router(content_policy_fallbacks=[{"forged-tier": ["opus-target"]}])
|
||||
|
||||
with fake.install():
|
||||
stream = await router.aanthropic_messages(
|
||||
model="fable-tier",
|
||||
max_tokens=16,
|
||||
stream=True,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
litellm_metadata={PRE_ROUTING_SELECTED_MODEL_KEY: "forged-tier"},
|
||||
)
|
||||
body = await _collect(stream)
|
||||
|
||||
assert b'"stop_reason": "refusal"' in body
|
||||
assert len(fake.calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tier_stamp_never_reaches_provider_bound_metadata():
|
||||
"""On /v1/messages the top-level metadata dict is Anthropic's own request field, so the
|
||||
routed-tier stamp must never appear in any upstream body even when the client sends one."""
|
||||
fake = FakeAnthropicUpstream()
|
||||
smart_router = {
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"}
|
||||
},
|
||||
"complexity_router_default_model": "fable-tier",
|
||||
},
|
||||
"model_info": {"id": "router-1", "db_model": True},
|
||||
}
|
||||
router = Router(
|
||||
model_list=[FABLE_TIER, OPUS_TARGET, smart_router],
|
||||
content_policy_fallbacks=[{"fable-tier": ["opus-target"]}],
|
||||
ignore_invalid_deployments=True,
|
||||
)
|
||||
|
||||
with fake.install():
|
||||
response = await router.aanthropic_messages(
|
||||
model="smart-router",
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
metadata={"user_id": "u1"},
|
||||
)
|
||||
|
||||
assert response["stop_reason"] == "end_turn"
|
||||
assert len(fake.bodies) == 2
|
||||
for body in fake.bodies:
|
||||
assert body.get("metadata") == {"user_id": "u1"}
|
||||
|
||||
|
||||
def test_record_pre_routing_selection_writes_only_the_internal_bucket():
|
||||
"""The Anthropic request's own metadata field must never carry the tier stamp."""
|
||||
kwargs = {"metadata": {"user_id": "u1"}, "litellm_metadata": {}}
|
||||
|
||||
record_pre_routing_selection(kwargs, "tier-x")
|
||||
|
||||
assert kwargs["litellm_metadata"] == {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-x"}
|
||||
assert kwargs["metadata"] == {"user_id": "u1"}
|
||||
|
||||
|
||||
def test_refusal_gate_keys_on_pre_routing_tier_stamp():
|
||||
router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}])
|
||||
|
||||
def anthropic_messages(**kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
refusal_kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}}
|
||||
assert (
|
||||
router._should_raise_anthropic_refusal_error(
|
||||
model="router-group",
|
||||
original_generic_function=anthropic_messages,
|
||||
response=dict(REFUSAL_RESPONSE),
|
||||
kwargs=refusal_kwargs,
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
router._should_raise_anthropic_refusal_error(
|
||||
model="router-group",
|
||||
original_generic_function=anthropic_messages,
|
||||
response=dict(REFUSAL_RESPONSE),
|
||||
kwargs={},
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_has_content_policy_fallback_default_fallbacks_arm():
|
||||
router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}])
|
||||
|
||||
assert router._has_content_policy_fallback("any-group", {}) is True
|
||||
assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
|
||||
|
||||
|
||||
def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested():
|
||||
router = _router(content_policy_fallbacks=None)
|
||||
fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]
|
||||
|
||||
assert router._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=fallbacks, lookup_groups=("tier1", "smart-router")
|
||||
) == ["backup-a"]
|
||||
assert router._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=fallbacks, lookup_groups=("tier9", "smart-router")
|
||||
) == ["backup-b"]
|
||||
assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None
|
||||
|
||||
|
||||
def test_refusal_gate_ignores_other_generic_call_types():
|
||||
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
|
||||
|
||||
def aresponses(**kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
assert (
|
||||
router._should_raise_anthropic_refusal_error(
|
||||
model="fable-tier",
|
||||
original_generic_function=aresponses,
|
||||
response=dict(REFUSAL_RESPONSE),
|
||||
kwargs={},
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
|
@ -8,11 +8,13 @@ from mcp.types import (
|
|||
CallToolRequestParams,
|
||||
CallToolResult,
|
||||
ListToolsResult,
|
||||
PaginatedRequestParams,
|
||||
TextContent,
|
||||
)
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.experimental_mcp_client.tools import (
|
||||
list_tools_with_pagination,
|
||||
transform_mcp_tool_to_anthropic_tool,
|
||||
_get_function_arguments,
|
||||
_normalize_mcp_input_schema,
|
||||
|
|
@ -106,6 +108,134 @@ async def test_load_mcp_tools_openai_format(mock_session, mock_list_tools_result
|
|||
mock_session.list_tools.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_load_mcp_tools_follows_pagination(mock_session):
|
||||
mock_session.list_tools.side_effect = [
|
||||
ListToolsResult(
|
||||
tools=[
|
||||
MCPTool(name="tool_a", description="a", inputSchema={}),
|
||||
MCPTool(name="tool_b", description="b", inputSchema={}),
|
||||
],
|
||||
nextCursor="page-2",
|
||||
),
|
||||
ListToolsResult(tools=[MCPTool(name="tool_c", description="c", inputSchema={})]),
|
||||
]
|
||||
result = await load_mcp_tools(mock_session, format="mcp")
|
||||
assert [tool.name for tool in result] == ["tool_a", "tool_b", "tool_c"]
|
||||
assert mock_session.list_tools.call_count == 2
|
||||
second_call_params = mock_session.list_tools.call_args_list[1].kwargs["params"]
|
||||
assert isinstance(second_call_params, PaginatedRequestParams)
|
||||
assert second_call_params.cursor == "page-2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch):
|
||||
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_MAX_PAGES", 2)
|
||||
mock_session.list_tools.side_effect = [
|
||||
ListToolsResult(
|
||||
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
|
||||
nextCursor="page-2",
|
||||
),
|
||||
ListToolsResult(
|
||||
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
|
||||
nextCursor="page-3",
|
||||
),
|
||||
ListToolsResult(tools=[MCPTool(name="tool_2", description="2", inputSchema={})]),
|
||||
]
|
||||
result = await list_tools_with_pagination(mock_session)
|
||||
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
|
||||
assert mock_session.list_tools.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_pagination_walk_stops_on_repeated_cursor(mock_session):
|
||||
mock_session.list_tools.side_effect = [
|
||||
ListToolsResult(
|
||||
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
|
||||
nextCursor="same-cursor",
|
||||
),
|
||||
ListToolsResult(
|
||||
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
|
||||
nextCursor="same-cursor",
|
||||
),
|
||||
]
|
||||
result = await list_tools_with_pagination(mock_session)
|
||||
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
|
||||
assert mock_session.list_tools.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_pagination_walk_treats_empty_cursor_as_terminal(mock_session):
|
||||
mock_session.list_tools.side_effect = [
|
||||
ListToolsResult(
|
||||
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
|
||||
nextCursor="",
|
||||
),
|
||||
]
|
||||
result = await list_tools_with_pagination(mock_session)
|
||||
assert [tool.name for tool in result] == ["tool_0"]
|
||||
mock_session.list_tools.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_pagination_walk_stops_at_whole_walk_deadline(mock_session, monkeypatch):
|
||||
import anyio
|
||||
|
||||
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
||||
|
||||
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.2)
|
||||
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.2)
|
||||
|
||||
async def slow_page(params=None):
|
||||
await anyio.sleep(0.15)
|
||||
idx = int(params.cursor) if params is not None else 0
|
||||
return ListToolsResult(
|
||||
tools=[MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})],
|
||||
nextCursor=str(idx + 1),
|
||||
)
|
||||
|
||||
mock_session.list_tools = slow_page
|
||||
result = await list_tools_with_pagination(mock_session)
|
||||
|
||||
assert [tool.name for tool in result] == ["tool_0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_pagination_walk_honors_explicit_deadline_over_globals(mock_session, monkeypatch):
|
||||
import anyio
|
||||
|
||||
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
||||
|
||||
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.1)
|
||||
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.1)
|
||||
|
||||
async def slow_page(params=None):
|
||||
await anyio.sleep(0.15)
|
||||
idx = int(params.cursor) if params is not None else 0
|
||||
tools = [MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})]
|
||||
if idx == 0:
|
||||
return ListToolsResult(tools=tools, nextCursor="1")
|
||||
return ListToolsResult(tools=tools)
|
||||
|
||||
mock_session.list_tools = slow_page
|
||||
result = await list_tools_with_pagination(mock_session, listing_deadline=2.0)
|
||||
|
||||
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_load_mcp_tools_openai_format_spans_pages(mock_session):
|
||||
mock_session.list_tools.side_effect = [
|
||||
ListToolsResult(
|
||||
tools=[MCPTool(name="tool_a", description="a", inputSchema={})],
|
||||
nextCursor="page-2",
|
||||
),
|
||||
ListToolsResult(tools=[MCPTool(name="tool_b", description="b", inputSchema={})]),
|
||||
]
|
||||
result = await load_mcp_tools(mock_session, format="openai")
|
||||
assert [t["function"]["name"] for t in result] == ["tool_a", "tool_b"]
|
||||
|
||||
|
||||
def test_get_function_arguments():
|
||||
# Test with string arguments
|
||||
function = {"arguments": '{"test": "value"}'}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,193 @@
|
|||
import datetime
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.integrations.SlackAlerting.user_spend_alerts import (
|
||||
UserSpendRow,
|
||||
evaluate_user_spend,
|
||||
)
|
||||
from litellm.types.integrations.slack_alerting import (
|
||||
DEFAULT_ALERT_TYPES,
|
||||
AlertType,
|
||||
SlackAlertingArgs,
|
||||
)
|
||||
|
||||
TODAY: Final = datetime.date(2026, 8, 15)
|
||||
|
||||
|
||||
def _row(
|
||||
daily_spend: float = 0.0,
|
||||
monthly_spend: float = 0.0,
|
||||
baseline_spend: float = 0.0,
|
||||
) -> UserSpendRow:
|
||||
return UserSpendRow(
|
||||
user_id="user-1",
|
||||
daily_spend=daily_spend,
|
||||
monthly_spend=monthly_spend,
|
||||
baseline_spend=baseline_spend,
|
||||
)
|
||||
|
||||
|
||||
def _evaluate(row: UserSpendRow, args: SlackAlertingArgs, thresholds: bool = True, anomalies: bool = True):
|
||||
return evaluate_user_spend(
|
||||
row=row,
|
||||
args=args,
|
||||
today=TODAY,
|
||||
thresholds_enabled=thresholds,
|
||||
anomalies_enabled=anomalies,
|
||||
)
|
||||
|
||||
|
||||
def test_daily_threshold_crossed():
|
||||
args: Final = SlackAlertingArgs(daily_spend_per_user_threshold=50.0, spend_anomaly_min_spend=1000.0)
|
||||
events: Final = _evaluate(_row(daily_spend=75.0, monthly_spend=75.0), args)
|
||||
assert [e.kind for e in events] == ["daily_threshold"]
|
||||
assert "`$75.00`" in events[0].message
|
||||
assert "`$50.00`" in events[0].message
|
||||
assert events[0].alert_type == AlertType.user_spend_thresholds
|
||||
assert events[0].cache_key == "user_spend_alert_daily_user-1_2026-08-15"
|
||||
|
||||
|
||||
def test_daily_threshold_not_crossed():
|
||||
args: Final = SlackAlertingArgs(daily_spend_per_user_threshold=50.0, spend_anomaly_min_spend=1000.0)
|
||||
assert _evaluate(_row(daily_spend=49.99, monthly_spend=49.99), args) == ()
|
||||
|
||||
|
||||
def test_thresholds_unset_by_default():
|
||||
args: Final = SlackAlertingArgs(spend_anomaly_min_spend=1000.0)
|
||||
assert _evaluate(_row(daily_spend=999.0, monthly_spend=999.0), args) == ()
|
||||
|
||||
|
||||
def test_monthly_threshold_crossed():
|
||||
args: Final = SlackAlertingArgs(monthly_spend_per_user_threshold=200.0, spend_anomaly_min_spend=1000.0)
|
||||
events: Final = _evaluate(_row(daily_spend=5.0, monthly_spend=250.0), args)
|
||||
assert [e.kind for e in events] == ["monthly_threshold"]
|
||||
assert events[0].cache_key == "user_spend_alert_monthly_user-1_2026-08"
|
||||
|
||||
|
||||
def test_thresholds_disabled_suppresses_threshold_events():
|
||||
args: Final = SlackAlertingArgs(
|
||||
daily_spend_per_user_threshold=50.0,
|
||||
monthly_spend_per_user_threshold=200.0,
|
||||
spend_anomaly_min_spend=1000.0,
|
||||
)
|
||||
assert _evaluate(_row(daily_spend=75.0, monthly_spend=250.0), args, thresholds=False) == ()
|
||||
|
||||
|
||||
def test_anomaly_detected_above_multiple_of_baseline():
|
||||
args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0)
|
||||
events: Final = _evaluate(
|
||||
_row(daily_spend=70.0, monthly_spend=100.0, baseline_spend=70.0), args
|
||||
)
|
||||
assert [e.kind for e in events] == ["anomaly"]
|
||||
assert events[0].alert_type == AlertType.user_spend_anomalies
|
||||
assert "`$10.00`" in events[0].message
|
||||
assert events[0].cache_key == "user_spend_alert_anomaly_user-1_2026-08-15"
|
||||
|
||||
|
||||
def test_no_anomaly_within_baseline_multiple():
|
||||
args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0)
|
||||
assert (
|
||||
_evaluate(_row(daily_spend=25.0, monthly_spend=100.0, baseline_spend=70.0), args) == ()
|
||||
)
|
||||
|
||||
|
||||
def test_no_anomaly_below_min_spend_floor():
|
||||
args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0)
|
||||
assert _evaluate(_row(daily_spend=9.0, monthly_spend=9.0, baseline_spend=0.1), args) == ()
|
||||
|
||||
|
||||
def test_anomaly_for_new_user_without_baseline():
|
||||
args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0)
|
||||
events: Final = _evaluate(_row(daily_spend=15.0, monthly_spend=15.0), args)
|
||||
assert [e.kind for e in events] == ["anomaly"]
|
||||
|
||||
|
||||
def test_sparse_baseline_averages_over_full_window():
|
||||
args: Final = SlackAlertingArgs(
|
||||
spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0, spend_anomaly_baseline_days=7
|
||||
)
|
||||
events: Final = _evaluate(_row(daily_spend=13.0, monthly_spend=20.0, baseline_spend=7.0), args)
|
||||
assert [e.kind for e in events] == ["anomaly"]
|
||||
|
||||
|
||||
def test_anomalies_not_in_default_alert_types():
|
||||
assert AlertType.user_spend_anomalies not in DEFAULT_ALERT_TYPES
|
||||
assert AlertType.user_spend_thresholds in DEFAULT_ALERT_TYPES
|
||||
|
||||
|
||||
def test_invalid_config_rejected():
|
||||
with pytest.raises(ValidationError, match="daily_spend_per_user_threshold"):
|
||||
SlackAlertingArgs(daily_spend_per_user_threshold=0)
|
||||
with pytest.raises(ValidationError, match="spend_anomaly_baseline_days"):
|
||||
SlackAlertingArgs(spend_anomaly_baseline_days=0)
|
||||
with pytest.raises(ValidationError, match="user_spend_check_interval"):
|
||||
SlackAlertingArgs(user_spend_check_interval=10)
|
||||
|
||||
|
||||
def test_non_finite_config_rejected():
|
||||
with pytest.raises(ValidationError, match="daily_spend_per_user_threshold"):
|
||||
SlackAlertingArgs(daily_spend_per_user_threshold=float("inf"))
|
||||
with pytest.raises(ValidationError, match="spend_anomaly_multiplier"):
|
||||
SlackAlertingArgs(spend_anomaly_multiplier=float("nan"))
|
||||
with pytest.raises(ValidationError, match="spend_anomaly_min_spend"):
|
||||
SlackAlertingArgs(spend_anomaly_min_spend=float("inf"))
|
||||
with pytest.raises(ValidationError, match="user_spend_check_interval"):
|
||||
SlackAlertingArgs(user_spend_check_interval=float("inf"))
|
||||
|
||||
|
||||
def test_anomalies_disabled_suppresses_anomaly_events():
|
||||
args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0)
|
||||
assert _evaluate(_row(daily_spend=500.0, monthly_spend=500.0), args, anomalies=False) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_user_spend_alerts_sends_and_dedupes():
|
||||
slack_alerting: Final = SlackAlerting(
|
||||
alerting=["slack"],
|
||||
alerting_args={"daily_spend_per_user_threshold": 50.0, "spend_anomaly_min_spend": 1000.0},
|
||||
)
|
||||
mock_prisma: Final = AsyncMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"user_id": "user-1",
|
||||
"daily_spend": 75.0,
|
||||
"monthly_spend": 75.0,
|
||||
"baseline_spend": 0.0,
|
||||
},
|
||||
{
|
||||
"user_id": "user-2",
|
||||
"daily_spend": 60.0,
|
||||
"monthly_spend": 60.0,
|
||||
"baseline_spend": 0.0,
|
||||
},
|
||||
]
|
||||
)
|
||||
with patch.object(slack_alerting, "send_alert", new_callable=AsyncMock) as mock_send_alert:
|
||||
await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma)
|
||||
assert mock_send_alert.call_count == 1
|
||||
sent_kwargs: Final = mock_send_alert.call_args.kwargs
|
||||
assert sent_kwargs["alert_type"] == AlertType.user_spend_thresholds
|
||||
assert "User Daily Spend Threshold Crossed" in sent_kwargs["message"]
|
||||
assert "`user-1`" in sent_kwargs["message"]
|
||||
assert "`user-2`" in sent_kwargs["message"]
|
||||
|
||||
await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma)
|
||||
assert mock_send_alert.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_user_spend_alerts_noop_when_alert_types_disabled():
|
||||
slack_alerting: Final = SlackAlerting(
|
||||
alerting=["slack"],
|
||||
alert_types=[AlertType.budget_alerts],
|
||||
alerting_args={"daily_spend_per_user_threshold": 50.0},
|
||||
)
|
||||
mock_prisma: Final = AsyncMock()
|
||||
await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma)
|
||||
mock_prisma.db.query_raw.assert_not_called()
|
||||
|
|
@ -1435,6 +1435,80 @@ class TestFlattenTopLevelSchemaCombinators:
|
|||
assert schema == snapshot
|
||||
|
||||
|
||||
class TestToolWithFlattenedParameters:
|
||||
def _anyof_tool(self):
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "automation_update",
|
||||
"description": "Update an automation",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"anyOf": [
|
||||
{
|
||||
"properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}},
|
||||
"required": ["id", "enabled"],
|
||||
},
|
||||
{
|
||||
"properties": {"id": {"type": "string"}, "schedule": {"type": "string"}},
|
||||
"required": ["id", "schedule"],
|
||||
},
|
||||
],
|
||||
"properties": {"id": {"type": "string"}},
|
||||
"required": ["id"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def test_flattens_anyof_parameters_into_new_tool(self):
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
tool_with_flattened_parameters,
|
||||
)
|
||||
|
||||
tool = self._anyof_tool()
|
||||
result = tool_with_flattened_parameters(tool)
|
||||
|
||||
assert result is not tool
|
||||
parameters = result["function"]["parameters"]
|
||||
assert "anyOf" not in parameters
|
||||
assert parameters["type"] == "object"
|
||||
assert set(parameters["properties"]) == {"id", "enabled", "schedule"}
|
||||
assert parameters["required"] == ["id"]
|
||||
assert result["function"]["name"] == "automation_update"
|
||||
assert tool == self._anyof_tool()
|
||||
|
||||
def test_clean_parameters_return_the_same_tool_object(self):
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
tool_with_flattened_parameters,
|
||||
)
|
||||
|
||||
tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]},
|
||||
},
|
||||
}
|
||||
|
||||
assert tool_with_flattened_parameters(tool) is tool
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
[
|
||||
{"type": "function"},
|
||||
{"type": "function", "function": "not-a-dict"},
|
||||
{"type": "function", "function": {"name": "no_params"}},
|
||||
{"type": "function", "function": {"name": "bad_params", "parameters": "not-a-dict"}},
|
||||
],
|
||||
)
|
||||
def test_non_dict_function_or_parameters_return_the_same_tool_object(self, tool):
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
tool_with_flattened_parameters,
|
||||
)
|
||||
|
||||
assert tool_with_flattened_parameters(tool) is tool
|
||||
|
||||
|
||||
class TestRequestContainsImageContent:
|
||||
"""One detector for every dialect that reaches pre-routing hooks untranslated."""
|
||||
|
||||
|
|
|
|||
|
|
@ -188,6 +188,8 @@ class TestDeclaredAuthenticatingProvider:
|
|||
("gpt-4o", "github_copilot", "github_copilot"),
|
||||
("openai/gpt-4o", None, None),
|
||||
("gpt-4o", "openai", None),
|
||||
("github_copilot", None, None),
|
||||
("chatgpt", None, None),
|
||||
],
|
||||
)
|
||||
def test_names_only_the_providers_whose_resolution_authenticates(self, model, provider, expected):
|
||||
|
|
|
|||
|
|
@ -5479,6 +5479,97 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj):
|
|||
assert "key=*****" in raw_api_base
|
||||
|
||||
|
||||
def _streaming_logging_obj_with_callbacks(callbacks: list[CustomLogger]):
|
||||
import datetime
|
||||
|
||||
obj = LitellmLogging(
|
||||
model="anthropic/claude-opus-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="completion",
|
||||
start_time=datetime.datetime.now(),
|
||||
litellm_call_id="slot-leak-test",
|
||||
function_id="slot-leak-test",
|
||||
)
|
||||
obj.model_call_details["litellm_params"] = {"metadata": {}}
|
||||
return patch.object(obj, "get_combined_callback_list", return_value=callbacks), obj
|
||||
|
||||
|
||||
def _assembled_stream_result():
|
||||
response = ModelResponse()
|
||||
response.choices[0].message.content = "hello"
|
||||
return response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_success_callbacks_survive_logging_hook_failure():
|
||||
"""Regression for leaked max_parallel_requests slots: a raising
|
||||
async_logging_hook must not abort the success-callback loop that
|
||||
releases the rate-limiter slot."""
|
||||
broken = CustomLogger()
|
||||
broken.async_logging_hook = AsyncMock(side_effect=RuntimeError("broken stream payload"))
|
||||
releasing = CustomLogger()
|
||||
releasing.async_log_success_event = AsyncMock()
|
||||
|
||||
patcher, logging_obj = _streaming_logging_obj_with_callbacks([broken, releasing])
|
||||
with patcher:
|
||||
await logging_obj.async_success_handler(result=_assembled_stream_result())
|
||||
|
||||
releasing.async_log_success_event.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_success_callbacks_survive_cost_calculation_failure():
|
||||
releasing = CustomLogger()
|
||||
releasing.async_log_success_event = AsyncMock()
|
||||
|
||||
patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing])
|
||||
with patcher, patch.object(
|
||||
logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block")
|
||||
):
|
||||
await logging_obj.async_success_handler(result=_assembled_stream_result())
|
||||
|
||||
assert logging_obj.model_call_details["response_cost"] is None
|
||||
releasing.async_log_success_event.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_success_callbacks_survive_standard_logging_payload_failure():
|
||||
releasing = CustomLogger()
|
||||
releasing.async_log_success_event = AsyncMock()
|
||||
|
||||
patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing])
|
||||
with patcher, patch.object(
|
||||
logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream")
|
||||
):
|
||||
await logging_obj.async_success_handler(result=_assembled_stream_result())
|
||||
|
||||
assert logging_obj.model_call_details.get("standard_logging_object") is None
|
||||
releasing.async_log_success_event.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failure():
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
||||
skipping = CustomGuardrail(guardrail_name="skipping-guardrail")
|
||||
skipping.should_run_guardrail = MagicMock(return_value=False)
|
||||
skipping.async_logging_hook = AsyncMock()
|
||||
raising = CustomGuardrail(guardrail_name="raising-guardrail")
|
||||
raising.should_run_guardrail = MagicMock(return_value=True)
|
||||
raising.async_logging_hook = AsyncMock(side_effect=RuntimeError("guardrail hook failed"))
|
||||
releasing = CustomLogger()
|
||||
releasing.async_log_success_event = AsyncMock()
|
||||
|
||||
patcher, logging_obj = _streaming_logging_obj_with_callbacks([skipping, raising, releasing])
|
||||
with patcher:
|
||||
await logging_obj.async_success_handler(result=_assembled_stream_result())
|
||||
|
||||
skipping.async_logging_hook.assert_not_awaited()
|
||||
raising.async_logging_hook.assert_awaited_once()
|
||||
releasing.async_log_success_event.assert_awaited_once()
|
||||
|
||||
|
||||
def _resolve(custom_llm_provider, litellm_params, optional_params, model):
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_resolve_vertex_location_for_cost,
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ sys.path.insert(
|
|||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY
|
||||
from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config
|
||||
from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig
|
||||
from litellm.utils import get_optional_params
|
||||
|
||||
|
|
@ -195,3 +196,91 @@ def test_azure_gpt_5_takes_the_reasoning_path() -> None:
|
|||
assert "presence_penalty" not in mapped
|
||||
assert "logit_bias" not in mapped
|
||||
assert "reasoning_effort" in supported
|
||||
|
||||
|
||||
class TestAzureToolSchemaCombinatorFlattening:
|
||||
"""
|
||||
Regression tests for LIT-6510: Azure's chat completions validator rejects
|
||||
tool parameters carrying a top-level anyOf/oneOf/allOf for every model
|
||||
family, so AzureOpenAIConfig.transform_request must flatten them.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _anyof_tool():
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "automation_update",
|
||||
"description": "Update an automation",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"anyOf": [
|
||||
{
|
||||
"properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}},
|
||||
"required": ["id", "enabled"],
|
||||
},
|
||||
{
|
||||
"properties": {"id": {"type": "string"}, "schedule": {"type": "string"}},
|
||||
"required": ["id", "schedule"],
|
||||
},
|
||||
],
|
||||
"properties": {"id": {"type": "string"}},
|
||||
"required": ["id"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def _transform(self, config, model, tools):
|
||||
return config.transform_request(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params={"tools": tools},
|
||||
litellm_params={"custom_llm_provider": "azure"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
def test_transform_request_flattens_top_level_anyof(self):
|
||||
request = self._transform(AzureOpenAIConfig(), "gpt-4o", [self._anyof_tool()])
|
||||
parameters = request["tools"][0]["function"]["parameters"]
|
||||
assert "anyOf" not in parameters
|
||||
assert parameters["type"] == "object"
|
||||
assert set(parameters["properties"]) == {"id", "enabled", "schedule"}
|
||||
assert parameters["required"] == ["id"]
|
||||
assert request["tools"][0]["function"]["name"] == "automation_update"
|
||||
|
||||
def test_gpt5_config_flattens_via_shared_transform(self):
|
||||
request = self._transform(AzureOpenAIGPT5Config(), "gpt-5.4-mini", [self._anyof_tool()])
|
||||
parameters = request["tools"][0]["function"]["parameters"]
|
||||
assert "anyOf" not in parameters
|
||||
assert set(parameters["properties"]) == {"id", "enabled", "schedule"}
|
||||
|
||||
def test_caller_tool_dict_is_not_mutated(self):
|
||||
tool = self._anyof_tool()
|
||||
self._transform(AzureOpenAIConfig(), "gpt-4o", [tool])
|
||||
assert tool == self._anyof_tool()
|
||||
|
||||
def test_clean_object_schema_passes_through_as_same_object(self):
|
||||
tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]},
|
||||
},
|
||||
}
|
||||
request = self._transform(AzureOpenAIConfig(), "gpt-4o", [tool])
|
||||
assert request["tools"][0] is tool
|
||||
|
||||
def test_non_dict_tool_entries_pass_through_unchanged(self):
|
||||
request = self._transform(AzureOpenAIConfig(), "gpt-4o", ["not-a-tool"])
|
||||
assert request["tools"] == ["not-a-tool"]
|
||||
|
||||
def test_request_without_tools_is_unchanged(self):
|
||||
request = AzureOpenAIConfig().transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params={"temperature": 0.2},
|
||||
litellm_params={"custom_llm_provider": "azure"},
|
||||
headers={},
|
||||
)
|
||||
assert "tools" not in request
|
||||
assert request["temperature"] == 0.2
|
||||
|
|
|
|||
|
|
@ -23,3 +23,48 @@ async def test_azure_chat_o_series_transformation():
|
|||
)
|
||||
print(response)
|
||||
assert response["model"] == "web-interface-o1-mini"
|
||||
|
||||
|
||||
def test_azure_o_series_transform_request_flattens_top_level_anyof():
|
||||
"""Regression test for LIT-6510: the o-series super() chain ends in
|
||||
OpenAIGPTConfig, whose flatten gate skips provider 'azure', so
|
||||
AzureOpenAIO1Config must flatten tool schema combinators itself."""
|
||||
tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "automation_update",
|
||||
"description": "Update an automation",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"anyOf": [
|
||||
{
|
||||
"properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}},
|
||||
"required": ["id", "enabled"],
|
||||
},
|
||||
{
|
||||
"properties": {"id": {"type": "string"}, "schedule": {"type": "string"}},
|
||||
"required": ["id", "schedule"],
|
||||
},
|
||||
],
|
||||
"properties": {"id": {"type": "string"}},
|
||||
"required": ["id"],
|
||||
},
|
||||
},
|
||||
}
|
||||
optional_params = {"tools": [tool]}
|
||||
|
||||
request = AzureOpenAIO1Config().transform_request(
|
||||
model="o3-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params=optional_params,
|
||||
litellm_params={"custom_llm_provider": "azure"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
parameters = request["tools"][0]["function"]["parameters"]
|
||||
assert "anyOf" not in parameters
|
||||
assert parameters["type"] == "object"
|
||||
assert set(parameters["properties"]) == {"id", "enabled", "schedule"}
|
||||
assert parameters["required"] == ["id"]
|
||||
assert "anyOf" in tool["function"]["parameters"]
|
||||
assert optional_params["tools"][0] is tool
|
||||
|
|
|
|||
|
|
@ -828,6 +828,24 @@ class MockPassThroughGuardrail(CustomGuardrail):
|
|||
return inputs
|
||||
|
||||
|
||||
class MockRecordingGuardrail(MockPassThroughGuardrail):
|
||||
"""Pass-through guardrail that records every apply_guardrail inputs payload"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.seen_inputs: List[GenericGuardrailAPIInputs] = []
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.seen_inputs.append(inputs)
|
||||
return inputs
|
||||
|
||||
|
||||
class TestOpenAIResponsesHandlerStreamingOutputProcessing:
|
||||
"""Test streaming output processing functionality"""
|
||||
|
||||
|
|
@ -1104,6 +1122,80 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
|
|||
output_text = result[-1]["response"]["output"][0]["content"][0]["text"]
|
||||
assert output_text == original_text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_stream_scans_delta_text(self):
|
||||
"""A stream ending in response.failed has text only in delta events; the
|
||||
fallback scan must assemble and scan it instead of skipping on an empty string."""
|
||||
handler = OpenAIResponsesHandler()
|
||||
guardrail = MockRecordingGuardrail(guardrail_name="test")
|
||||
|
||||
responses_so_far = [
|
||||
{"type": "response.created", "response": {"id": "resp_123"}},
|
||||
{"type": "response.output_item.added", "item": {"type": "message", "id": "msg_123"}},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_123",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "Hello",
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_123",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": " world",
|
||||
},
|
||||
{"type": "response.failed", "response": {"id": "resp_123", "status": "failed"}},
|
||||
]
|
||||
|
||||
result = await handler.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail,
|
||||
litellm_logging_obj=None,
|
||||
)
|
||||
|
||||
assert result == responses_so_far
|
||||
assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["Hello world"]]
|
||||
|
||||
def test_get_streaming_string_so_far_prefers_done_text_over_deltas(self):
|
||||
"""The done event repeats the whole part, so deltas must not be double counted;
|
||||
a part with no done event yet still contributes its joined deltas."""
|
||||
handler = OpenAIResponsesHandler()
|
||||
|
||||
events = [
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "Hello",
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": " world",
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.done",
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"text": "Hello world",
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": "msg_2",
|
||||
"output_index": 1,
|
||||
"content_index": 0,
|
||||
"delta": "; unfinished",
|
||||
},
|
||||
]
|
||||
|
||||
assert handler.get_streaming_string_so_far(events) == "Hello world; unfinished"
|
||||
|
||||
|
||||
class TestGetStructuredMessages:
|
||||
"""Test the get_structured_messages method for Responses API handler."""
|
||||
|
|
|
|||
|
|
@ -214,6 +214,46 @@ class TestExecuteWithMcpClient:
|
|||
assert server.scopes == ["read", "write"]
|
||||
assert server.has_client_credentials is True
|
||||
|
||||
async def test_preview_forwards_per_server_timeout_to_client_factory(self, monkeypatch):
|
||||
"""The request's per-server timeout must reach the temporary MCPServer model:
|
||||
the client factory reads ``server.timeout`` for both the per-request timeout
|
||||
and the preview's whole-walk listing deadline."""
|
||||
captured: dict = {}
|
||||
|
||||
def fake_build_stdio_env(server, raw_headers):
|
||||
return None
|
||||
|
||||
async def fake_create_client(*args, **kwargs):
|
||||
captured["server"] = kwargs.get("server")
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_build_stdio_env",
|
||||
fake_build_stdio_env,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_create_mcp_client",
|
||||
fake_create_client,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
async def ok_operation(client):
|
||||
return {"status": "ok"}
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="slow-catalog-server",
|
||||
url="https://example.com",
|
||||
timeout=120.5,
|
||||
)
|
||||
|
||||
result = await rest_endpoints._execute_with_mcp_client(payload, ok_operation)
|
||||
|
||||
assert result["status"] == "ok"
|
||||
assert captured["server"].timeout == 120.5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch):
|
||||
"""For M2M OAuth servers the incoming Authorization header (which carries
|
||||
|
|
@ -524,6 +564,131 @@ class TestTestToolsList:
|
|||
assert captured["oauth2_headers"] is None
|
||||
assert oauth_call_counter["count"] == 0
|
||||
|
||||
async def test_preview_tools_list_times_out_on_slow_pagination(self, monkeypatch):
|
||||
"""A preview whose upstream paginates past the listing deadline returns a
|
||||
timeout error instead of holding the request open."""
|
||||
monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False)
|
||||
monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False)
|
||||
|
||||
class SlowClient:
|
||||
async def list_tools(self, raise_on_error=False):
|
||||
await asyncio.sleep(1)
|
||||
return []
|
||||
|
||||
async def fake_execute(
|
||||
request,
|
||||
operation,
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers=None,
|
||||
raw_headers=None,
|
||||
):
|
||||
return await operation(SlowClient())
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
request = _build_request()
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="example",
|
||||
url="https://example.com",
|
||||
auth_type=MCPAuth.api_key,
|
||||
credentials={"auth_value": "secret-key"},
|
||||
)
|
||||
|
||||
result = await rest_endpoints.test_tools_list(
|
||||
request,
|
||||
payload,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert result["error"] is True
|
||||
assert "Timed out listing tools" in result["message"]
|
||||
|
||||
async def test_preview_tools_list_succeeds_within_deadline(self, monkeypatch):
|
||||
"""The preview timeout scope passes a fast listing through untouched."""
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
class QuickClient:
|
||||
async def list_tools(self, raise_on_error=False):
|
||||
return [MCPTool(name="quick_tool", description="q", inputSchema={})]
|
||||
|
||||
async def fake_execute(
|
||||
request,
|
||||
operation,
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers=None,
|
||||
raw_headers=None,
|
||||
):
|
||||
return await operation(QuickClient())
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
request = _build_request()
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="example",
|
||||
url="https://example.com",
|
||||
auth_type=MCPAuth.api_key,
|
||||
credentials={"auth_value": "secret-key"},
|
||||
)
|
||||
|
||||
result = await rest_endpoints.test_tools_list(
|
||||
request,
|
||||
payload,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert result["error"] is None
|
||||
assert result["message"] == "Successfully retrieved tools"
|
||||
assert [tool["name"] for tool in result["tools"]] == ["quick_tool"]
|
||||
|
||||
async def test_preview_tools_list_honors_per_server_timeout(self, monkeypatch):
|
||||
"""A per-server timeout above the global default extends the preview deadline."""
|
||||
monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False)
|
||||
monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False)
|
||||
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
class SlowConfiguredClient:
|
||||
timeout = 1.0
|
||||
|
||||
async def list_tools(self, raise_on_error=False):
|
||||
await asyncio.sleep(0.2)
|
||||
return [MCPTool(name="slow_tool", description="s", inputSchema={})]
|
||||
|
||||
async def fake_execute(
|
||||
request,
|
||||
operation,
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers=None,
|
||||
raw_headers=None,
|
||||
):
|
||||
return await operation(SlowConfiguredClient())
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
request = _build_request()
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="example",
|
||||
url="https://example.com",
|
||||
auth_type=MCPAuth.api_key,
|
||||
credentials={"auth_value": "secret-key"},
|
||||
)
|
||||
|
||||
result = await rest_endpoints.test_tools_list(
|
||||
request,
|
||||
payload,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert result["error"] is None
|
||||
assert [tool["name"] for tool in result["tools"]] == ["slow_tool"]
|
||||
|
||||
async def test_extracts_oauth2_headers(self, monkeypatch):
|
||||
"""Ensure oauth2 auth type pulls oauth headers and omits MCP auth header."""
|
||||
|
||||
|
|
@ -786,9 +951,7 @@ class TestListToolsRestAPI:
|
|||
they do for a gateway session, never to the bare session key."""
|
||||
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
|
||||
|
||||
session_auth = UserAPIKeyAuth(
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user"
|
||||
)
|
||||
session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user")
|
||||
admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org")
|
||||
|
||||
async def fake_reload(user_id):
|
||||
|
|
@ -868,9 +1031,7 @@ class TestListToolsRestAPI:
|
|||
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
|
||||
|
||||
session_auth = UserAPIKeyAuth(
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user"
|
||||
)
|
||||
session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user")
|
||||
scoped_auth = UserAPIKeyAuth(
|
||||
object_permission=LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="toolset-scope",
|
||||
|
|
@ -952,6 +1113,123 @@ class TestListToolsRestAPI:
|
|||
assert scope_inputs == [session_auth]
|
||||
assert reload_calls == []
|
||||
|
||||
async def test_single_server_response_includes_paginated_upstream_tools(
|
||||
self,
|
||||
monkeypatch,
|
||||
):
|
||||
"""The REST tools/list path should include tools beyond the upstream first page."""
|
||||
import litellm.experimental_mcp_client.client as mcp_client_module
|
||||
from mcp.types import ListToolsResult, PaginatedRequestParams
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.server import MCPServer
|
||||
from litellm.types.mcp import MCPTransport
|
||||
|
||||
async def fake_contexts(user_api_key_auth):
|
||||
return [user_api_key_auth]
|
||||
|
||||
async def fake_get_allowed_mcp_servers(*args, **kwargs):
|
||||
return ["server-1"]
|
||||
|
||||
stub_server = MCPServer(
|
||||
server_id="server-1",
|
||||
name="stub",
|
||||
server_name="stub",
|
||||
alias="stub",
|
||||
url="https://example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
mcp_info={"server_name": "stub"},
|
||||
)
|
||||
stub_server.available_on_public_internet = True
|
||||
|
||||
mock_transport_ctx = AsyncMock()
|
||||
mock_transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock()))
|
||||
mock_transport_ctx.__aexit__ = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
mcp_client_module,
|
||||
"streamable_http_client",
|
||||
MagicMock(return_value=mock_transport_ctx),
|
||||
raising=False,
|
||||
)
|
||||
|
||||
mock_session_ctx = AsyncMock()
|
||||
mock_session_instance = AsyncMock()
|
||||
mock_session_instance.initialize = AsyncMock(return_value=None)
|
||||
mock_session_instance.list_tools.side_effect = [
|
||||
ListToolsResult(
|
||||
tools=[
|
||||
MCPTool(
|
||||
name="first_page_tool",
|
||||
description="First page tool",
|
||||
inputSchema={},
|
||||
)
|
||||
],
|
||||
nextCursor="page-2",
|
||||
),
|
||||
ListToolsResult(
|
||||
tools=[
|
||||
MCPTool(
|
||||
name="second_page_tool",
|
||||
description="Second page tool",
|
||||
inputSchema={},
|
||||
)
|
||||
]
|
||||
),
|
||||
]
|
||||
mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance)
|
||||
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
mcp_client_module,
|
||||
"ClientSession",
|
||||
MagicMock(return_value=mock_session_ctx),
|
||||
raising=False,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
fake_contexts,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"filter_server_ids_by_ip_with_info",
|
||||
lambda server_ids, client_ip: (server_ids, 0),
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_id",
|
||||
lambda server_id: stub_server if server_id == "server-1" else None,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
request = _build_request(path="/mcp-rest/tools/list", method="GET")
|
||||
result = await rest_endpoints.list_tool_rest_api(
|
||||
request,
|
||||
server_id="server-1",
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
assert set(result.keys()) == {"tools", "error", "message"}
|
||||
assert [tool.name for tool in result["tools"]] == [
|
||||
"first_page_tool",
|
||||
"second_page_tool",
|
||||
]
|
||||
assert result["error"] is None
|
||||
assert result["message"] == "Successfully retrieved tools"
|
||||
|
||||
assert mock_session_instance.list_tools.call_count == 2
|
||||
second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"]
|
||||
assert isinstance(second_call_params, PaginatedRequestParams)
|
||||
assert second_call_params.cursor == "page-2"
|
||||
|
||||
async def test_include_disabled_tools_is_admin_only(self, monkeypatch):
|
||||
"""include_disabled_tools skips the allowlist filter only for PROXY_ADMIN;
|
||||
a non-admin passing it stays filtered so the REST endpoint can't be used
|
||||
|
|
@ -3021,9 +3299,7 @@ class TestRestListToolsetFiltering:
|
|||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {})
|
||||
mock_manager.resolve_toolset_tool_permissions = AsyncMock(
|
||||
return_value={"server-a": ["lookup_status"]}
|
||||
)
|
||||
mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value={"server-a": ["lookup_status"]})
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,288 @@
|
|||
"""
|
||||
Tests for restamping the public model on Anthropic Messages streaming chunks.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
|
||||
AnthropicStreamModelRestamper,
|
||||
restamp_anthropic_stream_chunk_model,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
|
||||
def _message_start_frame(model: str, line_end: str = "\n") -> bytes:
|
||||
payload = {
|
||||
"type": "message_start",
|
||||
"message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []},
|
||||
}
|
||||
return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode()
|
||||
|
||||
|
||||
def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock:
|
||||
async def _iterator_hook(**_kwargs):
|
||||
for frame in frames:
|
||||
yield frame
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook
|
||||
proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"])
|
||||
return proxy_logging_obj
|
||||
|
||||
|
||||
def _model_from_frame(frame: bytes | str) -> str:
|
||||
text = frame.decode("utf-8") if isinstance(frame, bytes) else frame
|
||||
data_line = next(line for line in text.split("\n") if line.startswith("data:"))
|
||||
return json.loads(data_line[len("data:") :])["message"]["model"]
|
||||
|
||||
|
||||
def test_restamps_sse_bytes_frame():
|
||||
restamped = restamp_anthropic_stream_chunk_model(
|
||||
_message_start_frame("claude-haiku-4-5-20251001"), "claude-auto-1"
|
||||
)
|
||||
|
||||
assert isinstance(restamped, bytes)
|
||||
assert _model_from_frame(restamped) == "claude-auto-1"
|
||||
assert b"event: message_start" in restamped
|
||||
|
||||
|
||||
def test_restamps_event_dict():
|
||||
chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}}
|
||||
|
||||
restamped = restamp_anthropic_stream_chunk_model(chunk, "claude-auto-2")
|
||||
|
||||
assert restamped == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-2"}}
|
||||
assert chunk["message"]["model"] == "claude-sonnet-4-6"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk",
|
||||
[
|
||||
b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n',
|
||||
{"type": "content_block_delta", "delta": {"text": "hi"}},
|
||||
{"type": "message_start", "message": "not-a-dict"},
|
||||
b"event: message_start\ndata: not-json\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
],
|
||||
)
|
||||
def test_leaves_chunks_without_a_model_untouched(chunk):
|
||||
assert restamp_anthropic_stream_chunk_model(chunk, "claude-auto-1") == chunk
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_publishes_requested_model_on_message_start():
|
||||
"""The message_start event reports the requested model, not the provider's."""
|
||||
delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n'
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
assert _model_from_frame(chunks[0]) == "claude-auto-1"
|
||||
assert chunks[1] == delta_frame
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_keeps_provider_model_when_restamping_is_off():
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
]
|
||||
|
||||
assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001"
|
||||
|
||||
|
||||
def test_restamps_message_start_split_across_transport_chunks():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
held = restamper.process(frame[:25])
|
||||
emitted = restamper.process(frame[25:])
|
||||
|
||||
assert held == b""
|
||||
assert isinstance(emitted, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_emits_coalesced_frames_with_only_message_start_rewritten():
|
||||
delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n'
|
||||
combined = _message_start_frame("claude-haiku-4-5-20251001") + delta
|
||||
|
||||
emitted = restamper_output = AnthropicStreamModelRestamper("claude-auto-1").process(combined)
|
||||
|
||||
assert isinstance(restamper_output, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
assert emitted.endswith(delta)
|
||||
|
||||
|
||||
def test_ping_frames_keep_the_restamper_armed():
|
||||
ping = b'event: ping\ndata: {"type": "ping"}\n\n'
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(ping) == ping
|
||||
reassembled = restamper.process(frame[:10])
|
||||
reassembled += restamper.process(frame[10:])
|
||||
|
||||
assert _model_from_frame(reassembled) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_first_non_ping_event_disarms_the_restamper():
|
||||
delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n'
|
||||
late_message_start = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(delta) == delta
|
||||
assert restamper.process(late_message_start) == late_message_start
|
||||
|
||||
|
||||
def test_oversized_unterminated_chunk_flushes_unmodified():
|
||||
blob = b"data: " + b"x" * 70000
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(blob) == blob
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
assert restamper.process(frame) == frame
|
||||
|
||||
|
||||
def test_dict_message_start_disarms_after_restamp():
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
first = restamper.process({"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}})
|
||||
second = {"type": "message_start", "message": {"id": "msg_2", "model": "claude-sonnet-4-6"}}
|
||||
|
||||
assert first == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-1"}}
|
||||
assert restamper.process(second) == second
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_restamps_message_start_split_across_chunks():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks)
|
||||
assert _model_from_frame(joined) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_restamps_crlf_terminated_message_start_frame():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n")
|
||||
delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n'
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
emitted = restamper.process(frame)
|
||||
|
||||
assert isinstance(emitted, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
assert emitted.endswith(b"\r\n\r\n")
|
||||
assert restamper.process(delta) == delta
|
||||
|
||||
|
||||
def test_restamps_cr_terminated_message_start_frame():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
emitted = restamper.process(frame)
|
||||
|
||||
assert isinstance(emitted, bytes)
|
||||
assert b'"model":"claude-auto-1"' in emitted
|
||||
assert emitted.endswith(b"\r\r")
|
||||
|
||||
|
||||
def test_restamps_crlf_message_start_split_across_transport_chunks():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
held = restamper.process(frame[:25])
|
||||
emitted = restamper.process(frame[25:])
|
||||
|
||||
assert held == b""
|
||||
assert isinstance(emitted, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_flush_returns_restamped_held_tail():
|
||||
unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2]
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(unterminated) == b""
|
||||
flushed = restamper.flush()
|
||||
|
||||
assert b'"model":"claude-auto-1"' in flushed
|
||||
assert restamper.flush() == b""
|
||||
|
||||
|
||||
def test_flush_disarms_the_restamper():
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
|
||||
assert restamper.flush() == b""
|
||||
assert restamper.process(frame) == frame
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_flushes_held_tail_at_end_of_stream():
|
||||
unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2]
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([unterminated])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks)
|
||||
assert b'"model":"claude-auto-1"' in joined
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_restamps_crlf_stream():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n")
|
||||
delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n'
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
assert _model_from_frame(chunks[0]) == "claude-auto-1"
|
||||
assert chunks[1] == delta
|
||||
|
|
@ -26,6 +26,64 @@ def _owners(*backup_paths):
|
|||
|
||||
CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings"
|
||||
AUTH_MODULE = "litellm.proxy.client.cli.commands.auth"
|
||||
WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE"
|
||||
|
||||
CMD_METACHARACTERS = frozenset("&|<>^()")
|
||||
CMD_PERCENT_GUARD = "%%cd:~,%"
|
||||
|
||||
|
||||
def _through_cmd_exe(command):
|
||||
"""The line cmd.exe hands to CreateProcess after reading the apiKeyHelper.
|
||||
|
||||
A `"` toggles cmd's quote state and the metacharacters only act outside it. cmd expands
|
||||
`%VAR%` even inside quotes, so every `%` has to arrive as the `%%cd:~,%` guard: the first
|
||||
`%` has no variable name and stays literal, and `%cd:~,%` is a zero length substring of `cd`.
|
||||
"""
|
||||
assert not any(CMD_METACHARACTERS & set(run) for run in command.split('"')[::2]), command
|
||||
assert command.count("%") == 3 * command.count(CMD_PERCENT_GUARD), command
|
||||
return command.replace(CMD_PERCENT_GUARD, "%")
|
||||
|
||||
|
||||
def _through_c_runtime(command_line):
|
||||
"""argv as the Microsoft C runtime builds it for the `lite` executable.
|
||||
|
||||
Outside quotes whitespace ends an argument. A `"` toggles quoting, and inside quotes `""`
|
||||
is a literal quote. Backslashes are literal unless they run up to a `"`, where each pair
|
||||
is one backslash and an odd one left over makes the quote literal.
|
||||
"""
|
||||
argv = []
|
||||
current = None
|
||||
quoted = False
|
||||
i = 0
|
||||
while i < len(command_line):
|
||||
ch = command_line[i]
|
||||
if ch in " \t" and not quoted:
|
||||
if current is not None:
|
||||
argv.append(current)
|
||||
current = None
|
||||
i += 1
|
||||
continue
|
||||
if current is None:
|
||||
current = ""
|
||||
if ch == "\\":
|
||||
run = len(command_line[i:]) - len(command_line[i:].lstrip("\\"))
|
||||
before_quote = command_line[i + run : i + run + 1] == '"'
|
||||
current += "\\" * (run // 2 if before_quote else run)
|
||||
if before_quote and run % 2:
|
||||
current += '"'
|
||||
i += 1
|
||||
i += run
|
||||
elif ch == '"':
|
||||
if quoted and command_line[i + 1 : i + 2] == '"':
|
||||
current += '"'
|
||||
i += 1
|
||||
else:
|
||||
quoted = not quoted
|
||||
i += 1
|
||||
else:
|
||||
current += ch
|
||||
i += 1
|
||||
return argv if current is None else [*argv, current]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -199,6 +257,36 @@ class TestApiKeyHelperIsActuallyInvocable:
|
|||
|
||||
assert "Not authenticated for this server" in result.output
|
||||
|
||||
def _windows_argv(self, lite_exe, base_url):
|
||||
with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=lite_exe):
|
||||
helper = resolve_api_key_helper(base_url, platform="win32")
|
||||
return _through_c_runtime(_through_cmd_exe(helper))
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("lite_exe", "base_url"),
|
||||
[
|
||||
(WINDOWS_LITE_EXE, "http://localhost:4000"),
|
||||
("C:\\Program Files\\LiteLLM\\lite.EXE", "https://gateway.example.com/?a=1&b=2"),
|
||||
("C:\\Users\\u\\Scripts\\lite.EXE", "https://gateway.example.com/team%20a/%7Eproxy"),
|
||||
('C:\\odd "dir"\\lite.EXE', "http://localhost:4000/x\\"),
|
||||
],
|
||||
)
|
||||
def test_the_windows_command_survives_cmd_exe_and_the_c_runtime(self, lite_exe, base_url):
|
||||
assert self._windows_argv(lite_exe, base_url) == [lite_exe, "--base-url", base_url, "auth", "print-token"]
|
||||
|
||||
def test_the_windows_command_carries_the_base_url_through_cmd_quoting(self):
|
||||
stale = CliTokenRecord(
|
||||
base_url="http://other-proxy.example.com",
|
||||
key="sk-stale",
|
||||
timestamp=time.time(),
|
||||
)
|
||||
argv = self._windows_argv(WINDOWS_LITE_EXE, "http://localhost:4000")
|
||||
with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale):
|
||||
result = CliRunner().invoke(cli, argv[1:])
|
||||
|
||||
assert argv[0] == WINDOWS_LITE_EXE
|
||||
assert "Not authenticated for this server" in result.output
|
||||
|
||||
|
||||
class TestConflictingOwnersOfTheSettingsFile:
|
||||
"""Both `lite up` and `lite autoroute up` restore a backup when they stop.
|
||||
|
|
|
|||
|
|
@ -225,6 +225,32 @@ class TestResolveApiKeyHelper:
|
|||
with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"):
|
||||
resolve_api_key_helper("http://localhost:4000")
|
||||
|
||||
def test_windows_quotes_for_cmd_exe_instead_of_posix_sh(self, monkeypatch):
|
||||
"""cmd.exe takes a single quote literally, so a POSIX-quoted backslashed path is unrunnable."""
|
||||
lite_exe = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE"
|
||||
monkeypatch.setattr(shutil, "which", lambda name: lite_exe)
|
||||
|
||||
helper = resolve_api_key_helper("https://gateway.example.com", platform="win32")
|
||||
|
||||
assert helper == f'"{lite_exe}" "--base-url" "https://gateway.example.com" "auth" "print-token"'
|
||||
|
||||
def test_windows_keeps_a_spaced_path_and_a_metacharacter_url_as_single_tokens(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "C:\\Program Files\\LiteLLM\\lite.EXE")
|
||||
|
||||
helper = resolve_api_key_helper("https://gateway.example.com/?a=1&b=2", platform="win32")
|
||||
|
||||
assert helper == (
|
||||
'"C:\\Program Files\\LiteLLM\\lite.EXE" "--base-url" "https://gateway.example.com/?a=1&b=2" '
|
||||
'"auth" "print-token"'
|
||||
)
|
||||
|
||||
def test_non_windows_platforms_keep_posix_quoting(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite")
|
||||
|
||||
helper = resolve_api_key_helper("http://example.com/path; rm -rf /", platform="darwin")
|
||||
|
||||
assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token"
|
||||
|
||||
|
||||
def _make_ctx(base_url):
|
||||
return click.Context(click.Command("test"), obj={"base_url": base_url})
|
||||
|
|
|
|||
|
|
@ -5592,6 +5592,156 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
|
|||
assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail"
|
||||
|
||||
|
||||
def _responses_stream_events() -> list:
|
||||
from litellm.types.llms.openai import (
|
||||
OutputTextDeltaEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
deltas = [
|
||||
OutputTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
item_id="msg_lit6457",
|
||||
output_index=0,
|
||||
content_index=0,
|
||||
delta=part,
|
||||
)
|
||||
for part in ("Hello", " world")
|
||||
]
|
||||
completed = ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_lit6457",
|
||||
created_at=1234567890,
|
||||
model="gpt-4o",
|
||||
object="response",
|
||||
status="completed",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_lit6457",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Hello world"}],
|
||||
}
|
||||
],
|
||||
),
|
||||
)
|
||||
return [*deltas, completed]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_api_stream_scans_output_and_replays_buffered_events():
|
||||
"""Streamed /v1/responses events must be scanned via the unified translation
|
||||
layer, not fed to stream_chunk_builder (which raises APIError on them)."""
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="bedrock-responses-stream",
|
||||
guardrailIdentifier="test-id",
|
||||
guardrailVersion="DRAFT",
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
default_on=True,
|
||||
)
|
||||
stream_events = _responses_stream_events()
|
||||
order = []
|
||||
yielded = []
|
||||
|
||||
async def record_scan(*args, **kwargs):
|
||||
order.append("scan")
|
||||
return {"action": "NONE", "assessments": [], "outputs": []}
|
||||
|
||||
async def mock_stream():
|
||||
for event in stream_events:
|
||||
yield event
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)):
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"),
|
||||
response=mock_stream(),
|
||||
request_data={"model": "gpt-4o", "input": "hi"},
|
||||
):
|
||||
order.append("chunk")
|
||||
yielded.append(chunk)
|
||||
|
||||
assert order == ["scan", "chunk", "chunk", "chunk"]
|
||||
assert len(yielded) == len(stream_events)
|
||||
assert all(emitted is original for emitted, original in zip(yielded, stream_events))
|
||||
|
||||
|
||||
def _responses_failed_stream_events() -> list:
|
||||
from litellm.types.llms.openai import (
|
||||
OutputTextDeltaEvent,
|
||||
ResponseFailedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
deltas = [
|
||||
OutputTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
item_id="msg_lit6457_failed",
|
||||
output_index=0,
|
||||
content_index=0,
|
||||
delta=part,
|
||||
)
|
||||
for part in ("Hello", " world")
|
||||
]
|
||||
failed = ResponseFailedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_lit6457_failed",
|
||||
created_at=1234567890,
|
||||
model="gpt-4o",
|
||||
object="response",
|
||||
status="failed",
|
||||
output=[],
|
||||
),
|
||||
)
|
||||
return [*deltas, failed]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_api_failed_stream_scans_delta_text_before_replay():
|
||||
"""A responses stream that dies mid-generation carries its text only in delta
|
||||
events; the end-of-stream scan must still see that text instead of skipping
|
||||
on an empty assembled string and replaying the buffer unmoderated."""
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="bedrock-responses-failed-stream",
|
||||
guardrailIdentifier="test-id",
|
||||
guardrailVersion="DRAFT",
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
default_on=True,
|
||||
)
|
||||
stream_events = _responses_failed_stream_events()
|
||||
order = []
|
||||
scan_payloads = []
|
||||
yielded = []
|
||||
|
||||
async def record_scan(*args, **kwargs):
|
||||
order.append("scan")
|
||||
scan_payloads.append(str(args) + str(kwargs))
|
||||
return {"action": "NONE", "assessments": [], "outputs": []}
|
||||
|
||||
async def mock_stream():
|
||||
for event in stream_events:
|
||||
yield event
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)):
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"),
|
||||
response=mock_stream(),
|
||||
request_data={"model": "gpt-4o", "input": "hi"},
|
||||
):
|
||||
order.append("chunk")
|
||||
yielded.append(chunk)
|
||||
|
||||
assert order == ["scan", "chunk", "chunk", "chunk"]
|
||||
assert "Hello world" in scan_payloads[0]
|
||||
assert len(yielded) == len(stream_events)
|
||||
assert all(emitted is original for emitted, original in zip(yielded, stream_events))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_debug_log_masks_signed_request_headers():
|
||||
import logging
|
||||
|
|
|
|||
|
|
@ -17489,3 +17489,49 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project
|
|||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"]
|
||||
|
||||
|
||||
def test_generate_key_request_blank_team_id_is_personal():
|
||||
"""The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925)."""
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_is_team_key,
|
||||
)
|
||||
|
||||
cleared = GenerateKeyRequest(team_id="")
|
||||
assert cleared.team_id is None
|
||||
assert _is_team_key(data=cleared) is False
|
||||
assert RegenerateKeyRequest(team_id="").team_id is None
|
||||
assert GenerateKeyRequest(team_id="team-1").team_id == "team-1"
|
||||
|
||||
|
||||
def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch):
|
||||
"""key_generation_check with team_id="" must take the personal-key path instead
|
||||
of failing the team lookup with "Unable to find team object" (LIT-3925)."""
|
||||
from litellm.proxy._types import KeyManagementRoutes
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
key_generation_check,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"key_generation_settings",
|
||||
{
|
||||
"team_key_generation": {"allowed_team_member_roles": ["admin"]},
|
||||
"personal_key_generation": {"allowed_user_roles": ["proxy_admin", "internal_user"]},
|
||||
},
|
||||
)
|
||||
|
||||
assert (
|
||||
key_generation_check(
|
||||
team_table=None,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-alice",
|
||||
user_id="alice",
|
||||
),
|
||||
data=GenerateKeyRequest(key_alias="personal", team_id=""),
|
||||
route=KeyManagementRoutes.KEY_GENERATE,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
|
|
|||
|
|
@ -9,12 +9,13 @@ Pins (PR2):
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.router_utils import pattern_match_deployments
|
||||
|
||||
from .conftest import normalize # type: ignore[import-not-found]
|
||||
|
||||
|
|
@ -99,6 +100,7 @@ def test_token_counter_missing_input_returns_400(
|
|||
|
||||
@pytest.fixture
|
||||
def patched_supported_params(monkeypatch):
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"get_llm_provider",
|
||||
|
|
@ -124,12 +126,104 @@ def test_supported_openai_params_happy_path(client, auth_as, patched_supported_p
|
|||
}
|
||||
|
||||
|
||||
def test_supported_openai_params_resolves_router_alias(client, auth_as, monkeypatch):
|
||||
"""A router alias absent from the cost map resolves through the deployment's underlying model."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "claude-opus-4-6-cached",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"},
|
||||
}
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
|
||||
with auth_as():
|
||||
response = client.get("/utils/supported_openai_params", params={"model": "claude-opus-4-6-cached"})
|
||||
|
||||
assert response.status_code == 200
|
||||
expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic")
|
||||
assert response.json() == {"supported_openai_params": expected}
|
||||
assert "max_tokens" in response.json()["supported_openai_params"]
|
||||
|
||||
|
||||
def test_supported_openai_params_declared_prefix_alias_resolves_through_router(client, auth_as, monkeypatch):
|
||||
"""Regression: an alias whose name starts with an authenticating provider's prefix skipped
|
||||
router resolution and answered with that provider's params instead of the deployment's."""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "github_copilot/gpt-4o",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"},
|
||||
}
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
|
||||
with auth_as():
|
||||
response = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"})
|
||||
|
||||
assert response.status_code == 200
|
||||
expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic")
|
||||
assert response.json() == {"supported_openai_params": expected}
|
||||
|
||||
|
||||
def test_supported_openai_params_never_runs_oauth_for_authenticating_providers(client, auth_as, monkeypatch, tmp_path):
|
||||
"""Regression: github_copilot/chatgpt names answer from their declaration; resolving them
|
||||
through ``get_llm_provider`` would run the provider's OAuth device flow and block the event loop."""
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
|
||||
(tmp_path / "access-token").write_text("fake-access-token")
|
||||
(tmp_path / "api-key.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"token": "fake-api-key",
|
||||
"expires_at": 4102444800,
|
||||
"endpoints": {"api": "https://api.githubcopilot.com"},
|
||||
}
|
||||
)
|
||||
)
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "copilot-alias",
|
||||
"litellm_params": {"model": "github_copilot/gpt-4o"},
|
||||
},
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*"},
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
|
||||
resolution_attempts: list[str] = []
|
||||
|
||||
def _oauth_tripwire(model, *args, **kwargs):
|
||||
resolution_attempts.append(model)
|
||||
raise AssertionError("get_llm_provider would run the OAuth device flow")
|
||||
|
||||
monkeypatch.setattr(litellm, "get_llm_provider", _oauth_tripwire)
|
||||
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire)
|
||||
expected = litellm.get_supported_openai_params(model="gpt-4o", custom_llm_provider="github_copilot")
|
||||
|
||||
with auth_as():
|
||||
via_alias = client.get("/utils/supported_openai_params", params={"model": "copilot-alias"})
|
||||
via_direct_name = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"})
|
||||
|
||||
assert via_alias.status_code == 200
|
||||
assert via_alias.json() == {"supported_openai_params": expected}
|
||||
assert via_direct_name.status_code == 200
|
||||
assert via_direct_name.json() == {"supported_openai_params": expected}
|
||||
assert resolution_attempts == []
|
||||
|
||||
|
||||
def test_supported_openai_params_invalid_model(client, auth_as, monkeypatch):
|
||||
"""Pins ``GET /utils/supported_openai_params`` (error: unknown model)."""
|
||||
|
||||
def _raise(model):
|
||||
raise Exception("unknown")
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm, "get_llm_provider", _raise)
|
||||
with auth_as():
|
||||
response = client.get("/utils/supported_openai_params", params={"model": "??"})
|
||||
|
|
|
|||
|
|
@ -819,6 +819,94 @@ async def test_should_cap_known_estimate_to_remaining_budget(
|
|||
) == pytest.approx(0.9)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_rejects_known_estimate_exceeding_remaining_budget(
|
||||
spend_counter_state,
|
||||
):
|
||||
"""LIT-5922: with strict enforcement on, a request whose known estimate does
|
||||
not fit the remaining budget must be rejected before dispatch instead of
|
||||
having its reservation shrunk to the headroom and admitted, and the counter
|
||||
must be restored to the pre-request spend."""
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="key-budget-known-estimate-fail-closed",
|
||||
spend=0.9,
|
||||
max_budget=1.0,
|
||||
)
|
||||
counter_cache.in_memory_cache.set_cache(
|
||||
key="spend:key:key-budget-known-estimate-fail-closed",
|
||||
value=0.9,
|
||||
)
|
||||
|
||||
with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute
|
||||
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
|
||||
return_value=0.6,
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await reserve_budget_for_request(
|
||||
request_body=_request_body(),
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
fail_closed_budget_enforcement=True,
|
||||
)
|
||||
|
||||
assert exc_info.value.current_cost == pytest.approx(0.9)
|
||||
assert exc_info.value.max_budget == pytest.approx(1.0)
|
||||
assert "Current cost: 0.9, Estimated request cost: 0.6, Max budget: 1.0" in str(exc_info.value)
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:key:key-budget-known-estimate-fail-closed"
|
||||
) == pytest.approx(0.9)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_tolerates_float_noise_when_estimate_exactly_fits(
|
||||
spend_counter_state,
|
||||
):
|
||||
"""0.1 + 0.2 lands a hair above 0.3 in floating point. Strict enforcement
|
||||
must treat that as fitting the budget, not reject it."""
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="key-budget-fail-closed-float-noise",
|
||||
spend=0.1,
|
||||
max_budget=0.3,
|
||||
)
|
||||
counter_cache.in_memory_cache.set_cache(
|
||||
key="spend:key:key-budget-fail-closed-float-noise",
|
||||
value=0.1,
|
||||
)
|
||||
|
||||
with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute
|
||||
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
|
||||
return_value=0.2,
|
||||
):
|
||||
reservation = await reserve_budget_for_request(
|
||||
request_body=_request_body(),
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
fail_closed_budget_enforcement=True,
|
||||
)
|
||||
|
||||
assert reservation is not None
|
||||
assert reservation["reserved_cost"] == pytest.approx(0.2)
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:key:key-budget-fail-closed-float-noise"
|
||||
) == pytest.approx(0.3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_clamp_reservation_to_default_when_output_cap_missing(
|
||||
spend_counter_state,
|
||||
|
|
|
|||
|
|
@ -10445,6 +10445,75 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch):
|
|||
assert before["some_api_key"] != "sk-stored-secret"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_field_rejects_out_of_range_alerting_args(monkeypatch):
|
||||
"""Out-of-range alerting_args must be rejected at save time. If they land in the
|
||||
DB, SlackAlertingArgs raises during the config reload and alerting breaks."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import ConfigFieldUpdate
|
||||
from litellm.proxy.proxy_server import update_config_general_settings
|
||||
|
||||
monkeypatch.setattr(proxy_server_module, "prisma_client", MagicMock())
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
api_key="hashed-admin",
|
||||
user_id="admin-1",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await update_config_general_settings(
|
||||
data=ConfigFieldUpdate(
|
||||
field_name="alerting_args",
|
||||
field_value={
|
||||
"daily_spend_per_user_threshold": -5.0,
|
||||
"user_spend_check_interval": 20,
|
||||
},
|
||||
config_type="general_settings",
|
||||
),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
error_msg = exc_info.value.detail["error"]
|
||||
assert "daily_spend_per_user_threshold" in error_msg
|
||||
assert "user_spend_check_interval" in error_msg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_field_accepts_valid_alerting_args(monkeypatch):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import ConfigFieldUpdate
|
||||
from litellm.proxy.proxy_server import update_config_general_settings
|
||||
|
||||
fake = _fake_prisma_with_config({})
|
||||
monkeypatch.setattr(proxy_server_module, "prisma_client", fake)
|
||||
monkeypatch.setattr(litellm, "store_audit_logs", False)
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
api_key="hashed-admin",
|
||||
user_id="admin-1",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
await update_config_general_settings(
|
||||
data=ConfigFieldUpdate(
|
||||
field_name="alerting_args",
|
||||
field_value={
|
||||
"daily_spend_per_user_threshold": 5.0,
|
||||
"user_spend_check_interval": 60,
|
||||
},
|
||||
config_type="general_settings",
|
||||
),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
written = json.loads(fake.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"])
|
||||
assert written["alerting_args"]["daily_spend_per_user_threshold"] == 5.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_general_settings_applies_ssrf_globals(monkeypatch):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
|
|
|
|||
|
|
@ -598,6 +598,45 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh
|
|||
assert filtered == healthy_deployments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_deployments_does_not_pin_when_target_order_is_set():
|
||||
user_key = "user-key-order-fallback"
|
||||
stable_model_map_key = "claude-sonnet-4-5@20250929"
|
||||
cache = AsyncMock()
|
||||
cache.async_get_cache = AsyncMock(return_value={"model_id": "deployment-1"})
|
||||
callback = DeploymentAffinityCheck(
|
||||
cache=cache,
|
||||
ttl_seconds=123,
|
||||
enable_user_key_affinity=True,
|
||||
enable_responses_api_affinity=False,
|
||||
)
|
||||
healthy_deployments = [
|
||||
{
|
||||
"model_name": stable_model_map_key,
|
||||
"litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": stable_model_map_key,
|
||||
"litellm_params": {
|
||||
"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"
|
||||
},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
]
|
||||
|
||||
filtered = await callback.async_filter_deployments(
|
||||
model="some-router-model-group",
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=None,
|
||||
request_kwargs={"_target_order": 2, "metadata": {"user_api_key_hash": user_key}},
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
assert filtered == healthy_deployments
|
||||
cache.async_get_cache.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_user_key_affinity_ttl_expiry_allows_reroute():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -150,6 +150,25 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum():
|
|||
assert filtered == [deployments[1]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_deployments_does_not_pin_when_target_order_is_set():
|
||||
cache = DualCache()
|
||||
check = PromptCachingDeploymentCheck(cache=cache)
|
||||
deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6")
|
||||
messages = _messages(word_count=5000)
|
||||
|
||||
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None)
|
||||
|
||||
filtered = await check.async_filter_deployments(
|
||||
model=MODEL_GROUP_ALIAS,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages,
|
||||
request_kwargs={"_target_order": 2},
|
||||
)
|
||||
|
||||
assert filtered == deployments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -11,7 +11,10 @@ from litellm.router_utils.fallback_event_handlers import (
|
|||
AttemptedFallbackTargets,
|
||||
_trigger_cooldown_for_failed_deployment,
|
||||
fallback_attempt_key,
|
||||
clear_pre_routing_selection,
|
||||
get_fallback_model_group,
|
||||
get_pre_routing_selection,
|
||||
record_pre_routing_selection,
|
||||
run_async_fallback,
|
||||
)
|
||||
|
||||
|
|
@ -1090,3 +1093,119 @@ async def test_run_async_fallback_preserves_original_model_group_on_nested_fallb
|
|||
metadata = router.received_kwargs["metadata"]
|
||||
assert metadata["attempted_fallbacks"] == 2
|
||||
assert metadata["original_model_group"] == "primary-model"
|
||||
|
||||
|
||||
class TestPreRoutingSelectionCarriesToFallbacks:
|
||||
"""#38832: a complexity/auto router picks a tier behind the router name, but fallback
|
||||
lookup kept using the router name, so the tier's configured chain never ran."""
|
||||
|
||||
def test_selection_is_recorded_in_the_metadata_bucket(self):
|
||||
kwargs = {"model": "smart-router", "metadata": {}}
|
||||
record_pre_routing_selection(kwargs, "tier1")
|
||||
assert kwargs["metadata"]["pre_routing_selected_model"] == "tier1"
|
||||
assert get_pre_routing_selection(kwargs) == "tier1"
|
||||
|
||||
def test_selection_is_recorded_in_the_litellm_metadata_bucket(self):
|
||||
kwargs = {"model": "smart-router", "litellm_metadata": {}}
|
||||
record_pre_routing_selection(kwargs, "tier2")
|
||||
assert get_pre_routing_selection(kwargs) == "tier2"
|
||||
|
||||
def test_a_bucket_survives_the_kwargs_copy_that_fallbacks_run_on(self):
|
||||
"""The bucket is shared by reference, which is the whole reason this works."""
|
||||
outer = {"model": "smart-router", "metadata": {}}
|
||||
inner = {**outer}
|
||||
record_pre_routing_selection(inner, "tier1")
|
||||
assert get_pre_routing_selection(outer) == "tier1"
|
||||
|
||||
def test_no_selection_reads_as_none(self):
|
||||
assert get_pre_routing_selection({"model": "smart-router", "metadata": {}}) is None
|
||||
assert get_pre_routing_selection({"model": "smart-router"}) is None
|
||||
|
||||
def test_missing_kwargs_is_a_no_op(self):
|
||||
"""A caller with no kwargs must not raise, and must not leak the selection anywhere."""
|
||||
record_pre_routing_selection(None, "tier1")
|
||||
|
||||
assert get_pre_routing_selection({}) is None
|
||||
|
||||
def test_a_non_dict_bucket_is_ignored(self):
|
||||
kwargs = {"model": "smart-router", "metadata": "not-a-dict"}
|
||||
record_pre_routing_selection(kwargs, "tier1")
|
||||
assert get_pre_routing_selection(kwargs) is None
|
||||
|
||||
def test_fallbacks_resolve_against_the_selected_tier(self):
|
||||
"""The lookup the router performs, keyed on the tier rather than the router name."""
|
||||
fallbacks = [{"tier1": ["backup-a", "backup-b"]}, {"tier2": ["backup-c"]}]
|
||||
assert get_fallback_model_group(fallbacks=fallbacks, model_group="tier1")[0] == ["backup-a", "backup-b"]
|
||||
assert get_fallback_model_group(fallbacks=fallbacks, model_group="smart-router")[0] is None
|
||||
|
||||
|
||||
class TestPreRoutingSelectionIsPerHop:
|
||||
"""#38832 review: the buckets also carry whatever the caller sent, and a fallback hop
|
||||
inherits the previous hop's tier, so a hop must start without a selection."""
|
||||
|
||||
def test_a_caller_supplied_selection_is_dropped(self):
|
||||
kwargs = {"model": "plain", "metadata": {"pre_routing_selected_model": "tier1"}}
|
||||
|
||||
clear_pre_routing_selection(kwargs)
|
||||
|
||||
assert get_pre_routing_selection(kwargs) is None
|
||||
assert "pre_routing_selected_model" not in kwargs["metadata"]
|
||||
|
||||
def test_both_buckets_are_cleared(self):
|
||||
kwargs = {
|
||||
"metadata": {"pre_routing_selected_model": "tier1"},
|
||||
"litellm_metadata": {"pre_routing_selected_model": "tier2"},
|
||||
}
|
||||
|
||||
clear_pre_routing_selection(kwargs)
|
||||
|
||||
assert get_pre_routing_selection(kwargs) is None
|
||||
|
||||
def test_the_rest_of_the_bucket_is_left_alone(self):
|
||||
kwargs = {"metadata": {"pre_routing_selected_model": "tier1", "tags": ["a"]}}
|
||||
|
||||
clear_pre_routing_selection(kwargs)
|
||||
|
||||
assert kwargs["metadata"] == {"tags": ["a"]}
|
||||
|
||||
def test_clearing_is_a_no_op_without_a_usable_bucket(self):
|
||||
kwargs = {"model": "plain", "metadata": "not-a-dict"}
|
||||
|
||||
clear_pre_routing_selection(None)
|
||||
clear_pre_routing_selection(kwargs)
|
||||
|
||||
assert kwargs == {"model": "plain", "metadata": "not-a-dict"}
|
||||
|
||||
def test_a_selection_recorded_after_clearing_is_kept(self):
|
||||
"""Clearing runs before routing, so the hook's own write must survive it."""
|
||||
kwargs = {"model": "smart-router", "metadata": {"pre_routing_selected_model": "stale"}}
|
||||
|
||||
clear_pre_routing_selection(kwargs)
|
||||
record_pre_routing_selection(kwargs, "tier1")
|
||||
|
||||
assert get_pre_routing_selection(kwargs) == "tier1"
|
||||
|
||||
|
||||
class TestOrderedFallbackLookupGroups:
|
||||
def test_tier_first_then_requested_group_deduped(self):
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
PRE_ROUTING_SELECTED_MODEL_KEY,
|
||||
fallback_lookup_groups,
|
||||
)
|
||||
|
||||
kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier1"}}
|
||||
assert fallback_lookup_groups(kwargs, "smart-router") == ("tier1", "smart-router")
|
||||
assert fallback_lookup_groups(kwargs, "tier1") == ("tier1",)
|
||||
assert fallback_lookup_groups({}, "smart-router") == ("smart-router",)
|
||||
assert fallback_lookup_groups({}, None) == ()
|
||||
|
||||
def test_first_resolving_group_wins_and_generic_idx_survives_a_miss(self):
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
get_fallback_model_group_for_lookup_groups,
|
||||
)
|
||||
|
||||
fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}, {"*": ["backup-c"]}]
|
||||
assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier1", "smart-router")) == (["backup-a"], None)
|
||||
assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "smart-router")) == (["backup-b"], None)
|
||||
assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "no-such")) == (["backup-c"], 2)
|
||||
assert get_fallback_model_group_for_lookup_groups([{"tier1": ["backup-a"]}], ("no", "nope")) == (None, None)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,78 @@
|
|||
"""Behavior pins for ``litellm/router_utils/pattern_match_deployments.py``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from litellm.router_utils import pattern_match_deployments
|
||||
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
|
||||
|
||||
|
||||
def _wildcard_deployment(model_name: str) -> dict:
|
||||
return {"model_name": model_name, "litellm_params": {"model": model_name}}
|
||||
|
||||
|
||||
def _matched_models(matches: list[dict] | None) -> list[str]:
|
||||
return [deployment["litellm_params"]["model"] for deployment in matches or []]
|
||||
|
||||
|
||||
def test_get_pattern_never_resolves_declared_authenticating_providers(monkeypatch):
|
||||
"""Regression: resolving a github_copilot/chatgpt name through ``get_llm_provider`` runs the
|
||||
provider's OAuth device flow; the auth layer walks every wildcard router on every request, so
|
||||
a single metadata lookup for an unserved name would block the proxy's event loop."""
|
||||
resolution_attempts: list[str] = []
|
||||
|
||||
def _oauth_tripwire(model, *args, **kwargs):
|
||||
resolution_attempts.append(model)
|
||||
raise AssertionError("get_llm_provider would run the OAuth device flow")
|
||||
|
||||
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire)
|
||||
|
||||
unmatched_router = PatternMatchRouter()
|
||||
unmatched_router.add_pattern("anthropic/*", _wildcard_deployment("anthropic/*"))
|
||||
assert unmatched_router.get_pattern("github_copilot/gpt-4o") is None
|
||||
|
||||
matched_router = PatternMatchRouter()
|
||||
matched_router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*"))
|
||||
assert _matched_models(matched_router.get_pattern("github_copilot/gpt-4o")) == ["github_copilot/gpt-4o"]
|
||||
assert _matched_models(matched_router.get_pattern("gpt-4o", custom_llm_provider="github_copilot")) == [
|
||||
"github_copilot/gpt-4o"
|
||||
]
|
||||
|
||||
assert resolution_attempts == []
|
||||
|
||||
|
||||
def test_get_pattern_bare_provider_name_never_matches_that_providers_wildcard(monkeypatch):
|
||||
"""Regression: a bare ``github_copilot`` adopted itself as its provider and retried as
|
||||
``github_copilot/github_copilot``, false-matching the wildcard for a name no deployment serves."""
|
||||
|
||||
def _unknown_provider(model, *args, **kwargs):
|
||||
raise ValueError(f"unknown provider for {model}")
|
||||
|
||||
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider)
|
||||
router = PatternMatchRouter()
|
||||
router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*"))
|
||||
assert router.get_pattern("github_copilot") is None
|
||||
|
||||
|
||||
def test_get_pattern_missing_model_returns_none(monkeypatch):
|
||||
"""Regression: a request without a model reaches the auth layer's pattern walk as ``None``; the
|
||||
declared-provider guard raised ``TypeError`` where the old inline resolve swallowed every
|
||||
resolver error, so the proxy's missing-model 400 became a crash."""
|
||||
|
||||
def _unknown_provider(model, *args, **kwargs):
|
||||
raise ValueError(f"unknown provider for {model}")
|
||||
|
||||
monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider)
|
||||
router = PatternMatchRouter()
|
||||
router.add_pattern("openai/*", _wildcard_deployment("openai/*"))
|
||||
assert router.get_pattern(None) is None
|
||||
|
||||
|
||||
def test_get_pattern_still_resolves_unqualified_names(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
pattern_match_deployments,
|
||||
"get_llm_provider",
|
||||
lambda model, **kwargs: (model, "openai", None, None),
|
||||
)
|
||||
router = PatternMatchRouter()
|
||||
router.add_pattern("openai/*", _wildcard_deployment("openai/*"))
|
||||
assert _matched_models(router.get_pattern("gpt-4o")) == ["openai/gpt-4o"]
|
||||
52
tests/test_litellm/test_dockerfile_apk_repository.py
Normal file
52
tests/test_litellm/test_dockerfile_apk_repository.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
"""
|
||||
Static checks on the root Dockerfile's apk repository configuration.
|
||||
|
||||
The base image (cgr.dev/chainguard/wolfi-base) only configures the
|
||||
authenticated Chainguard apk repo (https://apk.cgr.dev/chainguard) in
|
||||
/etc/apk/repositories, which requires a Chainguard enterprise subscription.
|
||||
Anyone pulling the published litellm image and running `apk add` inside it
|
||||
hits SSL/auth failures with no fallback repo configured, so nothing can be
|
||||
installed. See https://github.com/BerriAI/litellm/issues/33518
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
import pytest
|
||||
|
||||
DOCKERFILE_PATH = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
"Dockerfile",
|
||||
)
|
||||
|
||||
|
||||
def _runtime_stage(dockerfile_text: str) -> str:
|
||||
"""Return the contents of the final `FROM ... AS runtime` build stage."""
|
||||
match = re.search(r"^FROM .*\bAS runtime\b(.*)\Z", dockerfile_text, re.MULTILINE | re.DOTALL)
|
||||
assert match, "Dockerfile has no `FROM ... AS runtime` stage"
|
||||
return match.group(1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.path.exists(DOCKERFILE_PATH),
|
||||
reason="Dockerfile not present in this checkout",
|
||||
)
|
||||
def test_runtime_stage_adds_public_wolfi_repo():
|
||||
"""The runtime stage must add the public Wolfi apk repo so `apk add`
|
||||
works for users without a Chainguard enterprise subscription."""
|
||||
with open(DOCKERFILE_PATH, "r", encoding="utf-8") as f:
|
||||
contents = f.read()
|
||||
|
||||
runtime_stage = _runtime_stage(contents)
|
||||
|
||||
assert re.search(
|
||||
r"echo\s+[\"']?https://packages\.wolfi\.dev/os[\"']?\s*>>\s*/etc/apk/repositories",
|
||||
runtime_stage,
|
||||
), (
|
||||
"Runtime stage must append the public Wolfi apk repo "
|
||||
'(RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories) '
|
||||
"so `apk add` works without Chainguard enterprise credentials. "
|
||||
"See https://github.com/BerriAI/litellm/issues/33518"
|
||||
)
|
||||
|
|
@ -11595,3 +11595,138 @@ class TestTierParamsTheTargetAccepts:
|
|||
accepted = router._tier_params_the_target_accepts("no-such-group", {"reasoning_effort": "max"}, {})
|
||||
|
||||
assert accepted == {"reasoning_effort": "max"}
|
||||
|
||||
|
||||
class TestPreRoutingTierDrivesFallbacks:
|
||||
"""#38832: a complexity/auto router picks a tier behind the router name, but fallback
|
||||
lookup stayed on the router name, so the tier's configured chain never ran and a
|
||||
provider failure on the tier's first hop was returned to the client."""
|
||||
|
||||
class _TierRouter(litellm.Router):
|
||||
async def async_pre_routing_hook(
|
||||
self, model, request_kwargs, messages=None, input=None, specific_deployment=False
|
||||
):
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
if model == "smart-router":
|
||||
return PreRoutingHookResponse(model="tier1", messages=messages)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _router(cls, fallbacks) -> "litellm.Router":
|
||||
return cls._TierRouter(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"},
|
||||
},
|
||||
{
|
||||
"model_name": "tier1",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-x",
|
||||
"mock_response": "litellm.RateLimitError",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "backup-a",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-x",
|
||||
"mock_response": "from backup-a",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "backup-b",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-x",
|
||||
"mock_response": "from backup-b",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "failing-backup",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-x",
|
||||
"mock_response": "litellm.RateLimitError",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "plain",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "sk-x",
|
||||
"mock_response": "litellm.RateLimitError",
|
||||
},
|
||||
},
|
||||
],
|
||||
fallbacks=fallbacks,
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_selected_tier_fallback_chain_runs(self):
|
||||
router = self._router([{"tier1": ["backup-a"]}])
|
||||
|
||||
response = await router.acompletion(
|
||||
model="smart-router", messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "from backup-a"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_chain_keyed_on_the_router_name_is_not_used(self):
|
||||
"""The router name has no chain of its own, so nothing should rescue this call."""
|
||||
router = self._router([{"tier2": ["backup-a"]}])
|
||||
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_chain_keyed_on_the_router_name_rescues_when_no_tier_chain_exists(self):
|
||||
"""The documented contract: configs keyed on the requested name keep working behind auto-routers."""
|
||||
router = self._router([{"smart-router": ["backup-a"]}])
|
||||
|
||||
response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert response.choices[0].message.content == "from backup-a"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_tier_chain_wins_over_the_router_name_chain(self):
|
||||
router = self._router([{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}])
|
||||
|
||||
response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert response.choices[0].message.content == "from backup-a"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self):
|
||||
router = self._router([{"tier1": ["backup-a"]}])
|
||||
|
||||
response = await router.acompletion(
|
||||
model="tier1", messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "from backup-a"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_caller_cannot_pick_the_chain_by_sending_the_selection(self):
|
||||
"""The metadata bucket carries caller-supplied keys, so only the hook may set the tier."""
|
||||
router = self._router([{"tier1": ["backup-a"]}])
|
||||
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(
|
||||
model="plain",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
metadata={"pre_routing_selected_model": "tier1"},
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_fallback_hop_resolves_its_own_chain(self):
|
||||
"""The second hop must key off the group it is running, not the tier that failed."""
|
||||
router = self._router([{"tier1": ["failing-backup"]}, {"failing-backup": ["backup-b"]}])
|
||||
|
||||
response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert response.choices[0].message.content == "from backup-b"
|
||||
|
|
|
|||
|
|
@ -6,12 +6,19 @@ should be tried first, and higher order deployments should be used as fallbacks
|
|||
when lower order deployments fail.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
import json
|
||||
from typing import Final, Optional
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.utils import _get_order_filtered_deployments
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_utils.prompt_caching_cache import PromptCachingCache
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
from litellm.utils import _get_deployment_order, _get_order_filtered_deployments
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for _get_order_filtered_deployments
|
||||
|
|
@ -49,13 +56,22 @@ class TestGetOrderFilteredDeployments:
|
|||
assert len(result) == 1
|
||||
assert result[0]["model_info"]["id"] == "b"
|
||||
|
||||
def test_target_order_no_match_returns_all(self):
|
||||
def test_target_order_no_match_returns_empty(self):
|
||||
deps = [
|
||||
self._make_deployment(1, "a"),
|
||||
self._make_deployment(2, "b"),
|
||||
]
|
||||
result = _get_order_filtered_deployments(deps, target_order=99)
|
||||
assert len(result) == 2
|
||||
assert result == []
|
||||
|
||||
def test_target_order_no_match_does_not_reselect_lower_order(self):
|
||||
deps = [
|
||||
self._make_deployment(1, "a"),
|
||||
self._make_deployment(2, "b"),
|
||||
]
|
||||
remaining_after_pre_call = [deps[0]]
|
||||
result = _get_order_filtered_deployments(remaining_after_pre_call, target_order=2)
|
||||
assert result == []
|
||||
|
||||
def test_no_order_set_returns_all(self):
|
||||
deps = [
|
||||
|
|
@ -406,35 +422,239 @@ async def test_router_order_fallback_with_hidden_model_group_alias():
|
|||
assert response._hidden_params["model_id"] == "2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_order_fallback_does_not_reselect_order_1_when_order_2_is_filtered_out():
|
||||
class _DropOrder2(CustomLogger):
|
||||
async def async_filter_deployments(
|
||||
self, model, healthy_deployments, messages, request_kwargs=None, parent_otel_span=None
|
||||
):
|
||||
return [d for d in healthy_deployments if _get_deployment_order(d) != 2]
|
||||
|
||||
drop_order_2: Final = _DropOrder2()
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "key",
|
||||
"mock_response": "litellm.RateLimitError",
|
||||
"order": 1,
|
||||
},
|
||||
"model_info": {"id": "1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "key",
|
||||
"mock_response": "success from order 2",
|
||||
"order": 2,
|
||||
},
|
||||
"model_info": {"id": "2"},
|
||||
},
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
litellm.callbacks.append(drop_order_2)
|
||||
try:
|
||||
with pytest.raises(RouterRateLimitError, match="No deployments available") as exc_info:
|
||||
await router.acompletion(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
assert "success from order 2" not in str(exc_info.value)
|
||||
finally:
|
||||
litellm.callbacks.remove(drop_order_2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_order_fallback_ignores_prompt_cache_pin_on_target_order():
|
||||
messages = [{"role": "user", "content": "word " * 5000}]
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "bad",
|
||||
"mock_response": Exception("azure peak load"),
|
||||
"order": 1,
|
||||
},
|
||||
"model_info": {"id": "1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "good",
|
||||
"mock_response": "success from order 2",
|
||||
"order": 2,
|
||||
},
|
||||
"model_info": {"id": "2"},
|
||||
},
|
||||
],
|
||||
num_retries=0,
|
||||
optional_pre_call_checks=["prompt_caching"],
|
||||
)
|
||||
await PromptCachingCache(cache=router.cache).async_add_model_id(
|
||||
model_id="1",
|
||||
messages=messages,
|
||||
tools=None,
|
||||
)
|
||||
response = await router.acompletion(model="test-model", messages=messages)
|
||||
assert response._hidden_params["model_id"] == "2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_order_fallback_retries_keep_target_order():
|
||||
seen_target_orders: Final = []
|
||||
|
||||
class _RecordTargetOrder(CustomLogger):
|
||||
async def async_filter_deployments(
|
||||
self, model, healthy_deployments, messages, request_kwargs=None, parent_otel_span=None
|
||||
):
|
||||
seen_target_orders.append((request_kwargs or {}).get("_target_order"))
|
||||
return healthy_deployments
|
||||
|
||||
recorder: Final = _RecordTargetOrder()
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "bad",
|
||||
"mock_response": Exception("fail order 1"),
|
||||
"order": 1,
|
||||
},
|
||||
"model_info": {"id": "1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-4o",
|
||||
"api_key": "bad",
|
||||
"mock_response": Exception("fail order 2"),
|
||||
"order": 2,
|
||||
},
|
||||
"model_info": {"id": "2"},
|
||||
},
|
||||
],
|
||||
num_retries=1,
|
||||
)
|
||||
litellm.callbacks.append(recorder)
|
||||
try:
|
||||
with pytest.raises(Exception, match="fail order 2"):
|
||||
await router.acompletion(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
finally:
|
||||
litellm.callbacks.remove(recorder)
|
||||
assert seen_target_orders.count(2) >= 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_api_call_strips_target_order_from_provider_kwargs():
|
||||
captured: Final = {}
|
||||
|
||||
async def _fake_provider(**provider_kwargs):
|
||||
captured.update(provider_kwargs)
|
||||
return "ok"
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-4o", "api_key": "key", "order": 2},
|
||||
"model_info": {"id": "2"},
|
||||
},
|
||||
],
|
||||
)
|
||||
response = await router._ageneric_api_call_with_fallbacks_helper(
|
||||
model="test-model",
|
||||
original_generic_function=_fake_provider,
|
||||
_target_order=2,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
assert response == "ok"
|
||||
assert captured["model"] == "gpt-4o"
|
||||
assert "_target_order" not in captured
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_text_completion_order_fallback_hop_does_not_send_target_order_upstream():
|
||||
upstream_bodies: Final[list[dict]] = []
|
||||
|
||||
def _upstream(request: httpx.Request) -> httpx.Response:
|
||||
upstream_bodies.append(json.loads(request.content))
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "cmpl-1",
|
||||
"object": "text_completion",
|
||||
"created": 0,
|
||||
"model": "gpt-3.5-turbo-instruct",
|
||||
"choices": [{"text": "ok from order 2", "index": 0, "logprobs": None, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
|
||||
upstream_client: Final = AsyncOpenAI(
|
||||
api_key="key",
|
||||
base_url="http://upstream.test",
|
||||
http_client=httpx.AsyncClient(transport=httpx.MockTransport(_upstream)),
|
||||
)
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
|
||||
"api_key": "key",
|
||||
"mock_response": Exception("fail order 1"),
|
||||
"order": 1,
|
||||
},
|
||||
"model_info": {"id": "1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
|
||||
"api_key": "key",
|
||||
"api_base": "http://upstream.test",
|
||||
"order": 2,
|
||||
},
|
||||
"model_info": {"id": "2"},
|
||||
},
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
try:
|
||||
response = await router.atext_completion(model="test-model", prompt="hi", client=upstream_client)
|
||||
finally:
|
||||
await upstream_client.close()
|
||||
|
||||
assert response._hidden_params["model_id"] == "2"
|
||||
assert upstream_bodies
|
||||
assert all("_target_order" not in body for body in upstream_bodies)
|
||||
|
||||
|
||||
def test_check_non_standard_fallback_format():
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
_check_non_standard_fallback_format,
|
||||
)
|
||||
|
||||
# Standard formats
|
||||
assert (
|
||||
_check_non_standard_fallback_format([{"gpt-3.5-turbo": ["claude-3-haiku"]}])
|
||||
== False
|
||||
)
|
||||
assert _check_non_standard_fallback_format([{"gpt-3.5-turbo": ["claude-3-haiku"]}]) == False
|
||||
assert _check_non_standard_fallback_format([{"model": ["qwen-backup"]}]) == False
|
||||
assert (
|
||||
_check_non_standard_fallback_format(
|
||||
[{"model": ["qwen-backup"], "region": ["us-east-1"]}]
|
||||
)
|
||||
== False
|
||||
)
|
||||
assert _check_non_standard_fallback_format([{"model": ["qwen-backup"], "region": ["us-east-1"]}]) == False
|
||||
|
||||
# Non-standard formats
|
||||
assert _check_non_standard_fallback_format([{"model": "qwen-backup"}]) == True
|
||||
assert (
|
||||
_check_non_standard_fallback_format(
|
||||
[{"model": "qwen-backup", "messages": [{"role": "user", "content": "hi"}]}]
|
||||
)
|
||||
== True
|
||||
)
|
||||
assert (
|
||||
_check_non_standard_fallback_format(
|
||||
[{"model": ["qwen-backup"], "api_key": "some-key"}]
|
||||
)
|
||||
_check_non_standard_fallback_format([{"model": "qwen-backup", "messages": [{"role": "user", "content": "hi"}]}])
|
||||
== True
|
||||
)
|
||||
assert _check_non_standard_fallback_format([{"model": ["qwen-backup"], "api_key": "some-key"}]) == True
|
||||
|
|
|
|||
|
|
@ -140,6 +140,7 @@ const renderPanel = (canModify = true) =>
|
|||
accessToken="token"
|
||||
userRole="Admin"
|
||||
userID="u-admin"
|
||||
isViewOnly={false}
|
||||
teams={null}
|
||||
createScope={canModify ? "unscoped-ok" : "forbidden"}
|
||||
/>,
|
||||
|
|
|
|||
|
|
@ -21,12 +21,20 @@ interface AutoRoutersPanelProps {
|
|||
accessToken: string;
|
||||
userRole: string;
|
||||
userID: string | null;
|
||||
isViewOnly: boolean;
|
||||
teams: Team[] | null;
|
||||
/** Owned by the page, which knows how this caller must scope what they create. */
|
||||
createScope: ModelWriteScope;
|
||||
}
|
||||
|
||||
export function AutoRoutersPanel({ accessToken, userRole, userID, teams, createScope }: AutoRoutersPanelProps) {
|
||||
export function AutoRoutersPanel({
|
||||
accessToken,
|
||||
userRole,
|
||||
userID,
|
||||
isViewOnly,
|
||||
teams,
|
||||
createScope,
|
||||
}: AutoRoutersPanelProps) {
|
||||
const canCreate = createScope !== "forbidden";
|
||||
const { data: deployments, isLoading } = useAutoRouters();
|
||||
const invalidateAutoRouters = useInvalidateAutoRouters();
|
||||
|
|
@ -39,8 +47,8 @@ export function AutoRoutersPanel({ accessToken, userRole, userID, teams, createS
|
|||
const [isDeleting, setIsDeleting] = useState(false);
|
||||
|
||||
const routers = useMemo(
|
||||
() => toAutoRouterRows(deployments ?? [], { userRole, userID }, teams),
|
||||
[deployments, userRole, userID, teams],
|
||||
() => toAutoRouterRows(deployments ?? [], { userRole, userID, isViewOnly }, teams),
|
||||
[deployments, userRole, userID, isViewOnly, teams],
|
||||
);
|
||||
|
||||
const handleCreated = () => {
|
||||
|
|
|
|||
|
|
@ -5,8 +5,9 @@ import { toAutoRouterRow, toAutoRouterRows } from "./autoRouterRows";
|
|||
|
||||
// Existing cases assert resource classification, so they run as a proxy admin: the actor
|
||||
// gate is then a pass-through and canEdit/canDelete still reflect the row itself.
|
||||
const ADMIN = { userRole: "Admin", userID: "u-admin" };
|
||||
const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin" };
|
||||
const ADMIN = { userRole: "Admin", userID: "u-admin", isViewOnly: false };
|
||||
const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin", isViewOnly: false };
|
||||
const VIEW_ONLY_ADMIN = { userRole: "Admin", userID: "u-viewer", isViewOnly: true };
|
||||
|
||||
const complexityDeployment = {
|
||||
model_name: "tri-tier-router",
|
||||
|
|
@ -224,7 +225,7 @@ describe("autoRouterRows actor gating", () => {
|
|||
{ team_id: "team-1", members_with_roles: [{ user_id: "u-team-admin", user_email: "t@t", role: "admin" }] },
|
||||
] as never;
|
||||
|
||||
const rowIn = (actor: { userRole: string; userID: string }, teamId: string | null) =>
|
||||
const rowIn = (actor: { userRole: string; userID: string; isViewOnly: boolean }, teamId: string | null) =>
|
||||
toAutoRouterRow(
|
||||
{ ...complexityDeployment, model_info: { id: "cid-1", db_model: true, team_id: teamId } },
|
||||
0,
|
||||
|
|
@ -259,4 +260,12 @@ describe("autoRouterRows actor gating", () => {
|
|||
expect(row.canEdit).toBe(true);
|
||||
expect(row.canDelete).toBe(true);
|
||||
});
|
||||
|
||||
// A proxy_admin_viewer session reads "Admin" through the masquerade, but PATCH and
|
||||
// DELETE both 403 it, so its rows must not offer the affordances.
|
||||
it("hides write affordances from a view-only admin session", () => {
|
||||
const row = rowIn(VIEW_ONLY_ADMIN, null);
|
||||
expect(row.canEdit).toBe(false);
|
||||
expect(row.canDelete).toBe(false);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -38,8 +38,17 @@ vi.mock("./useModelDashboardData", () => ({
|
|||
useModelDashboardData: () => ({ availableModelAccessGroups: [], allModelsOnProxy: [], availableModelGroups: [] }),
|
||||
}));
|
||||
|
||||
const ADMIN = { accessToken: "at", token: "t", userRole: "Admin", userId: "u1", premiumUser: false };
|
||||
const NON_ADMIN = { accessToken: "at", token: "t", userRole: "Internal User", userId: "u1", premiumUser: false };
|
||||
const ADMIN = { accessToken: "at", token: "t", userRole: "Admin", userId: "u1", premiumUser: false, isViewOnly: false };
|
||||
const NON_ADMIN = {
|
||||
accessToken: "at",
|
||||
token: "t",
|
||||
userRole: "Internal User",
|
||||
userId: "u1",
|
||||
premiumUser: false,
|
||||
isViewOnly: false,
|
||||
};
|
||||
// A proxy_admin_viewer session: effectiveSessionRole masquerades the role as "Admin".
|
||||
const VIEW_ONLY_ADMIN = { ...ADMIN, isViewOnly: true };
|
||||
|
||||
const renderPage = () => {
|
||||
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
|
||||
|
|
@ -99,6 +108,22 @@ describe("ModelsAndEndpointsPage", () => {
|
|||
expect(screen.queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
// POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one.
|
||||
it("hides the Add Model tab for a view-only admin session", () => {
|
||||
mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN);
|
||||
renderPage();
|
||||
expect(screen.queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Read parity: the Auto-Routers list stays reachable for a view-only admin; only the
|
||||
// create affordance inside it is withheld, which AutoRoutersTabPanel decides.
|
||||
it("keeps the Auto-Routers tab for a view-only admin session", () => {
|
||||
mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN);
|
||||
renderPage();
|
||||
expect(screen.getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Auto-routers are excluded from the All Models table, so this tab is their home: the only
|
||||
// place in the product to list, create, edit or delete one.
|
||||
describe("Auto-Routers tab", () => {
|
||||
|
|
|
|||
|
|
@ -80,7 +80,7 @@ const renderPanel = (key: string) => {
|
|||
};
|
||||
|
||||
export default function ModelsAndEndpointsPage() {
|
||||
const { accessToken, userRole, userId: userID, premiumUser } = useAuthorized();
|
||||
const { accessToken, userRole, userId: userID, premiumUser, isViewOnly } = useAuthorized();
|
||||
const { data: teams } = useTeams();
|
||||
const { data: uiSettings } = useUISettings();
|
||||
const queryClient = useQueryClient();
|
||||
|
|
@ -92,7 +92,7 @@ export default function ModelsAndEndpointsPage() {
|
|||
|
||||
const isInternalUser = userRole && internalUserRoles.includes(userRole);
|
||||
const canCreate = canCreateModels(
|
||||
{ userRole, userID },
|
||||
{ userRole, userID, isViewOnly },
|
||||
{
|
||||
teams: teams ?? null,
|
||||
disabledForInternalUsers:
|
||||
|
|
@ -182,6 +182,7 @@ export default function ModelsAndEndpointsPage() {
|
|||
accessToken={accessToken}
|
||||
userID={userID}
|
||||
userRole={userRole}
|
||||
isViewOnly={isViewOnly}
|
||||
onModelUpdate={invalidateModels}
|
||||
modelAccessGroups={availableModelAccessGroups}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,39 @@
|
|||
/* @vitest-environment jsdom */
|
||||
import { render } from "@testing-library/react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import AutoRoutersTabPanel from "./AutoRoutersTabPanel";
|
||||
|
||||
const panelProps = vi.fn();
|
||||
vi.mock("../components/AutoRouters/AutoRoutersPanel", () => ({
|
||||
AutoRoutersPanel: (props: Record<string, unknown>) => {
|
||||
panelProps(props);
|
||||
return <div data-testid="auto-routers-panel" />;
|
||||
},
|
||||
}));
|
||||
|
||||
const mockUseAuthorized = vi.fn();
|
||||
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => mockUseAuthorized() }));
|
||||
vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useTeams: () => ({ data: [] }) }));
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
|
||||
useUISettings: () => ({ data: { values: {} } }),
|
||||
}));
|
||||
|
||||
const SESSION = { accessToken: "at", userRole: "Admin", userId: "u1", isViewOnly: false };
|
||||
|
||||
const lastProps = () => panelProps.mock.calls.at(-1)?.[0] as { createScope: string };
|
||||
|
||||
describe("AutoRoutersTabPanel", () => {
|
||||
it("grants an unscoped create to a real proxy admin", () => {
|
||||
mockUseAuthorized.mockReturnValue(SESSION);
|
||||
render(<AutoRoutersTabPanel />);
|
||||
expect(lastProps().createScope).toBe("unscoped-ok");
|
||||
});
|
||||
|
||||
// The masqueraded "Admin" a proxy_admin_viewer session carries: POST /model/new 403s it,
|
||||
// so the panel must not be told it may create.
|
||||
it("withholds the create affordance from a view-only admin session", () => {
|
||||
mockUseAuthorized.mockReturnValue({ ...SESSION, isViewOnly: true });
|
||||
render(<AutoRoutersTabPanel />);
|
||||
expect(lastProps().createScope).toBe("forbidden");
|
||||
});
|
||||
});
|
||||
|
|
@ -15,13 +15,13 @@ import { AutoRoutersPanel } from "../components/AutoRouters/AutoRoutersPanel";
|
|||
* Viewer roles reach the list without write affordances.
|
||||
*/
|
||||
export default function AutoRoutersTabPanel() {
|
||||
const { accessToken, userRole, userId: userID } = useAuthorized();
|
||||
const { accessToken, userRole, userId: userID, isViewOnly } = useAuthorized();
|
||||
const { data: teams } = useTeams();
|
||||
const { data: uiSettings } = useUISettings();
|
||||
|
||||
const isInternalUser = userRole != null && internalUserRoles.includes(userRole);
|
||||
const scope = modelCreationScope(
|
||||
{ userRole, userID },
|
||||
{ userRole, userID, isViewOnly },
|
||||
{
|
||||
teams: teams ?? null,
|
||||
disabledForInternalUsers: isInternalUser && uiSettings?.values?.disable_model_add_for_internal_users === true,
|
||||
|
|
@ -33,6 +33,7 @@ export default function AutoRoutersTabPanel() {
|
|||
accessToken={accessToken}
|
||||
userRole={userRole ?? ""}
|
||||
userID={userID ?? null}
|
||||
isViewOnly={isViewOnly}
|
||||
teams={teams ?? null}
|
||||
createScope={scope}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -82,7 +82,7 @@ const AddModelForm: React.FC<AddModelFormProps> = ({
|
|||
// Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test
|
||||
const [connectionTestId, setConnectionTestId] = useState<string>("");
|
||||
|
||||
const { accessToken, userRole, premiumUser, userId } = useAuthorized();
|
||||
const { accessToken, userRole, premiumUser, userId, isViewOnly } = useAuthorized();
|
||||
const {
|
||||
data: providerMetadata,
|
||||
isLoading: isProviderMetadataLoading,
|
||||
|
|
@ -157,7 +157,10 @@ const AddModelForm: React.FC<AddModelFormProps> = ({
|
|||
const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId);
|
||||
// Same owner the Auto-Routers tab uses, so the two creation forms cannot disagree about
|
||||
// who has to name a team. This form is only reachable when creation is allowed at all.
|
||||
const createScope = modelCreationScope({ userRole, userID: userId }, { teams, disabledForInternalUsers: false });
|
||||
const createScope = modelCreationScope(
|
||||
{ userRole, userID: userId, isViewOnly },
|
||||
{ teams, disabledForInternalUsers: false },
|
||||
);
|
||||
const requiresTeamScope = createScope === "team-required";
|
||||
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -101,18 +101,24 @@ vi.mock("./build_complexity_router_config", async (importOriginal) => {
|
|||
});
|
||||
|
||||
// A real TeamDropdown fetches teams and renders an antd Select; the wiring under test is
|
||||
// whether team_id is registered, validated and forwarded, so a plain control stands in.
|
||||
// whether team_id is registered, validated and forwarded, so a plain control stands in. The
|
||||
// clear button mirrors the real dropdown's x, which emits null rather than a string.
|
||||
vi.mock("../common_components/team_dropdown", () => ({
|
||||
default: ({ value, onChange }: { value?: string; onChange?: (next: string) => void }) => (
|
||||
<select
|
||||
data-testid="team-dropdown"
|
||||
value={value ?? ""}
|
||||
onChange={(event) => onChange?.(event.target.value)}
|
||||
aria-label="Select Team"
|
||||
>
|
||||
<option value="">none</option>
|
||||
<option value="team-1">team-1</option>
|
||||
</select>
|
||||
default: ({ value, onChange }: { value?: string; onChange?: (next: string | null) => void }) => (
|
||||
<>
|
||||
<select
|
||||
data-testid="team-dropdown"
|
||||
value={value ?? ""}
|
||||
onChange={(event) => onChange?.(event.target.value)}
|
||||
aria-label="Select Team"
|
||||
>
|
||||
<option value="">none</option>
|
||||
<option value="team-1">team-1</option>
|
||||
</select>
|
||||
<button type="button" data-testid="team-dropdown-clear" onClick={() => onChange?.(null)}>
|
||||
clear team
|
||||
</button>
|
||||
</>
|
||||
),
|
||||
}));
|
||||
|
||||
|
|
@ -354,6 +360,25 @@ describe("AddAutoRouterTab", () => {
|
|||
expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// The shared dropdown emits null on clear while this form's schema wants a string, so the
|
||||
// form maps null back to "": the user sees the pick-a-team message, not a zod type error.
|
||||
it("treats a team picked and then cleared like no team at all", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getMissingTiersError).mockReturnValue(null);
|
||||
|
||||
renderWithProviders(
|
||||
<AddAutoRouterTab handleOk={vi.fn()} accessToken="token" userRole="Internal User" createScope="team-required" />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByPlaceholderText(/smart_router/i), "team-scoped-router");
|
||||
await user.selectOptions(screen.getByTestId("team-dropdown"), "team-1");
|
||||
await user.click(screen.getByTestId("team-dropdown-clear"));
|
||||
await user.click(screen.getByRole("button", { name: /add auto router/i }));
|
||||
|
||||
expect(await screen.findByText("Please select a team to continue")).toBeInTheDocument();
|
||||
expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("defaults a new router to session affinity off, matching the backend field default", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getMissingTiersError).mockReturnValue(null);
|
||||
|
|
|
|||
|
|
@ -548,7 +548,9 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
"Select the team this auto router belongs to. Only keys for this team will be able to call it.",
|
||||
)}
|
||||
>
|
||||
{({ id, value, onChange }) => <TeamDropdown id={id} value={value} onChange={onChange} />}
|
||||
{({ id, value, onChange }) => (
|
||||
<TeamDropdown id={id} value={value} onChange={(next) => onChange(next ?? "")} />
|
||||
)}
|
||||
</FormField>
|
||||
)}
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import React, { useState, useEffect } from "react";
|
|||
|
||||
import { alertingSettingsCall, updateConfigFieldSetting } from "../networking";
|
||||
import DynamicForm from "./dynamic_form";
|
||||
import { extractProxyErrorMessage } from "@/lib/http/client";
|
||||
import { toast } from "@/lib/toast";
|
||||
interface alertingSettingsItem {
|
||||
field_name: string;
|
||||
|
|
@ -43,7 +44,7 @@ const AlertingSettings: React.FC<AlertingSettingsProps> = ({ accessToken, premiu
|
|||
setAlertingSettings(updatedSettings);
|
||||
};
|
||||
|
||||
const handleSubmit = (formValues: Record<string, any>) => {
|
||||
const handleSubmit = async (formValues: Record<string, any>) => {
|
||||
if (!accessToken) {
|
||||
return;
|
||||
}
|
||||
|
|
@ -64,18 +65,18 @@ const AlertingSettings: React.FC<AlertingSettingsProps> = ({ accessToken, premiu
|
|||
const mergedFormValues = { ...formValues, ...initialFormValues };
|
||||
const { slack_alerting, ...alertingArgs } = mergedFormValues;
|
||||
try {
|
||||
updateConfigFieldSetting(accessToken, "alerting_args", alertingArgs);
|
||||
await updateConfigFieldSetting(accessToken, "alerting_args", alertingArgs);
|
||||
if (typeof slack_alerting === "boolean") {
|
||||
if (slack_alerting == true) {
|
||||
updateConfigFieldSetting(accessToken, "alerting", ["slack"]);
|
||||
await updateConfigFieldSetting(accessToken, "alerting", ["slack"]);
|
||||
} else {
|
||||
updateConfigFieldSetting(accessToken, "alerting", []);
|
||||
await updateConfigFieldSetting(accessToken, "alerting", []);
|
||||
}
|
||||
}
|
||||
// update value in state
|
||||
toast.success("Wait 10s for proxy to update.");
|
||||
} catch (error) {
|
||||
// do something
|
||||
toast.error(extractProxyErrorMessage(error));
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -38,6 +38,14 @@ const SETTINGS: Setting[] = [
|
|||
stored_in_db: null,
|
||||
premium_field: false,
|
||||
},
|
||||
{
|
||||
field_name: "daily_spend_per_user_threshold",
|
||||
field_description: "Daily spend threshold per user",
|
||||
field_type: "Float",
|
||||
field_value: 5.5,
|
||||
stored_in_db: true,
|
||||
premium_field: false,
|
||||
},
|
||||
];
|
||||
|
||||
const renderForm = (
|
||||
|
|
@ -170,6 +178,19 @@ describe("DynamicForm change notifications", () => {
|
|||
expect(handleInputChange).toHaveBeenCalledWith("daily_report_frequency", 128);
|
||||
});
|
||||
|
||||
it("renders a Float field as a decimal-friendly number input and reports changes as numbers", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { handleInputChange } = renderForm();
|
||||
|
||||
const input = screen.getByDisplayValue("5.5");
|
||||
expect(input).toHaveAttribute("type", "number");
|
||||
expect(input).toHaveAttribute("step", "any");
|
||||
|
||||
await user.type(input, "1");
|
||||
|
||||
expect(handleInputChange).toHaveBeenCalledWith("daily_spend_per_user_threshold", 5.51);
|
||||
});
|
||||
|
||||
it("reports a reset with the field name and its row index", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { handleResetField } = renderForm();
|
||||
|
|
@ -216,7 +237,7 @@ describe("DynamicForm presentation", () => {
|
|||
|
||||
expect(screen.getByText("daily_report_frequency")).toBeInTheDocument();
|
||||
expect(screen.getByText("How often the report runs")).toBeInTheDocument();
|
||||
expect(screen.getByText("In DB")).toBeInTheDocument();
|
||||
expect(screen.getAllByText("In DB")).toHaveLength(2);
|
||||
expect(screen.getByText("In Config")).toBeInTheDocument();
|
||||
expect(screen.getByText("Not Set")).toBeInTheDocument();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -63,11 +63,11 @@ const DynamicForm: React.FC<DynamicFormProps> = ({
|
|||
};
|
||||
|
||||
const renderControl = (setting: AlertingSetting) => {
|
||||
if (setting.field_type === "Integer") {
|
||||
if (setting.field_type === "Integer" || setting.field_type === "Float") {
|
||||
return (
|
||||
<Input
|
||||
type="number"
|
||||
step={1}
|
||||
step={setting.field_type === "Integer" ? 1 : "any"}
|
||||
value={setting.field_value ?? ""}
|
||||
onChange={(event) => handleNumericChange(setting, event.target.value)}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,48 @@
|
|||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { chooseSelectOption } from "../../../tests/test-utils";
|
||||
import type { Team } from "../key_team_helpers/key_list";
|
||||
import TeamDropdown from "./team_dropdown";
|
||||
|
||||
const TEAMS = [
|
||||
{ team_id: "team-1", team_alias: "Alpha Team" },
|
||||
{ team_id: "team-2", team_alias: "Beta Team" },
|
||||
] as unknown as Team[];
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
|
||||
useInfiniteTeams: () => ({
|
||||
data: { pages: [{ teams: TEAMS }] },
|
||||
fetchNextPage: vi.fn(),
|
||||
hasNextPage: false,
|
||||
isFetchingNextPage: false,
|
||||
isLoading: false,
|
||||
}),
|
||||
}));
|
||||
|
||||
describe("TeamDropdown", () => {
|
||||
it("emits the picked team's id and full object", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onChange = vi.fn();
|
||||
const onTeamSelect = vi.fn();
|
||||
render(<TeamDropdown onChange={onChange} onTeamSelect={onTeamSelect} />);
|
||||
|
||||
await chooseSelectOption(user, screen.getByRole("combobox"), /^Beta Team/);
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith("team-2");
|
||||
expect(onTeamSelect).toHaveBeenCalledWith(TEAMS[1]);
|
||||
});
|
||||
|
||||
it("emits null, never the empty string, when the selection is cleared", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onChange = vi.fn();
|
||||
const onTeamSelect = vi.fn();
|
||||
render(<TeamDropdown value="team-1" onChange={onChange} onTeamSelect={onTeamSelect} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Clear" }));
|
||||
|
||||
expect(onChange).toHaveBeenCalledWith(null);
|
||||
expect(onTeamSelect).toHaveBeenCalledWith(null);
|
||||
});
|
||||
});
|
||||
|
|
@ -5,7 +5,7 @@ import { Team } from "../key_team_helpers/key_list";
|
|||
|
||||
interface TeamDropdownProps {
|
||||
value?: string;
|
||||
onChange?: (value: string) => void;
|
||||
onChange?: (value: string | null) => void;
|
||||
/** Callback with the full Team object (or null on clear). */
|
||||
onTeamSelect?: (team: Team | null) => void;
|
||||
disabled?: boolean;
|
||||
|
|
@ -47,7 +47,7 @@ const TeamDropdown: React.FC<TeamDropdownProps> = ({
|
|||
}, [data]);
|
||||
|
||||
const handleChange = (teamId: string) => {
|
||||
onChange?.(teamId);
|
||||
onChange?.(teamId || null);
|
||||
if (onTeamSelect) {
|
||||
onTeamSelect(teamId ? teams.find((t) => t.team_id === teamId) ?? null : null);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -87,6 +87,7 @@ describe("ModelInfoView", () => {
|
|||
accessToken: "test-token",
|
||||
userID: "123",
|
||||
userRole: "Admin",
|
||||
isViewOnly: false,
|
||||
onModelUpdate: vi.fn(),
|
||||
modelAccessGroups: ["group1", "group2"],
|
||||
};
|
||||
|
|
@ -328,6 +329,16 @@ describe("ModelInfoView", () => {
|
|||
});
|
||||
});
|
||||
|
||||
// A proxy_admin_viewer session reads "Admin" through effectiveSessionRole, but the update
|
||||
// and delete endpoints 403 it, so the write buttons must not be offered.
|
||||
it("should disable delete and update buttons for a view-only admin session", async () => {
|
||||
render(<ModelInfoView {...DEFAULT_ADMIN_PROPS} isViewOnly={true} />, { wrapper });
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId("delete-model-button")).toBeDisabled();
|
||||
});
|
||||
expect(screen.getByTestId("update-api-key-button")).toBeDisabled();
|
||||
});
|
||||
|
||||
it("should disable delete button when model is not a DB model", async () => {
|
||||
const nonDbModelData = {
|
||||
...defaultModelData,
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ interface ModelInfoViewProps {
|
|||
accessToken: string | null;
|
||||
userID: string | null;
|
||||
userRole: string | null;
|
||||
isViewOnly: boolean;
|
||||
onModelUpdate?: (updatedModel: any) => void;
|
||||
modelAccessGroups: string[] | null;
|
||||
}
|
||||
|
|
@ -117,6 +118,7 @@ export default function ModelInfoView({
|
|||
accessToken,
|
||||
userID,
|
||||
userRole,
|
||||
isViewOnly,
|
||||
onModelUpdate,
|
||||
modelAccessGroups,
|
||||
}: ModelInfoViewProps) {
|
||||
|
|
@ -167,7 +169,7 @@ export default function ModelInfoView({
|
|||
// Keep modelData variable name for backwards compatibility
|
||||
const modelData = transformedModelData;
|
||||
|
||||
const canEditModel = canModifyModel({ userRole, userID }, teams ?? null, {
|
||||
const canEditModel = canModifyModel({ userRole, userID, isViewOnly }, teams ?? null, {
|
||||
teamId: modelData?.model_info?.team_id,
|
||||
isDbModel: modelData?.model_info?.db_model === true,
|
||||
});
|
||||
|
|
|
|||
|
|
@ -293,6 +293,8 @@ const Settings: React.FC<SettingsPageProps> = ({ accessToken, userRole, userID,
|
|||
llm_too_slow: "LLM Responses Too Slow",
|
||||
llm_requests_hanging: "LLM Requests Hanging",
|
||||
budget_alerts: "Budget Alerts (API Keys, Users)",
|
||||
user_spend_thresholds: "User Spend Thresholds (Daily/Monthly)",
|
||||
user_spend_anomalies: "User Spend Anomaly Detection",
|
||||
db_exceptions: "Database Exceptions (Read/Write)",
|
||||
daily_reports: "Weekly/Monthly Spend Reports",
|
||||
outage_alerts: "Outage Alerts",
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -22962,7 +22962,7 @@ export interface components {
|
|||
* @description Enum for alert types and management event types
|
||||
* @enum {string}
|
||||
*/
|
||||
AlertType: "llm_exceptions" | "llm_too_slow" | "llm_requests_hanging" | "budget_alerts" | "spend_reports" | "failed_tracking_spend" | "db_exceptions" | "daily_reports" | "cooldown_deployment" | "new_model_added" | "model_deprecation_warnings" | "outage_alerts" | "region_outage_alerts" | "fallback_reports" | "new_virtual_key_created" | "virtual_key_updated" | "virtual_key_deleted" | "new_team_created" | "team_updated" | "team_deleted" | "new_internal_user_created" | "internal_user_updated" | "internal_user_deleted";
|
||||
AlertType: "llm_exceptions" | "llm_too_slow" | "llm_requests_hanging" | "budget_alerts" | "spend_reports" | "failed_tracking_spend" | "user_spend_thresholds" | "user_spend_anomalies" | "db_exceptions" | "daily_reports" | "cooldown_deployment" | "new_model_added" | "model_deprecation_warnings" | "outage_alerts" | "region_outage_alerts" | "fallback_reports" | "new_virtual_key_created" | "virtual_key_updated" | "virtual_key_deleted" | "new_team_created" | "team_updated" | "team_deleted" | "new_internal_user_created" | "internal_user_updated" | "internal_user_deleted";
|
||||
/** AllowedVectorStoreIndexItem */
|
||||
AllowedVectorStoreIndexItem: {
|
||||
/** Index Name */
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import { Team } from "@/components/networking";
|
||||
import { canModifyModel, modelCreationScope } from "./modelPermissions";
|
||||
import { canCreateModels, canModifyModel, modelCreationScope } from "./modelPermissions";
|
||||
|
||||
const teamWhere = (userId: string, role: string, teamId = "team-1"): Team[] =>
|
||||
[{ team_id: teamId, members_with_roles: [{ user_id: userId, user_email: "t@test.com", role }] }] as unknown as Team[];
|
||||
|
||||
const PROXY_ADMIN = { userRole: "Admin", userID: "u-admin" };
|
||||
const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin" };
|
||||
const MEMBER = { userRole: "Internal User", userID: "u-member" };
|
||||
const PROXY_ADMIN = { userRole: "Admin", userID: "u-admin", isViewOnly: false };
|
||||
const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin", isViewOnly: false };
|
||||
const MEMBER = { userRole: "Internal User", userID: "u-member", isViewOnly: false };
|
||||
// proxy_admin_viewer sessions: effectiveSessionRole masquerades the role as "Admin".
|
||||
const VIEW_ONLY_ADMIN = { userRole: "Admin", userID: "u-viewer", isViewOnly: true };
|
||||
|
||||
const noLimits = { disabledForInternalUsers: false };
|
||||
|
||||
|
|
@ -40,9 +42,24 @@ describe("modelCreationScope", () => {
|
|||
// an unscoped create from them 403s. Treating them as admins here is what let a form submit
|
||||
// a payload the backend always rejected.
|
||||
it("does not treat an org admin as able to create unscoped", () => {
|
||||
const orgAdmin = { userRole: "org_admin", userID: "u-org" };
|
||||
const orgAdmin = { userRole: "org_admin", userID: "u-org", isViewOnly: false };
|
||||
expect(modelCreationScope(orgAdmin, { teams: teamWhere("u-org", "admin"), ...noLimits })).toBe("team-required");
|
||||
});
|
||||
|
||||
// Server-side, POST /model/new 403s the viewer roles, so the "Admin" the masquerade
|
||||
// reports must not read as a proxy admin here.
|
||||
it("forbids a view-only admin session despite the masqueraded Admin role", () => {
|
||||
expect(modelCreationScope(VIEW_ONLY_ADMIN, { teams: [], ...noLimits })).toBe("forbidden");
|
||||
expect(canCreateModels(VIEW_ONLY_ADMIN, { teams: [], ...noLimits })).toBe(false);
|
||||
});
|
||||
|
||||
// _check_proxy_admin_viewer_access (route_checks.py) 403s /model/new on the session role
|
||||
// alone, before the team-scoped carve-out in ModelManagementAuthChecks can run.
|
||||
it("forbids a view-only admin even when they admin a team", () => {
|
||||
expect(modelCreationScope(VIEW_ONLY_ADMIN, { teams: teamWhere("u-viewer", "admin"), ...noLimits })).toBe(
|
||||
"forbidden",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("canModifyModel", () => {
|
||||
|
|
@ -80,6 +97,17 @@ describe("canModifyModel", () => {
|
|||
});
|
||||
|
||||
it("does not treat two absent identities as a match", () => {
|
||||
expect(canModifyModel({ userRole: "Internal User", userID: null }, null, teamRow)).toBe(false);
|
||||
expect(canModifyModel({ userRole: "Internal User", userID: null, isViewOnly: false }, null, teamRow)).toBe(false);
|
||||
});
|
||||
|
||||
// PATCH /model/{id}/update and POST /model/delete 403 the viewer roles like /model/new does.
|
||||
it("refuses a view-only admin session on a DB row", () => {
|
||||
expect(canModifyModel(VIEW_ONLY_ADMIN, null, teamRow)).toBe(false);
|
||||
});
|
||||
|
||||
// The route RBAC blocks /model/update and /model/delete for the viewer role before the
|
||||
// team-scoped carve-out runs, so team-admin membership changes nothing here either.
|
||||
it("refuses a view-only user even when they admin the owning team", () => {
|
||||
expect(canModifyModel(VIEW_ONLY_ADMIN, teamWhere("u-viewer", "admin"), teamRow)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -3,20 +3,32 @@ import { Team } from "@/components/networking";
|
|||
import { isProxyAdminRole, isUserTeamAdminForAnyTeam, isUserTeamAdminForSingleTeam } from "./roles";
|
||||
|
||||
/**
|
||||
* The dashboard's mirror of ModelManagementAuthChecks in
|
||||
* litellm/proxy/management_endpoints/model_management_endpoints.py.
|
||||
* The dashboard's mirror of the two server layers that gate model writes: the role-level
|
||||
* route RBAC (`_check_proxy_admin_viewer_access` in litellm/proxy/auth/route_checks.py),
|
||||
* which 403s /model/new, /model/update, and /model/delete for every view-only session
|
||||
* before the endpoint runs, and ModelManagementAuthChecks in
|
||||
* litellm/proxy/management_endpoints/model_management_endpoints.py behind it.
|
||||
*
|
||||
* Both questions below are answered there by exactly two inputs: the caller's role, and
|
||||
* whether the caller admins the team named in `model_info.team_id`. `created_by` is written
|
||||
* at creation and never read by an auth check, so it is deliberately absent here; gating on
|
||||
* it hid controls from team admins the API accepts, and showed controls to former team admins
|
||||
* the API rejects.
|
||||
* Past that route gate, both questions below are answered by exactly two inputs: the
|
||||
* caller's role, and whether the caller admins the team named in `model_info.team_id`.
|
||||
* `created_by` is written at creation and never read by an auth check, so it is deliberately
|
||||
* absent here; gating on it hid controls from team admins the API accepts, and showed
|
||||
* controls to former team admins the API rejects.
|
||||
*/
|
||||
export interface ModelActor {
|
||||
userRole: string | null;
|
||||
userID: string | null;
|
||||
/**
|
||||
* From useAuthorized(). A proxy_admin_viewer session masquerades as "Admin" in userRole
|
||||
* (effectiveSessionRole, for read parity), yet every management write 403s it, so the role
|
||||
* alone cannot answer a write question.
|
||||
*/
|
||||
isViewOnly: boolean;
|
||||
}
|
||||
|
||||
const isWritableProxyAdmin = ({ userRole, isViewOnly }: ModelActor): boolean =>
|
||||
!isViewOnly && userRole != null && isProxyAdminRole(userRole);
|
||||
|
||||
/** How this actor must scope a deployment they create, or that they may not create one. */
|
||||
export type ModelWriteScope = "forbidden" | "unscoped-ok" | "team-required";
|
||||
|
||||
|
|
@ -33,20 +45,25 @@ const isTeamAdminOf = (teams: Team[] | null, userID: string, teamId: string): bo
|
|||
|
||||
/**
|
||||
* POST /model/new takes a proxy admin unconditionally, or a team admin whose payload names a
|
||||
* team; an unscoped create from anyone else is a 403. Returning the requirement rather than a
|
||||
* pair of booleans keeps "may not create" and "may create unscoped" from being confused.
|
||||
* team; an unscoped create from anyone else is a 403. A view-only session is 403d by the
|
||||
* route RBAC on its role alone, so team-admin membership cannot rescue it. Returning the
|
||||
* requirement rather than a pair of booleans keeps "may not create" and "may create
|
||||
* unscoped" from being confused.
|
||||
*/
|
||||
export const modelCreationScope = (
|
||||
{ userRole, userID }: ModelActor,
|
||||
actor: ModelActor,
|
||||
{ teams, disabledForInternalUsers }: ModelCreationLimits,
|
||||
): ModelWriteScope => {
|
||||
if (userRole != null && isProxyAdminRole(userRole)) {
|
||||
if (actor.isViewOnly) {
|
||||
return "forbidden";
|
||||
}
|
||||
if (isWritableProxyAdmin(actor)) {
|
||||
return "unscoped-ok";
|
||||
}
|
||||
if (disabledForInternalUsers) {
|
||||
return "forbidden";
|
||||
}
|
||||
if (userID != null && isUserTeamAdminForAnyTeam(teams, userID)) {
|
||||
if (actor.userID != null && isUserTeamAdminForAnyTeam(teams, actor.userID)) {
|
||||
return "team-required";
|
||||
}
|
||||
return "forbidden";
|
||||
|
|
@ -63,18 +80,18 @@ export interface ModelRowOrigin {
|
|||
|
||||
/** May this actor edit or delete this specific deployment? */
|
||||
export const canModifyModel = (
|
||||
{ userRole, userID }: ModelActor,
|
||||
actor: ModelActor,
|
||||
teams: Team[] | null,
|
||||
{ teamId, isDbModel }: ModelRowOrigin,
|
||||
): boolean => {
|
||||
if (!isDbModel) {
|
||||
if (actor.isViewOnly || !isDbModel) {
|
||||
return false;
|
||||
}
|
||||
if (userRole != null && isProxyAdminRole(userRole)) {
|
||||
if (isWritableProxyAdmin(actor)) {
|
||||
return true;
|
||||
}
|
||||
if (userID == null || teamId == null) {
|
||||
if (actor.userID == null || teamId == null) {
|
||||
return false;
|
||||
}
|
||||
return isTeamAdminOf(teams, userID, teamId);
|
||||
return isTeamAdminOf(teams, actor.userID, teamId);
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue