diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c8765eb0bd0..32b8eb3d4d0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 8d589e6207c..d8318962633 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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] diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql new file mode 100644 index 00000000000..81b1cbc7ec3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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") diff --git a/litellm/constants.py b/litellm/constants.py index 6f0e9e7afe2..30d3bb1f26e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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))) diff --git a/litellm/integrations/otel/__init__.py b/litellm/integrations/otel/__init__.py index 94442e96adb..9c1205bb277 100644 --- a/litellm/integrations/otel/__init__.py +++ b/litellm/integrations/otel/__init__.py @@ -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", diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index af56734bec1..032441535e0 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -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, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 02499010624..aba9cc80240 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -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 diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 24f0b947b08..3d585c36b67 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -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``.""" diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 71388134e98..9748db2dcd2 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -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 diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 795810a7c40..199ab020559 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -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 { diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 314cfef6d84..9aa5a4f465f 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -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: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 3ef298aa336..c4b5cc628e2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -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, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b7dc3a7c73..1fc05ac4653 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3da899a5610..d07ac0c5586 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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, diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index a6e5eb2a0a0..ff9211742f3 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -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: diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index da1652cdb61..b1f074c26b7 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -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) """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index e9e729fb118..eecbce57468 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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(), diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 190d19f3d52..0d23e19f88d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -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, diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 88e5ae1b36d..b24c5ac298a 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -4,7 +4,7 @@ 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 @@ -264,6 +264,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 @@ -288,6 +289,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, @@ -310,10 +328,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 """ @@ -370,6 +389,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), @@ -447,6 +467,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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fe5a0e06d2e..0b99879f9fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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, diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index fce57e514c1..9f3c9d9db45 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -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, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 605455a2f73..e59c6adaf22 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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), diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f923702119c..7b02c1b8023 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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, ) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a69509fc37a..f6ced0bb9d2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -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, ), ) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 6b354a39101..bbb6d758814 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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): diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 05537da67d7..90713b270be 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -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): diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 056ff0cc5a2..61ca72d0f64 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -2,6 +2,7 @@ Types for auto-router management endpoints """ +from collections.abc import Mapping from typing import Final 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): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d97bdc3532f..8d66b624341 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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): diff --git a/schema.prisma b/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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") diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index e83897025a3..605dd8fb7e5 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -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: diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 34770bc596b..277478eebaf 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -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) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 93861d19922..85964529ada 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -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( diff --git a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py index b6d90b7f6a2..5e9c9f614e5 100644 --- a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py @@ -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: diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index 7053bd1dfd2..d76f7b356b2 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -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: diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 33ec557c339..73453478e5a 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -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( diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..2627bdb8038 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -70,6 +70,7 @@ from e2e_config import ( POLL_TIMEOUT, PROXY_BASE_URL, REQUEST_TIMEOUT, + settle_propagation, ) from transport import HttpTransport, SplitTransport, Transport @@ -161,13 +162,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. + + 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", @@ -180,7 +186,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: diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 50af6465d06..e1e5cc6c532 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -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 diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index aa734ee22cc..65b70f13a3b 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -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) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 41c02501acc..2573ad5a375 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -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 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 612ac1e5113..19d0cfc0b18 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -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 diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py index 544abab8dcf..c859f9b2f55 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py @@ -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"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index df3db3d2c57..f11324ca376 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 6ab0f2c08ab..9df72108332 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -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, diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 3b8b4af78d9..76bb11cc26d 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f56b8e113a9..e6c0eaee3c4 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -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 == [] diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index aa7d01bc880..0df11f224a2 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -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): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 76a695ce3fd..837fb93d331 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -7,6 +7,7 @@ import os import sys from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException @@ -15,12 +16,18 @@ sys.path.insert(0, os.path.abspath("../../../../../..")) import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentChunkResult, BedrockGuardrail, _redact_pii_matches, ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentItem, + BedrockTextContent, +) from litellm.types.utils import CallTypes, ModelResponse @@ -53,9 +60,7 @@ async def test__redact_pii_matches_function(): redacted_response = _redact_pii_matches(response_with_pii) # Verify that PII matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" @@ -173,12 +178,8 @@ async def test__redact_pii_matches_multiple_assessments(): redacted_response = _redact_pii_matches(response_multiple_assessments) # Verify all PII in all assessments are redacted - assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ - "piiEntities" - ] + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"]["piiEntities"] assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" @@ -199,9 +200,7 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII mock_bedrock_response = MagicMock() @@ -239,20 +238,11 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" - ) as mock_debug, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug") as mock_debug, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method that should log the redacted response @@ -275,37 +265,23 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): bedrock_response_log_call = call break - assert ( - bedrock_response_log_call is not None - ), "Should have logged Bedrock AI response" + assert bedrock_response_log_call is not None, "Should have logged Bedrock AI response" # Extract the logged response data - logged_response = bedrock_response_log_call[0][ - 1 - ] # Second argument to debug call + logged_response = bedrock_response_log_call[0][1] # Second argument to debug call # Verify that the logged response has redacted PII assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] - == "[REDACTED]" + logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) # Verify other fields are preserved assert logged_response["action"] == "GUARDRAIL_INTERVENED" - assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["type"] - == "PHONE" - ) + assert logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["type"] == "PHONE" slg_list = request_data["metadata"]["standard_logging_guardrail_information"] assert ( - slg_list[0]["guardrail_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + slg_list[0]["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) @@ -319,9 +295,7 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII original_response_data = { @@ -361,17 +335,10 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method @@ -385,19 +352,12 @@ async def test_bedrock_guardrail_original_response_not_modified(): # (The json() method should return the original data) original_data = mock_bedrock_response.json() assert ( - original_data["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + original_data["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" ) # Verify that the returned BedrockGuardrailResponse contains original data - assert ( - result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "+1 412 555 1212" - ) + assert result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" print("Original response not modified test passed") @@ -454,18 +414,14 @@ async def test__redact_pii_matches_preserves_non_pii_entities(): redacted_response = _redact_pii_matches(response_with_mixed_data) # Verify that PII entity matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" # Verify that regex matches are also redacted (updated behavior) - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" @@ -496,9 +452,7 @@ async def test_pii_redaction_matches_debug_output_format(): "assessments": [ { "invocationMetrics": { - "guardrailCoverage": { - "textCharacters": {"guarded": 84, "total": 84} - }, + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, "guardrailProcessingLatency": 322, "usage": { "contentPolicyImageUnits": 0, @@ -553,9 +507,7 @@ async def test_pii_redaction_matches_debug_output_format(): redacted_response = _redact_pii_matches(original_response) # Verify the redacted response matches your expected debug output - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] # All PII matches should be redacted assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" @@ -570,34 +522,19 @@ async def test_pii_redaction_matches_debug_output_format(): assert pii_entities[0]["detected"] == True # Verify that the original response is unchanged - original_pii_entities = original_response["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"] - assert ( - original_pii_entities[0]["match"] == "John Smith" - ), "Original should be unchanged" - assert ( - original_pii_entities[1]["match"] == "324-12-3212" - ), "Original should be unchanged" - assert ( - original_pii_entities[2]["match"] == "607-456-7890" - ), "Original should be unchanged" + original_pii_entities = original_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert original_pii_entities[0]["match"] == "John Smith", "Original should be unchanged" + assert original_pii_entities[1]["match"] == "324-12-3212", "Original should be unchanged" + assert original_pii_entities[2]["match"] == "607-456-7890", "Original should be unchanged" # Verify all other metadata is preserved in redacted response assert redacted_response["action"] == "GUARDRAIL_INTERVENED" assert redacted_response["actionReason"] == "Guardrail blocked." assert redacted_response["blockedResponse"] == "Input blocked by PII policy" - assert ( - redacted_response["assessments"][0]["invocationMetrics"][ - "guardrailProcessingLatency" - ] - == 322 - ) + assert redacted_response["assessments"][0]["invocationMetrics"]["guardrailProcessingLatency"] == 322 print("PII redaction matches debug output format test passed") - print( - f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" - ) + print(f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}") print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") @@ -632,14 +569,10 @@ async def test__redact_pii_matches_with_regex_matches(): redacted_response = _redact_pii_matches(response_with_regex) # Verify that regex matches are redacted - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" - assert ( - regexes[1]["match"] == "[REDACTED]" - ), "Credit card regex match should be redacted" + assert regexes[1]["match"] == "[REDACTED]", "Credit card regex match should be redacted" # Verify other fields are preserved assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" @@ -648,13 +581,9 @@ async def test__redact_pii_matches_with_regex_matches(): assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" # Verify original response is unchanged - original_regexes = response_with_regex["assessments"][0][ - "sensitiveInformationPolicy" - ]["regexes"] + original_regexes = response_with_regex["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" - assert ( - original_regexes[1]["match"] == "4111-1111-1111-1111" - ), "Original should be unchanged" + assert original_regexes[1]["match"] == "4111-1111-1111-1111", "Original should be unchanged" print("Regex matches redaction test passed") @@ -690,31 +619,17 @@ async def test__redact_pii_matches_with_custom_words(): # Verify that custom word matches are redacted custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "First custom word match should be redacted" - assert ( - custom_words[1]["match"] == "[REDACTED]" - ), "Second custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "First custom word match should be redacted" + assert custom_words[1]["match"] == "[REDACTED]", "Second custom word match should be redacted" # Verify other fields are preserved - assert ( - custom_words[0]["action"] == "BLOCKED" - ), "Custom word action should be preserved" - assert ( - custom_words[1]["action"] == "ANONYMIZED" - ), "Custom word action should be preserved" + assert custom_words[0]["action"] == "BLOCKED", "Custom word action should be preserved" + assert custom_words[1]["action"] == "ANONYMIZED", "Custom word action should be preserved" # Verify original response is unchanged - original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ - "customWords" - ] - assert ( - original_custom_words[0]["match"] == "confidential_data" - ), "Original should be unchanged" - assert ( - original_custom_words[1]["match"] == "secret_information" - ), "Original should be unchanged" + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"]["customWords"] + assert original_custom_words[0]["match"] == "confidential_data", "Original should be unchanged" + assert original_custom_words[1]["match"] == "secret_information", "Original should be unchanged" print("Custom words redaction test passed") @@ -750,41 +665,21 @@ async def test__redact_pii_matches_with_managed_words(): redacted_response = _redact_pii_matches(response_with_managed_words) # Verify that managed word matches are redacted - managed_words = redacted_response["assessments"][0]["wordPolicy"][ - "managedWordLists" - ] + managed_words = redacted_response["assessments"][0]["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "First managed word match should be redacted" - assert ( - managed_words[1]["match"] == "[REDACTED]" - ), "Second managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "First managed word match should be redacted" + assert managed_words[1]["match"] == "[REDACTED]", "Second managed word match should be redacted" # Verify other fields are preserved - assert ( - managed_words[0]["action"] == "BLOCKED" - ), "Managed word action should be preserved" - assert ( - managed_words[0]["type"] == "PROFANITY" - ), "Managed word type should be preserved" - assert ( - managed_words[1]["action"] == "ANONYMIZED" - ), "Managed word action should be preserved" - assert ( - managed_words[1]["type"] == "HATE_SPEECH" - ), "Managed word type should be preserved" + assert managed_words[0]["action"] == "BLOCKED", "Managed word action should be preserved" + assert managed_words[0]["type"] == "PROFANITY", "Managed word type should be preserved" + assert managed_words[1]["action"] == "ANONYMIZED", "Managed word action should be preserved" + assert managed_words[1]["type"] == "HATE_SPEECH", "Managed word type should be preserved" # Verify original response is unchanged - original_managed_words = response_with_managed_words["assessments"][0][ - "wordPolicy" - ]["managedWordLists"] - assert ( - original_managed_words[0]["match"] == "inappropriate_word" - ), "Original should be unchanged" - assert ( - original_managed_words[1]["match"] == "offensive_term" - ), "Original should be unchanged" + original_managed_words = response_with_managed_words["assessments"][0]["wordPolicy"]["managedWordLists"] + assert original_managed_words[0]["match"] == "inappropriate_word", "Original should be unchanged" + assert original_managed_words[1]["match"] == "offensive_term", "Original should be unchanged" print("Managed words redaction test passed") @@ -841,9 +736,7 @@ async def test__redact_pii_matches_comprehensive_coverage(): # PII entities pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] - assert ( - pii_entities[0]["match"] == "[REDACTED]" - ), "PII entity match should be redacted" + assert pii_entities[0]["match"] == "[REDACTED]", "PII entity match should be redacted" # Regex matches regexes = assessment["sensitiveInformationPolicy"]["regexes"] @@ -851,15 +744,11 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Custom words custom_words = assessment["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "Custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "Custom word match should be redacted" # Managed words managed_words = assessment["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "Managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "Managed word match should be redacted" # Verify all other fields are preserved assert pii_entities[0]["type"] == "EMAIL" @@ -868,21 +757,10 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Verify original response is unchanged original_assessment = comprehensive_response["assessments"][0] - assert ( - original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] - == "user@example.com" - ) - assert ( - original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] - == "555-123-4567" - ) - assert ( - original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" - ) - assert ( - original_assessment["wordPolicy"]["managedWordLists"][0]["match"] - == "inappropriate" - ) + assert original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "user@example.com" + assert original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] == "555-123-4567" + assert original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + assert original_assessment["wordPolicy"]["managedWordLists"][0]["match"] == "inappropriate" print("Comprehensive coverage redaction test passed") @@ -914,9 +792,7 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method to avoid actual AWS credential loading - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -926,10 +802,12 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + ) print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") @@ -944,9 +822,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) # Create guardrail without explicit aws_bedrock_runtime_endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -960,9 +836,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -972,10 +846,10 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint from environment is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, f"Expected URL to contain env endpoint. Got: {prepped_request.url}" print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") @@ -988,9 +862,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) # Create guardrail without any custom endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -1004,9 +876,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey aws_region_name = "us-west-2" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1017,9 +887,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey # Verify that the default endpoint is used expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected default URL. Got: {prepped_request.url}" + assert prepped_request.url == expected_url, f"Expected default URL. Got: {prepped_request.url}" print(f"Default endpoint test passed. URL: {prepped_request.url}") @@ -1057,9 +925,7 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1069,10 +935,12 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch ) # Verify that the parameter takes precedence over environment variable - expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + expected_url = ( + f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + ) print(f"Parameter precedence test passed. URL: {prepped_request.url}") @@ -1081,14 +949,10 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" # Create a BedrockGuardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the make_bedrock_api_request method - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api_request: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api_request: # Test the apply_guardrail method with tool_calls in response inputs = { "texts": [], @@ -1115,14 +979,9 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): assert guardrailed_inputs is not None assert "tool_calls" in guardrailed_inputs assert len(guardrailed_inputs["tool_calls"]) == 1 - assert ( - guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" - ) + assert guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" - assert ( - guardrailed_inputs["tool_calls"][0]["function"]["arguments"] - == '{"location":"São Paulo"}' - ) + assert guardrailed_inputs["tool_calls"][0]["function"]["arguments"] == '{"location":"São Paulo"}' # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") @@ -1136,14 +995,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): policies (e.g. PII on model output) then returned action=NONE for non-streaming completions that go through unified_guardrail -> process_output_response. """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1168,14 +1023,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): @pytest.mark.asyncio async def test_bedrock_apply_guardrail_request_uses_INPUT_source(): """input_type='request' must call Bedrock with source=INPUT and user messages.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1258,12 +1109,8 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): # Mock AWS-related methods with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -1431,9 +1278,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" def _create_guardrail(self) -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") @pytest.mark.asyncio async def test_should_handle_all_null_policy_sub_lists(self): @@ -1554,9 +1399,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: { "sensitiveInformationPolicy": { "piiEntities": None, - "regexes": [ - {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} - ], + "regexes": [{"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"}], }, } ], @@ -1611,18 +1454,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_none_texts_in_inputs(self): """inputs[\"texts\"] is explicitly None — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {"texts": None} # Explicit None mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1645,18 +1484,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_missing_texts_key(self): """inputs has no \"texts\" key at all — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {} # No "texts" key mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1677,9 +1512,7 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Test 1: ANONYMIZED action should NOT raise exception anonymized_response = { @@ -1700,9 +1533,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): ], } - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) + should_raise = guardrail._should_raise_guardrail_blocked_exception(anonymized_response) assert should_raise is False, "ANONYMIZED actions should not raise exceptions" # Test 2: BLOCKED action should raise exception @@ -1710,13 +1541,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "I can't provide that information."}], "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } + {"topicPolicy": {"topics": [{"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}]}} ], } @@ -1738,19 +1563,13 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): } ] }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"}]}, } ], } should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" + assert should_raise is True, "Mixed actions with any BLOCKED should raise exceptions" # Test 4: NONE action should not raise exception none_response = { @@ -1782,9 +1601,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): When logging_event_type is set, it must be forwarded to standard guardrail logging. When omitted, INPUT maps to pre_call (legacy). """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1795,13 +1612,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): mock_bedrock_response.json.return_value = { "action": "NONE", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ], } @@ -1811,12 +1622,8 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -1831,15 +1638,13 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): request_data=request_data, logging_event_type=GuardrailEventHooks.during_call, ) - assert ( - mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call - ) + assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call # Raw Bedrock JSON is forwarded; redaction runs once in # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. assert ( - mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] == "GG" ) @@ -1855,9 +1660,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): @pytest.mark.asyncio async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1873,15 +1676,9 @@ async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): prepared_request.headers = {} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), - patch.object( - guardrail, "_prepare_request", return_value=prepared_request - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=prepared_request) as mock_prepare_request, patch.object( guardrail, "get_guardrail_dynamic_request_body_params", @@ -1933,9 +1730,7 @@ async def test_during_call_hook_invokes_bedrock_async_moderation_hook(): "model": "gpt-4", "messages": [{"role": "user", "content": "test"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="test_key", user_id="test_user" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), call_type="completion", ) finally: @@ -1991,11 +1786,7 @@ def test_extract_blocked_assessments_multiple_policies(): "action": "GUARDRAIL_INTERVENED", "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Investment", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Investment", "type": "DENY", "action": "BLOCKED"}]}, "contentPolicy": { "filters": [ { @@ -2006,9 +1797,7 @@ def test_extract_blocked_assessments_multiple_policies(): } ] }, - "wordPolicy": { - "customWords": [{"match": "forbidden", "action": "BLOCKED"}] - }, + "wordPolicy": {"customWords": [{"match": "forbidden", "action": "BLOCKED"}]}, } ], } @@ -2023,13 +1812,7 @@ def test_extract_blocked_assessments_only_anonymized_returns_empty(): response = { "action": "GUARDRAIL_INTERVENED", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } assert g._extract_blocked_assessments(response) == [] @@ -2049,23 +1832,14 @@ def test_get_http_exception_includes_assessments_and_identifier(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "BLOCKED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "BLOCKED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) assert isinstance(exc, HTTPException) assert exc.status_code == 400 assert exc.detail["error"] == "Violated guardrail policy" - assert ( - exc.detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question." - ) + assert exc.detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question." assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r" assert exc.detail["guardrailVersion"] == "1" assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy" @@ -2088,15 +1862,11 @@ def test_extract_violation_category_names_mixed_policies(): {"name": "Tax Advice", "action": "BLOCKED"}, ] }, - "contentPolicy": { - "filters": [{"type": "VIOLENCE", "action": "BLOCKED"}] - }, + "contentPolicy": {"filters": [{"type": "VIOLENCE", "action": "BLOCKED"}]}, "wordPolicy": { "managedWordLists": [{"type": "PROFANITY", "action": "BLOCKED"}], }, - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}] - }, + "sensitiveInformationPolicy": {"piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}]}, } ], } @@ -2120,13 +1890,9 @@ def test_extract_violation_category_names_does_not_leak_user_input(): "assessments": [ { "wordPolicy": { - "customWords": [ - {"match": "secret-codeword-abc-123", "action": "BLOCKED"} - ], - }, - "sensitiveInformationPolicy": { - "regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}] + "customWords": [{"match": "secret-codeword-abc-123", "action": "BLOCKED"}], }, + "sensitiveInformationPolicy": {"regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}]}, } ], } @@ -2166,13 +1932,7 @@ def test_extract_violation_category_names_skips_anonymized(): g = _make_guardrail() response = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}] - } - } - ], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}]}}], } assert g._extract_violation_category_names(response) == [] @@ -2190,9 +1950,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): ``tracing_detail`` so downstream loggers (OTEL, ...) can surface the raw provider verdict as a queryable attribute without re-parsing the redacted guardrail_response blob.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2202,13 +1960,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}]}}], } request_data = { @@ -2217,12 +1969,8 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2253,9 +2001,7 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): """If the Bedrock response omits ``action`` (older / partial payloads), the field must be left off ``tracing_detail`` rather than written as ``None`` — downstream code expects strings or absence, not nulls.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2266,12 +2012,8 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): mock_bedrock_response.json.return_value = {"assessments": []} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2300,13 +2042,7 @@ def test_get_http_exception_no_blocked_assessments_omits_field(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "blocked"}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) @@ -2370,9 +2106,7 @@ async def test_streaming_post_call_only_runs_output_scan(): yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: out = [] async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), @@ -2382,18 +2116,11 @@ async def test_streaming_post_call_only_runs_output_scan(): out.append(chunk) assert len(out) >= 1 - output_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT" - ] + output_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT"] assert len(output_calls) == 1 assert output_calls[0].kwargs.get("request_data") is request_data - assert ( - output_calls[0].kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) - input_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT" - ] + assert output_calls[0].kwargs.get("logging_event_type") == GuardrailEventHooks.post_call + input_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT"] assert len(input_calls) == 0 @@ -2432,9 +2159,7 @@ async def test_streaming_post_call_output_only_path_passes_request_data_to_make_ yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: async for _ in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), response=mock_stream(), @@ -2488,9 +2213,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): ) minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=UserAPIKeyAuth(), @@ -2499,10 +2222,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): sources = [c.kwargs.get("source") for c in mock_make.call_args_list] assert sources == ["OUTPUT"] - assert ( - mock_make.call_args.kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) + assert mock_make.call_args.kwargs.get("logging_event_type") == GuardrailEventHooks.post_call # --------------------------------------------------------------------------- @@ -2522,9 +2242,7 @@ _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo." def _grounding_guardrail() -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") def _grounding_messages() -> list: @@ -2556,27 +2274,19 @@ def _model_response(content: str) -> ModelResponse: # Expected OUTPUT content blocks, keyed by their grounding qualifier, so the # per-test assertions read as the block sequence they expect. -_GROUNDING_SOURCE_BLOCK = { - "text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]} -} +_GROUNDING_SOURCE_BLOCK = {"text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]}} _QUERY_BLOCK = {"text": {"text": _GROUNDING_QUERY_TEXT, "qualifiers": ["query"]}} -_GUARD_BLOCK = { - "text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]} -} +_GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]}} def _input_request(messages: list) -> dict: """Arrange a guardrail and act: build the Bedrock INPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="INPUT", messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) def _output_request(messages: list, response=None) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="OUTPUT", response=response, messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2600,9 +2310,7 @@ def test_grounding_input_leaves_existing_guarded_text_unqualified(): """An existing guarded_text input block keeps its legacy unqualified payload.""" expected_request = {"source": "INPUT", "content": [{"text": {"text": "policy"}}]} - actual_request = _input_request( - [{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}] - ) + actual_request = _input_request([{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}]) assert actual_request == expected_request @@ -2615,9 +2323,7 @@ def test_grounding_output_assembles_source_query_and_response(): "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], } - actual_request = _output_request( - _grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(_grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2629,9 +2335,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): "content": [{"text": {"text": "Hi there."}}], } - actual_request = _output_request( - [{"role": "user", "content": "hello"}], _model_response("Hi there.") - ) + actual_request = _output_request([{"role": "user", "content": "hello"}], _model_response("Hi there.")) assert actual_request == expected_request @@ -2639,9 +2343,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_combines_multiple_sources(): """Every grounding_source block is emitted; Bedrock combines them into one corpus.""" uk_source_text = "London is the capital of UK." - uk_source_block = { - "text": {"text": uk_source_text, "qualifiers": ["grounding_source"]} - } + uk_source_block = {"text": {"text": uk_source_text, "qualifiers": ["grounding_source"]}} messages = [ { "role": "system", @@ -2662,9 +2364,7 @@ def test_grounding_output_combines_multiple_sources(): ], } - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2709,9 +2409,7 @@ def test_grounding_source_trusted_only_from_app_roles(role, is_trusted): if is_trusted: expected_content = [_GROUNDING_SOURCE_BLOCK, *expected_content] - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == {"source": "OUTPUT", "content": expected_content} @@ -2748,12 +2446,8 @@ async def test_grounding_output_blocked_raises_400(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -2791,13 +2485,7 @@ def _blocked_bedrock_httpx_response() -> MagicMock: response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}]}}], } return response @@ -2820,12 +2508,8 @@ async def test_make_bedrock_api_request_block_raises_modify_response_when_flag_s mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2857,12 +2541,8 @@ async def test_make_bedrock_api_request_block_raises_http_400_when_flag_unset(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2903,12 +2583,8 @@ async def test_async_pre_call_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2953,12 +2629,8 @@ async def test_async_moderation_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2998,12 +2670,8 @@ async def test_async_post_call_success_hook_attaches_original_response_on_block( mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -3032,9 +2700,7 @@ async def test_apply_guardrail_propagates_modify_response_on_block(): disable_exception_on_block=True, ) - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.side_effect = ModifyResponseException( message="Sorry, the model cannot answer this question.", model="bedrock-nova-micro", @@ -3276,6 +2942,1160 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n assert response is not None +def _too_large_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = { + "message": "Input is too long. Content size exceeds the maximum input size in text units." + } + response.text = json.dumps(response.json.return_value) + return response + + +def _other_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = {"message": "guardrailIdentifier is not valid"} + response.text = json.dumps(response.json.return_value) + return response + + +def _throttling_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 429 + response.json.return_value = {"message": "Rate exceeded"} + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _too_large_throttling_httpx_response() -> MagicMock: + """The shape AWS actually returns for an oversized ApplyGuardrail request when + the guardrail has an active content-filter policy: a 429 ThrottlingException, + not the documented 400 ValidationException. Message taken from a live call.""" + response = MagicMock() + response.status_code = 429 + response.json.return_value = { + "message": ( + "Input text size (3273 text units) exceeds the maximum allowed " + "(1000 text units) for the content filter policy (Classic tier)." + ) + } + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _passing_bedrock_httpx_response(marker: str) -> MagicMock: + """A successful ApplyGuardrail response tagged with `marker` so tests can + verify which chunk produced which output/usage after merging.""" + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "NONE", + "outputs": [{"text": marker}], + "assessments": [], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _blocking_bedrock_httpx_response(marker: str) -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": marker}], + "assessments": [{"topicPolicy": {"topics": [{"name": marker, "type": "DENY", "action": "BLOCKED"}]}}], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _bedrock_guardrail_for_chunk_tests() -> "BedrockGuardrail": + return BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunks_on_too_large_validation_error(): + """A too-large 400 on the whole-content call must trigger a bisect-and-retry, + and the two chunk responses must be merged (assessments concatenated, usage + summed, outputs concatenated) rather than losing either half's result.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first half of a very long message"}, + {"role": "user", "content": "second half of a very long message"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _blocking_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + detail = exc_info.value.detail + assert exc_info.value.status_code == 400 + assert "chunk-2" in detail["bedrock_guardrail_response"] + assert detail["assessments"][0]["matches"][0]["name"] == "chunk-2" + + +@pytest.mark.asyncio +async def test_apply_guardrail_merges_usage_and_outputs_across_chunks_when_both_pass(): + """When both chunks pass clean, the merged response must still carry both + chunks' outputs/usage forward (needed for accurate logging/telemetry) and + must not itself raise.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + assert result.get("usage", {}).get("contentPolicyUnits") == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_recurses_past_first_bisection_into_four_chunks(): + """A payload that is still too large after one bisection must keep splitting + -- chunking is not capped at two pieces. Four messages where both the + whole-content call AND both first-level halves are too large must recurse + one level deeper into four chunks that all fit, not give up after the + first split.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "message one"}, + {"role": "user", "content": "message two"}, + {"role": "user", "content": "message three"}, + {"role": "user", "content": "message four"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole content: [1,2,3,4] + _too_large_validation_httpx_response(), # first half: [1,2] + _passing_bedrock_httpx_response("message one"), + _passing_bedrock_httpx_response("message two"), + _too_large_validation_httpx_response(), # second half: [3,4] + _passing_bedrock_httpx_response("message three"), + _passing_bedrock_httpx_response("message four"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["message one", "message two", "message three", "message four"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_when_grounding_present(): + """Contextual-grounding requests are scored holistically against the whole + source; chunking them would silently produce misleading grounding scores. + A too-large error on a grounded request must propagate unchanged, not be + bisected.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_on_non_size_validation_error(): + """A 400 for an unrelated validation problem (e.g. a bad guardrail id) must + not trigger chunking -- retrying a bad-config error split into pieces would + just fail twice more and mask the real problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _other_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + assert "not valid" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_unsplittable_text_propagates_original_error(): + """A too-large error on content that has been bisected down to text too + short to split further (< 2 characters) must propagate the original error + rather than looping or crashing.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "a"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_single_item_splits_by_text_and_succeeds(): + """A too-large error on content that is already down to a single content + item must be bisected by that item's own text (not abandoned), so an + oversized single message can still be scanned successfully in halves.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "one giant single block of text"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + response = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 3 + assert response.get("action") == "NONE" + + +def _raised_bedrock_error(status_code: int, message: str) -> httpx.HTTPStatusError: + """A non-200 the way `AsyncHTTPHandler.post` actually surfaces it. + + That handler calls `response.raise_for_status()`, so in production a non-200 from + Bedrock arrives as a raised `httpx.HTTPStatusError` carrying the response, never + as a returned response object. Tests that return the response instead exercise a + branch real traffic never reaches. A real `httpx.Response` is used rather than a + MagicMock because the transport helper branches on + `isinstance(err_response, httpx.Response)`.""" + response = httpx.Response( + status_code=status_code, + json={"message": message}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail"), + ) + return httpx.HTTPStatusError(message, request=response.request, response=response) + + +_TOO_LARGE_MESSAGE = "Input is too long. Content size exceeds the maximum input size in text units." + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_once_when_client_raises_for_status(): + """The too-large attempt recovered by chunking must still produce exactly one + telemetry entry when the HTTP client raises for status, which is what really + happens: `AsyncHTTPHandler.post` calls `raise_for_status()`. + + Regression for per-attempt `guardrail_failed_to_respond` entries leaking out of + the transport helper on a request that ultimately succeeded, which made a + recovered request look like several failures plus a success.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["success"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_still_logs_once_when_client_raises(): + """Suppressing the transport helper's per-attempt logging must not swallow the only + record of a genuine failure: an unsplittable too-large request still has to produce + exactly one `guardrail_failed_to_respond` entry, not zero.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "x"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _post_side_effect(*_args, **_kwargs): + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["guardrail_failed_to_respond"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_single_item_split_twice_still_yields_one_output_per_item(): + """One oversized content item that needs two levels of text bisection ends up + as four text fragments, and all four must still collapse back into exactly + ONE output entry, because they all came from one original content item. + + Downstream masking (`_apply_masking_to_messages`) walks the merged outputs by + a running index across the original, unchunked message list, so emitting more + than one entry for a single message shifts every later message's masked text + onto the wrong message and drops the surplus. Regression for fragment + grouping assuming fragments only ever arrive as adjacent sibling *pairs*, + which holds for one bisection level but not for two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "aaaa bbbb cccc dddd eeee ffff gggg hhhh"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole single item + _too_large_validation_httpx_response(), # first half + _passing_bedrock_httpx_response("q1"), + _passing_bedrock_httpx_response("q2"), + _too_large_validation_httpx_response(), # second half + _passing_bedrock_httpx_response("q3"), + _passing_bedrock_httpx_response("q4"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["q1q2q3q4"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_retries_after_throttling_then_succeeds(): + """A chunk call throttled with a 429 must be retried with backoff and + eventually succeed, rather than surfacing the 429 to the caller.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _throttling_httpx_response() + if call_count == 3: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 4 + mock_sleep.assert_awaited() + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_exactly_once_as_success(): + """A too-large 400 that is recovered by chunking must not leave behind a + 'guardrail_failed_to_respond' telemetry entry for the initial oversized + attempt: the whole logical request (1 too-large attempt + 2 chunk + attempts here) must produce exactly one standard-logging entry, and it + must reflect the eventual success, not the transient too-large failure.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_logs_exactly_once_as_failed(): + """A too-large error that cannot be recovered (chunking disabled by + contextual grounding) must still log exactly once, as a failure -- not be + silently dropped by the chunking telemetry consolidation.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_merge_preserves_masking_position(): + """An earlier chunk that comes back clean (empty `outputs`) must not + shift a later chunk's masked text onto the wrong message. Regression for: + flattening outputs without positional metadata let a later chunk's PII + redaction get applied to the first message while the actual PII-bearing + message (in a later chunk) was forwarded unmasked.""" + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + mask_request_content=True, + ) + + request_data = { + "model": "bedrock-nova-micro", + "messages": [ + {"role": "user", "content": "clean chunk with nothing to mask"}, + {"role": "user", "content": "chunk with PII: John Doe"}, + ], + } + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + def _clean_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"action": "NONE", "assessments": []} + return response + + def _masked_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "chunk with PII: [NAME]"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "NAME", "match": "John Doe", "action": "ANONYMIZED"}] + } + } + ], + } + return response + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _clean_httpx_response() + return _masked_httpx_response() + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=request_data, + call_type="acompletion", + ) + + assert call_count == 3 + updated_messages = request_data["messages"] + assert updated_messages[0]["content"] == "clean chunk with nothing to mask" + assert updated_messages[1]["content"] == "chunk with PII: [NAME]" + + +@pytest.mark.asyncio +async def test_apply_guardrail_accepted_content_costs_exactly_one_call(): + """Content AWS accepts must cost exactly one ApplyGuardrail call, however far over + the chunk budget it is. Chunking is a recovery path, not something every request + pays for. Regression for: bin-packing eagerly on every request, which split + conversations AWS was happy to take whole and multiplied billed calls and guardrail + latency on traffic that never had a size problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [{"role": "user", "content": item_text} for _ in range(3)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + return _passing_bedrock_httpx_response(f"batch-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 1 + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_small_content_makes_exactly_one_call(): + """Content that fits entirely within the budget in a single batch must + make exactly one ApplyGuardrail call -- confirms bin-packing does not + introduce an extra probe call for the common (small-request) case.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "short message one"}, + {"role": "user", "content": "short message two"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _passing_bedrock_httpx_response("single-batch") + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_batch_under_budget_still_rejected_falls_back_to_bisection(): + """A batch that fits the budget guess but is still rejected by AWS as too large + (a lower real per-account/region/policy cap) must fall back to bisection for that + batch only, and any other batch from the same request that AWS already accepted + must not be re-sent. + + Three half-budget items pack into two batches once the whole-payload probe is + rejected, so the sequence is probe, batch one (rejected), its two halves, batch + two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [ + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count in (1, 2): + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 5 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-3", "chunk-4", "chunk-5"] + + +def test_split_bedrock_content_single_item_splits_on_whitespace_not_mid_word(): + """A single content item whose raw character midpoint would fall inside a + word must instead split at the nearest whitespace, so neither fragment + ends or begins mid-token. Regression for the Veria AI review finding: a + denied word/PII pattern straddling a raw character-midpoint cut could be + truncated on both fragments and scan clean on each, then reassemble into + the original unmasked text -- a detection bypass.""" + text = ("a" * 20) + " " + ("b" * 30) + raw_midpoint = len(text) // 2 + assert text[raw_midpoint] == "b" + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + + assert first_text + second_text == text + assert first_text == ("a" * 20) + " " + assert second_text == "b" * 30 + + +def test_split_bedrock_content_single_item_with_no_whitespace_falls_back_to_midpoint(): + """A single giant token with no whitespace anywhere has no safe split + point, so the split must fall back to the raw character midpoint rather + than failing or looping.""" + text = "a" * 40 + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + assert first_text + second_text == text + assert len(first_text) == 20 + assert len(second_text) == 20 + + +def test_bin_pack_bedrock_content_packs_minimal_batches_within_budget(): + """Many medium items should pack into the minimal number of in-order + batches that each stay within budget, not one batch per item.""" + items = [BedrockContentItem(text=BedrockTextContent(text="x" * 30)) for _ in range(10)] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert sum(len(batch) for batch in batches) == 10 + for batch in batches: + combined_len = sum(len(item["text"]["text"]) for item in batch) + assert combined_len <= 100 + assert len(batches) == 4 + + +def test_bin_pack_bedrock_content_oversized_single_item_becomes_its_own_batch(): + """An item whose own text already exceeds the budget must not be + pre-split here -- it becomes its own oversized batch, and only the + reactive bisection fallback (on an AWS rejection) may split it later.""" + small_item = BedrockContentItem(text=BedrockTextContent(text="short")) + oversized_item = BedrockContentItem(text=BedrockTextContent(text="x" * 200)) + items = [small_item, oversized_item, small_item] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert batches == ((small_item,), (oversized_item,), (small_item,)) + + +def test_bin_pack_bedrock_content_empty_content_makes_exactly_one_empty_batch(): + """Empty content must still pack into exactly one (empty) batch, matching + pre-bin-packing behavior of sending the content list as-is in one call -- + bin-packing must not turn an empty request into zero ApplyGuardrail calls.""" + assert BedrockGuardrail._bin_pack_bedrock_content([], budget=100) == ((),) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_reported_as_429_bisects_without_burning_retries(): + """AWS reports an oversized ApplyGuardrail request as a 429 ThrottlingException + (not the documented 400 ValidationException) when the guardrail has an active + content-filter policy. That is not a transient throttle -- re-posting the same + oversized content can never succeed -- so it must bisect immediately instead of + consuming the exponential-backoff retry budget first. + + Regression for a bug found against a live guardrail: because the throttle retry + only keyed off status 429, every oversized chunk burned all + _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES attempts (each a billed AWS call, + each preceded by a backoff sleep) before bisection got a chance, at every level + of the recursion.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_throttling_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_sleep.assert_not_awaited() + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["half-2", "half-3"] + + +def test_chunk_budget_defaults_to_apply_guardrail_per_second_quota(): + """The default budget must track ApplyGuardrail's default quota of 25 text units + (about 1,000 characters each) per second. Packing to that size and posting + sequentially is what stops chunking from trading a size error for a throttle, so + this default is a deliberate match to AWS behaviour rather than an arbitrary + number.""" + assert BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS == 25_000 + assert BedrockGuardrail(guardrailIdentifier="g", guardrailVersion="DRAFT").chunk_budget_chars == 25_000 + + +@pytest.mark.asyncio +async def test_configured_chunk_budget_changes_how_content_is_packed(): + """An account with raised quotas can set a larger `chunk_budget_chars` and have it + actually drive packing once AWS has rejected a payload, spending fewer + ApplyGuardrail calls for the same content instead of being pinned to the + conservative default. + + Four 20,000-character messages are 80,000 characters total, and every call here is + preceded by the one whole-payload probe AWS rejects. At the 25,000 default only one + message fits per batch, so it is the probe plus four; at 50,000 two fit per batch, + so it is the probe plus two.""" + messages = [{"role": "user", "content": "x" * 20_000} for _ in range(4)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _calls_made_with_budget(budget: int) -> int: + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + chunk_budget_chars=budget, + ) + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + posted = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal posted + posted += 1 + if posted == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response("ok") + + mock_post.side_effect = _post_side_effect + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + return mock_post.await_count + + assert await _calls_made_with_budget(BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS) == 5 + assert await _calls_made_with_budget(50_000) == 3 + + +def test_split_index_never_produces_an_empty_fragment(): + """Both fragments must be non-empty for every splittable text, so bisection always + makes progress. + + A text whose only qualifying whitespace is its final character is the dangerous + shape: taking that boundary puts the split at len(text), leaving the first fragment + identical to the input that was just rejected and the second empty. The recursion + would then resubmit the unchanged fragment forever and exhaust the stack instead of + scanning or surfacing Bedrock's error.""" + for text in ("ab ", "xxxx ", ("x" * 40) + " ", " ab", "a b", "ab", " "): + split_at = BedrockGuardrail._nearest_whitespace_split_index(text) + assert 0 < split_at < len(text), f"degenerate split {split_at} for {text!r}" + assert text[:split_at] and text[split_at:], f"empty fragment for {text!r}" + assert text[:split_at] + text[split_at:] == text + + +@pytest.mark.asyncio +async def test_oversized_single_item_with_trailing_space_gives_up_instead_of_recursing(): + """An oversized single item whose only space is trailing must bottom out and + surface Bedrock's error, not recurse forever. + + AWS is modelled the way it really behaves, rejecting every attempt, because the + danger is a fragment identical to the input that was just rejected: AWS would + reject it again, and each retry would split it into the same unchanged fragment. + A split that always shrinks the text terminates and re-raises; one that can return + the whole text raises RecursionError instead. The call-count bound is generous: + halving 41 characters down to unsplittable is a handful of attempts, nowhere near + a stack limit.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + messages = [{"role": "user", "content": ("x" * 40) + " "}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = lambda *_a, **_k: _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as excinfo: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert excinfo.value.status_code == 400 + assert mock_post.await_count < 200 + + class TestBedrockOnlyScanNewMessages: """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. @@ -3559,14 +4379,10 @@ class TestBedrockIncrementalFlagInteractions: session = {"litellm_session_id": "sess-flags-mask"} with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1 mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once" @pytest.mark.asyncio @@ -3586,9 +4402,7 @@ class TestBedrockIncrementalFlagInteractions: assert mock_api.call_count == 2, "incremental attempt + full-scan fallback" assert result["texts"] == ["MASKED q1"], "masked content must be applied" mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats" @pytest.mark.asyncio @@ -3647,9 +4461,7 @@ async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should } with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.return_value = MagicMock( - action="NONE", output=[], outputs=[], assessments=[] - ) + mock_api.return_value = MagicMock(action="NONE", output=[], outputs=[], assessments=[]) await guardrail.async_moderation_hook( data=data, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u"), @@ -3707,3 +4519,146 @@ class TestScanOnlyToolResultsWithLatestRoleFilter: assert result["texts"] == ["TOOL-RESULT"] warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) assert "scan_only_tool_results" in warning_text + + +@pytest.mark.parametrize("separator", ["\n", "\t", "\r\n", "　"]) +def test_split_bedrock_content_splits_on_any_whitespace_not_just_space(separator): + """Regression: the midpoint split must land on any Unicode whitespace, not only an + ASCII space. + + Matching only " " left the boundary unguarded for exactly the payloads that grow + large enough to need splitting: JSON lines, source code, logs and transcripts are + newline or tab delimited. A deny-listed word sitting at the midpoint of one was cut + in half, scanned clean on both fragments, and reassembled intact, which is the + single-token detection bypass the whitespace split exists to close.""" + text = separator.join(["aaaaaaa"] * 4) + separator + "BADWORDXYZ" + separator + separator.join(["bbbbbbb"] * 4) + + first, second = BedrockGuardrail._split_bedrock_content([BedrockContentItem(text=BedrockTextContent(text=text))]) + + first_text = first[0]["text"]["text"] + second_text = second[0]["text"]["text"] + assert first_text + second_text == text, "split must stay lossless" + assert "BADWORDXYZ" in first_text or "BADWORDXYZ" in second_text, "split severed the token" + + +def test_merge_bedrock_responses_preserves_fields_the_merge_has_no_opinion_on(): + """Regression: merging must not drop AWS response fields it does not itself merge. + + The merged response used to be rebuilt from an empty dict holding only action, + outputs, assessments and usage, so actionReason, guardrailCoverage and anything AWS + adds later vanished from the guardrail_json_response the Admin UI renders, on every + ApplyGuardrail request rather than only chunked ones.""" + chunk = BedrockContentChunkResult( + response={ + "action": "NONE", + "actionReason": "No action.", + "guardrailCoverage": {"textCharacters": {"guarded": 41, "total": 41}}, + "usage": {"contentPolicyUnits": 1}, + }, + content=[BedrockContentItem(text=BedrockTextContent(text="hello"))], + fragment_group_size=1, + ) + + merged = BedrockGuardrail._merge_bedrock_guardrail_responses([chunk]) + + assert merged["actionReason"] == "No action." + assert merged["guardrailCoverage"] == {"textCharacters": {"guarded": 41, "total": 41}} + + +def test_merge_bedrock_usage_sums_counters_not_on_the_known_list(): + """Regression: usage counters were summed from a hardcoded list of six keys, so the + ones AWS also returns (contentPolicyImageUnits, the automatedReasoning pair) were + reported as absent no matter what the chunks actually used.""" + chunks = [ + BedrockContentChunkResult( + response={"action": "NONE", "usage": {"contentPolicyImageUnits": units, "contentPolicyUnits": 1}}, + content=[BedrockContentItem(text=BedrockTextContent(text="x"))], + fragment_group_size=1, + ) + for units in (3, 4) + ] + + usage = BedrockGuardrail._merge_bedrock_guardrail_responses(chunks)["usage"] + + assert usage["contentPolicyImageUnits"] == 7 + assert usage["contentPolicyUnits"] == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_exception_inside_200_logs_failure_and_proceeds(): + """Regression: AWS can report a failure inside an HTTP 200 body via Output.__type, + and that must be logged as guardrail_failed_to_respond rather than success. + + Real AWS does this: an unrecognised operation path on bedrock-runtime answers + HTTP 200 with {"Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}}. + Consolidating telemetry had replaced the derived status with a hardcoded "success", + which reported a failed scan as a clean one. The request itself still proceeds, which + is the behaviour of the code before chunking existed.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + exception_response = MagicMock() + exception_response.status_code = 200 + exception_response.json.return_value = { + "Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}, + "Version": "1.0", + } + exception_response.text = json.dumps(exception_response.json.return_value) + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.return_value = exception_response + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + assert result is not None, "the request proceeds, as it did before chunking existed" + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_failure_logs_a_dict_not_a_bare_string(): + """Regression: the consolidated failure logger must log guardrail_json_response as a + dict, the shape the pre-chunking code and the InvokeGuardrailChecks path both use. + + Consolidating telemetry had changed it to a bare string on the ApplyGuardrail path + only, which breaks any consumer that reads it as a mapping and leaves the two paths + in this file inconsistent.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.side_effect = _raised_bedrock_error(400, "guardrailIdentifier is not valid") + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_log.assert_called_once() + logged = mock_log.call_args.kwargs["guardrail_json_response"] + assert isinstance(logged, dict), f"expected a dict, got {type(logged).__name__}" + assert "error" in logged diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 71e775842e3..8edb56ce25e 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 9c22575f8a5..ef32d0f7f4a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -292,6 +292,7 @@ class TestAutoRouterBenchmarks: ROW = _SessionAggRow( router_name="live-auto", router_type="complexity", + tier_turns={}, sessions=4, turns=40, unordered_turns=1, @@ -375,6 +376,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 @@ -426,6 +442,29 @@ class TestAutoRouterBenchmarks: 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 + class TestAutoRouterQualitySignals: """The endpoint's cohort split and its refusal to report a baseline it cannot stand behind. diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index eedffa1ea5f..1e47010b57c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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(): """ diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py new file mode 100644 index 00000000000..4a021627c3e --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -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", + ) diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 59268e1427b..5fb4392eec6 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -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"] diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/test_litellm/responses/test_responses_api_bridge_flag.py index 463af6562f1..f94c31831bf 100644 --- a/tests/test_litellm/responses/test_responses_api_bridge_flag.py +++ b/tests/test_litellm/responses/test_responses_api_bridge_flag.py @@ -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" diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8e9e32f5898..356556f3563 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -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" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 991f8eaa934..0a0cfe9a617 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -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 } } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 2bda0f72cec..e8331294972 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -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(); + + 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", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 95642b93019..42ff8bcc4c2 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -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 = ({ 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 = ({ accessToken, userID, userRole, premiumUser label={ Models{" "} - + } - rules={[ - { - required: true, - message: "Please select at least one model", - }, - ]} name="models" > = new Set([ "disable_global_guardrails", ]); +const TEAM_MODEL_BADGE_COLORS: Record = { + "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; guardrails?: string[]; policies?: string[]; @@ -483,7 +497,7 @@ const TeamInfoView: React.FC = ({ 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 = ({ Models
- {info.models.length === 0 || info.models.includes("all-proxy-models") ? ( - All proxy models - ) : ( - <> - {info.models.map((model: string, index: number) => ( - - {model} - - ))} - {(info.access_group_models || []).map((model: string, index: number) => ( - - {model} - - ))} - + {computeTeamModelBadges(info.models, info.access_group_models || [], info.access_group_details).map( + (badge, index) => ( + + + {badge.label} + + + ), )}
@@ -982,7 +989,7 @@ const TeamInfoView: React.FC = ({ { + 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" }, + ]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts new file mode 100644 index 00000000000..91ddf0f4045 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts @@ -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)}`, + }), + ), + ]; +} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0fcf206fb08..3a1b142059a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21420,6 +21420,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; };