scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

Merge origin/litellm_internal_staging into shadow-eval-pre-adoption

Browse source 
Resolves two real conflicts (the PR's actual base branch is
litellm_internal_staging, not main):

- litellm/types/management_endpoints/auto_router_endpoints.py: kept both
  the Mapping and Literal imports, both used by pre-existing types.
- tests/e2e/proxy_client.py: kept upstream's more detailed create_model()
  docstring covering multi-replica propagation.

Everything else auto-merged cleanly. Regenerated the OpenAPI schema
(schema.d.ts) to pick up upstream's tier_turns addition to the auto-router
benchmarks response.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Abhimanyu Kapur 2026-08-07 21:18:17 -07:00
parent 13fe386d50
commit 2092b809d3
67 changed files with 4609 additions and 818 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
basedpyright-code-budget.json
View file

@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45110
"limit": 45098
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 39838
"limit": 39826
},
"reportUnknownParameterType": {
"limit": 20237
},
"reportUnknownVariableType": {
"limit": 31383
"limit": 31371
},
"reportUnnecessaryCast": {
"limit": 122

112
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -3,6 +3,7 @@
import base64
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast
from uuid import NAMESPACE_URL, uuid5
@ -32,9 +33,12 @@ from litellm.proxy._types import (
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
apply_unified_file_ids,
ensure_batch_response_managed_file_ids,
get_batch_id_from_unified_batch_id,
get_content_type_from_file_object,
get_model_id_from_unified_batch_id,
map_raw_file_ids_to_unified,
normalize_mime_type_for_provider,
resolve_managed_output_file_model_name,
)
@ -62,6 +66,9 @@ if TYPE_CHECKING:
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from prisma.models import (
LiteLLM_ManagedObjectTable as PrismaManagedObjectRow,
)
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
@ -75,6 +82,30 @@ else:
PrismaClient = Any
def _sanitized_parse_error(e: Exception) -> str:
return (
str(e.errors(include_input=False, include_url=False, include_context=False))
if isinstance(e, ValidationError)
else type(e).__name__
)
def _decode_json_blob(blob: object) -> object:
return json.loads(blob) if isinstance(blob, str) else blob
def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMBatch]:
try:
batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object))
except Exception as e:
verbose_logger.warning(
f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}"
)
return None
batch_obj.id = row.unified_object_id
return batch_obj
def _parse_managed_file_object(
raw_file_object: object, unified_file_id: str
) -> Optional[OpenAIFileObject]:
@ -82,15 +113,9 @@ def _parse_managed_file_object(
return None
try:
return OpenAIFileObject.model_validate(raw_file_object)
except ValidationError as e:
verbose_logger.warning(
f"Failed to parse managed file object {unified_file_id}: "
f"{e.errors(include_input=False, include_url=False, include_context=False)}"
)
return None
except Exception as e:
verbose_logger.warning(
f"Failed to parse managed file object {unified_file_id}: {type(e).__name__}"
f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}"
)
return None
@ -349,7 +374,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
detail=f"Invalid 'after' cursor: no batch found with id '{after}'.",
)
page_size = limit or 20
page_size: Final = min(limit or 20, 100)
cursor_args: Dict[str, Any] = (
{"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
)
@ -363,25 +388,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
has_more = len(batches) > page_size
batch_objects: List[LiteLLMBatch] = []
for batch in batches[:page_size]:
try:
batch_data = (
json.loads(batch.file_object)
if isinstance(batch.file_object, str)
else batch.file_object
)
batch_obj = LiteLLMBatch.model_validate(batch_data)
batch_obj.id = batch.unified_object_id
batch_objects.append(batch_obj)
parsed_rows: Final = tuple(
(row, batch_obj)
for row in batches[:page_size]
if (batch_obj := _parse_managed_batch_row(row)) is not None
)
unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified(
raw_file_ids=frozenset(
file_id
for _, batch_obj in parsed_rows
for file_id in (batch_obj.input_file_id, batch_obj.output_file_id, batch_obj.error_file_id)
if file_id and not _is_base64_encoded_unified_file_id(file_id)
),
prisma_client=self.prisma_client,
)
resolved_batches: Final = [
await self._resolve_listed_batch(
row=row,
batch_obj=batch_obj,
unified_id_by_raw_id=unified_id_by_raw_id,
user_api_key_dict=user_api_key_dict,
)
for row, batch_obj in parsed_rows
]
return build_list_page(
[batch_obj for batch_obj in resolved_batches if batch_obj is not None],
has_more=has_more,
)
except Exception as e:
verbose_logger.warning(
f"Failed to parse batch object {batch.unified_object_id}: {e}"
)
continue
return build_list_page(batch_objects, has_more=has_more)
async def _resolve_listed_batch(
self,
row: "PrismaManagedObjectRow",
batch_obj: LiteLLMBatch,
unified_id_by_raw_id: Mapping[str, str],
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[LiteLLMBatch]:
apply_unified_file_ids(batch_obj, unified_id_by_raw_id)
try:
await ensure_batch_response_managed_file_ids(
response=batch_obj,
managed_files_obj=self,
prisma_client=self.prisma_client,
verbose_proxy_logger=verbose_logger,
user_api_key_dict=user_api_key_dict,
db_batch_object=row,
unified_batch_id=_is_base64_encoded_unified_file_id(
row.unified_object_id
),
)
except Exception as e:
verbose_logger.warning(
f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}"
)
return None
return batch_obj
async def get_user_created_file_ids(
self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str]

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}';

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
tier_turns Json @default("{}")
@@id([api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")

1
litellm/constants.py
View file

@ -280,6 +280,7 @@ TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECO
GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
)
BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
# Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
# Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)))

2
litellm/integrations/otel/__init__.py
View file

@ -57,6 +57,7 @@ from litellm.integrations.otel.model.semconv import (
Metric,
Network,
NetworkTransport,
RpcSystem,
Server,
resolve_operation,
resolve_provider,
@ -102,6 +103,7 @@ __all__ = [
"ProxyRequestSpanData",
"RequestContext",
"RequestIdentity",
"RpcSystem",
"Server",
"ServerInfo",
"ServiceSpanData",

5
litellm/integrations/otel/mappers/genai.py
View file

@ -31,7 +31,9 @@ from litellm.integrations.otel.model.semconv import (
MCP,
Error,
GenAI,
JsonRpc,
LiteLLM,
RpcSystem,
Server,
)
from litellm.integrations.otel.model.spans import db_system
@ -94,11 +96,14 @@ class GenAIMapper:
_MCP_ATTRS: dict[str, Callable[[MCPToolCallSpanData], AttrValue | None]] = {
GenAI.OPERATION_NAME: lambda d: d.operation.value,
JsonRpc.SYSTEM: lambda d: RpcSystem.JSONRPC.value if d.server_address and d.server_port else None,
MCP.METHOD_NAME: lambda d: d.method,
MCP.SESSION_ID: lambda d: d.session_id,
GenAI.TOOL_NAME: lambda d: d.tool_name or None,
GenAI.TOOL_CALL_ARGUMENTS: lambda d: d.arguments_json,
GenAI.TOOL_CALL_RESULT: lambda d: d.result_json,
Server.ADDRESS: lambda d: d.server_address,
Server.PORT: lambda d: d.server_port,
LiteLLM.MCP_SERVER_NAME: lambda d: d.server_name,
LiteLLM.CALL_ID: lambda d: d.identity.call_id or None,
f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost,

33
litellm/integrations/otel/model/payloads.py
View file

@ -364,6 +364,34 @@ class LLMCallSpanData:
# --- the MCP tool-call model ------------------------------------------------- #
def _upstream_address_port(resource: str | None) -> tuple[str | None, int | None]:
"""Split a redacted MCP server origin into ``server.address`` / ``server.port``.
``mcp_server_resource`` is a scheme + host + port origin with userinfo, path,
query and fragment already stripped. The port falls back to the scheme default
when the origin omits it, because a consumer that keys a downstream dependency
off the address renders a missing port as ``0``.
The origin is rebuilt without its IPv6 brackets upstream, so reading the port can
raise on an address the host check still admits: a zone-scoped ``fe80::1%25eth0``
leaves a truthy hostname of ``fe80`` behind. Both halves are read inside the guard
so an unparseable origin yields no address rather than propagating out of span
construction, matching how the redactor guards the same split.
"""
if not resource:
return None, None
try:
parsed: Final = urlsplit(resource)
hostname: Final = parsed.hostname
port: Final = parsed.port
except ValueError:
return None, None
if not hostname:
return None, None
default_port: Final = 443 if parsed.scheme == "https" else 80 if parsed.scheme == "http" else None
return hostname, port or default_port
@dataclass(frozen=True)
class MCPToolCallSpanData:
"""One MCP ``tools/call`` execution, parsed from a closed request's payload.
@ -378,6 +406,8 @@ class MCPToolCallSpanData:
method: str
tool_name: str
server_name: str | None
server_address: str | None
server_port: int | None
session_id: str | None
arguments_json: str | None
result_json: str | None
@ -390,11 +420,14 @@ class MCPToolCallSpanData:
cls, payload: StandardLoggingPayload, capture_content: bool = False
) -> MCPToolCallSpanData:
meta: Final = _mcp_tool_call_metadata(cast(Mapping[str, object], payload))
address, port = _upstream_address_port(as_str(meta.get("mcp_server_resource")) or None)
return cls(
operation=resolve_operation(as_str(payload.get("call_type"))),
method=MCPMethod.TOOLS_CALL.value,
tool_name=as_str(meta.get("name")) or "",
server_name=as_str(meta.get("mcp_server_name")),
server_address=address,
server_port=port,
session_id=as_str(meta.get("mcp_session_id")),
arguments_json=(
_json_or_none(meta.get("arguments")) if capture_content and meta.get("arguments") is not None else None

14
litellm/integrations/otel/model/semconv.py
View file

@ -130,11 +130,25 @@ class JsonRpc:
"""JSON-RPC keys carried on MCP spans. The error/status code lives in the
``rpc.*`` namespace per semconv, not ``jsonrpc.*``."""
SYSTEM: Final = "rpc.system"
REQUEST_ID: Final = "jsonrpc.request.id"
PROTOCOL_VERSION: Final = "jsonrpc.protocol.version"
RESPONSE_STATUS_CODE: Final = "rpc.response.status_code"
class RpcSystem(str, Enum):
"""Well-known values for ``rpc.system``. MCP frames every message as JSON-RPC 2.0.
Naming the system also classifies the span: a CLIENT span carrying none of the
``rpc.*``/``http.*``/``db.*``/``messaging.*`` families records no span type or
subtype in backends that derive those from the attribute family. It is emitted
only alongside ``server.address``/``server.port``, since a backend that reads it
as a downstream dependency names that dependency from the server address.
"""
JSONRPC = "jsonrpc"
class NetworkTransport(str, Enum):
"""Well-known values for ``network.transport``."""

56
litellm/integrations/websearch_interception/handler.py
View file

@ -9,7 +9,7 @@ server-side using litellm router's search tools.
import asyncio
import math
import uuid
from collections.abc import AsyncIterator, Mapping
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, cast
import litellm
@ -37,6 +37,8 @@ from litellm.types.integrations.custom_logger import (
AgenticLoopRequestPatch,
)
from litellm.types.integrations.websearch_interception import (
AnthropicSearchQuery,
AnthropicServerToolUseBlock,
WebSearchInterceptionConfig,
)
from litellm.types.llms.openai import AllMessageValues
@ -833,22 +835,48 @@ class WebSearchInterceptionLogger(CustomLogger):
def _build_native_result_blocks(
tool_calls: list[dict],
structured_results: list[SearchResponse | None],
) -> list[dict[str, object]]:
"""Build one ``web_search_tool_result`` block per tool_call."""
blocks: Final[list[dict[str, object]]] = []
for i, tool_call in enumerate(tool_calls):
tool_use_id = tool_call.get("id") or ""
structured = structured_results[i] if i < len(structured_results) else None
blocks.append(
WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id=tool_use_id,
search_response=structured,
)
) -> tuple[Mapping[str, object], ...]:
"""
Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call.
The pair is what Anthropic's spec requires: a bare result block, or one
keyed by the model's ``toolu_...`` id instead of a ``srvtoolu_...`` one,
is rejected on replay ("String should match pattern '^srvtoolu_'") and
leaves native clients without a search to attach the sources to.
"""
return tuple(
block
for i, tool_call in enumerate(tool_calls)
for block in WebSearchInterceptionLogger._native_result_pair(
query=WebSearchInterceptionLogger._tool_call_query(tool_call),
search_response=structured_results[i] if i < len(structured_results) else None,
)
return blocks
)
@staticmethod
def _inject_native_blocks(response: Any, native_blocks: list[dict[str, object]]) -> Any:
def _tool_call_query(tool_call: Mapping[str, object]) -> str:
tool_input: Final = tool_call.get("input")
if not isinstance(tool_input, Mapping):
return ""
query: Final = tool_input.get("query")
return query if isinstance(query, str) else ""
@staticmethod
def _native_result_pair(
query: str,
search_response: SearchResponse | None,
) -> tuple[Mapping[str, object], Mapping[str, object]]:
tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
return (
AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(),
WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id=tool_use_id,
search_response=search_response,
),
)
@staticmethod
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
"""Prepend native blocks to response content, dict or object form."""
if not native_blocks:
return response

10
litellm/integrations/websearch_interception/transformation.py
View file

@ -412,6 +412,15 @@ class WebSearchTransformation:
block that should accompany the model's text reply when the original
request used a native ``web_search_*`` tool.
The spec'd shape carries page text only in ``encrypted_content``, an
opaque server-issued blob that we cannot mint. Emitting the four spec
fields alone would drop the snippet entirely, leaving the client (and
the model, on any replayed follow-up turn) with URLs and titles but no
evidence to answer from, forcing a fetch per result. So the snippet is
carried in an additive ``snippet`` key alongside the spec fields.
``encrypted_content`` stays empty rather than holding plaintext, which
would assert encryption semantics that do not hold.
Spec reference:
https://docs.anthropic.com/en/api/web-search-tool
@ -438,6 +447,7 @@ class WebSearchTransformation:
"title": title,
"page_age": page_age,
"encrypted_content": "",
"snippet": getattr(r, "snippet", "") or "",
}
)
return {

151
litellm/llms/anthropic/common_utils.py
View file

@ -4,9 +4,12 @@ This file contains common utils for anthropic calls.
import copy
import re
from typing import Any, Final
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final, Literal
import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -1057,6 +1060,152 @@ def sanitize_tool_use_ids_in_anthropic_messages(messages: list[Any]) -> list[Any
return out
class _ReplayedSearchQuery(BaseModel):
model_config = ConfigDict(extra="allow")
query: str = ""
class _ReplayedWebSearchResult(BaseModel):
model_config = ConfigDict(extra="allow")
type: Literal["web_search_result"]
url: str = ""
title: str = ""
snippet: str = ""
encrypted_content: str = ""
class _ReplayedWebSearchToolResult(BaseModel):
model_config = ConfigDict(extra="allow")
type: Literal["web_search_tool_result"]
tool_use_id: str
content: tuple[_ReplayedWebSearchResult, ...]
class _ReplayedServerToolUse(BaseModel):
model_config = ConfigDict(extra="allow")
type: Literal["server_tool_use"]
id: str
input: _ReplayedSearchQuery = _ReplayedSearchQuery()
class _TextBlock(BaseModel):
type: Literal["text"] = "text"
text: str
_WEB_SEARCH_TOOL_RESULT_ADAPTER: Final = TypeAdapter(_ReplayedWebSearchToolResult)
_SERVER_TOOL_USE_ADAPTER: Final = TypeAdapter(_ReplayedServerToolUse)
def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchToolResult | None:
"""
The parsed block when it is a ``web_search_tool_result`` carrying no
``encrypted_content``, else None for anything Anthropic itself issued.
An empty ``content`` list is flattenable too. It is what the interceptor emits
when a search legitimately returns nothing and when a search raises, and it
carries neither evidence to preserve nor an ``encrypted_content`` to respect,
so leaving it in place only buys the 400 this whole function exists to avoid.
"""
try:
parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block)
except ValidationError:
return None
if any(result.encrypted_content for result in parsed.content):
return None
return parsed
def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
try:
return _SERVER_TOOL_USE_ADAPTER.validate_python(block)
except ValidationError:
return None
def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str:
header: Final = f"Web search results for '{query}':" if query else "Web search results:"
if not results:
return f"{header}\n\nNo results were returned."
body: Final = "\n\n".join(
"\n".join(
line
for line in (
f"Title: {result.title}" if result.title else "",
f"URL: {result.url}" if result.url else "",
f"Snippet: {result.snippet}" if result.snippet else "",
)
if line
)
for result in results
)
return f"{header}\n\n{body}" if body else header
def _rewrite_replayed_web_search_block(
block: object,
flattenable: Mapping[str, _ReplayedWebSearchToolResult],
queries: Mapping[str, str],
) -> object | None:
parsed_result: Final = _flattenable_web_search_tool_result(block)
if parsed_result is not None:
return _TextBlock(
text=_render_web_search_results(queries.get(parsed_result.tool_use_id, ""), parsed_result.content)
).model_dump()
parsed_use: Final = _replayed_server_tool_use(block)
if parsed_use is not None and parsed_use.id in flattenable:
return None
return block
def _flatten_web_search_results_in_message(message: object) -> object:
if not isinstance(message, Mapping) or not isinstance(message.get("content"), Sequence):
return message
content: Final = message["content"]
if isinstance(content, str):
return message
flattenable: Final = MappingProxyType(
{
parsed.tool_use_id: parsed
for parsed in (_flattenable_web_search_tool_result(block) for block in content)
if parsed is not None
}
)
if not flattenable:
return message
queries: Final = MappingProxyType(
{
parsed.id: parsed.input.query
for parsed in (_replayed_server_tool_use(block) for block in content)
if parsed is not None
}
)
rewritten: Final = tuple(_rewrite_replayed_web_search_block(block, flattenable, queries) for block in content)
return {**message, "content": [b for b in rewritten if b is not None]} # mutable-ok: JSON wire format
def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: as sibling sanitizers
messages: list[Any],
) -> list[Any]:
"""
Return a new message list with replayed ``web_search_tool_result`` blocks that
carry no ``encrypted_content`` rewritten into plain ``text`` blocks holding the
same title / url / snippet evidence.
``encrypted_content`` is an opaque blob only Anthropic's own search backend can
mint, so blocks synthesized by LiteLLM (websearch interception against a search
provider) are rejected with ``Invalid encrypted_content in search_result block``
when a native client loops them back as history. Flattening them keeps the
evidence in the conversation instead of 400ing the follow-up turn, and leaves
genuine Anthropic-issued blocks untouched.
"""
return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format
def process_anthropic_headers(headers: httpx.Headers | dict) -> dict:
openai_headers: Final = {}
if "anthropic-ratelimit-requests-limit" in headers:

3
litellm/llms/anthropic/experimental_pass_through/messages/handler.py
View file

@ -14,6 +14,7 @@ from typing import Any, Final, cast
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.common_utils import (
flatten_unencrypted_web_search_results_in_anthropic_messages,
sanitize_tool_use_ids_in_anthropic_messages,
strip_empty_text_blocks_from_anthropic_messages,
)
@ -222,6 +223,7 @@ async def anthropic_messages(
# Replay of cross-provider tool history (e.g. kimi -> Anthropic) may carry
# ids like ``functions.Bash:0`` that violate Anthropic's id pattern.
messages = sanitize_tool_use_ids_in_anthropic_messages(messages)
messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages)
from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
@ -413,6 +415,7 @@ def anthropic_messages_handler(
if not kwargs.pop("_litellm_messages_presanitized", False):
messages = strip_empty_text_blocks_from_anthropic_messages(messages)
messages = sanitize_tool_use_ids_in_anthropic_messages(messages)
messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages)
from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,

7
litellm/proxy/_types.py
View file

@ -3894,12 +3894,19 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse):
##########################################
class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase):
access_group_id: str
access_group_name: str
models: tuple[str, ...]
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
team_member_budget_table: LiteLLM_BudgetTableFull | None = None
# Resources inherited from access groups (separate from direct assignments)
access_group_models: list[str] | None = None
access_group_mcp_server_ids: list[str] | None = None
access_group_agent_ids: list[str] | None = None
access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None
class TeamInfoResponseObject(TypedDict):

5
litellm/proxy/auth/auth_checks.py
View file

@ -13,6 +13,7 @@ import asyncio
import math
import re
import time
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
from fastapi import HTTPException, Request, status
@ -2834,7 +2835,7 @@ async def get_org_object(
async def _get_resources_from_access_groups(
access_group_ids: list[str],
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
prisma_client: PrismaClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
@ -2893,7 +2894,7 @@ async def _get_resources_from_access_groups(
async def _get_models_from_access_groups(
access_group_ids: list[str],
access_group_ids: Sequence[str],
prisma_client: PrismaClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,

7
litellm/proxy/auth/model_checks.py
View file

@ -1,6 +1,7 @@
# What is this?
## Common checks for /v1/models and `/model/info`
import copy
from collections.abc import Sequence
from typing import Any, Final
import litellm
@ -178,8 +179,8 @@ def get_team_models(
def get_complete_model_list(
key_models: list[str],
team_models: list[str],
key_models: Sequence[str],
team_models: Sequence[str],
proxy_model_list: list[str],
user_model: str | None,
infer_model_from_keys: bool | None,
@ -203,7 +204,7 @@ def get_complete_model_list(
def append_unique(models):
for model in models:
if model not in unique_models:
if model not in unique_models and model != SpecialModelNames.no_default_models.value:
unique_models.append(model)
if key_models:

13
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -49,6 +49,7 @@ class AutoRouterTurnTransaction:
cache_hit: bool
cache_ttl_seconds: int | None
cache_touched: bool
tier: str | None = None
class TurnCacheFacts(NamedTuple):
@ -152,11 +153,13 @@ def build_autorouter_turn_transaction(
return None
usage_object_raw: Final = metadata.get("usage_object")
cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None)
tier_raw: Final = routing_decision.get("tier")
return AutoRouterTurnTransaction(
api_key=api_key,
session_id=_bounded_session_id(session_id),
router_name=router_name,
router_type=str(routing_decision.get("router_type") or "unknown"),
tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None,
model=model,
turn_at=turn_at,
total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
@ -184,6 +187,8 @@ _COVERED: Final = _p("covered")
_CACHE_HIT: Final = _p("cache_hit")
_CACHE_TTL: Final = _p("cache_ttl_seconds")
_TOUCHED: Final = _p("cache_touched")
_TIER: Final = f"{_p('tier')}::text"
_TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)"
_IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at"
_SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}"
@ -201,7 +206,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t (
last_model, models, turns, unordered_turns, covered_turns, cache_hits,
same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns
)
VALUES (
{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
@ -211,7 +216,8 @@ VALUES (
0, 0, 0, 0,
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
{_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8
{_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
{_TIER_DELTA}
)
ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
turns = t.turns + 1,
@ -242,6 +248,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
ELSE COALESCE((t.models -> {_MODEL} ->> 'ttl')::int, {_CACHE_TTL}::int) END)
)),
last_model = (CASE WHEN {_IN_ORDER} THEN {_MODEL} ELSE t.last_model END),
tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")}
THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1)
ELSE t.tier_turns END),
first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at),
last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at)
"""

857
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -9,11 +9,14 @@ import os
import sys
sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path
import asyncio
import copy
import json
import re
import sys
from collections.abc import AsyncGenerator, Mapping
from collections.abc import AsyncGenerator, Mapping, Sequence
from datetime import datetime, timezone
from itertools import accumulate, groupby
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast
import httpx
@ -23,6 +26,7 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
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.core_helpers import redact_nested_match_and_regex_keys
@ -46,6 +50,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrailOutput,
BedrockGuardrailQualifier,
BedrockGuardrailResponse,
BedrockGuardrailUsage,
BedrockRequest,
BedrockTextContent,
)
@ -53,6 +58,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from botocore.awsrequest import AWSPreparedRequest
from botocore.credentials import Credentials
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -71,6 +77,17 @@ from litellm.types.utils import (
GUARDRAIL_NAME: Final = "bedrock"
_BEDROCK_DYNAMIC_BODY_DENYLIST: Final = frozenset({"content", "source"})
_BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = (
"text unit",
"maximum input size",
"content size",
"too long",
"too large",
"exceeds the maximum",
)
_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3
_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5
_BEDROCK_WHITESPACE: Final = re.compile(r"\s")
# Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required).
_BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke"
# InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with
@ -118,6 +135,29 @@ class GuardrailMessageFilterResult(NamedTuple):
target_indices: list[int] | None
class BedrockContentChunkResult(NamedTuple):
"""One chunk's ApplyGuardrail response, paired with enough bookkeeping to
reconstruct global masked-output positions once every chunk is back.
`content` is the exact content items this chunk was called with -- needed
so an all-clear chunk (empty `outputs`) can still contribute one unmasked
placeholder per item it covers, keeping every later chunk's masked text
aligned to its original global position. `fragment_group_size` is 1 for an
ordinary chunk, and otherwise the total number of consecutive chunk results
that together make up ONE original content item's own text (split because a
list of length 1 could not be bisected by list length). All of them must be
concatenated back into that one item's masked output rather than treated as
separate items. It is a count rather than a boolean because one item can be
bisected more than once: two levels of splitting produce four fragments for
a single item, not two, and grouping them in fixed pairs would emit two
outputs for one message and shift every later message's masked text.
"""
response: BedrockGuardrailResponse
content: tuple[BedrockContentItem, ...]
fragment_group_size: int
class ApplyGuardrailMessageSelection(NamedTuple):
"""Messages selected for an apply_guardrail scan + write-back metadata."""
@ -168,12 +208,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
content_filter_threshold: float | None = 0.5,
prompt_attack_threshold: float | None = 0.5,
pii_confidence_threshold: float | None = 0.5,
chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
**kwargs,
):
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self.guardrailIdentifier = guardrailIdentifier
self.guardrailVersion = guardrailVersion
self.guardrail_provider = "bedrock"
self.chunk_budget_chars = chunk_budget_chars
self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
# Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
@ -759,12 +801,35 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data: dict | None = None,
logging_event_type: GuardrailEventHooks | None = None,
) -> BedrockGuardrailResponse:
"""Scan `messages`/`response` with ApplyGuardrail, chunking if it is too large.
Content is bin-packed into budget-sized batches and each batch posted
sequentially, every batch independently falling back to bisection if AWS
rejects it. The per-batch responses are merged so callers cannot tell whether
chunking happened.
Content using contextual grounding opts out of chunking entirely: grounding is
scored holistically against the whole reference source, so bisecting it would
fragment that evaluation and yield misleading scores. Such a request keeps the
old behavior of surfacing a too-large error rather than being split.
`logging_event_type` drives what UI and spend logs report. It is distinct from
Bedrock's `source`, which is INPUT vs OUTPUT for the API body and must not be
confused with the proxy hook (pre_call / during_call / post_call); when omitted,
the legacy source-derived mapping is kept for backward compatibility.
A guardrail *block* is logged where it happens, in
`_post_apply_guardrail_content`, because chunking stops immediately and there is
no later merged response to log instead. Everything else that fails out of the
chunking flow (an unrecoverable too-large error, a non-size validation error,
exhausted throttle retries) is a genuine end-to-end failure of this one logical
guardrail call and is logged exactly once here.
"""
start_time: Final = datetime.now(timezone.utc)
credentials, aws_region_name = self._load_credentials()
bedrock_request_data: Final[dict] = dict(
self.convert_to_bedrock_format(source=source, messages=messages, response=response)
)
bedrock_guardrail_response: BedrockGuardrailResponse = BedrockGuardrailResponse()
api_key: str | None = None
if request_data:
dynamic_request_body_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data)
@ -778,6 +843,257 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
if request_data.get("api_key") is not None:
api_key = request_data["api_key"]
event_type: Final = (
logging_event_type
if logging_event_type is not None
else (GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call)
)
content: Final[tuple[BedrockContentItem, ...]] = tuple(bedrock_request_data.get("content") or ())
allow_chunking: Final = not self._content_uses_contextual_grounding(content)
try:
responses: Final = await self._apply_guardrail_content_with_chunking(
content=content,
base_request_data=bedrock_request_data,
credentials=credentials,
aws_region_name=aws_region_name,
api_key=api_key,
request_data=request_data,
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
)
except HTTPException as exc:
if not isinstance(exc.detail, dict):
self._log_apply_guardrail_failure(
detail=exc.detail,
request_data=request_data,
event_type=event_type,
start_time=start_time,
)
raise
merged_response: Final = self._merge_bedrock_guardrail_responses(responses)
self._log_apply_guardrail_success(
merged_response=merged_response,
request_data=request_data,
event_type=event_type,
start_time=start_time,
)
return merged_response
async def _apply_guardrail_content_with_chunking(
self,
content: Sequence[BedrockContentItem],
base_request_data: Mapping[str, Any],
credentials: "Credentials",
aws_region_name: str,
api_key: str | None,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
allow_chunking: bool,
) -> tuple[BedrockContentChunkResult, ...]:
"""Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large.
Tries `content` as a single call first. AWS's per-request "maximum input
size in text units" quota is account/region/policy-dependent and cannot be
predicted ahead of time, so it is only ever discovered reactively: on an
error whose message indicates the input was too large (a ThrottlingException
in practice, a ValidationException per the docs -- see
``_is_input_too_large_error``), the content is re-sent in smaller pieces.
Probing with the whole payload first is what keeps a request AWS would have
accepted at exactly one call. Packing into fixed batches up front instead
would split conversations AWS was happy to take whole, multiplying billed
calls and guardrail latency on traffic that never had a size problem, and
no fixed budget can avoid that because the real cap is unknown here.
Once a rejection proves the payload is over the cap, a multi-item payload is
re-sent as ``chunk_budget_chars``-sized batches rather than bisected: that
reaches a working size in one step instead of paying an O(log n) ladder of
rejected calls. Bisection remains the fallback for anything bin-packing
cannot make smaller, which is what makes the recursion terminate: a batch
already inside the budget packs back to itself, so it falls through to the
split below. A single oversized
content item (one very long message) is split by its own text instead of
by list length, since a list of length 1 has no items left to bisect --
the resulting fragments all carry a ``fragment_group_size`` so the merge
step can recombine them into the one content item they came from, rather
than treating each fragment as its own item when reconstructing positions
for masking. That count covers however many fragments the item ended up
split into, not just two, since it can be bisected repeatedly: the
outermost single-item split stamps the total leaf count on every leaf
below it, overwriting any smaller count an inner split had set. A real
guardrail block on any (sub-)chunk raises immediately
-- callers must not lose that signal by continuing to post the remaining
chunks.
"""
try:
response: Final = await self._post_apply_guardrail_content_with_retry(
content=content,
base_request_data=base_request_data,
credentials=credentials,
aws_region_name=aws_region_name,
api_key=api_key,
request_data=request_data,
event_type=event_type,
start_time=start_time,
)
return (
BedrockContentChunkResult(
response=response,
content=tuple(content),
fragment_group_size=1,
),
)
except HTTPException as exc:
if allow_chunking and self._is_input_too_large_error(exc.detail):
batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars)
if len(batches) > 1:
verbose_proxy_logger.warning(
"Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; "
"re-sending as %d batches of at most %d characters",
len(content),
len(batches),
self.chunk_budget_chars,
)
batch_results: Final = [ # mutable-ok: await needs a list comprehension; frozen to a tuple below
await self._apply_guardrail_content_with_chunking(
content=batch,
base_request_data=base_request_data,
credentials=credentials,
aws_region_name=aws_region_name,
api_key=api_key,
request_data=request_data,
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
)
for batch in batches
]
return tuple(result for results in batch_results for result in results)
split_content: Final = self._split_bedrock_content(content)
if split_content is None:
raise
first_half, second_half = split_content
is_single_item_text_split: Final = len(content) == 1
verbose_proxy_logger.warning(
"Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; "
"splitting into %d + %d and retrying each",
len(content),
len(first_half),
len(second_half),
)
first_results: Final = await self._apply_guardrail_content_with_chunking(
content=first_half,
base_request_data=base_request_data,
credentials=credentials,
aws_region_name=aws_region_name,
api_key=api_key,
request_data=request_data,
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
)
second_results: Final = await self._apply_guardrail_content_with_chunking(
content=second_half,
base_request_data=base_request_data,
credentials=credentials,
aws_region_name=aws_region_name,
api_key=api_key,
request_data=request_data,
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
)
combined_results: Final = tuple(first_results) + tuple(second_results)
if is_single_item_text_split:
return tuple(
result._replace(fragment_group_size=len(combined_results)) for result in combined_results
)
return combined_results
raise
async def _post_apply_guardrail_content_with_retry(
self,
content: Sequence[BedrockContentItem],
base_request_data: Mapping[str, Any],
credentials: "Credentials",
aws_region_name: str,
api_key: str | None,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
) -> BedrockGuardrailResponse:
"""Post one ApplyGuardrail call for `content`, retrying with exponential
backoff on AWS ThrottlingException (HTTP 429).
Chunking already trades one oversized call for several smaller ones, so
retries here are capped low -- they must not multiply per-request latency
by an order of magnitude when the account's per-second text-unit quota is
the binding constraint rather than the per-request size quota.
A too-large rejection is deliberately excluded from the retry. AWS reports
it as a ThrottlingException (429), not only as a ValidationException, but
unlike a genuine throttle it is not transient: re-posting the same
oversized content can never succeed. Retrying it would burn every backoff
sleep and every (billed) attempt before the caller's bisection gets a
chance to split the content, at every level of the recursion.
"""
for attempt in range(_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + 1):
try:
return await self._post_apply_guardrail_content(
content=content,
base_request_data=base_request_data,
credentials=credentials,
aws_region_name=aws_region_name,
api_key=api_key,
request_data=request_data,
event_type=event_type,
start_time=start_time,
)
except HTTPException as exc:
if (
exc.status_code != 429
or self._is_input_too_large_error(exc.detail)
or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES
):
raise
await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt))
raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted")
async def _post_apply_guardrail_content(
self,
content: Sequence[BedrockContentItem],
base_request_data: Mapping[str, Any],
credentials: "Credentials",
aws_region_name: str,
api_key: str | None,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
) -> BedrockGuardrailResponse:
"""Make exactly one signed ApplyGuardrail HTTP call for `content` and
parse the result. Raises HTTPException on a guardrail block or any
non-200 response (including 429, handled by the retry wrapper above).
AWS also reports some failures inside a 200 body, tagging ``Output.__type``
with an Exception marker. Those deliberately do NOT raise: the request proceeds,
matching the behaviour of this code before chunking existed. The marker survives
the merge, so the one consolidated log entry still records
``guardrail_failed_to_respond`` rather than a success. Making that path fail
closed is a separate change, tracked apart from this PR, and belongs behind the
existing ``unreachable_fallback`` setting rather than a hardcoded status.
A block is logged here rather than by the caller: it ends the whole chunking
flow immediately, with no further chunks attempted, so there is no later
merged response for the caller to log instead.
"""
bedrock_request_data: Final = { # mutable-ok: outbound JSON request body
**base_request_data,
"content": content,
} # mutable-ok: outbound JSON request body
prepared_request: Final = self._prepare_request(
credentials=credentials,
data=bedrock_request_data,
@ -792,42 +1108,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prepared_request.headers,
)
# UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API
# body, which must not be confused with the proxy hook (pre_call / during_call /
# post_call). When omitted, keep legacy mapping for backward compatibility.
if logging_event_type is not None:
event_type = logging_event_type
else:
event_type = GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call
httpx_response: Final = await self._sign_and_post(
prepared_request=prepared_request,
request_data=request_data,
event_type=event_type,
start_time=start_time,
log_transport_failure=False,
)
#########################################################
# Add guardrail information to request trace
#########################################################
_json_response: Final = httpx_response.json()
tracing_detail: Final = self._build_tracing_detail(_json_response)
# Raw Bedrock JSON is passed here; match/regex redaction runs once inside
# CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=_json_response,
request_data=request_data or {},
guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response),
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
event_type=event_type,
tracing_detail=tracing_detail or None,
)
#########################################################
if httpx_response.status_code == 200:
_json_response: Final = httpx_response.json()
# check if the response was flagged
verbose_proxy_logger.debug(
"Bedrock AI response : %s",
@ -835,19 +1125,462 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response):
self._log_apply_guardrail_attempt(
httpx_response=httpx_response,
json_response=_json_response,
request_data=request_data,
event_type=event_type,
start_time=start_time,
)
raise self._get_http_exception_for_blocked_guardrail(
bedrock_guardrail_response, request_data=request_data
)
else:
status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
verbose_proxy_logger.error(
"Bedrock AI: error in response. Status code: %s, response: %s",
httpx_response.status_code,
httpx_response.text,
)
raise HTTPException(status_code=status_code, detail=detail_message)
return bedrock_guardrail_response
return bedrock_guardrail_response
status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
verbose_proxy_logger.error(
"Bedrock AI: error in response. Status code: %s, response: %s",
httpx_response.status_code,
httpx_response.text,
)
raise HTTPException(status_code=status_code, detail=detail_message)
def _log_apply_guardrail_attempt(
self,
httpx_response: httpx.Response,
json_response: dict, # mutable-ok: raw AWS JSON payload
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
) -> None:
"""Log a single ApplyGuardrail HTTP attempt as-is (its own status,
derived from its own response). Used only for the blocked-content
case, which ends the whole chunking flow immediately."""
tracing_detail: Final = self._build_tracing_detail(BedrockGuardrailResponse(**json_response))
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=json_response,
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response),
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
event_type=event_type,
tracing_detail=tracing_detail or None,
)
def _log_apply_guardrail_success(
self,
merged_response: BedrockGuardrailResponse,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
) -> None:
"""Log one logical ApplyGuardrail call -- possibly several chunk calls
under the hood -- using its final merged response, so a chunked
request produces exactly one telemetry entry, the same as an
unchunked one would.
AWS can report a failure inside an HTTP 200 body by tagging
``Output.__type`` with an exception marker. That marker survives the merge,
so the status is derived from the merged response rather than assumed to be
a success, which is what the pre-chunking code reported for that shape."""
tracing_detail: Final = self._build_tracing_detail(merged_response)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=dict(merged_response), # mutable-ok: logging helper requires a dict
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status=(
"guardrail_failed_to_respond"
if "Exception" in str((merged_response.get("Output") or {}).get("__type", ""))
else "success"
),
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
event_type=event_type,
tracing_detail=tracing_detail or None,
)
def _log_apply_guardrail_failure(
self,
detail: object,
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
) -> None:
"""Log one logical ApplyGuardrail call that failed end-to-end (an
unrecoverable too-large error, a non-size validation error, or
exhausted throttle retries) as a single failure, rather than logging
every failed attempt chunking made along the way."""
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": str(detail)}, # mutable-ok: logging helper requires a dict
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
event_type=event_type,
)
@staticmethod
def _content_uses_contextual_grounding(content: Sequence[BedrockContentItem]) -> bool:
"""True if any content item carries a contextual-grounding qualifier
(``grounding_source``, ``query``, or the ``guard_content`` the response
itself is tagged with once grounding is present)."""
for item in content:
if (item.get("text") or {}).get("qualifiers"): # mutable-ok: read-only empty fallback
return True
return False
@staticmethod
def _bin_pack_bedrock_content(
content: Sequence[BedrockContentItem],
budget: int,
) -> tuple[tuple[BedrockContentItem, ...], ...]:
"""Pack whole content items, in order, into batches whose combined text
length stays within `budget`, in a single pass that carries the running
total rather than re-summing the open batch per item.
This is the fast-path half of the hybrid chunking strategy: bin-packing
at a conservative fixed budget keeps the common case at O(n / budget)
ApplyGuardrail calls instead of the O(log n) round trips pure reactive
bisection pays on every oversized request. An item whose own text
already exceeds `budget` is not split here -- it becomes its own
(still oversized) batch and is sent as-is; if AWS rejects that batch as
too large, `_apply_guardrail_content_with_chunking`'s existing
recursive-bisection fallback takes over for that batch only.
`budget` comes from the guardrail's ``chunk_budget_chars`` setting and
defaults to 25,000, matching ApplyGuardrail's default quota of 25 text
units (roughly 1,000 characters each) per second. Packing to that size and
posting sequentially is what keeps chunking from tripping the rate quota
and trading a size error for a throttle. Accounts with raised quotas can
configure a larger budget to spend fewer calls.
The budget is not a correctness dependency either way. AWS's effective cap
varies by account, region, and policy, is not a fixed character count, and
cannot be read from config, so any batch it still rejects falls back to
bisection, which self-corrects however wrong the value was. An over-large
budget therefore costs one extra probe-and-bisect round trip rather than
failing the request.
"""
if not content:
return (tuple(content),)
lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content)
def assign(carried: tuple[int, int], length: int) -> tuple[int, int]:
batch_index, used = carried
if used + length <= budget:
return batch_index, used + length
return batch_index + 1, length
batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:])
return tuple(
tuple(item for _, item in group)
for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0])
)
@staticmethod
def _split_bedrock_content(
content: Sequence[BedrockContentItem],
) -> tuple[tuple[BedrockContentItem, ...], tuple[BedrockContentItem, ...]] | None:
"""Bisect `content` into two roughly-equal, non-empty halves.
When `content` already holds more than one item, it is split by list
length. When it holds exactly one item, that item's own text is split
instead (a list of length 1 has no items left to bisect, but one very
long message is still a single content item) -- at the whitespace
character nearest the midpoint rather than a raw character index, so
the cut never lands inside a word/token. This is a plain, lossless
cut with no overlap: concatenating the two fragments in order always
reproduces the original text exactly, so merging back at
``_merge_logical_unit_outputs`` needs no reconciliation step.
Known, accepted limitation: whitespace splitting only guards against
*accidentally* severing a single token (one denied word, one PII
pattern) across the cut. It does not, and cannot without an overlap
window, stop a *multi-word* denied phrase deliberately positioned to
straddle the boundary -- each fragment can scan clean on its own and
still reassemble into the flagged phrase. AWS's own guidance on this
API acknowledges the same gap for input chunking ("a critical piece of
text could span two (or more) chunks if not carefully divided") with
no documented resolution, and overlap-and-reconcile was evaluated and
rejected for this PR: AWS's masking output has no documented
length-preservation guarantee, so reconciling an overlap region against
masked text is not sound in general. Out of scope for this PR.
Returns None when there is nothing left to split -- a single item
whose text is too short to halve into two non-empty pieces -- so the
caller can give up and propagate the original too-large error instead
of recursing forever.
"""
if len(content) > 1:
midpoint: Final = max(1, len(content) // 2)
return tuple(content[:midpoint]), tuple(content[midpoint:])
text_content: Final = content[0].get("text") or BedrockTextContent()
text: Final = text_content.get("text") or ""
if len(text) < 2:
return None
split_at: Final = BedrockGuardrail._nearest_whitespace_split_index(text)
qualifiers: Final = text_content.get("qualifiers")
def fragment(piece: str) -> BedrockContentItem:
block: Final = (
BedrockTextContent(text=piece, qualifiers=qualifiers) if qualifiers else BedrockTextContent(text=piece)
)
return BedrockContentItem(text=block)
return (fragment(text[:split_at]),), (fragment(text[split_at:]),)
@staticmethod
def _nearest_whitespace_split_index(text: str) -> int:
"""Return the index nearest `text`'s midpoint that falls on a whitespace
boundary, so splitting `text[:i]` / `text[i:]` there never severs a word.
Any Unicode whitespace counts, not just an ASCII space. Matching only `" "`
would leave the boundary unguarded for exactly the payloads that get large
enough to need splitting: JSON lines, source code, logs and transcripts are
newline or tab delimited, so a deny-listed word sitting at the midpoint of
one would be cut in half, scan clean on both fragments, and reassemble
intact.
The returned index always leaves both sides non-empty, which is what makes
the caller's recursion terminate. A boundary that would put the split at 0
or at ``len(text)`` is discarded: it would hand back a fragment identical to
the text just rejected as too large, AWS would reject that again, and each
retry would re-split it into the same unchanged fragment until the stack ran
out. The dangerous shape is a text whose only space at or after the midpoint
is its final character.
Falls back to the raw midpoint when no usable whitespace boundary exists, either
because `text` has none at all (a single giant token) or because the only
candidates were degenerate. That is still a correct, lossless split, just no
longer guaranteed word-safe for those cases. `text` must be at least two
characters, which `_split_bedrock_content` guarantees, so the midpoint itself
is never degenerate.
"""
midpoint: Final = len(text) // 2
before: Final = max((found.end() for found in _BEDROCK_WHITESPACE.finditer(text, 0, midpoint)), default=None)
after_match: Final = _BEDROCK_WHITESPACE.search(text, midpoint)
candidates: Final = sorted(
(split for split in (before, after_match.end() if after_match else None) if split is not None),
key=lambda split: abs(split - midpoint),
)
return next((split for split in candidates if 0 < split < len(text)), midpoint)
@staticmethod
def _is_input_too_large_error(detail: object) -> bool:
"""True if `detail` is an AWS error message for input exceeding the
per-request text-unit quota.
Matched on the message rather than the status code on purpose: AWS is not
consistent about which error it raises for this. Observed against a live
guardrail with an active content-filter policy, an oversized request comes
back as a *ThrottlingException* (429) reading ``Input text size (3273 text
units) exceeds the maximum allowed (1000 text units) for the content filter
policy (Classic tier)``, while the documented failure mode is a
ValidationException (400). Keying off the message covers both.
A guardrail *block* is also raised as an HTTPException with status 400,
but its ``detail`` is always a dict (built by
``_get_http_exception_for_blocked_guardrail``); a non-200 API error's
``detail`` is always the plain string returned by
``_parse_bedrock_guardrail_error_response``. Checking ``isinstance(detail,
str)`` is therefore sufficient to never mistake a real block for a
too-large error.
"""
if not isinstance(detail, str):
return False
lowered: Final = detail.lower()
return any(substring in lowered for substring in _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS)
@staticmethod
def _merge_bedrock_guardrail_responses(
chunk_results: Sequence[BedrockContentChunkResult],
) -> BedrockGuardrailResponse:
"""Merge the per-chunk ApplyGuardrail responses of a chunked request into
one, so a caller cannot tell whether chunking happened.
Only ever called with responses that all passed (a block raises
immediately from ``_apply_guardrail_content_with_chunking`` and is never
added to this list). ``action`` is only set on the merged response when
at least one chunk's raw response included it, and left absent otherwise
-- mirroring a real single-call response and matching what
``_build_tracing_detail`` treats as "Bedrock didn't report an action".
Fields this merge has no opinion on (``actionReason``, ``guardrailCoverage``,
``blockedResponse``, anything AWS adds later) are carried over from the chunk
responses rather than dropped, so the response and the logged telemetry keep
the shape a single unchunked call returned. The merged keys below win.
Per AWS's documented ApplyGuardrail contract, a single call's ``outputs``
is positionally parallel to the ``content`` items *of that call*: an
entry per item when anything in the call was masked, or an empty list
when nothing in the whole call was masked. Downstream masking
(``_apply_masking_to_messages``) walks the merged ``outputs`` by a single
running index across the *original, unchunked* message list, so a later
chunk's masked text must land at the same global position it would have
if chunking had never happened. Naively concatenating each chunk's
``outputs`` breaks that whenever a chunk had nothing masked (its empty
list would otherwise silently swallow its items' slots, shifting every
later chunk's masked text left onto the wrong message). So every
item -- masked or not -- always contributes exactly one entry here,
falling back to that item's own original (unmasked) text when its
chunk returned no output for it; a wholly-untouched result is then
collapsed back to an empty ``outputs`` list to match a real single-call
no-op response. A chunk that returns a nonzero output count not equal
to its item count is passed through as-is instead of guessed at, since
AWS's docs don't cover partial masking within one multi-item call.
"""
logical_units: Final = BedrockGuardrail._group_fragment_units(chunk_results)
per_unit_outputs: Final = tuple(BedrockGuardrail._merge_logical_unit_outputs(unit) for unit in logical_units)
merged_outputs: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list
output for outputs, _ in per_unit_outputs for output in outputs
]
any_masked: Final = any(masked for _, masked in per_unit_outputs)
actions: Final = tuple(
chunk_result.response.get("action")
for chunk_result in chunk_results
if isinstance(chunk_result.response.get("action"), str)
)
merged_action: Final = (
"GUARDRAIL_INTERVENED" if "GUARDRAIL_INTERVENED" in actions else (actions[-1] if actions else None)
)
merged_assessments: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list
assessment
for chunk_result in chunk_results
for assessment in (chunk_result.response.get("assessments") or []) # mutable-ok: logged payload
]
any_usage_reported: Final = any(chunk_result.response.get("usage") for chunk_result in chunk_results)
merged: Final[BedrockGuardrailResponse] = cast( # cast-ok: TypedDict assembled from a comprehension
BedrockGuardrailResponse,
{ # mutable-ok: builds the TypedDict payload
key: value for chunk_result in chunk_results for key, value in chunk_result.response.items()
},
)
if merged_action is not None:
merged["action"] = merged_action
if merged_outputs and any_masked:
merged["outputs"] = merged_outputs
merged["output"] = merged_outputs
if merged_assessments:
merged["assessments"] = merged_assessments
if any_usage_reported:
merged["usage"] = BedrockGuardrail._sum_bedrock_guardrail_usage(chunk_results)
return merged
@staticmethod
def _sum_bedrock_guardrail_usage(
chunk_results: Sequence[BedrockContentChunkResult],
) -> BedrockGuardrailUsage:
"""Sum each chunk's ``usage`` counters field-by-field into one totals dict.
Keys are taken from the responses rather than from a fixed list, so a counter
this code does not know about (AWS has added several) is still summed and
reported instead of being silently dropped to zero."""
chunk_usages: Final = tuple(
chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback
for chunk_result in chunk_results
)
return cast( # cast-ok: TypedDict assembled from a comprehension
BedrockGuardrailUsage,
{ # mutable-ok: builds the TypedDict payload
key: sum(usage.get(key) or 0 for usage in chunk_usages)
for key in dict.fromkeys(key for usage in chunk_usages for key in usage)
},
)
@staticmethod
def _group_fragment_units(
chunk_results: Sequence[BedrockContentChunkResult],
) -> tuple[tuple[BedrockContentChunkResult, ...], ...]:
"""Group consecutive text-fragment chunk results back into the one content
item each group came from, leaving every ordinary chunk result as a unit of
one.
The group size is read off the results themselves rather than assumed,
because a single content item can be bisected repeatedly: two levels of
splitting yield four fragments for one item, not two. Assuming a fixed pair
here would emit two outputs for one message and shift every later message's
masked text onto the wrong message."""
def advance(carried: tuple[int, bool], result: BedrockContentChunkResult) -> tuple[int, bool]:
remaining, _ = carried
if remaining == 0:
return max(1, result.fragment_group_size) - 1, True
return remaining - 1, False
starts: Final = tuple(
index
for index, (_, starts_unit) in enumerate(tuple(accumulate(chunk_results, advance, initial=(0, False)))[1:])
if starts_unit
)
return tuple(tuple(chunk_results[start:end]) for start, end in zip(starts, starts[1:] + (len(chunk_results),)))
@staticmethod
def _merge_logical_unit_outputs(
unit: tuple[BedrockContentChunkResult, ...],
) -> tuple[tuple[BedrockGuardrailOutput, ...], bool]:
"""Reduce one logical unit (a fragment group of any size, or a single chunk
result) to the ``BedrockGuardrailOutput`` entries it contributes to the
merged response, plus whether any masking actually happened in it.
Per AWS's documented ApplyGuardrail contract, a single call's
``outputs`` is positionally parallel to the ``content`` items *of that
call*: an entry per item when anything in the call was masked, or an
empty list when nothing in the whole call was masked. Downstream
masking (``_apply_masking_to_messages``) walks the merged ``outputs``
by a single running index across the *original, unchunked* message
list, so a later chunk's masked text must land at the same global
position it would have if chunking had never happened. So every item
-- masked or not -- always contributes exactly one entry here, falling
back to that item's own original (unmasked) text when its chunk
returned no output for it. A chunk that returns a nonzero output count
not equal to its item count is passed through as-is instead of guessed
at, since AWS's docs don't cover partial masking within one multi-item
call.
A unit holding more than one result is a fragment group: every result in it
is one fragment of a single content item's text, so the group collapses to
one entry built from each fragment's masked text (or that fragment's own
original text where it came back unmasked), concatenated in order. This
holds for any group size, not only two.
"""
if len(unit) > 1:
def fragment_outputs(result: BedrockContentChunkResult) -> tuple[BedrockGuardrailOutput, ...]:
return tuple(result.response.get("outputs") or result.response.get("output") or ())
def fragment_text(result: BedrockContentChunkResult) -> str:
source: Final = (result.content[0].get("text") or {}).get( # mutable-ok: read-only fallback
"text"
) or ""
outputs: Final = fragment_outputs(result)
masked: Final = outputs[0].get("text") if outputs else None
return masked if masked is not None else source
merged_text: Final = "".join(fragment_text(result) for result in unit)
any_masked: Final = any(fragment_outputs(result) for result in unit)
return (BedrockGuardrailOutput(text=merged_text),), any_masked
(chunk_result,) = unit
chunk_outputs: Final = chunk_result.response.get("outputs") or chunk_result.response.get("output") or ()
if len(chunk_outputs) == len(chunk_result.content):
return tuple(chunk_outputs), bool(chunk_outputs)
if not chunk_outputs:
return tuple(
BedrockGuardrailOutput(
text=(item.get("text") or {}).get("text") or "" # mutable-ok: read-only fallback
)
for item in chunk_result.content
), False
return tuple(chunk_outputs), True
async def _sign_and_post(
self,
@ -855,6 +1588,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data: dict | None,
event_type: GuardrailEventHooks,
start_time: "datetime",
log_transport_failure: bool = True,
) -> httpx.Response:
"""POST a signed Bedrock request, logging+raising on network/HTTP errors.
@ -862,6 +1596,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
transport-error handling cannot drift. Returns the raw ``httpx.Response`` on
success (including non-2xx that httpx did not raise on); the 200-path logging,
status and tracing stay with each caller because the two APIs report differently.
``log_transport_failure=False`` suppresses the ``guardrail_failed_to_respond``
entry for a non-200 that is re-raised as an ``HTTPException``, for callers that
own consolidated per-request logging. The ApplyGuardrail path needs this:
``AsyncHTTPHandler.post`` calls ``raise_for_status()``, so every non-200 lands
in this handler, and one logical request can legitimately produce several of
them (a too-large probe, then each rejected bisection level) while still
succeeding overall. Logging per attempt would report a recovered request as
several failures plus a success.
The connection-level branch below (timeout, endpoint down) still logs
unconditionally: it re-raises the original exception rather than an
``HTTPException``, so no consolidating caller catches it, and suppressing it
would drop the only record of the failure.
"""
try:
return await self.async_handler.post(
@ -882,16 +1630,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
status_code,
detail_message,
) = self._parse_bedrock_guardrail_error_response(err_response)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": detail_message},
request_data=request_data or {},
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
event_type=event_type,
)
if log_transport_failure:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={ # mutable-ok: logging helper requires a dict
"error": detail_message
},
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
event_type=event_type,
)
raise HTTPException(status_code=status_code, detail=detail_message) from e
except HTTPException:
raise
@ -900,7 +1651,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": str(e)},
request_data=request_data or {},
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
@ -1027,7 +1778,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": detail_message},
request_data=request_data or {},
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
@ -1043,7 +1794,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": str(e)},
request_data=request_data or {},
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
@ -1061,7 +1812,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=self._sanitize_invoke_checks_response_for_logging(json_response),
request_data=request_data or {},
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status=self._get_invoke_checks_status(bool(violations)),
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),

1
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -20,6 +20,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
content_filter_threshold=litellm_params.content_filter_threshold,
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
pii_confidence_threshold=litellm_params.pii_confidence_threshold,
chunk_budget_chars=litellm_params.chunk_budget_chars,
default_on=litellm_params.default_on,
disable_exception_on_block=litellm_params.disable_exception_on_block,
mask_request_content=litellm_params.mask_request_content,

30
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -4,8 +4,9 @@ AUTO ROUTER MANAGEMENT ENDPOINTS
POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config
"""
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final
from pydantic import BaseModel, TypeAdapter
@ -265,6 +266,7 @@ async def preview_auto_router_routing(
class _SessionAggRow(BaseModel):
router_name: str
router_type: str
tier_turns: Mapping[str, int]
sessions: int
turns: int
unordered_turns: int
@ -289,6 +291,23 @@ class _SessionAggRow(BaseModel):
_SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow])
_BENCHMARKS_SQL: Final = """
WITH windowed AS (
SELECT * FROM "LiteLLM_AutoRouterSession"
WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
),
tier_maps AS (
SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns
FROM (
SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns
FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv
GROUP BY router_name, router_type, kv.key
) per_tier
GROUP BY router_name, router_type
)
SELECT
agg.*,
COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns
FROM (
SELECT
router_name,
router_type,
@ -311,10 +330,11 @@ SELECT
COALESCE(SUM(spend), 0)::float8 AS spend,
COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
FROM "LiteLLM_AutoRouterSession"
WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
FROM windowed
GROUP BY router_name, router_type
ORDER BY SUM(spend) DESC
) agg
LEFT JOIN tier_maps USING (router_name, router_type)
ORDER BY agg.spend DESC
"""
@ -371,6 +391,7 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
return _SessionAggRow(
router_name="",
router_type="",
tier_turns=MappingProxyType({}),
sessions=sum(row.sessions for row in rows),
turns=sum(row.turns for row in rows),
unordered_turns=sum(row.unordered_turns for row in rows),
@ -448,6 +469,7 @@ async def get_auto_router_benchmarks(
AutoRouterBenchmarkGroup(
router_name=row.router_name,
router_type=row.router_type,
tier_turns=row.tier_turns,
**_benchmark_totals(row).model_dump(),
)
for row in rows

88
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -57,6 +57,7 @@ from litellm.proxy._types import (
SpecialManagementEndpointEnums,
SpecialModelNames,
SpecialProxyStrings,
TeamAccessGroupModelGrant,
TeamAddMemberResponse,
TeamInfoResponseObject,
TeamInfoResponseObjectTeamTable,
@ -3829,7 +3830,7 @@ async def _add_team_member_budget_table(
) -> TeamInfoResponseObjectTeamTable:
try:
team_budget: Final = await _budget_db(prisma_client).find_unique(where={"budget_id": team_member_budget_id})
team_info_response_object.team_member_budget_table = team_budget
return team_info_response_object.model_copy(update={"team_member_budget_table": team_budget})
except Exception:
verbose_proxy_logger.info(
"Team member budget table not found, passed team_member_budget_id=%s", team_member_budget_id
@ -3838,21 +3839,34 @@ async def _add_team_member_budget_table(
return team_info_response_object
async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None:
"""Populate access_group_models / mcp_server_ids / agent_ids on the team
info response by resolving inherited resources from its access groups."""
async def _resolve_team_access_group_resources(
_team_info: TeamInfoResponseObjectTeamTable,
) -> TeamInfoResponseObjectTeamTable:
"""Return a copy of the team info with access_group_models / mcp_server_ids /
agent_ids / details resolved from its access groups."""
if not _team_info.access_group_ids:
return
return _team_info
ag_lookup: Final = await _batch_resolve_access_group_resources(_team_info.access_group_ids)
models, mcp_ids, agent_ids = set(), set(), set()
for ag_id in _team_info.access_group_ids:
if ag_id in ag_lookup:
models.update(ag_lookup[ag_id]["models"])
mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"])
agent_ids.update(ag_lookup[ag_id]["agent_ids"])
_team_info.access_group_models = list(models)
_team_info.access_group_mcp_server_ids = list(mcp_ids)
_team_info.access_group_agent_ids = list(agent_ids)
resolved_groups: Final = tuple(
ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup
)
return _team_info.model_copy(
update={
"access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}),
"access_group_mcp_server_ids": list(
{s for group in resolved_groups for s in (group.access_mcp_server_ids or [])}
),
"access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}),
"access_group_details": tuple(
TeamAccessGroupModelGrant(
access_group_id=group.access_group_id,
access_group_name=group.access_group_name,
models=tuple(group.access_model_names or ()),
)
for group in resolved_groups
),
}
)
@router.get("/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)])
@ -3958,11 +3972,11 @@ async def team_info(
)
# Resolve resources inherited from access groups
await _resolve_team_access_group_resources(_team_info)
resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info)
response_object: Final = TeamInfoResponseObject(
team_id=team_id,
team_info=_team_info,
team_info=resolved_team_info,
keys=keys,
team_memberships=returned_tm,
)
@ -4391,32 +4405,21 @@ async def _build_team_list_where_conditions(
async def _batch_resolve_access_group_resources(
all_access_group_ids: list[str],
) -> dict[str, dict[str, list[str]]]:
) -> dict[str, LiteLLM_AccessGroupTable]:
"""
Batch-fetch access groups in a single DB query and return a per-group
resource map.
Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}.
Missing/invalid groups are silently omitted.
Batch-fetch access groups in a single DB query and return them keyed by
access_group_id. Missing/invalid groups are silently omitted.
"""
from litellm.proxy.proxy_server import prisma_client as _prisma_client
if not all_access_group_ids or _prisma_client is None:
return {}
unique_ids: Final = list(set(all_access_group_ids))
unique_ids: Final = tuple(frozenset(all_access_group_ids))
rows: Final = await _access_group_db(_prisma_client).find_many(
where={"access_group_id": {"in": unique_ids}},
)
result: Final[dict[str, dict[str, list[str]]]] = {}
for row in rows:
result[row.access_group_id] = {
"models": list(row.access_model_names or []),
"mcp_server_ids": list(row.access_mcp_server_ids or []),
"agent_ids": list(row.access_agent_ids or []),
}
return result
return {row.access_group_id: row for row in rows}
def _convert_teams_to_response_models(
@ -4710,15 +4713,18 @@ async def list_team_v2(
all_ag_ids: Final = [ag_id for t in team_items_with_ag for ag_id in (t.access_group_ids or [])]
ag_lookup: Final = await _batch_resolve_access_group_resources(all_ag_ids)
for team_item in team_items_with_ag:
models, mcp_ids, agent_ids = set(), set(), set()
for ag_id in team_item.access_group_ids or []:
if ag_id in ag_lookup:
models.update(ag_lookup[ag_id]["models"])
mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"])
agent_ids.update(ag_lookup[ag_id]["agent_ids"])
team_item.access_group_models = list(models)
team_item.access_group_mcp_server_ids = list(mcp_ids)
team_item.access_group_agent_ids = list(agent_ids)
team_groups = tuple(
ag_lookup[ag_id] for ag_id in (team_item.access_group_ids or []) if ag_id in ag_lookup
)
team_item.access_group_models = list(
{m for group in team_groups for m in (group.access_model_names or [])}
)
team_item.access_group_mcp_server_ids = list(
{s for group in team_groups for s in (group.access_mcp_server_ids or [])}
)
team_item.access_group_agent_ids = list(
{a for group in team_groups for a in (group.access_agent_ids or [])}
)
return {
"teams": team_list,

30
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -1,6 +1,7 @@
import base64
import mimetypes
import re
from collections.abc import Mapping
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Optional
@ -16,6 +17,7 @@ if TYPE_CHECKING:
from prisma.models import LiteLLM_ManagedObjectTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
from litellm.types.utils import LiteLLMBatch
@ -1002,6 +1004,34 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None:
pass
async def map_raw_file_ids_to_unified(
raw_file_ids: frozenset[str], prisma_client: "PrismaClient | None"
) -> Mapping[str, str]:
if not raw_file_ids or not prisma_client:
return MappingProxyType({})
managed_files: Final = await ManagedFileRepository(prisma_client).table.find_many(
where={"flat_model_file_ids": {"hasSome": sorted(raw_file_ids)}} # mutable-ok: prisma where is a plain dict
)
return MappingProxyType(
{
raw_id: managed_file.unified_file_id
for managed_file in managed_files
for raw_id in managed_file.flat_model_file_ids
if raw_id in raw_file_ids
}
)
def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mapping[str, str]) -> None:
for file_attr, raw_id in (
("input_file_id", getattr(response, "input_file_id", None)),
("output_file_id", getattr(response, "output_file_id", None)),
("error_file_id", getattr(response, "error_file_id", None)),
):
if isinstance(raw_id, str) and raw_id in unified_id_by_raw_id:
setattr(response, file_attr, unified_id_by_raw_id[raw_id])
async def ensure_batch_response_managed_file_ids(
response,
managed_files_obj,

1
litellm/proxy/schema.prisma
View file

@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
tier_turns Json @default("{}")
@@id([api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")

136
litellm/proxy/utils.py
View file

@ -165,6 +165,7 @@ if TYPE_CHECKING:
from prisma.client import TransactionManager
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
@ -6419,6 +6420,74 @@ def construct_database_url_from_env_vars() -> str | None:
return None
async def _get_validated_team_object(
user_api_key_dict: "UserAPIKeyAuth",
team_id: str,
prisma_client: "PrismaClient",
user_api_key_cache: "UserApiKeyCache",
proxy_logging_obj: "ProxyLogging",
) -> "LiteLLM_TeamTableCachedObj":
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.management_endpoints.team_endpoints import validate_membership
team_object: Final = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object)
return team_object
async def _get_team_object_for_access_groups(
team_id: str | None,
prisma_client: Optional["PrismaClient"],
user_api_key_cache: Optional["UserApiKeyCache"],
proxy_logging_obj: Optional["ProxyLogging"],
) -> Optional["LiteLLM_TeamTableCachedObj"]:
from litellm.proxy.auth.auth_checks import get_team_object
if team_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None:
return None
try:
return await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
except HTTPException:
verbose_proxy_logger.debug("Could not fetch team %s while listing models", team_id)
return None
async def _get_access_group_models(
user_api_key_dict: "UserAPIKeyAuth",
team_object: Optional["LiteLLM_TeamTableCachedObj"],
prisma_client: Optional["PrismaClient"],
user_api_key_cache: Optional["UserApiKeyCache"],
proxy_logging_obj: Optional["ProxyLogging"],
) -> tuple[str, ...]:
from litellm.proxy.auth.auth_checks import (
_get_models_from_access_groups,
get_authorized_resources_from_key_access_groups,
)
team_group_models: Final = await _get_models_from_access_groups(
access_group_ids=(team_object.access_group_ids or ()) if team_object is not None else (),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
key_group_models: Final = await get_authorized_resources_from_key_access_groups(
valid_token=user_api_key_dict,
team_object=team_object,
resource_field="access_model_names",
)
return tuple(dict.fromkeys((*team_group_models, *key_group_models)))
async def get_available_models_for_user(
user_api_key_dict: "UserAPIKeyAuth",
llm_router: Optional["Router"],
@ -6450,13 +6519,11 @@ async def get_available_models_for_user(
Returns:
List of model names available to the user
"""
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.auth.model_checks import (
get_complete_model_list,
get_key_models,
get_team_models,
)
from litellm.proxy.management_endpoints.team_endpoints import validate_membership
# Get proxy model list and access groups
if llm_router is None:
@ -6466,31 +6533,33 @@ async def get_available_models_for_user(
proxy_model_list = llm_router.get_model_names()
model_access_groups = llm_router.get_model_access_groups()
# Get key models
key_models = get_key_models(
user_api_key_dict=user_api_key_dict,
proxy_model_list=proxy_model_list,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
)
# Get team models
team_models: list[str] = user_api_key_dict.team_models
# If specific team_id is provided, validate and get team models
if team_id and prisma_client and proxy_logging_obj and user_api_key_cache:
key_models = []
team_object: Final = await get_team_object(
requested_team_object: Final = (
await _get_validated_team_object(
user_api_key_dict=user_api_key_dict,
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object)
team_models = team_object.models
if team_id and prisma_client and proxy_logging_obj and user_api_key_cache
else None
)
team_models = get_team_models(
team_models=team_models,
key_models: Final[Sequence[str]] = (
()
if requested_team_object is not None
else get_key_models(
user_api_key_dict=user_api_key_dict,
proxy_model_list=proxy_model_list,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
)
)
team_models: Final = get_team_models(
team_models=(
requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models
),
proxy_model_list=proxy_model_list,
model_access_groups=model_access_groups,
include_model_access_groups=include_model_access_groups,
@ -6498,10 +6567,31 @@ async def get_available_models_for_user(
effective_team_id: Final = team_id or user_api_key_dict.team_id
access_group_models: Final = (
await _get_access_group_models(
user_api_key_dict=user_api_key_dict,
team_object=requested_team_object
or await _get_team_object_for_access_groups(
team_id=effective_team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if key_models or team_models
else ()
)
granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models
granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models
# Get complete model list
all_models: Final = get_complete_model_list(
key_models=key_models,
team_models=team_models,
key_models=granted_key_models,
team_models=granted_team_models,
proxy_model_list=proxy_model_list,
user_model=user_model,
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),

1
litellm/responses/main.py
View file

@ -1064,6 +1064,7 @@ def responses(
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout if timeout is not None else request_timeout,
allowed_openai_params=allowed_openai_params,
**kwargs,
)

6
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -1730,6 +1730,7 @@ class ComplexityRouter(CustomLogger):
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
tier=self._tier_for_model(routed_model),
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
conversation_continuing=conversation_continuing,
@ -1797,7 +1798,8 @@ class ComplexityRouter(CustomLogger):
if user_message is None:
verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model")
if not self.config.plugins and self.config.default_model:
default_model_first: Final = not self.config.plugins and self.config.default_model
if default_model_first:
# No plugins configured: preserve the pre-existing default_model-first
# priority exactly (changing it would be a silent behavior change for
# every non-plugin user, not just a security fix).
@ -1809,12 +1811,14 @@ class ComplexityRouter(CustomLogger):
routed_model = await self._pick_model_for_tier(
ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs
)
fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause="default_fallback",
tier=fallback_tier,
conversation_continuing=conversation_continuing,
),
)

10
litellm/types/guardrails.py
View file

@ -5,6 +5,7 @@ from typing import Any, Final, Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing_extensions import Required, TypedDict
from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
from litellm.types.proxy.guardrails.guardrail_hooks.akto import (
AktoConfigModel,
)
@ -525,6 +526,15 @@ class BedrockGuardrailConfigModel(BaseModel):
description="InvokeGuardrailChecks: block when any sensitiveInformation confidenceScore "
">= this value (scores are in [0,1]). Set to null to make PII detection detect-only.",
)
chunk_budget_chars: int = Field(
default=BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
gt=0,
description="ApplyGuardrail: batch size, in characters, used to re-send content after AWS "
"has rejected a request as too large. Requests AWS accepts are always sent in a single "
"call, so this has no effect until a rejection happens. Defaults to 25,000; a batch AWS "
"still rejects is bisected automatically, so this value only trades round trips against "
"batch size and cannot fail a request on its own.",
)
class LakeraV2GuardrailConfigModel(BaseModel):

23
litellm/types/integrations/websearch_interception.py
View file

@ -2,7 +2,28 @@
Type definitions for WebSearch Interception integration.
"""
from typing import TypedDict
from typing import Literal, TypedDict
from pydantic import BaseModel
class AnthropicSearchQuery(BaseModel):
"""``input`` of an Anthropic ``server_tool_use`` block for a web search."""
query: str
class AnthropicServerToolUseBlock(BaseModel):
"""
The ``server_tool_use`` block that must accompany a ``web_search_tool_result``.
Anthropic requires the pair, with a ``srvtoolu_``-prefixed id shared by both.
"""
type: Literal["server_tool_use"] = "server_tool_use"
id: str
name: Literal["web_search"] = "web_search"
input: AnthropicSearchQuery
class WebSearchInterceptionConfig(TypedDict, total=False):

11
litellm/types/management_endpoints/auto_router_endpoints.py
View file

@ -2,6 +2,7 @@
Types for auto-router management endpoints
"""
from collections.abc import Mapping
from typing import Final, Literal
from pydantic import BaseModel, Field, field_validator
@ -120,6 +121,16 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals):
router_name: str = Field(description="The auto-router alias requests were sent to")
router_type: str = Field(description="complexity, adaptive or quality")
tier_turns: Mapping[str, int] = Field(
default_factory=dict,
description="Turns per tier, keyed by the tier name the routing decision recorded at "
"request time (never re-derived at read time, since the tier-to-model mapping is "
"mutable config). Tier names are scoped to this group's router_type and are not "
"comparable across types: a complexity router reports 'simple'/'medium'/'complex'/"
"'reasoning', a quality router reports its numeric quality tier, and an adaptive router "
"records no tier at all. Turns no tier served (the classifier fell back to default_model) "
"are absent rather than pooled under a sentinel key, so the values may sum to less than turns",
)
class AutoRouterBenchmarksResponse(BaseModel):

3
litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -28,6 +28,9 @@ class BedrockGuardrailUsage(TypedDict, total=False):
sensitiveInformationPolicyUnits: int | None
sensitiveInformationPolicyFreeUnits: int | None
contextualGroundingPolicyUnits: int | None
contentPolicyImageUnits: int | None
automatedReasoningPolicyUnits: int | None
automatedReasoningPolicies: int | None
class BedrockGuardrailOutput(TypedDict, total=False):

18
litellm/types/router.py
View file

@ -17,7 +17,12 @@ from .completion import CompletionRequest
from .embedding import EmbeddingRequest
from .llms.openai import OpenAIFileObject
from .search import SearchProvider
from .utils import CustomPricingLiteLLMParams, ModelResponse, StandardLoggingRoutingDecision
from .utils import (
CustomPricingLiteLLMParams,
MirroredPricingParams,
ModelResponse,
StandardLoggingRoutingDecision,
)
class ConfigurableClientsideParamsCustomAuth(TypedDict):
@ -122,7 +127,7 @@ class UpdateRouterConfig(BaseModel):
model_config = ConfigDict(protected_namespaces=())
class ModelInfo(BaseModel):
class ModelInfo(MirroredPricingParams):
id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance
db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config.
updated_at: datetime.datetime | None = None
@ -424,14 +429,7 @@ class DeploymentTypedDict(TypedDict, total=False):
model_info: dict
SPECIAL_MODEL_INFO_PARAMS = [
"input_cost_per_token",
"output_cost_per_token",
"input_cost_per_character",
"output_cost_per_character",
"cache_read_input_token_cost",
"cache_creation_input_token_cost",
]
SPECIAL_MODEL_INFO_PARAMS = tuple(MirroredPricingParams.model_fields)
class Deployment(BaseModel):

21
litellm/types/utils.py
View file

@ -3245,10 +3245,23 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
litellm_disabled_callbacks: list[str] | None
class CustomPricingLiteLLMParams(BaseModel):
## CUSTOM PRICING ##
class MirroredPricingParams(BaseModel):
"""Pricing overrides that ``Deployment.__init__`` mirrors from ``litellm_params``
onto ``model_info``, so both blobs hold the same rate.
Declared once and inherited by both sides of that mirror, so the two can't drift.
"""
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
input_cost_per_character: float | None = None
output_cost_per_character: float | None = None
cache_read_input_token_cost: float | None = None
cache_creation_input_token_cost: float | None = None
class CustomPricingLiteLLMParams(MirroredPricingParams):
## CUSTOM PRICING ##
input_cost_per_second: float | None = None
output_cost_per_second: float | None = None
output_cost_per_second_1080p: float | None = None
@ -3259,7 +3272,6 @@ class CustomPricingLiteLLMParams(BaseModel):
# This allows any model_info parameter to be set in litellm_params
input_cost_per_token_flex: float | None = None
input_cost_per_token_priority: float | None = None
cache_creation_input_token_cost: float | None = None
cache_creation_input_token_cost_above_1hr: float | None = None
cache_creation_input_token_cost_above_200k_tokens: float | None = None
cache_creation_input_token_cost_above_272k_tokens: float | None = None
@ -3268,7 +3280,6 @@ class CustomPricingLiteLLMParams(BaseModel):
cache_creation_input_token_cost_flex: float | None = None
cache_creation_input_token_cost_priority: float | None = None
cache_creation_input_audio_token_cost: float | None = None
cache_read_input_token_cost: float | None = None
cache_read_input_token_cost_flex: float | None = None
cache_read_input_token_cost_priority: float | None = None
cache_read_input_token_cost_above_200k_tokens: float | None = None
@ -3276,7 +3287,6 @@ class CustomPricingLiteLLMParams(BaseModel):
cache_read_input_token_cost_above_272k_tokens_priority: float | None = None
cache_read_input_token_cost_above_272k_tokens_flex: float | None = None
cache_read_input_audio_token_cost: float | None = None
input_cost_per_character: float | None = None
input_cost_per_character_above_128k_tokens: float | None = None
input_cost_per_audio_token: float | None = None
input_cost_per_token_cache_hit: float | None = None
@ -3298,7 +3308,6 @@ class CustomPricingLiteLLMParams(BaseModel):
output_cost_per_token_batches: float | None = None
output_cost_per_token_flex: float | None = None
output_cost_per_token_priority: float | None = None
output_cost_per_character: float | None = None
output_cost_per_audio_token: float | None = None
output_cost_per_token_above_128k_tokens: float | None = None
output_cost_per_token_above_200k_tokens: float | None = None

1
schema.prisma
View file

@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
tier_turns Json @default("{}")
@@id([api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")

7
tests/e2e/a2a/a2a_client.py
View file

@ -17,6 +17,7 @@ from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import settle_propagation
from e2e_http import NoBody, Result, Success, get_external, is_ok
from proxy_client import ProxyClient
@ -298,7 +299,9 @@ class A2AClient:
the next DB reload. A card read or message/send issued the instant this
returns can therefore 404 on the agent it just created. Waiting here keeps
every caller from having to poll, the same way ProxyClient.create_model
waits for a new model to become servable.
waits for a new model to become servable -- including the settle that
covers the other replicas, since one successful card read only proves the
replica that answered it has the agent.
"""
result = self.proxy.transport.post(
"/v1/agents",
@ -307,7 +310,9 @@ class A2AClient:
response_type=AgentResponse,
)
if isinstance(result, Success):
written_at = time.monotonic()
self._await_agent_servable(result.data.agent_id)
settle_propagation(written_at)
return result
def _await_agent_servable(self, agent_id: str) -> None:

26
tests/e2e/e2e_config.py
View file

@ -7,6 +7,7 @@ environment so the same tests run against localhost or a deployed proxy.
from __future__ import annotations
import os
import time
import uuid
from pathlib import Path
@ -75,6 +76,18 @@ POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120"))
POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5"))
REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60"))
# How long a control-plane write (/model/new, /guardrails, /v1/agents) may take to
# reach EVERY replica. Distinct from POLL_TIMEOUT, which is sized for spend-row
# flush; this one is sized for the proxy's config reload
# (`proxy_config_reload_interval_seconds`, 30s by default and 7s on the e2e stack)
# plus margin.
#
# The barriers below wait this out instead of returning on first sight, because a
# single successful read only proves ONE replica converged: every request opens a
# fresh connection, so a load-balanced Service routes each one independently and
# the next call re-rolls. See ProxyClient._await_model_servable.
PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15"))
EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes")
# Deliberately modest concurrency. The suite shares its proxy with every other
@ -137,3 +150,16 @@ def unique_marker() -> str:
"""A short unique token per call/run, so concurrent runs and the shared
response cache never collide on prompts, tags, or customer ids."""
return uuid.uuid4().hex[:12]
def settle_propagation(written_at: float) -> None:
"""Block until PROPAGATION_TIMEOUT has elapsed since `written_at`, a
`time.monotonic()` stamp taken the moment a control-plane write returned.
Callers that already polled for the object still need this: the poll proves one
replica has it, not all of them. Waiting out the config-reload budget is what
makes the object safe to use on whichever replica the next request lands on.
"""
remaining = PROPAGATION_TIMEOUT - (time.monotonic() - written_at)
if remaining > 0:
time.sleep(remaining)

76
tests/e2e/guardrails/guardrails_client.py
View file

@ -11,7 +11,7 @@ from typing import Literal
from pydantic import BaseModel
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker
from e2e_http import NoBody, Result, Success, unwrap
from lifecycle import ResourceManager
from models import (
@ -104,25 +104,14 @@ class GuardrailsClient:
proxy: ProxyClient
def create_content_filter_guardrail(self, name: str, blocked_keyword: str) -> str:
return unwrap(
self.proxy.transport.post(
"/guardrails",
headers=self.proxy.transport.master,
json=GuardrailCreateBody(
guardrail=GuardrailSpecBody(
guardrail_name=name,
litellm_params=ContentFilterParamsBody(
mode="pre_call",
default_on=True,
blocked_words=[
BlockedWordBody(keyword=blocked_keyword, action="BLOCK")
],
),
)
),
response_type=GuardrailCreateResponse,
)
).guardrail_id
return self.register(
name,
ContentFilterParamsBody(
mode="pre_call",
default_on=True,
blocked_words=[BlockedWordBody(keyword=blocked_keyword, action="BLOCK")],
),
)
def create_bedrock_guardrail(
self,
@ -141,24 +130,15 @@ class GuardrailsClient:
test takes out whatever else is running. Callers select the guardrail
per-request instead, which keeps the blast radius to the test that wants it.
"""
return unwrap(
self.proxy.transport.post(
"/guardrails",
headers=self.proxy.transport.master,
json=GuardrailCreateBody(
guardrail=GuardrailSpecBody(
guardrail_name=name,
litellm_params=BedrockGuardrailParamsBody(
mode="pre_call",
default_on=default_on,
guardrailIdentifier=identifier,
guardrailVersion=version,
),
)
),
response_type=GuardrailCreateResponse,
)
).guardrail_id
return self.register(
name,
BedrockGuardrailParamsBody(
mode="pre_call",
default_on=default_on,
guardrailIdentifier=identifier,
guardrailVersion=version,
),
)
def create_backend_model(self, resources: ResourceManager, prefix: str = "e2e-guard-backend") -> str:
"""Register a gemini chat deployment for a guardrail test to run against
@ -174,11 +154,19 @@ class GuardrailsClient:
return model_name
def register(self, name: str, params: GuardrailParamsBody) -> str:
"""Register any guardrail via POST /guardrails and return its id. New
built-ins register with default_on=False and are opted into per request
via the chat body's `guardrails` list, so one guardrail under test never
intercepts unrelated traffic on the shared proxy."""
return unwrap(
"""Register any guardrail via POST /guardrails and return its id, once every
replica can be expected to serve it. New built-ins register with
default_on=False and are opted into per request via the chat body's
`guardrails` list, so one guardrail under test never intercepts unrelated
traffic on the shared proxy.
/guardrails is a control-plane route and guardrails reach the data plane on
the config reload, so a request naming this guardrail the instant the POST
returns can 404 with "Guardrail not found" on a replica that has not
reloaded. There is no data-plane read that lists guardrails, so unlike
ProxyClient.create_model this settles on the propagation budget alone with
nothing to poll first."""
guardrail_id = unwrap(
self.proxy.transport.post(
"/guardrails",
headers=self.proxy.transport.master,
@ -188,6 +176,8 @@ class GuardrailsClient:
response_type=GuardrailCreateResponse,
)
).guardrail_id
settle_propagation(time.monotonic())
return guardrail_id
def delete_guardrail(self, guardrail_id: str) -> None:
_ = self.proxy.transport.delete(

14
tests/e2e/llm_translation/test_vertex_passthrough_e2e.py
View file

@ -23,11 +23,12 @@ model, spend > 0), correlated by the x-litellm-call-id header.
"""
import os
import time
import pytest
from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_config import settle_propagation, unique_marker
from e2e_http import NoBody, require_successful_call, unwrap
from lifecycle import ResourceManager
from models import SpendLogRow
@ -90,7 +91,14 @@ class _ModelDeleteBody(BaseModel):
def _add_vertex_passthrough_model(
client: PassthroughClient, model_name: str, project: str, credentials: str
) -> str:
return unwrap(
"""Register the passthrough deployment and settle before the caller uses it.
This body carries `use_in_pass_through` and a pinned `model_info.id`, so it
cannot go through ProxyClient.create_model -- but it needs that helper's
propagation settle just the same, or the passthrough call can land on a replica
that has not reloaded yet.
"""
model_id = unwrap(
client.proxy.transport.post(
"/model/new",
headers=client.proxy.transport.master,
@ -108,6 +116,8 @@ def _add_vertex_passthrough_model(
response_type=_ModelNewResponse,
)
).model_id
settle_propagation(time.monotonic())
return model_id
def _delete_model(client: PassthroughClient, model_id: str) -> None:

3
tests/e2e/logging/logging_client.py
View file

@ -23,7 +23,7 @@ from typing import Callable, Literal
import pytest
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation
from proxy_client import ProxyClient
from e2e_http import (
URL,
@ -409,6 +409,7 @@ class LoggingClient:
)
guardrail_id = response.guardrail_id
assert guardrail_id, f"create guardrail returned no id: {response!r}"
settle_propagation(time.monotonic())
return guardrail_id
def delete_guardrail(self, guardrail_id: str) -> None:

5
tests/e2e/mcp/mcp_client.py
View file

@ -18,6 +18,7 @@ from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict, Field, RootModel
from e2e_config import settle_propagation
from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap
from models import KeyGenerateBody, ObjectPermission
from proxy_client import ProxyClient
@ -330,7 +331,7 @@ class McpClient:
tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is
unique per test, so default_on only ever intercepts this test's own
banned tool call on the shared proxy."""
return unwrap(
guardrail_id = unwrap(
self.proxy.transport.post(
"/guardrails",
headers=self.proxy.transport.master,
@ -345,6 +346,8 @@ class McpClient:
response_type=GuardrailCreateResponse,
)
).guardrail_id
settle_propagation(time.monotonic())
return guardrail_id
def delete_guardrail(self, guardrail_id: str) -> None:
_ = self.proxy.transport.delete(

25
tests/e2e/proxy_client.py
View file

@ -70,6 +70,7 @@ from e2e_config import (
POLL_TIMEOUT,
PROXY_BASE_URL,
REQUEST_TIMEOUT,
settle_propagation,
)
from transport import HttpTransport, SplitTransport, Transport
@ -279,18 +280,18 @@ class ProxyClient:
"""Register a deployment under `model_name` and return its proxy-assigned
model_id, once the model is actually servable on the data plane.
/model/new is a control-plane route; in a split control/data-plane
deployment the gateway (data plane, which serves /chat, /ocr, ...) only
picks the new model up on its next DB reload, so a call issued the instant
this returns can race the reload and 400 with "Invalid model name passed".
We therefore poll the data-plane /v1/models until the model appears before
handing back, so callers can invoke it immediately. In the monolithic case
it is already present on the first poll, so this adds one request.
/model/new is a control-plane route; the data plane (which serves /chat,
/ocr, ...) only picks the new model up on its next DB reload, so a call
issued the instant this returns can race the reload and 400 with "Invalid
model name passed". We poll the data-plane /v1/models until the model
appears, then settle for the remainder of the propagation budget.
First listing must arrive within `model_servable_timeout` (not the longer
spend `poll_timeout`). The model must then stay listed for
`model_servable_db_sync_seconds` (product default DB reload interval) so every
gateway worker has run add_deployment before callers use the model."""
Both steps are needed, and the second is the one that matters at >1 replica.
The poll proves *a* replica is serving the model; it cannot prove they all
are, because every request opens a fresh connection and a load-balanced
Service routes each one independently -- so the caller's next request
re-rolls and can land on a replica that has not reloaded yet. Waiting out
PROPAGATION_TIMEOUT is what makes the model safe to use anywhere."""
model_id = unwrap(
self.transport.post(
"/model/new",
@ -303,7 +304,9 @@ class ProxyClient:
response_type=ModelNewResponse,
)
).model_id
written_at = time.monotonic()
self._await_model_servable(model_name)
settle_propagation(written_at)
return model_id
def _await_model_servable(self, model_name: str) -> None:

184
tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py
View file

@ -1813,6 +1813,190 @@ def _create_unified_batch_id(model_id: str, batch_id: str) -> str:
return base64.urlsafe_b64encode(unified_str.encode()).decode().rstrip("=")
def _decode_unified_id(b64_id: str) -> str:
return base64.urlsafe_b64decode(b64_id + "=" * (-len(b64_id) % 4)).decode()
def _terminal_batch_record(
unified_batch_uid: str,
raw_input_file_id: str,
raw_output_file_id: str,
raw_error_file_id: str,
):
record = MagicMock()
record.unified_object_id = unified_batch_uid
record.created_by = "owner-user"
record.team_id = "owner-team"
record.status = "cancelled"
record.file_object = json.dumps(
{
"id": "batch-raw-456",
"object": "batch",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"status": "cancelled",
"created_at": 1234567890,
"input_file_id": raw_input_file_id,
"output_file_id": raw_output_file_id,
"error_file_id": raw_error_file_id,
}
)
return record
@pytest.mark.asyncio
async def test_list_batches_registers_and_returns_unified_output_file_ids():
"""A stored batch blob with raw provider file IDs (e.g. persisted by the cost
poller for a cancelled batch) must be listed with unified managed IDs, and the
output/error files must be registered in the managed file table so GET
/files/{id}/content can route them."""
from litellm.proxy._types import UserAPIKeyAuth
unified_batch_uid = _create_unified_batch_id("model-123", "batch-456")
raw_input_file_id = "file-list-in-1"
raw_output_file_id = "file-list-out-1"
raw_error_file_id = "file-list-err-1"
unified_input_file_id = base64.urlsafe_b64encode(
b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch"
).decode()
prisma_client = AsyncMock()
prisma_client.db.litellm_managedobjecttable.find_many.return_value = [
_terminal_batch_record(
unified_batch_uid, raw_input_file_id, raw_output_file_id, raw_error_file_id
)
]
input_file_row = MagicMock()
input_file_row.unified_file_id = unified_input_file_id
input_file_row.flat_model_file_ids = [raw_input_file_id]
prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(
return_value=[input_file_row]
)
prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
result = await proxy_managed_files.list_user_batches(
user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"),
limit=10,
)
listed = result["data"][0]
assert listed.id == unified_batch_uid
assert listed.input_file_id == unified_input_file_id
bulk_lookup = prisma_client.db.litellm_managedfiletable.find_many.await_args
assert set(bulk_lookup.kwargs["where"]["flat_model_file_ids"]["hasSome"]) == {
raw_input_file_id,
raw_output_file_id,
raw_error_file_id,
}
decoded_output = _decode_unified_id(listed.output_file_id)
assert decoded_output.startswith("litellm_proxy")
assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output
assert "llm_output_file_model_id,model-123" in decoded_output
assert "target_model_names,gpt-5-batch" in decoded_output
decoded_error = _decode_unified_id(listed.error_file_id)
assert f"llm_output_file_id,{raw_error_file_id}" in decoded_error
upsert_calls = prisma_client.db.litellm_managedfiletable.upsert.await_args_list
stored_raw_ids = {
c.kwargs["data"]["create"]["flat_model_file_ids"][0] for c in upsert_calls
}
assert stored_raw_ids == {raw_output_file_id, raw_error_file_id}
for c in upsert_calls:
assert c.kwargs["data"]["create"]["created_by"] == "owner-user"
assert c.kwargs["data"]["create"]["team_id"] == "owner-team"
@pytest.mark.asyncio
async def test_list_batches_resolves_existing_managed_rows_without_minting():
"""When the raw provider file IDs already have managed file rows, listing must
swap in the existing unified IDs via one bulk lookup for the whole page, with
no per-row queries and no duplicate upserts."""
from litellm.proxy._types import UserAPIKeyAuth
unified_input_file_id = base64.urlsafe_b64encode(
b"litellm_proxy:application/octet-stream;unified_id,in-9;target_model_names,gpt-5-batch"
).decode()
raw_output_file_ids = ["file-list-out-existing-1", "file-list-out-existing-2"]
existing_unified_output_ids = [
base64.urlsafe_b64encode(
f"litellm_proxy:application/json;unified_id,u-{i};llm_output_file_id,{raw_id}".encode()
).decode()
for i, raw_id in enumerate(raw_output_file_ids)
]
records = [
_terminal_batch_record(
_create_unified_batch_id("model-123", f"batch-{i}"),
unified_input_file_id,
raw_id,
"",
)
for i, raw_id in enumerate(raw_output_file_ids)
]
prisma_client = AsyncMock()
prisma_client.db.litellm_managedobjecttable.find_many.return_value = records
existing_rows = [
MagicMock(unified_file_id=unified_id, flat_model_file_ids=[raw_id])
for raw_id, unified_id in zip(raw_output_file_ids, existing_unified_output_ids)
]
prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(
return_value=existing_rows
)
prisma_client.db.litellm_managedfiletable.find_first = AsyncMock()
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
result = await proxy_managed_files.list_user_batches(
user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"),
limit=10,
)
assert [b.output_file_id for b in result["data"]] == existing_unified_output_ids
prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once()
prisma_client.db.litellm_managedfiletable.find_first.assert_not_awaited()
prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited()
@pytest.mark.asyncio
async def test_list_batches_caps_page_size_at_100():
"""The list page size must be capped at 100 rows (matching OpenAI's limit)
even when the caller asks for more, so one request cannot fan out into an
unbounded scan."""
from litellm.proxy._types import UserAPIKeyAuth
prisma_client = AsyncMock()
prisma_client.db.litellm_managedobjecttable.find_many.return_value = []
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
result = await proxy_managed_files.list_user_batches(
user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"),
limit=100000,
)
assert (
prisma_client.db.litellm_managedobjecttable.find_many.await_args.kwargs["take"]
== 101
)
assert result["data"] == []
@pytest.mark.asyncio
async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception():
from litellm.proxy._types import UserAPIKeyAuth

102
tests/proxy_behavior/spend/test_autorouter_session_rollup.py
View file

@ -38,11 +38,13 @@ async def _turn(
tokens: int = 100,
spend: float = 0.01,
saved: float = 0.02,
tier: "str | None" = None,
) -> None:
touched: Final = 1 if (hit or ttl is not None or not covered) else 0
await db.execute_raw(
UPSERT_AUTOROUTER_SESSION_SQL,
key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched,
tier,
)
@ -195,6 +197,106 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db
assert [(row["router_type"], row["sessions"]) for row in matching] == [("complexity", 1), ("quality", 1)]
async def test_tier_turns_count_each_tier_that_served_a_turn(db):
key = f"k-{uuid.uuid4()}"
await _turn(db, key, "A", T0, tier="simple")
await _turn(db, key, "B", T0 + timedelta(seconds=10), tier="complex")
await _turn(db, key, "A", T0 + timedelta(seconds=20), tier="simple")
assert (await _row(db, key))["tier_turns"] == {"simple": 2, "complex": 1}
async def test_an_untiered_turn_increments_no_tier_counter(db):
key = f"k-{uuid.uuid4()}"
await _turn(db, key, "A", T0, tier=None)
assert (await _row(db, key))["tier_turns"] == {}
await _turn(db, key, "A", T0 + timedelta(seconds=10), tier="medium")
await _turn(db, key, "A", T0 + timedelta(seconds=20), tier=None)
row = await _row(db, key)
assert row["tier_turns"] == {"medium": 1}
assert row["turns"] == 3
async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_the_map(db):
key = f"k-{uuid.uuid4()}"
await _turn(db, key, "A", T0, router_type="complexity", tier="medium")
await _turn(db, key, "A", T0 + timedelta(seconds=10), router_type="quality", tier="2")
await _turn(db, key, "A", T0 + timedelta(seconds=20), router_type="complexity", tier="medium")
row = await _row(db, key)
assert row["tier_turns"] == {"medium": 2}
assert row["turns"] == 3
async def test_an_out_of_order_turn_still_counts_toward_its_tier(db):
key = f"k-{uuid.uuid4()}"
await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple")
await _turn(db, key, "A", T0, tier="simple")
row = await _row(db, key)
assert row["tier_turns"] == {"simple": 2}
assert row["unordered_turns"] == 1
async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier="simple")
await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, tier="simple")
await _turn(db, key, "B", T0 + timedelta(seconds=20), session_id=f"s-{uuid.uuid4()}", router=router, tier="complex")
await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
rows = await db.query_raw(
_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)
grouped = next(row for row in rows if row["router_name"] == router)
assert grouped["tier_turns"] == {"simple": 2, "complex": 1}
assert grouped["turns"] == 4
async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
await _turn(
db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity", tier="medium"
)
await _turn(
db,
key,
"A",
T0 + timedelta(seconds=10),
session_id=f"s-{uuid.uuid4()}",
router=router,
router_type="quality",
tier="2",
)
rows = await db.query_raw(
_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)
by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router}
assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}}
async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
rows = await db.query_raw(
_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)
grouped = next(row for row in rows if row["router_name"] == router)
assert grouped["tier_turns"] == {}
async def test_a_miss_that_touched_no_cache_does_not_advance_the_ttl_clock(db):
key = f"k-{uuid.uuid4()}"
await _turn(db, key, "A", T0, ttl=300)

79
tests/test_litellm/integrations/otel/test_otel_v2_logger.py
View file

@ -356,6 +356,7 @@ def _mcp_payload(**overrides):
"arguments": {"city": "Paris"},
"result": {"temp_c": 21},
"mcp_server_name": "weather-mcp",
"mcp_server_resource": "https://weather.example.com",
"mcp_session_id": "sess-abc123",
},
},
@ -452,6 +453,8 @@ def test_mcp_tool_call_failure_marks_error():
assert span.name == "tools/call get_weather"
assert span.status.status_code is StatusCode.ERROR
assert span.attributes["error.type"] == "MCPError"
assert span.attributes["rpc.system"] == "jsonrpc"
assert span.attributes["server.address"] == "weather.example.com"
def test_mcp_tool_call_deduped_on_repeat():
@ -531,6 +534,82 @@ _MCP_SPAN_CASES = [
]
def test_mcp_tool_call_names_its_rpc_system_and_upstream():
"""A tool-call span names the RPC system, and always alongside the upstream it called.
A CLIENT span holding none of the ``rpc.*``/``http.*``/``db.*``/``messaging.*``
families is unclassifiable, so a backend deriving a span type from them has nothing
to derive from: Elastic APM indexed these spans as ``span.type=unknown`` with no
``span.subtype`` at all, and its span-links API then rejected the whole trace with
``Missing required fields (span.subtype)``.
The two assertions are one invariant, not two. Naming the RPC system makes a
consumer treat the span as a downstream dependency and key that dependency off
``server.address``/``server.port``; emitting the first without the second names the
dependency ``:0``, which is worse than leaving the span unclassified.
"""
logger, exporter = _logger()
asyncio.run(
logger.async_log_success_event(
{"standard_logging_object": _mcp_payload()}, None, None, None
)
)
(span,) = exporter.get_finished_spans()
assert span.attributes["rpc.system"] == "jsonrpc"
assert span.attributes["server.address"] == "weather.example.com"
assert span.attributes["server.port"] == 443
@pytest.mark.parametrize(
"resource",
[None, "mcp://weather.example.com", "ws://weather.example.com"],
ids=["no-resource", "scheme-with-no-default-port", "ws-scheme"],
)
def test_mcp_tool_call_omits_rpc_system_without_a_complete_upstream(resource):
"""A tool call drops the RPC system unless the full destination resolved.
``mcp_server_resource`` is absent whenever the tool name resolves to no registered
server, and it is ``None`` for a transport with no host to log at all (stdio). A
host-bearing scheme outside the HTTP(S) default-port map resolves an address but no
port, and the ``url`` field is not scheme-validated, so that state is reachable from
config. Each case names the dependency ``:0`` or ``host:0`` if ``rpc.system`` ships
on its own, so the pairing is enforced here rather than left to the extractors
happening to agree.
"""
logger, exporter = _logger()
payload = _mcp_payload()
if resource is None:
del payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_resource"]
else:
payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_resource"] = resource
asyncio.run(
logger.async_log_success_event({"standard_logging_object": payload}, None, None, None)
)
(span,) = exporter.get_finished_spans()
assert "rpc.system" not in span.attributes
assert "server.port" not in span.attributes
assert span.attributes["mcp.method.name"] == "tools/call"
def test_mcp_list_tools_omits_rpc_system_without_an_upstream():
"""The discovery span carries no upstream identity, so it must not claim to be RPC.
``tools/list`` reaches the callbacks with no ``mcp_tool_call_metadata``, so there is
no ``server.address`` to attach and a listing can span several upstreams anyway.
Naming ``rpc.system`` here would buy a ``span.subtype`` at the cost of a bogus ``:0``
dependency node in every consumer that aggregates on it.
"""
logger, exporter = _logger()
asyncio.run(
logger.async_log_success_event(
{"standard_logging_object": _mcp_list_payload()}, None, None, None
)
)
(span,) = exporter.get_finished_spans()
assert "rpc.system" not in span.attributes
assert "server.address" not in span.attributes
@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES)
def test_mcp_span_nests_under_transport_without_propagated_context(
make_payload, span_name

38
tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
View file

@ -24,7 +24,11 @@ from litellm.integrations.otel import (
resolve_provider,
)
from litellm.integrations.otel.model import spans as spans_mod
from litellm.integrations.otel.model.payloads import LLMCallSpanData, RequestIdentity
from litellm.integrations.otel.model.payloads import (
LLMCallSpanData,
RequestIdentity,
_upstream_address_port,
)
from litellm.integrations.otel.model.spans import (
SPAN_REGISTRY,
LiteLLMSpanKind,
@ -177,6 +181,7 @@ def test_mcp_attribute_vocabulary_is_complete():
"mcp.resource.uri",
"jsonrpc.request.id",
"jsonrpc.protocol.version",
"rpc.system",
"rpc.response.status_code",
"gen_ai.operation.name",
"gen_ai.tool.name",
@ -808,3 +813,34 @@ def test_promoted_baggage_is_bounded_allowlist():
# http.* is never a promoted key
assert HTTP.ROUTE not in promoted
assert HTTP.REQUEST_METHOD not in promoted
@pytest.mark.parametrize(
"resource, expected",
[
("https://weather.example.com", ("weather.example.com", 443)),
("http://weather.example.com", ("weather.example.com", 80)),
("https://weather.example.com:8443", ("weather.example.com", 8443)),
("mcp://weather.example.com", ("weather.example.com", None)),
("http://::1:8080", (None, None)),
("http://fe80::1%25eth0:80", (None, None)),
(None, (None, None)),
("", (None, None)),
],
ids=[
"https-default",
"http-default",
"explicit-port",
"no-default-port",
"ipv6-unbracketed",
"ipv6-zone-scoped",
"none",
"empty",
],
)
def test_upstream_address_port(resource, expected):
"""The redacted MCP origin resolves to the address and port a consumer names its
dependency from. A scheme outside the default-port map yields no port, and an IPv6
origin yields nothing at all because the redactor rebuilds it without its brackets;
both are why the mapper gates ``rpc.system`` on the complete pair."""
assert _upstream_address_port(resource) == expected

47
tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py
View file

@ -134,6 +134,30 @@ class TestBuildWebSearchToolResultBlock:
assert first["title"] == "LiteLLM Docs"
assert first["page_age"] == "2025-01-15"
assert first["encrypted_content"] == ""
assert first["snippet"] == "Unified interface for LLMs."
def test_snippet_carried_for_every_result(self):
# The snippet is the only field carrying page text. Losing it leaves the
# client and the model with nothing to answer from, forcing a fetch per
# result.
block = WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id="toolu_abc",
search_response=_make_search_response(),
)
assert [r["snippet"] for r in block["content"]] == [
"Unified interface for LLMs.",
"Pay-per-use pricing model.",
]
def test_missing_snippet_degrades_to_empty_string(self):
response = SearchResponse(
results=[SearchResult(title="T", url="https://x/", snippet="")]
)
block = WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id="toolu_abc",
search_response=response,
)
assert block["content"][0]["snippet"] == ""
def test_handles_none_search_response(self):
block = WebSearchTransformation.build_web_search_tool_result_block(
@ -223,11 +247,15 @@ class TestBuildPlanAttachesBlocks:
)
blocks = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY)
assert isinstance(blocks, list)
assert len(blocks) == 1
assert blocks[0]["type"] == "web_search_tool_result"
assert blocks[0]["tool_use_id"] == "toolu_one"
assert blocks[0]["content"][0]["url"] == "https://docs.litellm.ai/"
assert isinstance(blocks, tuple)
assert [b["type"] for b in blocks] == [
"server_tool_use",
"web_search_tool_result",
]
assert blocks[0]["id"].startswith("srvtoolu_")
assert blocks[0]["input"] == {"query": "what is litellm"}
assert blocks[1]["tool_use_id"] == blocks[0]["id"]
assert blocks[1]["content"][0]["url"] == "https://docs.litellm.ai/"
@pytest.mark.asyncio
async def test_metadata_does_not_carry_blocks_when_flag_absent(self):
@ -479,6 +507,9 @@ class TestLegacyPathMatchesNewPath:
kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True},
)
assert out["content"][0]["type"] == "web_search_tool_result"
assert out["content"][0]["tool_use_id"] == "toolu_legacy"
assert out["content"][1]["type"] == "text"
assert [b["type"] for b in out["content"]] == [
"server_tool_use",
"web_search_tool_result",
"text",
]
assert out["content"][1]["tool_use_id"] == out["content"][0]["id"]

55
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py
View file

@ -732,6 +732,61 @@ def test_handler_skips_strip_when_presanitized():
assert result is not None
def test_handler_flattens_replayed_unencrypted_web_search_results():
"""Synthesized search blocks replayed as history must reach the provider as text."""
from litellm.llms.anthropic.experimental_pass_through.messages import handler
captured = {}
def fake_base_handler(*args, **kwargs):
captured.update(kwargs)
return "stub"
with patch.object(
handler.base_llm_http_handler,
"anthropic_messages_handler",
side_effect=fake_base_handler,
):
handler.anthropic_messages_handler(
max_tokens=10,
messages=[
{"role": "user", "content": "latest litellm version?"},
{
"role": "assistant",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {"query": "latest litellm version"},
},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": [
{
"type": "web_search_result",
"url": "https://github.com/BerriAI/litellm/releases",
"title": "Releases",
"page_age": None,
"encrypted_content": "",
"snippet": "Latest release v1.95.0",
}
],
},
],
},
{"role": "user", "content": "which version?"},
],
model="anthropic/claude-3-5-sonnet-20241022",
custom_llm_provider="anthropic",
)
replayed = captured["messages"][1]["content"]
assert [b["type"] for b in replayed] == ["text"]
assert "Snippet: Latest release v1.95.0" in replayed[0]["text"]
def test_presanitized_flag_not_leaked_to_provider_params():
"""The private sentinel must be popped, never forwarded as a request param."""
from litellm.llms.anthropic.experimental_pass_through.messages import handler

186
tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py
View file

@ -10,8 +10,10 @@ Verifies that:
- ANTHROPIC_API_KEY / ANTHROPIC_API_BASE take precedence over their aliases.
"""
import json
import os
import sys
from types import SimpleNamespace
from unittest.mock import patch
import pytest
@ -1457,6 +1459,190 @@ class TestAnthropicThinkingSignatureSelfHeal:
out = strip_empty_text_blocks_from_anthropic_messages(msgs)
assert [b["type"] for b in out[0]["content"]] == ["tool_result"]
def test_flatten_unencrypted_web_search_results_keeps_snippet_evidence(self):
from litellm.llms.anthropic.common_utils import (
flatten_unencrypted_web_search_results_in_anthropic_messages,
)
msgs = [
{"role": "user", "content": "latest litellm version?"},
{
"role": "assistant",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {"query": "latest litellm version"},
},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": [
{
"type": "web_search_result",
"url": "https://github.com/BerriAI/litellm/releases",
"title": "Releases",
"page_age": None,
"encrypted_content": "",
"snippet": "Latest release v1.95.0",
}
],
},
{"type": "text", "text": "v1.95.0"},
],
},
]
out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs)
assert out[0] is msgs[0]
assert [b["type"] for b in out[1]["content"]] == ["text", "text"]
flattened = out[1]["content"][0]["text"]
assert "Web search results for 'latest litellm version':" in flattened
assert "URL: https://github.com/BerriAI/litellm/releases" in flattened
assert "Snippet: Latest release v1.95.0" in flattened
assert msgs[1]["content"][0]["type"] == "server_tool_use"
@pytest.mark.parametrize("results", [[], None], ids=["empty_list", "search_raised"])
def test_flatten_unencrypted_web_search_results_flattens_a_resultless_search(self, results):
"""A search that found nothing, or that raised, still has to be flattened.
Both cases reach the client as ``content: []``, and leaving that block in
place ships an unsupported tag to Bedrock on the next turn just as surely
as a populated one does.
"""
from litellm.integrations.websearch_interception.transformation import (
WebSearchTransformation,
)
from litellm.llms.anthropic.common_utils import (
flatten_unencrypted_web_search_results_in_anthropic_messages,
)
block = WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id="srvtoolu_1",
search_response=None if results is None else SimpleNamespace(results=results),
)
assert block["content"] == [], "fixture drifted from what the interceptor emits"
msgs = [
{
"role": "assistant",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {"query": "who won"},
},
block,
{"type": "text", "text": "I could not find that."},
],
}
]
out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs)
assert [b["type"] for b in out[0]["content"]] == ["text", "text"]
assert out[0]["content"][0]["text"] == ("Web search results for 'who won':\n\nNo results were returned.")
@pytest.mark.parametrize("results", [[SimpleNamespace(title="Rome", url="u", snippet="s", date=None)], []])
def test_flatten_unencrypted_web_search_results_is_idempotent(self, results):
"""Flattening twice must equal flattening once.
The agentic loop re-enters the same entry point for its follow-up call and
hands it the original history, so this runs again on already-flattened
messages once per iteration. A pass that appended instead of replacing
would duplicate the evidence on every loop.
"""
from litellm.integrations.websearch_interception.transformation import (
WebSearchTransformation,
)
from litellm.llms.anthropic.common_utils import (
flatten_unencrypted_web_search_results_in_anthropic_messages,
)
msgs = [
{
"role": "assistant",
"content": [
{"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "when"}},
WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id="srvtoolu_1",
search_response=SimpleNamespace(results=results),
),
{"type": "text", "text": "753 BC."},
],
}
]
once = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs)
twice = flatten_unencrypted_web_search_results_in_anthropic_messages(once)
assert [b["type"] for b in once[0]["content"]] == ["text", "text"]
assert json.dumps(twice) == json.dumps(once)
def test_flatten_unencrypted_web_search_results_preserves_real_anthropic_blocks(self):
from litellm.llms.anthropic.common_utils import (
flatten_unencrypted_web_search_results_in_anthropic_messages,
)
msgs = [
{
"role": "assistant",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {"query": "q"},
},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": [
{
"type": "web_search_result",
"url": "https://example.com",
"title": "Example",
"page_age": None,
"encrypted_content": "EqgfCioIARgBIiQ4",
}
],
},
],
}
]
out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs)
assert out[0] is msgs[0]
def test_flatten_unencrypted_web_search_results_leaves_error_blocks_alone(self):
from litellm.llms.anthropic.common_utils import (
flatten_unencrypted_web_search_results_in_anthropic_messages,
)
msgs = [
{
"role": "assistant",
"content": [
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": {
"type": "web_search_tool_result_error",
"error_code": "max_uses_exceeded",
},
}
],
}
]
out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs)
assert out[0] is msgs[0]
def test_sanitize_tool_use_ids_in_anthropic_messages(self):
from litellm.llms.anthropic.common_utils import (
sanitize_tool_use_ids_in_anthropic_messages,

65
tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
View file

@ -4,6 +4,7 @@ import json
import os
import sys
from datetime import datetime
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
@ -2515,3 +2516,67 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag(
assert thinking.get("type") == "enabled"
assert isinstance(thinking.get("budget_tokens"), int)
assert "output_config" not in flipped
@pytest.mark.parametrize(
"search_results, expected_evidence",
[
pytest.param(
[SimpleNamespace(title="Rome", url="https://ex.com/rome", snippet="Founded 753 BC.", date=None)],
"Snippet: Founded 753 BC.",
id="search_returned_results",
),
pytest.param([], "No results were returned.", id="search_returned_nothing"),
],
)
def test_replayed_intercepted_search_turn_leaves_no_unsupported_block_for_bedrock(search_results, expected_evidence):
"""A native client replaying an intercepted search turn must not 400 on Bedrock.
``websearch_interception`` hands Claude Desktop an Anthropic-native
``server_tool_use`` + ``web_search_tool_result`` pair, and Anthropic's protocol
obliges the client to replay that assistant turn verbatim on every later turn.
Bedrock's Anthropic schema defines neither tag, so both have to be gone from the
outbound body by the time it is signed, with the search evidence carried forward
as text instead. Built from the real builder rather than a hand-written fixture
so the two cannot drift apart.
"""
from litellm.integrations.websearch_interception.transformation import (
WebSearchTransformation,
)
from litellm.llms.anthropic.common_utils import (
flatten_unencrypted_web_search_results_in_anthropic_messages,
)
from litellm.types.router import GenericLiteLLMParams
replayed_turn = [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {"query": "when was Rome founded"},
},
WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id="srvtoolu_1",
search_response=SimpleNamespace(results=search_results),
),
{"type": "text", "text": "Rome was founded in 753 BC."},
]
messages = [
{"role": "user", "content": [{"type": "text", "text": "When was Rome founded?"}]},
{"role": "assistant", "content": replayed_turn},
{"role": "user", "content": [{"type": "text", "text": "Repeat the year."}]},
]
body = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request(
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=flatten_unencrypted_web_search_results_in_anthropic_messages(messages),
anthropic_messages_optional_request_params={"max_tokens": 64},
litellm_params=GenericLiteLLMParams(),
headers={},
)
serialized = json.dumps(body)
assert "web_search_tool_result" not in serialized
assert "server_tool_use" not in serialized
assert expected_evidence in serialized
assert "Rome was founded in 753 BC." in serialized

25
tests/test_litellm/proxy/auth/test_model_checks.py
View file

@ -735,3 +735,28 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion():
litellm.vertex_language_models.discard(fake_model)
litellm.add_known_models(model_cost_map={})
assert fake_model not in litellm.models_by_provider["vertex_ai"]
def test_get_complete_model_list_drops_no_default_models_sentinel():
from litellm.proxy.auth.model_checks import get_complete_model_list
result = get_complete_model_list(
key_models=["no-default-models", "model-a"],
team_models=[],
proxy_model_list=["model-a", "model-b"],
user_model=None,
infer_model_from_keys=False,
)
assert result == ["model-a"]
def test_get_complete_model_list_sentinel_only_grants_nothing():
from litellm.proxy.auth.model_checks import get_complete_model_list
result = get_complete_model_list(
key_models=["no-default-models"],
team_models=["no-default-models"],
proxy_model_list=["model-a", "model-b"],
user_model=None,
infer_model_from_keys=False,
)
assert result == []

22
tests/test_litellm/proxy/db/test_autorouter_session_rollup.py
View file

@ -93,6 +93,19 @@ class TestBuildTransaction:
def test_requests_without_a_routing_decision_are_skipped(self, metadata: dict):
assert _build(metadata=metadata) is None
def test_the_tier_the_decision_recorded_is_carried_onto_the_transaction(self):
transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": "reasoning"}))
assert transaction is not None and transaction.tier == "reasoning"
@pytest.mark.parametrize("tier", [None, "", 3, {"tier": "medium"}])
def test_a_decision_without_a_usable_tier_records_no_tier(self, tier: object):
transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": tier}))
assert transaction is not None and transaction.tier is None
def test_a_decision_that_never_mentions_tier_records_no_tier(self):
transaction = _build()
assert transaction is not None and transaction.tier is None
def test_router_name_falls_back_to_the_payload_model_group(self):
transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"}))
assert transaction is not None and transaction.router_name == "live-auto"
@ -167,7 +180,11 @@ class _FakeClient:
self.db = _FakeDB(failures, poison_session)
def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, 0, 0)) -> AutoRouterTurnTransaction:
def _transaction(
session_id: str = "s1",
at: datetime = datetime(2026, 8, 1, 12, 0, 0),
tier: str | None = "medium",
) -> AutoRouterTurnTransaction:
return AutoRouterTurnTransaction(
api_key="k1",
session_id=session_id,
@ -182,6 +199,7 @@ def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12,
cache_hit=False,
cache_ttl_seconds=None,
cache_touched=False,
tier=tier,
)
@ -201,7 +219,7 @@ class TestFlush:
assert sql == UPSERT_AUTOROUTER_SESSION_SQL
assert params == (
"k1", "s1", "live-auto", "complexity", "bedrock/haiku",
"2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0,
"2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, "medium",
)
def test_a_connect_error_retries_the_same_statement(self):

1989
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

File diff suppressed because it is too large Load diff

35
tests/test_litellm/proxy/guardrails/test_init_guardrails.py
View file

@ -38,6 +38,41 @@ def test_initialize_presidio_guardrail():
assert result["litellm_params"].mode == "pre_call"
def test_initialize_bedrock_forwards_chunk_budget_chars():
"""Regression: `chunk_budget_chars` set in config.yaml must reach the guardrail.
The field lives on BedrockGuardrailConfigModel, so LitellmParams parsed it and the
Admin UI rendered it, but initialize_bedrock enumerates its kwargs explicitly and
dropped it. The setting validated and then silently did nothing. Asserting through
initialize_guardrail rather than the constructor is the point: constructing
BedrockGuardrail directly bypasses the only path a user can actually reach.
"""
import litellm
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
test_guardrail = {
"guardrail_name": "test_bedrock_chunk_budget",
"litellm_params": {
"guardrail": SupportedGuardrailIntegrations.BEDROCK.value,
"mode": "pre_call",
"guardrailIdentifier": "test-guardrail",
"guardrailVersion": "DRAFT",
"chunk_budget_chars": 60_000,
},
}
guardrail_handler = InMemoryGuardrailHandler()
guardrail_handler.initialize_guardrail(guardrail=test_guardrail)
initialized = [
callback
for callback in litellm.callbacks
if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_chunk_budget"
]
assert initialized, "bedrock guardrail was not registered as a callback"
assert initialized[-1].chunk_budget_chars == 60_000
def test_initialize_guardrail_preserves_guardrail_info():
"""
Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the

39
tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
View file

@ -292,6 +292,7 @@ class TestAutoRouterBenchmarks:
ROW = _SessionAggRow(
router_name="live-auto",
router_type="complexity",
tier_turns={},
sessions=4,
turns=40,
unordered_turns=1,
@ -377,6 +378,21 @@ class TestAutoRouterBenchmarks:
assert totals.avg_turns_per_session == 10.0
assert totals.spend == 10.0
def test_tier_names_stay_scoped_to_the_router_type_that_recorded_them(self):
quality = self.ROW.model_copy(
update={"router_name": "quality-auto", "router_type": "quality", "tier_turns": {"2": 7}}
)
complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}})
assert complexity.tier_turns == {"medium": 7}
assert quality.tier_turns == {"2": 7}
def test_summed_totals_carry_no_tier_map_because_names_are_router_scoped(self):
from litellm.proxy.management_endpoints.auto_router_endpoints import _summed_agg_row
quality = self.ROW.model_copy(update={"router_type": "quality", "tier_turns": {"2": 7}})
complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}})
assert _summed_agg_row([complexity, quality]).tier_turns == {}
@pytest.mark.asyncio
async def test_non_admin_roles_cannot_read_benchmarks(self):
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
@ -427,3 +443,26 @@ class TestAutoRouterBenchmarks:
assert response.routers_in_scope == 1
assert response.groups[0].router_name == "live-auto"
assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0
@pytest.mark.asyncio
@pytest.mark.parametrize(
"wire_value, expected", [({"simple": 24, "complex": 16}, {"simple": 24, "complex": 16}), ({}, {})]
)
async def test_the_tier_map_reaches_the_response_as_the_jsonb_column_returns_it(
self, wire_value: dict, expected: dict, monkeypatch: pytest.MonkeyPatch
):
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
class _DB:
async def query_raw(self, sql: str, *params: object):
return [{**TestAutoRouterBenchmarks.ROW.model_dump(), "tier_turns": wire_value}]
monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})())
response = await get_auto_router_benchmarks(
user_api_key_dict=ADMIN,
start_date="2026-07-01",
end_date="2026-08-01",
)
assert response.groups[0].tier_turns == expected

84
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -932,8 +932,11 @@ async def test_add_team_member_budget_table_success():
)
# Verify the result
assert result == team_info_response
assert result.team_member_budget_table == mock_budget_record
assert result == team_info_response.model_copy(
update={"team_member_budget_table": mock_budget_record}
)
assert team_info_response.team_member_budget_table is None
# Verify database call was made correctly
mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with(
@ -8637,9 +8640,9 @@ class TestBatchResolveAccessGroupResources:
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma):
result = await _batch_resolve_access_group_resources(["ag-1"])
assert sorted(result["ag-1"]["models"]) == ["claude-3", "gpt-4"]
assert result["ag-1"]["mcp_server_ids"] == ["mcp-1"]
assert sorted(result["ag-1"]["agent_ids"]) == ["agent-1", "agent-2"]
assert sorted(result["ag-1"].access_model_names) == ["claude-3", "gpt-4"]
assert result["ag-1"].access_mcp_server_ids == ["mcp-1"]
assert sorted(result["ag-1"].access_agent_ids) == ["agent-1", "agent-2"]
@pytest.mark.asyncio
async def test_multiple_access_groups(self):
@ -8668,8 +8671,8 @@ class TestBatchResolveAccessGroupResources:
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma):
result = await _batch_resolve_access_group_resources(["ag-1", "ag-2"])
assert result["ag-1"]["models"] == ["gpt-4"]
assert result["ag-2"]["models"] == ["gemini"]
assert result["ag-1"].access_model_names == ["gpt-4"]
assert result["ag-2"].access_model_names == ["gemini"]
@pytest.mark.asyncio
async def test_missing_access_group_omitted(self):
@ -8735,6 +8738,75 @@ class TestBatchResolveAccessGroupResources:
assert "ag-1" in result
class TestResolveTeamAccessGroupResources:
"""Tests for the per-team access group resolution on /team/info."""
@pytest.mark.asyncio
async def test_populates_flat_lists_and_per_group_details(self):
"""access_group_details must attribute each model to the group granting it,
so the UI can show provenance on hover; flat lists stay for back-compat.
Duplicated ids must collapse to one entry (response amplification), and the
input object must stay untouched (resolution returns a copy)."""
from litellm.proxy._types import TeamInfoResponseObjectTeamTable
from litellm.proxy.management_endpoints.team_endpoints import (
_resolve_team_access_group_resources,
)
row1 = MagicMock()
row1.access_group_id = "ag-1"
row1.access_group_name = "shared-models"
row1.access_model_names = ["gpt-4", "claude-3"]
row1.access_mcp_server_ids = ["mcp-1"]
row1.access_agent_ids = []
row2 = MagicMock()
row2.access_group_id = "ag-2"
row2.access_group_name = "extra-models"
row2.access_model_names = ["claude-3", "gemini"]
row2.access_mcp_server_ids = []
row2.access_agent_ids = ["agent-1"]
fake_prisma = MagicMock()
fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(
return_value=[row1, row2]
)
team_info = TeamInfoResponseObjectTeamTable(
team_id="team-1", access_group_ids=["ag-1", "ag-2", "ag-1", "ag-missing"]
)
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma):
resolved = await _resolve_team_access_group_resources(team_info)
assert team_info.access_group_details is None
assert sorted(resolved.access_group_models or []) == [
"claude-3",
"gemini",
"gpt-4",
]
assert resolved.access_group_mcp_server_ids == ["mcp-1"]
assert resolved.access_group_agent_ids == ["agent-1"]
assert [
(d.access_group_id, d.access_group_name, d.models)
for d in (resolved.access_group_details or [])
] == [
("ag-1", "shared-models", ("gpt-4", "claude-3")),
("ag-2", "extra-models", ("claude-3", "gemini")),
]
@pytest.mark.asyncio
async def test_no_access_groups_leaves_details_unset(self):
from litellm.proxy._types import TeamInfoResponseObjectTeamTable
from litellm.proxy.management_endpoints.team_endpoints import (
_resolve_team_access_group_resources,
)
team_info = TeamInfoResponseObjectTeamTable(team_id="team-1", access_group_ids=[])
resolved = await _resolve_team_access_group_resources(team_info)
assert resolved.access_group_details is None
assert resolved.access_group_models is None
@pytest.mark.asyncio
async def test_verify_team_access_denies_unauthorized_user():
"""

97
tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py Normal file
View file

@ -0,0 +1,97 @@
import os
import sys
from types import MappingProxyType
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.proxy.openai_files_endpoints.common_utils import (
apply_unified_file_ids,
map_raw_file_ids_to_unified,
)
from litellm.types.utils import LiteLLMBatch
def _batch(input_file_id, output_file_id, error_file_id) -> LiteLLMBatch:
return LiteLLMBatch(
id="batch-1",
completion_window="24h",
created_at=1234567890,
endpoint="/v1/chat/completions",
input_file_id=input_file_id,
object="batch",
status="cancelled",
output_file_id=output_file_id,
error_file_id=error_file_id,
)
@pytest.mark.asyncio
async def test_map_raw_file_ids_to_unified_empty_ids_skips_db():
prisma_client = MagicMock()
assert await map_raw_file_ids_to_unified(frozenset(), prisma_client) == {}
prisma_client.db.litellm_managedfiletable.find_many.assert_not_called()
@pytest.mark.asyncio
async def test_map_raw_file_ids_to_unified_no_prisma_client_returns_empty():
assert await map_raw_file_ids_to_unified(frozenset({"file-raw-1"}), None) == {}
@pytest.mark.asyncio
async def test_map_raw_file_ids_to_unified_bulk_queries_and_filters_to_requested_ids():
row_a = MagicMock(
unified_file_id="unified-a",
flat_model_file_ids=["file-raw-a", "file-raw-other"],
)
row_b = MagicMock(unified_file_id="unified-b", flat_model_file_ids=["file-raw-b"])
prisma_client = MagicMock()
prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row_a, row_b])
mapping = await map_raw_file_ids_to_unified(
frozenset({"file-raw-b", "file-raw-a", "file-raw-missing"}), prisma_client
)
prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with(
where={"flat_model_file_ids": {"hasSome": ["file-raw-a", "file-raw-b", "file-raw-missing"]}}
)
assert dict(mapping) == {"file-raw-a": "unified-a", "file-raw-b": "unified-b"}
def test_apply_unified_file_ids_swaps_only_mapped_ids():
batch = _batch(input_file_id="file-raw-in", output_file_id="file-raw-out", error_file_id=None)
apply_unified_file_ids(batch, MappingProxyType({"file-raw-out": "unified-out"}))
assert batch.input_file_id == "file-raw-in"
assert batch.output_file_id == "unified-out"
assert batch.error_file_id is None
def test_apply_unified_file_ids_swaps_all_three_ids():
batch = _batch(
input_file_id="file-raw-in",
output_file_id="file-raw-out",
error_file_id="file-raw-err",
)
apply_unified_file_ids(
batch,
MappingProxyType(
{
"file-raw-in": "unified-in",
"file-raw-out": "unified-out",
"file-raw-err": "unified-err",
}
),
)
assert (batch.input_file_id, batch.output_file_id, batch.error_file_id) == (
"unified-in",
"unified-out",
"unified-err",
)

117
tests/test_litellm/proxy/utils/helpers/test_model_access.py
View file

@ -9,6 +9,7 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import (
create_model_info_response,
get_available_models_for_user,
hash_token,
is_known_model,
is_known_vector_store_index,
model_dump_with_preserved_fields,
@ -404,3 +405,119 @@ async def test_get_available_models_for_user_error_path_complete_list_raises(
general_settings={},
user_model=None,
)
@pytest.mark.asyncio
async def test_get_available_models_for_user_resolves_team_access_group_models(
monkeypatch,
):
from litellm.models.access_group import LiteLLM_AccessGroupTable
from litellm.models.team import LiteLLM_TeamTableCachedObj
team = LiteLLM_TeamTableCachedObj(
team_id="team-1",
models=["no-default-models"],
access_group_ids=["ag-1"],
)
access_group = LiteLLM_AccessGroupTable(
access_group_id="ag-1",
access_group_name="repro-group",
access_model_names=["model-a", "model-b"],
assigned_team_ids=["team-1"],
)
async def _get_team_object(**_kwargs):
return team
async def _get_access_object(**_kwargs):
return access_group
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object)
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object)
result = await get_available_models_for_user(
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-test-key",
user_id="user-1",
team_id="team-1",
models=["all-team-models"],
team_models=["no-default-models"],
),
llm_router=_router_with_models(["model-a", "model-b", "model-c"]),
general_settings={},
user_model=None,
prisma_client=MagicMock(),
proxy_logging_obj=MagicMock(),
user_api_key_cache=MagicMock(),
)
assert sorted(result) == ["model-a", "model-b"]
@pytest.mark.asyncio
async def test_get_available_models_for_user_without_access_groups_grants_nothing(
monkeypatch,
):
from litellm.models.team import LiteLLM_TeamTableCachedObj
async def _get_team_object(**_kwargs):
return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"])
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object)
result = await get_available_models_for_user(
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-test-key",
user_id="user-1",
team_id="team-1",
models=["all-team-models"],
team_models=["no-default-models"],
),
llm_router=_router_with_models(["model-a", "model-b"]),
general_settings={},
user_model=None,
prisma_client=MagicMock(),
proxy_logging_obj=MagicMock(),
user_api_key_cache=MagicMock(),
)
assert result == []
@pytest.mark.asyncio
async def test_get_available_models_for_user_resolves_key_access_group_models(
monkeypatch,
):
from litellm.models.access_group import LiteLLM_AccessGroupTable
from litellm.models.team import LiteLLM_TeamTableCachedObj
async def _get_team_object(**_kwargs):
return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"])
async def _get_access_object(**_kwargs):
return LiteLLM_AccessGroupTable(
access_group_id="ag-1",
access_group_name="key-group",
access_model_names=["model-b"],
assigned_key_ids=[hash_token("sk-test-key")],
)
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object)
monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
result = await get_available_models_for_user(
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-test-key",
user_id="user-1",
team_id="team-1",
models=["no-default-models"],
team_models=["no-default-models"],
access_group_ids=["ag-1"],
),
llm_router=_router_with_models(["model-a", "model-b"]),
general_settings={},
user_model=None,
prisma_client=MagicMock(),
proxy_logging_obj=MagicMock(),
user_api_key_cache=MagicMock(),
)
assert result == ["model-b"]

39
tests/test_litellm/responses/test_responses_api_bridge_flag.py
View file

@ -16,6 +16,7 @@ sys.path.insert(
import litellm
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.utils import Choices, Message, ModelResponse, Usage
class TestUseResponsesApiBridgeFlag:
@ -130,6 +131,44 @@ class TestUseResponsesApiBridgeFlag:
mock_bridge_handler.assert_called_once()
@patch("litellm.acompletion")
@patch(
"litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config"
)
async def test_allowed_openai_params_forwarded_through_bridge(
self, mock_get_config, mock_acompletion
):
"""allowed_openai_params is a named param of responses(), so it must be
explicitly forwarded to the bridge; otherwise litellm.acompletion raises
UnsupportedParamsError for params the caller explicitly allowed."""
mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig()
mock_acompletion.return_value = ModelResponse(
id="chatcmpl_123",
model="openai/my-custom-model",
choices=[
Choices(
index=0,
message=Message(role="assistant", content="Answer"),
finish_reason="stop",
)
],
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
await litellm.aresponses(
model="openai/my-custom-model",
input="Hello",
use_chat_completions_api=True,
allowed_openai_params=["reasoning_effort"],
reasoning={"effort": "high"},
litellm_logging_obj=MagicMock(),
)
mock_acompletion.assert_called_once()
assert mock_acompletion.call_args.kwargs.get("allowed_openai_params") == [
"reasoning_effort"
]
@patch("litellm.responses.file_search.emulated_handler._call_aresponses")
@patch(
"litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config"

56
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -3477,6 +3477,23 @@ class TestSessionAffinity:
# Pinned to the first turn's model, not re-classified down to SIMPLE.
assert second.model == "o1-preview"
@pytest.mark.asyncio
async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config):
mock_router_instance.cache = DualCache()
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=session_affinity_config,
)
request_kwargs = self._request_kwargs("session-1")
await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE
)
pinned = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert pinned.routing_decision["tier"] == "REASONING"
@pytest.mark.asyncio
async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config):
mock_router_instance.cache = DualCache()
@ -4362,7 +4379,24 @@ class TestRoutingDecisionContents:
assert decision is not None
assert decision["cause"] == "default_fallback"
assert decision["routed_model"] == response.model
assert "tier" not in decision
assert decision.get("tier") == "MEDIUM"
@pytest.mark.asyncio
async def test_a_default_model_fallback_claims_no_tier(self, mock_router_instance, basic_config):
router = ComplexityRouter(
model_name="test-complexity-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={**basic_config, "default_model": "gpt-4o"},
)
response = await router.async_pre_routing_hook(
model="test-complexity-router",
request_kwargs={},
messages=[{"role": "system", "content": "be nice"}],
)
assert response is not None
assert response.routing_decision is not None
assert response.routing_decision["cause"] == "default_fallback"
assert "tier" not in response.routing_decision
@pytest.mark.asyncio
async def test_session_pin_decision(self, mock_router_instance, basic_config):
@ -5907,11 +5941,21 @@ class TestConversationShapeDiscriminator:
)
builds = source.split("self._build_routing_decision(")[1:]
assert builds
missing = [
i
for i, block in enumerate(builds)
if "conversation_continuing=conversation_continuing" not in block.split("),")[0]
]
missing = []
for i, block in enumerate(builds):
depth = 0
end = 0
for j, char in enumerate(block):
if char == "(":
depth += 1
elif char == ")":
depth -= 1
if depth < 0:
end = j
break
extracted = block[:end]
if "conversation_continuing=conversation_continuing" not in extracted:
missing.append(i)
assert not missing, f"routing decisions {missing} do not carry the conversation shape"

73
tests/test_litellm/types/test_router.py Normal file
View file

@ -0,0 +1,73 @@
import pytest
from litellm.types.router import (
SPECIAL_MODEL_INFO_PARAMS,
Deployment,
LiteLLM_Params,
ModelInfo,
)
from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams
def test_model_info_declares_mirrored_pricing_fields():
"""The pricing keys Deployment mirrors onto model_info must be declared fields, not
extras that only survive because ModelInfo sets extra="allow"."""
for field in SPECIAL_MODEL_INFO_PARAMS:
assert field in ModelInfo.model_fields
info = ModelInfo(id="x", input_cost_per_token=1e-06)
assert info.__pydantic_extra__ == {}
assert info.input_cost_per_token == 1e-06
def test_special_model_info_params_cannot_drift_from_the_mirror():
assert SPECIAL_MODEL_INFO_PARAMS == tuple(MirroredPricingParams.model_fields)
assert set(SPECIAL_MODEL_INFO_PARAMS) <= set(CustomPricingLiteLLMParams.model_fields)
assert set(SPECIAL_MODEL_INFO_PARAMS) <= set(LiteLLM_Params.model_fields)
def test_custom_pricing_params_keeps_every_field_it_had():
"""The mirrored fields moved to a base class; none of them may go missing from
CustomPricingLiteLLMParams, whose model_fields drive custom-pricing detection."""
for field in (
"input_cost_per_token",
"output_cost_per_token",
"input_cost_per_character",
"output_cost_per_character",
"cache_read_input_token_cost",
"cache_creation_input_token_cost",
"input_cost_per_second",
"cache_read_input_token_cost_flex",
"input_cost_per_character_above_128k_tokens",
"output_cost_per_audio_token",
):
assert field in CustomPricingLiteLLMParams.model_fields
@pytest.mark.parametrize("field", SPECIAL_MODEL_INFO_PARAMS)
def test_deployment_mirrors_pricing_from_litellm_params_onto_model_info(field):
deployment = Deployment(
model_name="my-model",
litellm_params=LiteLLM_Params(model="gpt-4o", **{field: 3e-06}),
)
assert getattr(deployment.model_info, field) == 3e-06
assert deployment.model_info.model_dump(exclude_none=True)[field] == 3e-06
def test_unset_pricing_is_still_absent_from_dumps():
"""/model/info responses and DB writes dump model_info with exclude_none=True, so
declaring the pricing fields must not start emitting ~6 null keys per deployment."""
dumped = ModelInfo(id="x").model_dump(exclude_none=True)
assert [field for field in SPECIAL_MODEL_INFO_PARAMS if field in dumped] == []
def test_pricing_strings_are_coerced_to_float():
"""Cost values arrive from the DB and the Admin UI as strings; they must land as
floats so cost calculation doesn't multiply a str."""
info = ModelInfo(id="x", output_cost_per_token="0.000002")
assert info.output_cost_per_token == 2e-06
def test_invalid_pricing_is_rejected():
with pytest.raises(ValueError):
ModelInfo(id="x", input_cost_per_token="free")

6
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 23245
"limit": 23235
},
"LIT002": {
"limit": 27179
"limit": 27176
},
"LIT003": {
"limit": 269
@ -30,6 +30,6 @@
"limit": 16769
},
"LIT011": {
"limit": 5602
"limit": 5598
}
}

28
ui/litellm-dashboard/src/components/Teams.test.tsx
View file

@ -613,6 +613,34 @@ describe("Teams - access_group_ids in team create", () => {
);
});
});
it("creates a team with no models selected, sending the no-default-models sentinel instead of an empty list", async () => {
renderWithQueryClient(<Teams accessToken="test-token" userID="user-123" userRole="Admin" />);
const createButton = screen.getAllByRole("button", { name: /create team/i })[0];
act(() => {
fireEvent.click(createButton);
});
await waitFor(() => {
expect(screen.getByLabelText(/team name/i)).toBeInTheDocument();
});
fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Group Only Team" } });
const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i });
fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]);
await waitFor(() => {
expect(teamCreateCall).toHaveBeenCalledWith(
"test-token",
expect.objectContaining({
team_alias: "Group Only Team",
models: ["no-default-models"],
}),
);
});
});
});
describe("Teams - metadata key-value pairs in team create", () => {

11
ui/litellm-dashboard/src/components/Teams.tsx
View file

@ -42,6 +42,7 @@ interface TeamProps {
import DeleteResourceModal from "./common_components/DeleteResourceModal";
import { teamCreateCall } from "./networking";
import { normalizeTeamModelSelection } from "./team/teamModelAccess";
import { ModelSelect } from "./ModelSelect/ModelSelect";
const canCreateOrManageTeams = (
@ -351,7 +352,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
}
}
await teamCreateCall(accessToken, formValues);
await teamCreateCall(accessToken, { ...formValues, models: normalizeTeamModelSelection(formValues.models) });
NotificationsManager.success("Team created");
await refreshTeams();
form.resetFields();
@ -618,17 +619,11 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
label={
<span>
Models{" "}
<Tooltip title="These are the models that your selected team has access to">
<Tooltip title="These are the models that your selected team has access to. Leave empty to grant no models directly, e.g. when the team gets its models from access groups">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
}
rules={[
{
required: true,
message: "Please select at least one model",
},
]}
name="models"
>
<ModelSelect

41
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View file

@ -35,6 +35,12 @@ import { CheckIcon, CopyIcon } from "lucide-react";
import React, { useEffect, useMemo, useState } from "react";
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
import AccessGroupSelector from "../common_components/AccessGroupSelector";
import {
computeTeamModelBadges,
normalizeTeamModelSelection,
TeamAccessGroupModelGrant,
TeamModelBadgeKind,
} from "./teamModelAccess";
import MetadataKeyValueFields, {
metadataObjectToPairs,
metadataPairsToObject,
@ -82,6 +88,13 @@ const UI_MANAGED_METADATA_KEYS: ReadonlySet<string> = new Set([
"disable_global_guardrails",
]);
const TEAM_MODEL_BADGE_COLORS: Record<TeamModelBadgeKind, "red" | "gray" | "blue" | "green"> = {
"all-proxy": "red",
"no-default": "gray",
direct: "blue",
"access-group": "green",
};
export interface TeamMembership {
user_id: string;
team_id: string;
@ -132,6 +145,7 @@ export interface TeamData {
access_group_models?: string[];
access_group_mcp_server_ids?: string[];
access_group_agent_ids?: string[];
access_group_details?: TeamAccessGroupModelGrant[];
router_settings?: Record<string, any>;
guardrails?: string[];
policies?: string[];
@ -483,7 +497,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
const updateData: any = {
team_id: teamId,
team_alias: values.team_alias,
models: values.models,
models: normalizeTeamModelSelection(values.models),
tpm_limit: sanitizeNumeric(values.tpm_limit),
rpm_limit: sanitizeNumeric(values.rpm_limit),
model_tpm_limit: modelTpmLimit,
@ -764,21 +778,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<Card>
<Text>Models</Text>
<div className="mt-2 flex flex-wrap gap-2">
{info.models.length === 0 || info.models.includes("all-proxy-models") ? (
<Badge color="red">All proxy models</Badge>
) : (
<>
{info.models.map((model: string, index: number) => (
<Badge key={`direct-${index}`} color="blue">
{model}
</Badge>
))}
{(info.access_group_models || []).map((model: string, index: number) => (
<Badge key={`ag-${index}`} color="green" title="From access group">
{model}
</Badge>
))}
</>
{computeTeamModelBadges(info.models, info.access_group_models || [], info.access_group_details).map(
(badge, index) => (
<Tooltip key={`${badge.kind}-${badge.label}-${index}`} title={badge.tooltip}>
<span>
<Badge color={TEAM_MODEL_BADGE_COLORS[badge.kind]}>{badge.label}</Badge>
</span>
</Tooltip>
),
)}
</div>
</Card>
@ -982,7 +989,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<Form.Item
label="Models"
name="models"
rules={[{ required: true, message: "Please select at least one model" }]}
extra="Leave empty to grant no models directly. The team keeps any models granted through its access groups"
>
<ModelSelect
value={form.getFieldValue("models") || []}

86
ui/litellm-dashboard/src/components/team/teamModelAccess.test.ts Normal file
View file

@ -0,0 +1,86 @@
import { describe, expect, it } from "vitest";
import { computeTeamModelBadges, normalizeTeamModelSelection, TeamAccessGroupModelGrant } from "./teamModelAccess";
const GRANTS: TeamAccessGroupModelGrant[] = [
{ access_group_id: "ag-1", access_group_name: "shared", models: ["haiku", "gpt-4o-mini"] },
{ access_group_id: "ag-2", access_group_name: "extra", models: ["haiku", "sonnet"] },
];
describe("normalizeTeamModelSelection", () => {
it("substitutes the no-default-models sentinel for an empty selection", () => {
expect(normalizeTeamModelSelection([])).toEqual(["no-default-models"]);
expect(normalizeTeamModelSelection(undefined)).toEqual(["no-default-models"]);
});
it("passes a non-empty selection through untouched", () => {
expect(normalizeTeamModelSelection(["gpt-4o-mini"])).toEqual(["gpt-4o-mini"]);
expect(normalizeTeamModelSelection(["all-proxy-models"])).toEqual(["all-proxy-models"]);
});
});
describe("computeTeamModelBadges", () => {
it("attributes group-only models to the groups granting them", () => {
const badges = computeTeamModelBadges(["sonnet-direct"], [], GRANTS);
expect(badges).toEqual([
{
label: "sonnet-direct",
kind: "direct",
tooltip: "Granted directly in the team's model list",
},
{ label: "haiku", kind: "access-group", tooltip: "Granted via access groups shared, extra" },
{ label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" },
{ label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" },
]);
});
it("marks a model both direct and group-granted on the direct badge, without a duplicate badge", () => {
const badges = computeTeamModelBadges(["haiku"], [], GRANTS);
expect(badges).toEqual([
{
label: "haiku",
kind: "direct",
tooltip: "Granted directly in the team's model list, and also via access groups shared, extra",
},
{ label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" },
{ label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" },
]);
});
it("shows the no-default-models sentinel as its own badge and keeps group badges visible", () => {
const badges = computeTeamModelBadges(["no-default-models"], [], [GRANTS[0]]);
expect(badges.map((b) => [b.label, b.kind])).toEqual([
["No default models", "no-default"],
["haiku", "access-group"],
["gpt-4o-mini", "access-group"],
]);
});
it("still shows group badges when the empty model list grants everything", () => {
const badges = computeTeamModelBadges([], [], [GRANTS[0]]);
expect(badges[0]).toEqual({
label: "All proxy models",
kind: "all-proxy",
tooltip: "The team's model list is empty, so it can access every model on the proxy",
});
expect(badges.slice(1).map((b) => b.label)).toEqual(["haiku", "gpt-4o-mini"]);
});
it("distinguishes the all-proxy-models sentinel from an empty list in the tooltip", () => {
const badges = computeTeamModelBadges(["all-proxy-models"], [], []);
expect(badges).toEqual([
{
label: "All proxy models",
kind: "all-proxy",
tooltip: "Granted by the All Proxy Models entry in the team's model list",
},
]);
});
it("falls back to the flat access_group_models list when per-group details are absent", () => {
const badges = computeTeamModelBadges(["direct-model"], ["haiku"], undefined);
expect(badges).toEqual([
{ label: "direct-model", kind: "direct", tooltip: "Granted directly in the team's model list" },
{ label: "haiku", kind: "access-group", tooltip: "Granted via an access group" },
]);
});
});

82
ui/litellm-dashboard/src/components/team/teamModelAccess.ts Normal file
View file

@ -0,0 +1,82 @@
export const ALL_PROXY_MODELS = "all-proxy-models";
export const NO_DEFAULT_MODELS = "no-default-models";
export interface TeamAccessGroupModelGrant {
access_group_id: string;
access_group_name: string;
models: string[];
}
export type TeamModelBadgeKind = "all-proxy" | "no-default" | "direct" | "access-group";
export interface TeamModelBadge {
label: string;
kind: TeamModelBadgeKind;
tooltip: string;
}
export function normalizeTeamModelSelection(models: string[] | undefined): string[] {
return models && models.length > 0 ? models : [NO_DEFAULT_MODELS];
}
const describeGroups = (names: string[]): string =>
names.length > 1 ? `access groups ${names.join(", ")}` : `access group ${names[0]}`;
export function computeTeamModelBadges(
models: string[],
accessGroupModels: string[],
accessGroupDetails: TeamAccessGroupModelGrant[] | undefined,
): TeamModelBadge[] {
const grants = accessGroupDetails ?? [];
const groupNamesFor = (model: string): string[] =>
grants.filter((g) => g.models.includes(model)).map((g) => g.access_group_name);
const viaGroups = (model: string): string => {
const names = groupNamesFor(model);
return names.length > 0 ? describeGroups(names) : "an access group";
};
const allProxy = models.length === 0 || models.includes(ALL_PROXY_MODELS);
const directModels = allProxy ? [] : models.filter((m) => m !== NO_DEFAULT_MODELS);
const groupModels = [...new Set(grants.length > 0 ? grants.flatMap((g) => g.models) : accessGroupModels)].filter(
(m) => !directModels.includes(m),
);
const allProxyBadge: TeamModelBadge = {
label: "All proxy models",
kind: "all-proxy",
tooltip: models.includes(ALL_PROXY_MODELS)
? "Granted by the All Proxy Models entry in the team's model list"
: "The team's model list is empty, so it can access every model on the proxy",
};
const noDefaultBadge: TeamModelBadge = {
label: "No default models",
kind: "no-default",
tooltip: "No models are granted directly. Access comes only from access groups",
};
const headBadge = (): TeamModelBadge[] => {
if (allProxy) return [allProxyBadge];
if (models.includes(NO_DEFAULT_MODELS)) return [noDefaultBadge];
return [];
};
return [
...headBadge(),
...directModels.map(
(m): TeamModelBadge => ({
label: m,
kind: "direct",
tooltip:
groupNamesFor(m).length > 0
? `Granted directly in the team's model list, and also via ${viaGroups(m)}`
: "Granted directly in the team's model list",
}),
),
...groupModels.map(
(m): TeamModelBadge => ({
label: m,
kind: "access-group",
tooltip: `Granted via ${viaGroups(m)}`,
}),
),
];
}

19
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -21478,6 +21478,13 @@ export interface components {
* @description What the routed traffic actually cost
*/
spend: number;
/**
* Tier Turns
* @description Turns per tier, keyed by the tier name the routing decision recorded at request time (never re-derived at read time, since the tier-to-model mapping is mutable config). Tier names are scoped to this group's router_type and are not comparable across types: a complexity router reports 'simple'/'medium'/'complex'/'reasoning', a quality router reports its numeric quality tier, and an adaptive router records no tier at all. Turns no tier served (the classifier fell back to default_model) are absent rather than pooled under a sentinel key, so the values may sum to less than turns
*/
tier_turns?: {
[key: string]: number;
};
/** Turns */
turns: number;
};
@ -35510,6 +35517,10 @@ export interface components {
base_model?: string | null;
/** Blocked */
blocked?: boolean | null;
/** Cache Creation Input Token Cost */
cache_creation_input_token_cost?: number | null;
/** Cache Read Input Token Cost */
cache_read_input_token_cost?: number | null;
/** Created At */
created_at?: string | null;
/** Created By */
@ -35521,6 +35532,14 @@ export interface components {
db_model: boolean;
/** Id */
id: string | null;
/** Input Cost Per Character */
input_cost_per_character?: number | null;
/** Input Cost Per Token */
input_cost_per_token?: number | null;
/** Output Cost Per Character */
output_cost_per_character?: number | null;
/** Output Cost Per Token */
output_cost_per_token?: number | null;
/** Team Id */
team_id?: string | null;
/** Team Public Model Name */

2
ui/litellm-dashboard/tsconfig.tsbuildinfo
View file

File diff suppressed because one or more lines are too long