diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 57ca267e504..26e4e06a796 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -105,7 +105,7 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38269 }, "reportUnknownParameterType": { "limit": 19584 diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql new file mode 100644 index 00000000000..e5d48abcb52 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql @@ -0,0 +1,15 @@ +-- DropForeignKey +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN + ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey"; + END IF; +END $$; + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN + ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) diff --git a/litellm/__init__.py b/litellm/__init__.py index ede8a73453d..2a85a3d4200 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -475,6 +475,7 @@ prometheus_metrics_config: Optional[List] = None prometheus_exclude_metrics: Optional[List[str]] = None prometheus_exclude_labels: Optional[List[str]] = None prometheus_emit_stream_label: bool = False +prometheus_emit_input_sequence_length_label: bool = False prometheus_deployment_and_latency_caller_identity: Literal[ "api_key_alias", "user_email", @@ -546,7 +547,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None #### PII MASKING #### output_parse_pii: bool = False ############################################# -from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete model_cost = get_model_cost_map(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount @@ -2405,3 +2406,5 @@ def __getattr__(name: str) -> Any: # ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time + +mark_litellm_import_complete() diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index f856fe0f2b3..8132008731f 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -6,9 +6,33 @@ be settable from user input. Context variables are scoped to the current asyncio task and cannot be injected via HTTP request bodies. """ +from collections.abc import Generator +from contextlib import contextmanager from contextvars import ContextVar +from datetime import datetime, timezone from typing import Final # When True, suppresses async logging and billing for internal sub-calls # (e.g., emulated file-search steps that make nested LLM calls). is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False) + +# One request prices its totals, its per-token-type lines and the rates it reports on +# separate code paths. Each reads the clock for off-peak pricing, so without a pinned +# moment they can land on either side of a window boundary and disagree with each other. +_billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", default=None) + + +@contextmanager +def pinned_billing_time(moment: datetime) -> Generator[None]: + """Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read.""" + token: Final = _billing_time.set(moment) + try: + yield + finally: + _billing_time.reset(token) + + +def current_billing_time() -> datetime: + """The pinned billing moment, or now in UTC outside a pinned block.""" + pinned: Final = _billing_time.get() + return pinned if pinned is not None else datetime.now(timezone.utc) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index a4e6fa50901..306a8871b12 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -13,6 +13,7 @@ from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( BedrockAgentCoreA2ATransformation, ) +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler: Returns: A2A JSON-RPC response dict from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await run_aws_signing( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, @@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler: Yields: A2A streaming response events from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await run_aws_signing( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 4fe069b0b7d..c3c18e3d009 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -25,7 +25,7 @@ import litellm from litellm import ModelResponse from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( - responses_reasoning_item_from_thinking_blocks, + responses_reasoning_items_from_thinking_blocks, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( @@ -129,8 +129,8 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: return stored raw_blocks: Final = msg.get("thinking_blocks") or () blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json - from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload + replayed: Final = responses_reasoning_items_from_thinking_blocks(blocks) + return [dict(item) for item in replayed] # mutable-ok: API message payload def _build_reasoning_item( @@ -227,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False): provider_specific_fields: Mapping[str, object] -def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: +def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw string payload in ``input`` rather than ``arguments``; both map to @@ -755,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Tool calls accumulate into the single trailing tool_calls choice # like the typed branches above; a choice per call would hide every # call after choices[0] from chat clients - accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index)) + accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index)) tool_call_index += 1 elif handle_raw_dict_callback is not None: choice, index = handle_raw_dict_callback(item=raw_item, index=index) @@ -1409,7 +1409,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") in ("function_call", "custom_tool_call"): - converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) + converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) provider_specific_fields: Final = converted.get("provider_specific_fields") function_chunk: Final = ChatCompletionToolCallFunctionChunk( @@ -1484,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): index=0, delta=Delta( tool_calls=( - _tool_call_dict_from_output_item( + tool_call_dict_from_output_item( output_item, parsed_chunk.get("output_index", 0) ), ) diff --git a/litellm/constants.py b/litellm/constants.py index 108a914e9c1..0a7ef363e92 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -335,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv( ########### v2 Architecture constants for managing writing updates to the database ########### REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer" +REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer" @@ -397,6 +398,18 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range( minimum=1, maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS, ) +TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_EXACT_CHARS", + default=4_000_000, + minimum=1, + maximum=1_000_000_000, +) +TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", + default=4, + minimum=1, + maximum=256, +) MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) @@ -569,6 +582,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float( LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100) LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000) LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0 +AWS_SIGNING_MAX_THREADS: Final = 16 DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) @@ -1668,6 +1682,7 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3"))) RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2")) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) +PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25"))) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index d4c6c87efc8..814eaaf76f7 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import TranscriptionUsageObjectTransformation, ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + BilledTokenRates, CostCalculatorUtils, _generic_cost_per_character, _get_regional_uplift_multiplier, @@ -1125,6 +1126,7 @@ def _store_cost_breakdown_in_logging_obj( service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + billed_token_rates: BilledTokenRates | None = None, ) -> None: """ Helper function to store cost breakdown in the logging object. @@ -1169,6 +1171,7 @@ def _store_cost_breakdown_in_logging_obj( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + billed_token_rates=billed_token_rates, ) except Exception as breakdown_error: @@ -1737,6 +1740,7 @@ def completion_cost( _reasoning_cost: float | None = None _cache_read_cost: float | None = None _cache_creation_cost: float | None = None + _billed_token_rates: BilledTokenRates | None = None if cost_per_token_usage_object is not None and model: _breakdown_provider: str | None = ( custom_llm_provider if isinstance(custom_llm_provider, str) else None @@ -1748,10 +1752,12 @@ def completion_cost( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + custom_cost_per_token=custom_cost_per_token, ) _reasoning_cost = _token_type_breakdown.reasoning_cost _cache_read_cost = _token_type_breakdown.cache_read_cost _cache_creation_cost = _token_type_breakdown.cache_creation_cost + _billed_token_rates = _token_type_breakdown.rates _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar, @@ -1771,6 +1777,7 @@ def completion_cost( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + billed_token_rates=_billed_token_rates, ) return _final_cost diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 3503468c735..421a3c84d75 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -18,6 +18,7 @@ from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServ from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client from mcp.shared.message import SessionMessage +from mcp.shared.session import RequestResponder from typing_extensions import Unpack _TransportStreams: TypeAlias = tuple[ @@ -53,15 +54,22 @@ def missing_streamable_http_client_error() -> ImportError: ) -from mcp.types import CallToolRequestParams as MCPCallToolRequestParams -from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import ( + METHOD_NOT_FOUND, + ClientResult, GetPromptRequestParams, GetPromptResult, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, Prompt, ResourceTemplate, + ServerNotification, + ServerRequest, TextContent, ) +from mcp.types import CallToolRequestParams as MCPCallToolRequestParams +from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import Tool as MCPTool from pydantic import AnyUrl @@ -146,8 +154,8 @@ _SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT) otherwise carries JSON-RPC error codes.""" -def _as_read_timeout(exc: BaseException) -> TimeoutError | None: - """The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``. +def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: + """Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``. The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error @@ -442,6 +450,18 @@ class MCPClient: in_flight_error: BaseException | None = None try: read_stream, write_stream = transport[0], transport[1] + stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future() + + async def receive_message( + message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, + ) -> None: + if not isinstance(message, (ValueError, httpx.RequestError, OSError)): + return + if not stream_error.done(): + stream_error.set_result(message) + # The SDK closes pending requests when its message handler raises. + raise RuntimeError("MCP response stream failed") + # Build session kwargs with optional callbacks session_kwargs: Final[dict[str, Any]] = {} if self._sampling_callback is not None: @@ -456,6 +476,7 @@ class MCPClient: read_stream, write_stream, read_timeout_seconds=timedelta(seconds=self.timeout), + message_handler=receive_message, **session_kwargs, ) session: Final = await session_ctx.__aenter__() @@ -467,6 +488,10 @@ class MCPClient: if isinstance(ins, str) and ins.strip(): self._last_initialize_instructions = ins.strip() return await operation(session) + except McpError: + if stream_error.done(): + raise stream_error.result() + raise finally: try: await session_ctx.__aexit__(None, None, None) @@ -501,11 +526,10 @@ class MCPClient: transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) except Exception as e: - read_timeout: Final = _as_read_timeout(e) + read_timeout: Final = as_mcp_read_timeout(e) if read_timeout is not None: verbose_logger.warning( - "MCP client timed out after %ss waiting for %s to answer; the server accepted the " - "request and ended its response stream without a JSON-RPC reply", + "MCP client timed out after %ss waiting for a valid MCP response from %s", self.timeout, self.server_url or "stdio", ) @@ -757,8 +781,19 @@ class MCPClient: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_prompts_operation(session: ClientSession): - return await session.list_prompts() + async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.prompts is None: + return ListPromptsResult(prompts=[]) + try: + return await session.list_prompts() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_prompts is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListPromptsResult(prompts=[]) try: result: Final = await self.run_with_session(_list_prompts_operation) @@ -834,8 +869,19 @@ class MCPClient: """List available resources from the server.""" verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") - async def _list_resources_operation(session: ClientSession): - return await session.list_resources() + async def _list_resources_operation(session: ClientSession) -> ListResourcesResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.resources is None: + return ListResourcesResult(resources=[]) + try: + return await session.list_resources() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_resources is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListResourcesResult(resources=[]) try: result: Final = await self.run_with_session(_list_resources_operation) @@ -870,8 +916,19 @@ class MCPClient: """List available resource templates from the server.""" verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") - async def _list_resource_templates_operation(session: ClientSession): - return await session.list_resource_templates() + async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult: + capabilities: Final = session.get_server_capabilities() + if capabilities is not None and capabilities.resources is None: + return ListResourceTemplatesResult(resourceTemplates=[]) + try: + return await session.list_resource_templates() + except McpError as error: + if error.error.code != METHOD_NOT_FOUND: + raise + verbose_logger.debug( + "MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error + ) + return ListResourceTemplatesResult(resourceTemplates=[]) try: result: Final = await self.run_with_session(_list_resource_templates_operation) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index eba6c862f7a..db5f790615f 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -21,6 +21,8 @@ from types import MappingProxyType from typing import Final, TypeVar from urllib.parse import urlparse +import httpx + from litellm._logging import verbose_logger from litellm.integrations.batch_utils import ( BatchSendCancelled, @@ -418,7 +420,7 @@ class AzureSentinelLogger(CustomBatchLogger): "Content-Type": "application/json", } - async def _send_batch(batch: Sequence[_QueuedPayload]): + async def _send_batch(batch: Sequence[_QueuedPayload]) -> httpx.Response: body: Final = safe_dumps(batch) return await self.async_httpx_client.post( url=api_endpoint, diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 37d6a7e793d..77bf4820a1a 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -850,20 +850,24 @@ class CustomGuardrail(CustomLogger): if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: return None - # CHECK IF GUARDRAIL REJECTS THE REQUEST target: Final = self._deployment_hook_target() - hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data - result: Final = await target.async_post_call_success_hook( - user_api_key_dict=UserAPIKeyAuth( - user_id=request_data.get("user_api_key_user_id"), - team_id=request_data.get("user_api_key_team_id"), - end_user_id=request_data.get("user_api_key_end_user_id"), - api_key=request_data.get("user_api_key_hash"), - request_route=request_data.get("user_api_key_request_route"), - ), - data=hook_request_data, - response=response, - ) + try: + if target is not self: + request_data["guardrail_to_apply"] = self # rebind-ok: dispatch consumes this key + result: Final = await target.async_post_call_success_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id=request_data.get("user_api_key_user_id"), + team_id=request_data.get("user_api_key_team_id"), + end_user_id=request_data.get("user_api_key_end_user_id"), + api_key=request_data.get("user_api_key_hash"), + request_route=request_data.get("user_api_key_request_route"), + ), + data=request_data, + response=response, + ) + finally: + if target is not self: + request_data.pop("guardrail_to_apply", None) if not self._is_valid_response_type(result): return None diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 9576eabaa34..b75369965de 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -2,6 +2,7 @@ # On success, logs events to Langfuse import inspect import os +import re import traceback from collections.abc import Callable, Iterable, Mapping from datetime import datetime @@ -63,6 +64,44 @@ def _object_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, dict) else None +def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, object]]: + """Header pairs with the key type widened back to what a caller-supplied dict can actually hold.""" + return mapping.items() + + +def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool: + if not isinstance(trace_id, str) or not isinstance(session_id, str): + return False + request: Final = _object_mapping(proxy_server_request) + raw_headers: Final = _object_mapping(request.get("headers")) if request is not None else None + if raw_headers is None: + return False + headers: Final = MappingProxyType( + {key.lower(): value for key, value in _widened_items(raw_headers) if isinstance(key, str)} + ) + if headers.get("x-litellm-trace-id"): + return False + if headers.get("langfuse_trace_id") is not None: + return False + if trace_id != session_id and headers.get("langfuse_session_id") != session_id: + return False + if headers.get("x-litellm-session-id") == trace_id: + return True + if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", trace_id) is None: + return False + user_agent: Final = headers.get("user-agent") + codex: Final = isinstance(user_agent, str) and re.match(r"^codex[-_ /]", user_agent, re.IGNORECASE) is not None + return any( + value == trace_id + and ( + key == "x-session-id" + or re.fullmatch(r"x-.+-session-id", key) is not None + or (codex and key in ("session-id", "session_id", "thread-id", "conversation_id")) + ) + for key, value in headers.items() + ) + + class _UsageObject(Protocol): """Token-count surface the Langfuse logger reads off a response usage payload.""" @@ -609,6 +648,18 @@ class LangFuseLogger: # This allows continuing an existing trace while still returning the correct trace_id if existing_trace_id is not None: trace_id = existing_trace_id + resolved_trace_id: Final = ( + litellm_call_id or trace_id + if existing_trace_id is None + and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request")) + else trace_id + ) + if resolved_trace_id != trace_id: + verbose_logger.debug( + "Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace", + trace_id, + resolved_trace_id, + ) requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) update_trace_keys: Final = ( requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else () @@ -663,7 +714,7 @@ class LangFuseLogger: trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm" else: # don't overwrite an existing trace trace_params = { - "id": trace_id, + "id": resolved_trace_id, "name": trace_name, "session_id": session_id, "input": masked_input if not mask_input else "redacted-by-litellm", @@ -845,13 +896,13 @@ class LangFuseLogger: # Verify langfuse accepted our trace_id; if it differs, log a warning but still return our intended value # to match expected test behavior if hasattr(generation_client, "trace_id") and generation_client.trace_id: - if generation_client.trace_id != trace_id: + if generation_client.trace_id != resolved_trace_id: verbose_logger.warning( "Langfuse trace_id mismatch: set %s, but langfuse returned %s. Using our intended trace_id for consistency.", - trace_id, + resolved_trace_id, generation_client.trace_id, ) - return trace_id, generation_id + return resolved_trace_id, generation_id except Exception: verbose_logger.error("Langfuse Layer Error - %s", traceback.format_exc()) return None, None diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6766d246894..540ce6738fc 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -246,6 +246,7 @@ class PrometheusLogger(CustomLogger): # logger so toggling these flags only takes effect after a # restart, keeping init-time and runtime label sets in sync. self._cached_metric_labels: dict[str, list[str]] = {} + self._emit_input_sequence_length_label = litellm.prometheus_emit_input_sequence_length_label is True _custom_buckets: Final = litellm.prometheus_latency_buckets self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS @@ -1522,6 +1523,11 @@ class PrometheusLogger(CustomLogger): # 2. Pyright does not allow us to run isinstance(standard_logging_payload, StandardLoggingPayload) <- this would be ideal enum_values=enum_values, label_context=label_context, + input_sequence_length=( + self._get_input_sequence_length(standard_logging_payload, kwargs, response_obj) + if self._emit_input_sequence_length_label + else None + ), ) # set x-ratelimit headers @@ -2192,6 +2198,36 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(remaining_tokens) + @staticmethod + def _get_input_sequence_length( + standard_logging_payload: StandardLoggingPayload, + kwargs: Mapping[str, object], + response_obj: object, + ) -> str: + prompt_tokens: Final = standard_logging_payload.get("prompt_tokens") + if prompt_tokens: + return get_input_sequence_length_bucket(prompt_tokens) + combined_usage: Final = kwargs.get("combined_usage_object") + if ( + combined_usage is not None + and getattr(kwargs.get("_litellm_upstream_reported_usage"), "total_tokens", None) is not None + ): + return get_input_sequence_length_bucket(None) + reported_usage: Final = ( + response_obj.get("usage") if isinstance(response_obj, dict) else getattr(response_obj, "usage", None) + ) + if reported_usage is None and combined_usage is None: + return get_input_sequence_length_bucket(None) + usage_metadata: Final = standard_logging_payload["metadata"].get("usage_object") + if isinstance(usage_metadata, Mapping): + return get_input_sequence_length_bucket(usage_metadata.get("prompt_tokens")) + if combined_usage is None and isinstance(response_obj, dict): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + normalized_usage: Final[Mapping[str, object]] = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj) + return get_input_sequence_length_bucket(normalized_usage.get("prompt_tokens")) + return get_input_sequence_length_bucket(prompt_tokens) + def _set_latency_metrics( self, kwargs: dict, @@ -2202,7 +2238,16 @@ class PrometheusLogger(CustomLogger): user_api_team_alias: str | None, enum_values: UserAPIKeyLabelValues, label_context: PrometheusLabelFactoryContext | None = None, + input_sequence_length: str | None = None, ): + latency_enum_values: Final = ( + replace(enum_values, input_sequence_length=input_sequence_length) + if input_sequence_length is not None + else enum_values + ) + latency_label_context: Final = ( + PrometheusLabelFactoryContext(latency_enum_values) if input_sequence_length is not None else label_context + ) # latency metrics end_time: Final[datetime] = kwargs.get("end_time") or datetime.now() start_time: Final[datetime | None] = kwargs.get("start_time") @@ -2220,8 +2265,8 @@ class PrometheusLogger(CustomLogger): supported_enum_labels=self.get_labels_for_metric( metric_name="litellm_llm_api_time_to_first_token_metric" ), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_llm_api_time_to_first_token_metric.labels(**_ttft_labels).observe(time_to_first_token_seconds) self._track_end_user_metric_series( @@ -2241,8 +2286,8 @@ class PrometheusLogger(CustomLogger): if api_call_total_time_seconds is not None: _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_llm_api_latency_metric"), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_llm_api_latency_metric.labels(**_labels).observe(api_call_total_time_seconds) self._track_end_user_metric_series( @@ -2272,8 +2317,8 @@ class PrometheusLogger(CustomLogger): ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_request_total_latency_metric"), - enum_values=enum_values, - label_context=label_context, + enum_values=latency_enum_values, + label_context=latency_label_context, ) self.litellm_request_total_latency_metric.labels(**_labels).observe(_observed_total_time_seconds) self._track_end_user_metric_series( diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 712ce41d09e..972ac79e306 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -10,9 +10,11 @@ import asyncio import time from collections.abc import Mapping from datetime import datetime -from typing import Final, cast +from typing import TYPE_CHECKING, Final, cast from urllib.parse import quote +import httpx + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS @@ -24,7 +26,7 @@ from litellm.integrations.s3 import ( from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -35,6 +37,9 @@ from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger +if TYPE_CHECKING: + from botocore.credentials import Credentials + class S3Logger(CustomBatchLogger, BaseAWSLLM): def __init__( @@ -232,6 +237,26 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): f"{get_aws_dns_suffix(self.s3_region_name)}/{encoded_key}" ) + def _sign_put( + self, credentials: "Credentials", url: str, json_string: str, headers: Mapping[str, str] + ) -> dict[str, str]: # mutable-ok: [LIT001] AsyncHTTPHandler.put/HTTPHandler.put only accept dict headers + """ + ``RefreshableCredentials`` (IMDS roles) may refresh between the access key, secret and token + reads SigV4 performs, producing a mixed-generation signature that S3 rejects with 403. + Freezing first makes the three values one atomic snapshot. + """ + from botocore.auth import S3SigV4Auth + from botocore.awsrequest import AWSRequest + from botocore.credentials import RefreshableCredentials + + frozen: Final = ( + credentials.get_frozen_credentials() if isinstance(credentials, RefreshableCredentials) else credentials + ) + aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=dict(headers)) + aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) + S3SigV4Auth(frozen, "s3", aws_region_name).add_auth(aws_request) + return dict(aws_request.headers.items()) + def _sse_headers(self) -> Mapping[str, str]: candidates: Final = { "x-amz-server-side-encryption": self.s3_server_side_encryption, @@ -317,26 +342,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): try: import base64 import hashlib - - from botocore.auth import S3SigV4Auth - from botocore.awsrequest import AWSRequest except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") try: from litellm.litellm_core_utils.asyncify import asyncify asyncified_get_credentials: Final = asyncify(self.get_credentials) - credentials: Final = await asyncified_get_credentials( - aws_access_key_id=self.s3_aws_access_key_id, - aws_secret_access_key=self.s3_aws_secret_access_key, - aws_session_token=self.s3_aws_session_token, - aws_region_name=self.s3_region_name, - aws_session_name=self.s3_aws_session_name, - aws_profile_name=self.s3_aws_profile_name, - aws_role_name=self.s3_aws_role_name, - aws_web_identity_token=self.s3_aws_web_identity_token, - aws_sts_endpoint=self.s3_aws_sts_endpoint, - ) verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key) verbose_logger.debug("s3_v2 logger - s3_verify setting: %s", self.s3_verify) @@ -363,19 +374,28 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): **self._sse_headers(), } - # Sign the request - aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) - aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) + async def signed_put() -> httpx.Response: + credentials: Final = await asyncified_get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + aws_session_name=self.s3_aws_session_name, + aws_profile_name=self.s3_aws_profile_name, + aws_role_name=self.s3_aws_role_name, + aws_web_identity_token=self.s3_aws_web_identity_token, + aws_sts_endpoint=self.s3_aws_sts_endpoint, + ) + signed_headers: Final = await run_aws_signing(self._sign_put, credentials, url, json_string, headers) + try: + return await self.async_httpx_client.put(url, data=json_string, headers=signed_headers) + except httpx.HTTPStatusError as error: + return error.response - # Prepare the signed headers - signed_headers: Final = dict(aws_request.headers.items()) - - # Make the request with retry for transient S3 errors (500/503) max_retries: Final = 3 for attempt in range(max_retries): - response = await self.async_httpx_client.put(url, data=json_string, headers=signed_headers) - if response.status_code in (500, 503) and attempt < max_retries - 1: + response = await signed_put() + if response.status_code in (403, 500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", @@ -479,20 +499,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): try: import base64 import hashlib - - from botocore.auth import S3SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") try: verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key) - credentials: Final[Credentials] = self.get_credentials( - aws_access_key_id=self.s3_aws_access_key_id, - aws_secret_access_key=self.s3_aws_secret_access_key, - aws_session_token=self.s3_aws_session_token, - aws_region_name=self.s3_region_name, - ) url: Final = self._build_object_url(batch_logging_element.s3_object_key) @@ -516,22 +526,24 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): **self._sse_headers(), } - # Sign the request - aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) - aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) - - # Prepare the signed headers - signed_headers: Final = dict(aws_request.headers.items()) - httpx_client: Final = _get_httpx_client( params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None) ) - # Make the request with retry for transient S3 errors (500/503) + + def signed_put() -> httpx.Response: + credentials: Final = self.get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + ) + signed_headers: Final = self._sign_put(credentials, url, json_string, headers) + return httpx_client.put(url, data=json_string, headers=signed_headers) + max_retries: Final = 3 for attempt in range(max_retries): - response = httpx_client.put(url, data=json_string, headers=signed_headers) - if response.status_code in (500, 503) and attempt < max_retries - 1: + response = signed_put() + if response.status_code in (403, 500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s", @@ -597,7 +609,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="GET", url=url, headers=headers) - S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request) + await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 2b4c8c9928d..d787375ca3c 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -22,7 +22,7 @@ from litellm.constants import ( SQS_SEND_MESSAGE_ACTION, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request) + await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index cdc4810ff04..91a22144805 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -2,6 +2,7 @@ Pulls the cost + context window + provider route for known models from https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json This can be disabled by setting the LITELLM_LOCAL_MODEL_COST_MAP environment variable to True. +The ``lite`` and ``litellm-proxy`` CLI entry points also use the bundled map without fetching. ``` export LITELLM_LOCAL_MODEL_COST_MAP=True @@ -13,11 +14,14 @@ import hashlib import json import os import random +import sys +import threading import time from collections.abc import Awaitable, Callable from dataclasses import dataclass, replace from datetime import datetime, timezone from importlib.resources import files +from pathlib import Path from typing import Final, Protocol import httpx @@ -33,6 +37,12 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" +_CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) + + +def _is_cli_process() -> bool: + return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES + # Reserved top-level keys that are not model entries. They must be excluded # from the model-count integrity check so a real upstream shrink can't be masked. @@ -176,6 +186,11 @@ class GetModelCostMap: RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3 MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0 +_litellm_import_complete = threading.Event() + + +def mark_litellm_import_complete() -> None: + _litellm_import_complete.set() @dataclass(frozen=True, slots=True) @@ -314,12 +329,13 @@ async def _fetch_remote_model_cost_map_with_retry( def _fetch_remote_model_cost_map_with_retry_sync( url: str, timeout: int, - max_attempts: int, + attempts: range, sleep: Callable[[float], None], rng: random.Random, client: _SyncGetClient, ) -> ModelCostMapReloadResult: - for attempt in range(1, max_attempts + 1): + max_attempts: Final = attempts.stop - 1 + for attempt in attempts: outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout) if not isinstance(outcome, _FetchAttemptRetryable): return outcome @@ -520,6 +536,68 @@ def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMa return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) +def adopt_model_cost_map( + new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract +) -> int: + import litellm + from litellm import utils + + litellm.model_cost = new_model_cost_map + utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + litellm.add_known_models(model_cost_map=new_model_cost_map) + fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 + utils.reapply_runtime_model_cost_registrations() + return fetched_model_count + + +def _retry_remote_fetch_in_background( + url: str, + timeout: int, + max_attempts: int, + sleep: Callable[[float], None], + rng: random.Random, + client: _SyncGetClient, + first_outcome: _FetchAttemptRetryable, +) -> None: + try: + first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng) + if isinstance(first_wait, ModelCostMapReloadUnavailable): + return + sleep(first_wait) + result: Final = _fetch_remote_model_cost_map_with_retry_sync( + url=url, + timeout=timeout, + attempts=range(2, max_attempts + 1), + sleep=sleep, + rng=rng, + client=client, + ) + if isinstance(result, ModelCostMapReloadUnavailable): + verbose_logger.warning( + "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", + url, + max_attempts, + ) + return + _litellm_import_complete.wait() + if not GetModelCostMap.validate_model_cost_map( + fetched_map=result.model_cost_map, + backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) + return + finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map + _cost_map_source_info.source = "remote" + _cost_map_source_info.fallback_reason = None + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + adopt_model_cost_map(finalized) + except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays + verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) + + def get_model_cost_map( url: str, timeout: int = 5, @@ -531,10 +609,12 @@ def get_model_cost_map( """ Public entry point — returns the model cost map dict. - 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. + 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set or this is a ``lite`` / + ``litellm-proxy`` CLI process, uses the local backup only. 2. Otherwise fetches from ``url``, retrying transient HTTP errors - (429/5xx/transport) with Retry-After-aware backoff, validates - integrity, and falls back to the local backup on any failure. + (429/5xx/transport) with Retry-After-aware backoff in a background + thread, validates integrity, and falls back to the local backup on any + failure. Only the backup model count is cached (a single int) for validation. The full backup dict is only parsed when it must be *returned* as a @@ -543,7 +623,7 @@ def get_model_cost_map( _cost_map_source_info.loaded_at = datetime.now(timezone.utc) # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. - if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True @@ -553,24 +633,34 @@ def get_model_cost_map( _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False - result: Final = _fetch_remote_model_cost_map_with_retry_sync( - url=url, - timeout=timeout, - max_attempts=max_attempts, - sleep=sleep, - rng=rng if rng is not None else random.Random(), - client=client if client is not None else httpx, - ) - if isinstance(result, ModelCostMapReloadUnavailable): + fetch_client: Final = client if client is not None else httpx + fetch_rng: Final = rng if rng is not None else random.Random() + outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) + if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1: + threading.Thread( + target=_retry_remote_fetch_in_background, + kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping + "url": url, + "timeout": timeout, + "max_attempts": max_attempts, + "sleep": sleep, + "rng": fetch_rng, + "client": fetch_client, + "first_outcome": outcome, + }, + name="litellm-model-cost-map-retry", + daemon=True, + ).start() + if not isinstance(outcome, ModelCostMapReloaded): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, - result.reason, + outcome.reason, ) _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}" return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map - content: Final = result.model_cost_map + content: Final = outcome.model_cost_map # Validate using cached count (cheap int comparison, no file I/O) if not GetModelCostMap.validate_model_cost_map( @@ -587,4 +677,4 @@ def get_model_cost_map( _cost_map_source_info.source = "remote" _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(result).model_cost_map + return _finalize_loaded_model_cost_map(outcome).model_cost_map diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b0d6db20b31..de45414dbf3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -120,7 +120,6 @@ from litellm.types.utils import ( CachingDetails, CallTypes, CostBreakdown, - CostResponseTypes, CustomPricingLiteLLMParams, DynamicPromptManagementParamLiteral, EmbeddingResponse, @@ -203,7 +202,8 @@ if TYPE_CHECKING: from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config - from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, LoggedRelayResponse try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -590,6 +590,7 @@ class Logging(LiteLLMLoggingBaseClass): # Initialize cost breakdown field self.cost_breakdown: CostBreakdown | None = None + self.billed_token_rates: BilledTokenRates | None = None # Init Caching related details self.caching_details: CachingDetails | None = None @@ -1587,6 +1588,7 @@ class Logging(LiteLLMLoggingBaseClass): service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + billed_token_rates: "BilledTokenRates | None" = None, ) -> None: """ Helper method to store cost breakdown in the logging object. @@ -1606,8 +1608,10 @@ class Logging(LiteLLMLoggingBaseClass): service_tier: Tier the costs above were priced on, already resolved data_residency: Region uplift the costs above were priced on, already resolved vertex_location: Vertex AI location the costs above were priced on, already resolved + billed_token_rates: Per-token rates the costs above were billed at, already resolved """ + self.billed_token_rates = billed_token_rates self.cost_breakdown = CostBreakdown( input_cost=input_cost, output_cost=output_cost, @@ -2376,7 +2380,7 @@ class Logging(LiteLLMLoggingBaseClass): self, raw_bytes: list[bytes], provider_config: "BasePassthroughConfig", - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: all_chunks: Final = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) complete_streaming_response: Final = provider_config.handle_logging_collected_chunks( all_chunks=all_chunks, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index c05d4c29a5e..e5977ca4156 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -10,6 +10,7 @@ from typing import Any, Final, Literal, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import litellm +from litellm._internal_context import current_billing_time from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import ( select_tier_for_input, @@ -19,6 +20,7 @@ from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, CompletionTokensDetailsWrapper, + CostPerToken, DataResidency, ImageResponse, ModelInfo, @@ -305,7 +307,7 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_ than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(), or every window shifts by the host's offset. """ - reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference: Final = current_time if current_time is not None else current_billing_time() now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time() windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc for window in windows: @@ -392,7 +394,7 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose hours apply only on its weekdays. """ - reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference: Final = current_time if current_time is not None else current_billing_time() reference_utc: Final = ( reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc) ) @@ -1195,7 +1197,7 @@ def generic_cost_per_token( usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0 ) - billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) + billing_time: Final = current_time if current_time is not None else current_billing_time() ( prompt_base_cost, completion_base_cost, @@ -1309,42 +1311,90 @@ def _coerce_token_count(value: object) -> int: return value if isinstance(value, int) and value > 0 else 0 +@dataclass(frozen=True, slots=True) +class BilledTokenRates: + """Per-token rates one request's usage bills at, after token tiers, off-peak windows and the + regional multipliers the totals apply, so each cost line equals its token count times its rate.""" + + input_cost_per_token: float + output_cost_per_token: float + cache_read_input_token_cost: float + cache_creation_input_token_cost: float + cache_creation_input_token_cost_above_1hr: float + output_cost_per_reasoning_token: float + + def scaled(self, multiplier: float) -> "BilledTokenRates": + if multiplier == 1.0: + return self + return BilledTokenRates( + input_cost_per_token=self.input_cost_per_token * multiplier, + output_cost_per_token=self.output_cost_per_token * multiplier, + cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier, + cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier, + cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier, + output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier, + ) + + @dataclass(frozen=True, slots=True) class TokenTypeCostBreakdown: reasoning_cost: float cache_read_cost: float cache_creation_cost: float + rates: BilledTokenRates | None = None + """Rates these lines were billed at, so a caller reporting both cannot resolve them a second, + differently-argued way. None when the model's pricing could not be resolved.""" -def get_token_type_cost_breakdown( - model: str, - custom_llm_provider: str | None, +def _reasoning_token_count(usage: Usage) -> int: + parsed: Final = ( + parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 + ) + return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) + + +def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]: + """(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details + first, then the private top-level counters the Usage constructor mirrors cache tokens onto for + providers/callers that bypass the details.""" + parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None + parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0 + parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0 + return ( + parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)), + parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)), + parsed["cache_creation_token_details"] if parsed is not None else None, + ) + + +def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRates: + """Flat custom pricing has no tiers, uplifts or reasoning rate: cache tokens bill at the configured + cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does.""" + input_rate: Final = custom_cost_per_token["input_cost_per_token"] + output_rate: Final = custom_cost_per_token["output_cost_per_token"] + cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate) + return BilledTokenRates( + input_cost_per_token=input_rate, + output_cost_per_token=output_rate, + cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate), + cache_creation_input_token_cost=cache_creation_rate, + cache_creation_input_token_cost_above_1hr=cache_creation_rate, + output_cost_per_reasoning_token=output_rate, + ) + + +def _cost_map_billed_rates( + model_info: ModelInfo, usage: Usage, - service_tier: str | None = None, - data_residency: str | None = None, - vertex_location: str | None = None, - current_time: datetime | None = None, -) -> TokenTypeCostBreakdown: - """ - Provider-agnostic cost of reasoning and cache tokens, derived from the usage - object and model pricing alone. - - This works for every provider, including Perplexity/Cerebras/Dashscope whose - cost calculators bypass ``generic_cost_per_token``, because cache tokens always - land on ``prompt_tokens_details`` (via the Usage constructor and provider - transformations) and reasoning tokens on ``completion_tokens_details``. It reuses - the same rate-resolution primitives as the total-cost path so the breakdown can - never drift from the totals. Returns zeros (never raises) when the model or its - pricing cannot be resolved. - """ - try: - model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - return TokenTypeCostBreakdown(0.0, 0.0, 0.0) - - billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) + custom_llm_provider: str | None, + service_tier: str | None, + data_residency: str | None, + vertex_location: str | None, + current_time: datetime | None, +) -> BilledTokenRates: + billing_time: Final = current_time if current_time is not None else current_billing_time() ( - _prompt_base_cost, + prompt_base_cost, completion_base_cost, cache_creation_cost_rate, cache_creation_cost_above_1hr_rate, @@ -1356,13 +1406,6 @@ def get_token_type_cost_breakdown( current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) - - reasoning_tokens = ( - parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 - ) - if not reasoning_tokens: - reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) - reasoning_rate: Final = _resolve_billed_reasoning_rate( model_info=model_info, usage=usage, @@ -1370,57 +1413,103 @@ def get_token_type_cost_breakdown( completion_base_cost=completion_base_cost, current_time=billing_time, ) - reasoning_cost = float(reasoning_tokens) * reasoning_rate + multiplier: Final = ( + _get_regional_uplift_multiplier(model_info, data_residency) + * get_vertex_regional_endpoint_uplift(model_info, vertex_location) + * get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) + ) + return BilledTokenRates( + input_cost_per_token=prompt_base_cost, + output_cost_per_token=completion_base_cost, + cache_read_input_token_cost=cache_read_cost_rate, + cache_creation_input_token_cost=cache_creation_cost_rate, + cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate, + output_cost_per_reasoning_token=reasoning_rate, + ).scaled(multiplier) - cache_read_tokens = 0 - cache_creation_tokens = 0 - cache_creation_token_details: CacheCreationTokenDetails | None = None - if usage.prompt_tokens_details is not None: - prompt_tokens_details: Final = parse_prompt_tokens_details(usage) - cache_read_tokens = prompt_tokens_details["cache_hit_tokens"] - cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"] - cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"] - # Fall back to the private top-level counters the Usage constructor mirrors cache - # tokens onto, so providers/callers that bypass prompt_tokens_details are covered. - if not cache_read_tokens: - cache_read_tokens = _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)) - if not cache_creation_tokens: - cache_creation_tokens = _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)) - cache_read_cost = float(cache_read_tokens) * cache_read_cost_rate - cache_creation_cost = calculate_cache_writing_cost( - cache_creation_tokens=cache_creation_tokens, - cache_creation_token_details=cache_creation_token_details, - cache_creation_cost_above_1hr=cache_creation_cost_above_1hr_rate, - cache_creation_cost=cache_creation_cost_rate, +def get_billed_token_rates( + model: str, + custom_llm_provider: str | None, + usage: Usage, + service_tier: str | None = None, + data_residency: str | None = None, + vertex_location: str | None = None, + current_time: datetime | None = None, + custom_cost_per_token: CostPerToken | None = None, +) -> BilledTokenRates | None: + """Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type + breakdown resolve them. None when the model's pricing cannot be resolved.""" + if custom_cost_per_token is not None: + return _custom_pricing_rates(custom_cost_per_token) + try: + model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates + return None + return _cost_map_billed_rates( + model_info=model_info, + usage=usage, + custom_llm_provider=custom_llm_provider, + service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, + current_time=current_time, ) - # Apply the same flat regional-processing uplift the totals get, so per-type - # costs stay reconciled with input_cost/output_cost for regionalized OpenAI hosts. - uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) - if uplift != 1.0: - reasoning_cost *= uplift - cache_read_cost *= uplift - cache_creation_cost *= uplift - vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) - if vertex_uplift != 1.0: - reasoning_cost *= vertex_uplift - cache_read_cost *= vertex_uplift - cache_creation_cost *= vertex_uplift +def get_token_type_cost_breakdown( + model: str, + custom_llm_provider: str | None, + usage: Usage, + service_tier: str | None = None, + data_residency: str | None = None, + vertex_location: str | None = None, + current_time: datetime | None = None, + custom_cost_per_token: CostPerToken | None = None, +) -> TokenTypeCostBreakdown: + """ + Provider-agnostic cost of reasoning and cache tokens, derived from the usage + object and model pricing alone. - # Mirror the provider-specific geo uplift (e.g. Anthropic us: 1.1) the totals - # apply, so cache and reasoning line items stay reconciled with them. - geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) - if geo_multiplier != 1.0: - reasoning_cost *= geo_multiplier - cache_read_cost *= geo_multiplier - cache_creation_cost *= geo_multiplier + This works for every provider, including Perplexity/Cerebras/Dashscope whose + cost calculators bypass ``generic_cost_per_token``, because cache tokens always + land on ``prompt_tokens_details`` (via the Usage constructor and provider + transformations) and reasoning tokens on ``completion_tokens_details``. It reuses + the same rate resolution as the total-cost path (``get_billed_token_rates``) so the + breakdown can never drift from the totals. A deployment billed by + ``custom_cost_per_token`` is priced from those flat rates instead of the cost map and, + like its totals, bills cache writes flat rather than by their 5m/1h split. + Returns zeros (never raises) when the model or its pricing cannot be resolved. + """ + rates: Final = get_billed_token_rates( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + service_tier=service_tier, + data_residency=data_residency, + vertex_location=vertex_location, + current_time=current_time, + custom_cost_per_token=custom_cost_per_token, + ) + if rates is None: + return TokenTypeCostBreakdown(0.0, 0.0, 0.0) + cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage) + cache_creation_cost: Final = ( + float(cache_creation_tokens) * rates.cache_creation_input_token_cost + if custom_cost_per_token is not None + else calculate_cache_writing_cost( + cache_creation_tokens=cache_creation_tokens, + cache_creation_token_details=cache_creation_token_details, + cache_creation_cost_above_1hr=rates.cache_creation_input_token_cost_above_1hr, + cache_creation_cost=rates.cache_creation_input_token_cost, + ) + ) return TokenTypeCostBreakdown( - reasoning_cost=reasoning_cost, - cache_read_cost=cache_read_cost, + reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token, + cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost, cache_creation_cost=cache_creation_cost, + rates=rates, ) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index a6f10e1ede3..87524d86c61 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,7 @@ import json import re import time import traceback -from collections.abc import Iterable, Sequence +from collections.abc import Mapping, Sequence from typing import Final, Literal, cast import litellm @@ -151,6 +151,16 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: del choice.enhancements +def _invalid_choices_message(response_object: Mapping[str, object]) -> str: + raw_keys: Final = list(response_object.keys()) + if "choices" not in response_object: + return f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {raw_keys}" + return ( + f"LiteLLM: provider returned 'choices' that is not a list ({type(response_object['choices']).__name__}). " + f"Raw keys: {raw_keys}" + ) + + async def convert_to_streaming_response_async( response_object: dict | None = None, ): @@ -179,14 +189,12 @@ async def convert_to_streaming_response_async( choice_list: Final[list[StreamingChoices]] = [] - if not response_object.get("choices"): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( status_code=500, - message=( - f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) @@ -287,14 +295,12 @@ def convert_to_streaming_response( model_response_object: Final = ModelResponseStream() choice_list: Final[list[StreamingChoices]] = [] - if not response_object.get("choices"): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( status_code=500, - message=( - f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) @@ -623,15 +629,12 @@ def convert_to_model_response_object( return convert_to_streaming_response(response_object=response_object) choice_list: Final[list[Choices]] = [] - if not response_object.get("choices") or not isinstance(response_object["choices"], Iterable): + if not isinstance(response_object.get("choices"), list): from litellm.exceptions import APIError raise APIError( status_code=500, - message=( - "LiteLLM: provider returned a response with no 'choices'. " - f"Raw keys: {list(response_object.keys())}" - ), + message=_invalid_choices_message(response_object), llm_provider="", model="", ) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9b612993a69..31523c9309d 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -441,7 +441,15 @@ class LoggingCallbackManager: return result + def get_callback_objects(self) -> tuple[tuple[str, CustomLogger | Callable], ...]: + return tuple( + (self._get_callback_string(callback), callback) + for callback in self._get_all_callbacks() + if not isinstance(callback, str) + ) + def _get_callback_string(self, callback: CustomLogger | Callable | str) -> str: + from litellm.integrations.opentelemetry import OpenTelemetry from litellm.litellm_core_utils.custom_logger_registry import ( CustomLoggerRegistry, ) @@ -449,6 +457,8 @@ class LoggingCallbackManager: """Convert a callback to its string representation""" if isinstance(callback, str): return callback + elif isinstance(callback, OpenTelemetry) and callback.callback_name is not None: + return callback.callback_name elif isinstance(callback, CustomLogger): # Try to get the string representation from the registry callback_str: Final = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback)) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 00b80839dde..2485896184e 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,8 +6,8 @@ import io import json import mimetypes import re -from collections.abc import Iterable, Mapping, Sequence -from itertools import groupby +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from itertools import groupby, islice from os import PathLike from pathlib import Path from types import MappingProxyType @@ -1320,17 +1320,128 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo -def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]: +_SUBSCHEMA_KEYWORDS: Final = frozenset( + { + "additionalItems", + "additionalProperties", + "contains", + "else", + "if", + "items", + "not", + "propertyNames", + "then", + "unevaluatedItems", + "unevaluatedProperties", + } +) +_SUBSCHEMA_LIST_KEYWORDS: Final = frozenset({"allOf", "anyOf", "items", "oneOf", "prefixItems"}) +_SUBSCHEMA_MAP_KEYWORDS: Final = frozenset( + {"$defs", "definitions", "dependentSchemas", "patternProperties", "properties"} +) + +_MAX_SCHEMA_NESTING: Final = 1024 + + +def drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: + """Drop every regex in a schema position that Python's ``re`` cannot compile. + + OpenAI validates tool ``parameters`` against the 2020-12 metaschema with + ``jsonschema``'s format checker, which hands each ``pattern`` value and each + ``patternProperties`` key to ``re.compile``, so a regex written for an + ECMA-262 engine (Unicode property escapes such as ``\\p{Cc}``, as in Claude + Code's ``Artifact`` tool) is refused with "'...' is not a 'regex'" by every + model family on both the chat and Responses wires. Only schema positions are + walked (properties, items, combinators, ``$defs`` and the other applicators), + so a ``pattern`` key inside ``default``, ``examples``, ``const`` or vendor + extensions is data and stays. Outside strict mode the keyword is only a + hint, so dropping it costs the model a constraint and the caller nothing. + Compilable regexes and everything else pass through, the input is never + mutated, and the same object comes back when nothing was dropped. The walk + is level-order rather than recursive, rebuilt deepest level first, and stops + at more schema levels than a JSON parser admits, so a cyclic schema built in + code cannot spin it. + """ + rebuilt: dict[int, Mapping[str, object]] = {} # mutable-ok: per-call memo of rewritten nodes, deepest level first + for level in reversed(tuple(islice(_schema_levels(schema), _MAX_SCHEMA_NESTING))): + rebuilt.update( + (id(node), rewritten) + for node in level + if (rewritten := _node_without_non_python_regex(node, rebuilt)) is not node + ) + return rebuilt.get(id(schema), schema) + + +def _schema_levels(schema: Mapping[str, object]) -> Iterator[tuple[Mapping[str, object], ...]]: + frontier: tuple[Mapping[str, object], ...] = (schema,) # rebind-ok: level-order cursor, one level a round + while frontier: + yield frontier + frontier = tuple(child for node in frontier for child in _subschemas(node)) + + +def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]: + for key, value in node.items(): + if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): + yield from (sub for sub in value.values() if isinstance(sub, dict)) + elif key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list): + yield from (sub for sub in value if isinstance(sub, dict)) + elif key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict): + yield value + + +def _node_without_non_python_regex( + node: Mapping[str, object], rebuilt: Mapping[int, Mapping[str, object]] +) -> Mapping[str, object]: + kept: Final = { # mutable-ok: tool parameters are JSON dicts + key: _keyword_value_rebuilt(key, value, rebuilt) + for key, value in node.items() + if key != "pattern" or not isinstance(value, str) or _is_python_regex(value) + } + return node if len(kept) == len(node) and all(kept[key] is node[key] for key in kept) else kept + + +def _keyword_value_rebuilt(key: str, value: object, rebuilt: Mapping[int, Mapping[str, object]]) -> object: + if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): + kept: Final = { # mutable-ok: tool parameters are JSON dicts + name: rebuilt.get(id(sub), sub) + for name, sub in value.items() + if key != "patternProperties" or not isinstance(name, str) or _is_python_regex(name) + } + return value if len(kept) == len(value) and all(kept[name] is value[name] for name in kept) else kept + if key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list): + items: Final = [rebuilt.get(id(sub), sub) for sub in value] # mutable-ok: tool parameters are JSON lists + return value if all(new is old for new, old in zip(items, value, strict=True)) else items + if key in _SUBSCHEMA_KEYWORDS and isinstance(value, dict): + return rebuilt.get(id(value), value) + return value + + +def _is_python_regex(pattern: str) -> bool: + try: + re.compile(pattern) + except (re.error, RecursionError): + return False + return True + + +def flatten_combinators_and_drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: + return flatten_top_level_schema_combinators(drop_non_python_regex_patterns(schema)) + + +def tool_with_sanitized_parameters( + tool: Mapping[str, object], + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], +) -> Mapping[str, object]: function: Final = tool.get("function") if not isinstance(function, dict): return tool parameters: Final = function.get("parameters") if not isinstance(parameters, dict): return tool - flattened: Final = flatten_top_level_schema_combinators(parameters) - if flattened is parameters: + sanitized: Final = sanitize(parameters) + if sanitized is parameters: return tool - return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts + return {**tool, "function": {**function, "parameters": sanitized}} # mutable-ok: request tools are JSON dicts def _get_image_mime_type_from_url(url: str) -> str | None: @@ -1823,14 +1934,11 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]: return None, message_content -def _readable_thinking_text( - block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, -) -> str: +def _readable_thinking_text(block: Mapping[str, object]) -> str: """The text a chat model can read back, empty for redacted blocks and malformed ones.""" if block.get("type") != "thinking": return "" - thinking: Final = cast(ChatCompletionThinkingBlock, block).get("thinking") # cast-ok: narrowed by the type tag - return str(thinking or "") + return str(block.get("thinking") or "") def reasoning_content_from_thinking_blocks( @@ -1843,24 +1951,125 @@ def reasoning_content_from_thinking_blocks( return "\n".join(text for block in thinking_blocks if (text := _readable_thinking_text(block))) -def responses_reasoning_item_from_thinking_blocks( - thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], -) -> ChatCompletionReasoningItem | None: - """Build a Responses API `reasoning` input item from Anthropic thinking blocks. +ENCRYPTED_REASONING_SIGNATURE_PREFIX: Final = "litellm_encrypted_reasoning:" - The item carries no `id`: the Responses API rejects an empty one and 404s on any id it - did not mint itself, while an item without an id is always accepted. + +def encrypted_reasoning_signature(encrypted_content: str) -> str: + """The opaque value a Responses API reasoning item's `encrypted_content` travels in. + + Anthropic clients echo a thinking block's `signature` and a redacted block's `data` + back verbatim, so either field can carry the encrypted reasoning across turns; the + prefix tells the two apart from a signature Anthropic minted. """ + return f"{ENCRYPTED_REASONING_SIGNATURE_PREFIX}{encrypted_content}" + + +def _carries_encrypted_reasoning(signature: object) -> bool: + return isinstance(signature, str) and signature.startswith(ENCRYPTED_REASONING_SIGNATURE_PREFIX) + + +def encrypted_content_from_signature(signature: object) -> str | None: + if not isinstance(signature, str) or not _carries_encrypted_reasoning(signature): + return None + return signature.removeprefix(ENCRYPTED_REASONING_SIGNATURE_PREFIX) or None + + +def _encrypted_reasoning_field(block: Mapping[str, object]) -> object: + match block.get("type"): + case "thinking": + return block.get("signature") + case "redacted_thinking": + return block.get("data") + case _: + return None + + +def encrypted_content_of_block(block: Mapping[str, object]) -> str | None: + return encrypted_content_from_signature(_encrypted_reasoning_field(block)) + + +def is_encrypted_reasoning_block(block: object) -> bool: + """A thinking or redacted_thinking block carrying Responses API encrypted reasoning. + + Only the Responses API that minted the content can read it back, so an Anthropic + backend has to drop such a block rather than fail signature verification on it. + """ + if not isinstance(block, Mapping): + return False + mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance + return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping)) + + +def strip_encrypted_reasoning_from_messages(messages: object) -> None: + """Drop the bridge-tagged reasoning blocks a routed deployment cannot decrypt from + Anthropic-shaped history. + + The whole block goes, the way #40280 drops undecryptable Responses ``input`` items: a + provider that did not mint the block rejects it signed (a foreign signature) and unsigned + (a missing signature) alike, so keeping its text as an unsigned thinking block only moves + the 400 from the router to the provider. + + Mutates the content lists in place: the router's fallback snapshot shares these + message objects, so a rebound list would replay the stripped blocks on the fallback hop. + """ + if not isinstance(messages, list): + return + for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json + _strip_encrypted_reasoning_from_blocks(content) + + +def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: + return ( + cast(list[object], content) # cast-ok: narrowed by isinstance + for message in messages + if isinstance(message, Mapping) + for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance + if isinstance(content, list) + ) + + +def _strip_encrypted_reasoning_from_blocks(content: object) -> None: + blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance + kept: Final = tuple(block for block in blocks if not is_encrypted_reasoning_block(block)) + blocks[:] = kept # rebind-ok: shared with fallback snapshot + + +def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: + index, block = indexed_block + return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary" + + +def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) -> ChatCompletionReasoningItem | None: summary: Final[list[ChatCompletionReasoningSummaryTextBlock]] = [ # mutable-ok: API message payload ChatCompletionReasoningSummaryTextBlock(type="summary_text", text=text) - for block in thinking_blocks + for block in group if (text := _readable_thinking_text(block)) ] + encrypted_content: Final = encrypted_content_of_block(group[0]) + if encrypted_content is not None: + return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content) if not summary: return None return ChatCompletionReasoningItem(type="reasoning", summary=summary) +def responses_reasoning_items_from_thinking_blocks( + thinking_blocks: Iterable[Mapping[str, object]], +) -> tuple[ChatCompletionReasoningItem, ...]: + """Build Responses API `reasoning` input items from Anthropic thinking blocks. + + A block carrying encrypted reasoning replays the item it came from byte for byte; + a run of plain thinking blocks collapses into one summary-only item. No item carries + an `id`: the Responses API 404s on any id it did not mint itself and rejects an empty + one, while an item without an id is always accepted. + """ + return tuple( + item + for _, group in groupby(enumerate(thinking_blocks), key=_reasoning_replay_group_key) + if (item := _reasoning_item_from_block_group(tuple(block for _, block in group))) is not None + ) + + def _parse_content_for_reasoning( message_text: str | None, ) -> tuple[str | None, str | None]: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 56c1d605700..ece619e3883 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -46,6 +46,7 @@ from litellm.types.utils import GenericImageParsingChunk from .common_utils import ( convert_content_list_to_str, infer_content_type_from_url_and_content, + is_encrypted_reasoning_block, is_non_content_values_set, parse_tool_call_arguments, ) @@ -2299,13 +2300,16 @@ def sanitize_messages_for_tool_calling( def _is_unsignable_thinking_block(block: object) -> bool: - """A `thinking` block that Anthropic cannot accept on input. + """A thinking block that Anthropic cannot accept on input. Anthropic verifies the thinking signature cryptographically, so a block whose signature is null, empty, or missing (e.g. from an open-source reasoning model) - is rejected with a 400 and must be dropped rather than blanked or repaired. - `redacted_thinking` blocks carry no signature and are always kept. + is rejected with a 400 and must be dropped rather than blanked or repaired, and + so is a block whose signature or data carries another provider's encrypted + reasoning. A `redacted_thinking` block Anthropic minted is always kept. """ + if is_encrypted_reasoning_block(block): + return True if not isinstance(block, dict) or block.get("type") != "thinking": return False signature: Final = block.get("signature") diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index cf9604a0fd5..90698296142 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,6 +1,6 @@ import base64 import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast @@ -210,7 +210,7 @@ def apply_grounding_request_counts( class ChunkProcessor: - def __init__(self, chunks: list, messages: list | None = None): + def __init__(self, chunks: list, messages: Sequence | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages self.first_chunk = chunks[0] @@ -1004,8 +1004,9 @@ class ChunkProcessor: chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, - messages: list | None = None, + messages: Sequence | None = None, reasoning_tokens: int | None = None, + count_prompt_tokens: Callable[[], int] | None = None, ) -> Usage: """ Calculate usage for the given chunks. @@ -1030,7 +1031,9 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages) + returned_usage.prompt_tokens = prompt_tokens or ( + count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages) + ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ae17bac6ff..db23929e0c3 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1473,17 +1473,14 @@ class CustomStreamWrapper: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider == "cached_response": cached_chunk: Final = cast(ModelResponseStream, chunk) - chunk_finish_reason: Final = cached_chunk.choices[0].finish_reason + cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None + chunk_finish_reason: Final = cached_choice.finish_reason if cached_choice is not None else None response_obj = { - "text": cached_chunk.choices[0].delta.content, + "text": cached_choice.delta.content if cached_choice is not None else None, "is_finished": chunk_finish_reason is not None, "finish_reason": chunk_finish_reason, "original_chunk": cached_chunk, - "tool_calls": ( - cached_chunk.choices[0].delta.tool_calls - if hasattr(cached_chunk.choices[0].delta, "tool_calls") - else None - ), + "tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None), } completion_obj["content"] = response_obj["text"] diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3732ffd734c..1de2533d514 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,11 +3,15 @@ import base64 import io import struct -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast +import anyio +import anyio.lowlevel import httpx import tiktoken +from tokenizers import Tokenizer +from typing_extensions import ParamSpec, TypeVar import litellm from litellm import verbose_logger @@ -21,7 +25,10 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, + TOKEN_COUNTER_MAX_CONCURRENT_COUNTS, + TOKEN_COUNTER_MAX_EXACT_CHARS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -172,6 +179,13 @@ def calculate_tiles_needed( return total_tiles +def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int: + largest_tile_count: Final = calculate_tiles_needed( + MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES + ) + return base_tokens + (base_tokens * 2) * largest_tile_count + + def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: return struct.unpack(fmt, buffer) @@ -317,6 +331,32 @@ TokenCounterFunction = Callable[[str], int] Type for a function that counts tokens in a string. """ +EXTRAPOLATION_SAMPLES: Final = 16 +T_ParamSpec: Final = ParamSpec("T_ParamSpec") +T_Retval = TypeVar("T_Retval") +_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") + + +def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter: + existing: Final = _COUNT_OFFLOAD_LIMITER.get(None) + if existing is not None: + return existing + created: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + _COUNT_OFFLOAD_LIMITER.set(created) + return created + + +def offload_token_count( + function: Callable[T_ParamSpec, T_Retval], +) -> Callable[T_ParamSpec, Awaitable[T_Retval]]: + async def offloaded( + *args: T_ParamSpec.args, + **kwargs: T_ParamSpec.kwargs, # kwargs-ok: ParamSpec keeps the wrapped function's own keyword contract + ) -> T_Retval: + return await asyncify(function, limiter=_count_offload_limiter_for_this_loop())(*args, **kwargs) + + return offloaded + def _get_tiktoken_count_function( encode_length: Callable[[str], int], @@ -538,9 +578,40 @@ def _count_extra( return num_tokens +def _get_extrapolating_count_function( + count_exactly: TokenCounterFunction, + max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS, +) -> TokenCounterFunction: + def count_tokens(text: str) -> int: + if len(text) <= max_exact_chars: + return count_exactly(text) + samples: Final = _evenly_spaced_samples(text, max_exact_chars) + sampled_chars: Final = sum(len(sample) for sample in samples) + return round(sum(count_exactly(sample) for sample in samples) * len(text) / sampled_chars) + + return count_tokens + + +def _evenly_spaced_samples(text: str, total_chars: int) -> tuple[str, ...]: + sample_count: Final = min(EXTRAPOLATION_SAMPLES, total_chars) + sample_chars: Final = total_chars // sample_count + last_start: Final = len(text) - sample_chars + return tuple( + text[start : start + sample_chars] + for start in (last_start * index // max(sample_count - 1, 1) for index in range(sample_count)) + ) + + def _get_count_function( model: str | None, custom_tokenizer: dict | SelectTokenizerResponse | None = None, +) -> TokenCounterFunction: + return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer)) + + +def _get_exact_count_function( + model: str | None, + custom_tokenizer: dict | SelectTokenizerResponse | None = None, ) -> TokenCounterFunction: """ Get the function to count tokens based on the model and custom tokenizer.""" @@ -549,10 +620,10 @@ def _get_count_function( if model is not None or custom_tokenizer is not None: tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model) if tokenizer_json["type"] == "huggingface_tokenizer": + tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - enc: Final = tokenizer_json["tokenizer"].encode(text) - return len(enc.ids) + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e486be12fe2..9d50345d70d 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,10 +13,11 @@ Pattern Overview: """ import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Mapping, MutableSequence, Sequence from copy import deepcopy from dataclasses import dataclass from itertools import chain, repeat +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable from typing_extensions import ReadOnly, TypedDict, assert_never @@ -41,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_field, stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -153,6 +155,46 @@ class ExtractedInput: EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) +@dataclass(frozen=True, slots=True) +class _ToolCallShape: + name: str | None + arguments: str + + +@dataclass(frozen=True, slots=True) +class _SSEFieldRewrite: + """One field of one nested section of a buffered SSE event, rewritten.""" + + section: str + field: str + value: object + + +class _SSEEventRewriter(Protocol): + def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ... + + +def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]: + rewrite: Final = rewrite_event(event) + section: Final = None if rewrite is None else event.get(rewrite.section) + if rewrite is None or not isinstance(section, Mapping): + return event + return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} # mutable-ok: json.dumps needs a dict + + +def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: + """The guardrail-visible shape of each tool call, whether the guardrail handed + back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts.""" + functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls) + return tuple( + _ToolCallShape( + name=name if isinstance(name := stream_item_field(function, "name"), str) else None, + arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "", + ) + for function in functions + ) + + class _AnthropicSSEDelta(TypedDict, total=False): type: ReadOnly[str] text: ReadOnly[str] @@ -170,12 +212,18 @@ class AnthropicMessagesHandler(BaseTranslation): them through guardrail rewrites; downstream provider handling is out of scope. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True + assembles_streamed_response = True def __init__(self): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() + def post_call_hook_response(self, response: object) -> object: + if not isinstance(response, ModelResponse): + return response + return self.adapter.translate_openai_response_to_anthropic(response) + @staticmethod def _build_streaming_usage_response( responses_so_far: Sequence[object], @@ -1050,6 +1098,7 @@ class AnthropicMessagesHandler(BaseTranslation): first_choice.message.tool_calls, ) string_so_far = first_choice.message.content + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ()) guardrail_inputs: Final = GenericGuardrailAPIInputs() if string_so_far: guardrail_inputs["texts"] = [string_so_far] @@ -1084,6 +1133,19 @@ class AnthropicMessagesHandler(BaseTranslation): and guardrailed_texts[0] != string_so_far ): self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0]) + if deliver_ended_stream_rewrites: + returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls") + self._write_ended_stream_tool_call_rewrites( + responses_so_far, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=_tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) + and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tool_calls_list or () + ), + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) else: verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far @@ -1206,44 +1268,124 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _write_ended_stream_text_rewrite( - responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place rewritten_text: str, ) -> None: """Deliver an ended-stream guardrail text rewrite by rewriting the buffered chunks in place: the first ``text_delta`` carries the full rewritten text and every later one is blanked, leaving the surrounding - message and content-block framing untouched. Handles both chunk formats - this stream carries (parsed event dicts and raw SSE bytes).""" + message and content-block framing untouched.""" replacements: Final = chain((rewritten_text,), repeat("")) - for idx, item in enumerate(responses_so_far): - if isinstance(item, dict): - delta = item.get("delta") - if item.get("type") == "content_block_delta" and isinstance(delta, dict): - if delta.get("type") == "text_delta": - delta["text"] = next(replacements) - elif isinstance(item, (bytes, bytearray)): - responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer - AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements) - ) + + def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = event.get("delta") + if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping): + return None + if delta.get("type") != "text_delta": + return None + return _SSEFieldRewrite("delta", "text", next(replacements)) + + AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) + + @classmethod + def _write_ended_stream_tool_call_rewrites( + cls, + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + *, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Deliver ended-stream guardrail tool-call rewrites by rewriting the + buffered chunks in place: the rebuilt response lists tool calls in the + order of the stream's ``tool_use`` blocks, so the nth rewritten call lands + on the nth block, its first ``input_json_delta`` carrying the full rewritten + arguments, every later one blanked, and ``content_block_start`` carrying the + rewritten name. Blocks that do not line up with the rebuilt tool calls make + the rewrite undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + block_indices: Final = tuple( + index + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + and isinstance(index := event.get("index"), int) + ) + if len(block_indices) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + rewrites_by_block: Final = MappingProxyType( + { + index: after + for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + } + ) + argument_replacements: Final = MappingProxyType( + {index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()} + ) + + def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + index: Final = event.get("index") + if not isinstance(index, int) or index not in rewrites_by_block: + return None + match event.get("type"): + case "content_block_start": + name: Final = rewrites_by_block[index].name + if name is None: + return None + return _SSEFieldRewrite("content_block", "name", name) + case "content_block_delta": + delta: Final = event.get("delta") + if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta": + return None + return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index])) + case _: + return None + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use) @staticmethod - def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes: - """Rewrite every ``text_delta`` data line in one SSE chunk with the next - replacement text, leaving all other events and framing byte-identical.""" + def _rewrite_ended_stream_events( + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + rewrite_event: _SSEEventRewriter, + ) -> None: + """Replace every buffered event ``rewrite_event`` returns a rewrite for, in + both chunk formats this stream carries (parsed event dicts and raw SSE + bytes), leaving every other event and the framing untouched.""" + rewritten_items: Final = tuple( + AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far + ) + responses_so_far[:] = rewritten_items # rebind-ok: delivers the rewrites into the caller's buffer + + @staticmethod + def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object: + if isinstance(item, dict): + return _rewritten_event(_as_str_mapping(item), rewrite_event) + if isinstance(item, (bytes, bytearray)): + return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event) + return item + + @staticmethod + def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes: + """Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites, + leaving all other events and framing byte-identical.""" try: decoded: Final = sse_bytes.decode("utf-8") except UnicodeDecodeError: return sse_bytes return "\n\n".join( - AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n") + "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n")) + for block in decoded.split("\n\n") ).encode("utf-8") @staticmethod - def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str: - return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n")) - - @staticmethod - def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str: + def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str: if not line.startswith("data:"): return line try: @@ -1252,14 +1394,10 @@ class AnthropicMessagesHandler(BaseTranslation): ) except json.JSONDecodeError: return line - if not isinstance(data, dict) or data.get("type") != "content_block_delta": + if not isinstance(data, dict): return line - delta: Final = data.get("delta") - if not isinstance(delta, dict) or delta.get("type") != "text_delta": - return line - return "data: " + json.dumps( - {**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts - ) + rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event) + return line if rewritten is data else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 2b57883cc13..87c4ec8938e 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -21,6 +21,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, + is_encrypted_reasoning_block, ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, @@ -72,8 +73,13 @@ _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") + + def is_claude_code_user_agent(user_agent: str) -> bool: - return user_agent.startswith("claude-cli/") + """Claude Code sends its API calls through the Anthropic SDK as `claude-cli/` and its own + fetches, such as gateway model discovery, as `claude-code/`""" + return user_agent.startswith(_CLAUDE_CODE_USER_AGENT_PREFIXES) def _validated_claude_code_mapping(value: object) -> dict[object, object] | None: @@ -1201,6 +1207,32 @@ def strip_thinking_blocks_from_anthropic_messages(messages: list[Any]) -> list[A return out +def _without_encrypted_reasoning_blocks(message: dict) -> dict | None: # mutable-ok: Anthropic message payload shape + if not isinstance(message, Mapping): + return message + content: Final = message.get("content") + if not isinstance(content, list): + return message + kept: Final = [b for b in content if not is_encrypted_reasoning_block(b)] # mutable-ok: API message payload + if len(kept) == len(content): + return message + if not kept: + return None + return {**message, "content": kept} # mutable-ok: API message payload + + +def strip_encrypted_reasoning_blocks_from_anthropic_messages( + messages: Sequence[dict], # mutable-ok: Anthropic message payload shape +) -> list[dict]: # mutable-ok: AnthropicMessagesRequest.messages is typed list[dict] + """ + Drop thinking / redacted_thinking blocks that carry another provider's encrypted + reasoning (a turn the Responses API bridge served) before the request reaches + Anthropic, which cannot verify them. Anthropic's own signed blocks are kept. + """ + stripped: Final = (_without_encrypted_reasoning_blocks(m) for m in messages) + return [m for m in stripped if m is not None] # mutable-ok: API message payload + + def strip_thinking_blocks_from_anthropic_messages_request_dict( data: dict[str, Any], ) -> None: @@ -1629,11 +1661,16 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: def _anthropic_model_entry( - model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str] + model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str], listed_ids: Mapping[str, str] ) -> Mapping[str, object]: + listed_id: Final = listed_ids.get(model["id"]) + source: Final[Mapping[str, object]] = ( + MappingProxyType({"source_model": model["id"]}) if listed_id is not None else MappingProxyType({}) + ) return { # mutable-ok: JSON response body, serialized by the route and never mutated "type": "model", - "id": model["id"], + "id": listed_id or model["id"], + **source, "display_name": display_names.get(model["id"], model["id"]), "created_at": created_at, "max_input_tokens": model.get("max_input_tokens"), @@ -1644,6 +1681,7 @@ def _anthropic_model_entry( def create_anthropic_model_list_response( models: Sequence[ModelInfoResponse], display_names: Mapping[str, str] = MappingProxyType({}), + listed_ids: Mapping[str, str] = MappingProxyType({}), ) -> Mapping[str, object]: """Build the Anthropic-native /v1/models envelope. @@ -1653,17 +1691,19 @@ def create_anthropic_model_list_response( over from the OpenAI-shaped listing, named as the Messages API names them, and are always present because the vendor shape declares them nullable, not optional. display_names maps a listed model id to a configured human-readable name; ids - without an entry fall back to the id itself, matching the vendor behavior + without an entry fall back to the id itself, matching the vendor behavior. + listed_ids maps a model id to the id the caller should see it under (the Claude + Code view); ids without an entry are listed as they are """ created_at: Final = ( datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z") ) data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated - _anthropic_model_entry(model, created_at, display_names) for model in models + _anthropic_model_entry(model, created_at, display_names, listed_ids) for model in models ] return { # mutable-ok: JSON response body, serialized by the route and never mutated "data": data, "has_more": False, - "first_id": models[0]["id"] if models else None, - "last_id": models[-1]["id"] if models else None, + "first_id": data[0]["id"] if data else None, + "last_id": data[-1]["id"] if data else None, } diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index db890662132..9b3a57cc422 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -113,6 +113,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( from litellm.llms.anthropic.common_utils import ( is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, + strip_encrypted_reasoning_blocks_from_anthropic_messages, ) from litellm.llms.anthropic.experimental_pass_through.context_management import ( PolyfillResult, @@ -417,7 +418,8 @@ class LiteLLMAnthropicMessagesAdapter: model: str | None = None, ) -> list: new_messages: Final[list[AllMessageValues]] = [] - for m in messages: + replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) + for m in replayable_messages: user_message: ChatCompletionUserMessage | None = None tool_message_list: list[ChatCompletionToolMessage] = [] new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = [] @@ -1487,8 +1489,9 @@ class LiteLLMAnthropicMessagesAdapter: anthropic_content.insert(0, polyfill_result.compaction_block) ## extract finish reason + openai_finish_reason: Final = response.choices[0].finish_reason if response.choices else "stop" translated_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic( - openai_finish_reason=response.choices[0].finish_reason + openai_finish_reason=openai_finish_reason ) anthropic_finish_reason: Final = ( "refusal" diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 8267da157ad..9f9346fad4d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -25,6 +25,7 @@ from ...common_utils import ( AnthropicModelInfo, optionally_handle_anthropic_oauth, strip_advisor_blocks_from_messages, + strip_encrypted_reasoning_blocks_from_anthropic_messages, ) DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -613,7 +614,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): messages = strip_advisor_blocks_from_messages(messages) anthropic_messages_request: Final[AnthropicMessagesRequest] = AnthropicMessagesRequest( - messages=messages, + messages=strip_encrypted_reasoning_blocks_from_anthropic_messages(messages), max_tokens=max_tokens, model=model, **anthropic_messages_optional_request_params, diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 0445c23ed8c..7731c883d9f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -19,6 +19,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.utils import ProviderConfigManager from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper @@ -34,6 +35,15 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, return extra_kwargs or {} +def _provider_returns_encrypted_reasoning(model: str, custom_llm_provider: object) -> bool: + provider: Final = ( + custom_llm_provider if isinstance(custom_llm_provider, str) else litellm.get_llm_provider(model=model)[1] + ) + provider_model: Final = local_model_name(model, provider) + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config(provider, provider_model) + return responses_config is not None and "include" in responses_config.get_supported_openai_params(provider_model) + + def _build_responses_kwargs( *, max_tokens: int, @@ -85,8 +95,13 @@ def _build_responses_kwargs( request_data["output_format"] = output_format anthropic_request: Final = AnthropicMessagesRequest(**request_data) - responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request) forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) + responses_kwargs: Final = _ADAPTER.translate_request( + anthropic_request, + include_encrypted_reasoning=_provider_returns_encrypted_reasoning( + model, forwarded_kwargs.get("custom_llm_provider") + ), + ) # Normalize reasoning effort based on model capabilities # (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported) @@ -111,7 +126,7 @@ def _build_responses_kwargs( responses_kwargs["stream"] = True # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) - excluded: Final = {"anthropic_messages"} + excluded: Final = frozenset(("anthropic_messages",)) for key, value in forwarded_kwargs.items(): if key == "litellm_logging_obj" and value is not None: from litellm.litellm_core_utils.litellm_logging import ( @@ -132,6 +147,14 @@ def _build_responses_kwargs( if explicit_prompt_cache_key is not None: responses_kwargs["prompt_cache_key"] = explicit_prompt_cache_key + deployment_include: Final = forwarded_kwargs.get("include") + bridge_include: Final = responses_kwargs.get("include") + if isinstance(deployment_include, list) and isinstance(bridge_include, list): + responses_kwargs["include"] = [ + *bridge_include, + *(item for item in deployment_include if item not in bridge_include), + ] + return responses_kwargs diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 2e0a6a9df8f..f753e87fee3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -9,13 +9,19 @@ from typing import TYPE_CHECKING, Any, Final from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( refusal_stop_details, responses_output_refusal_text, ) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage -from .transformation import LiteLLMAnthropicToResponsesAPIAdapter +from .transformation import ( + REASONING_SUMMARY_PART_SEPARATOR, + LiteLLMAnthropicToResponsesAPIAdapter, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject @@ -29,9 +35,10 @@ class AnthropicResponsesStreamWrapper: response.created -> message_start response.output_item.added -> content_block_start (if message/function_call) response.output_text.delta -> content_block_delta (text_delta) + response.reasoning_summary_part.added -> content_block_delta (thinking_delta separator) response.reasoning_summary_text.delta -> content_block_delta (thinking_delta) response.function_call_arguments.delta -> content_block_delta (input_json_delta) - response.output_item.done -> content_block_stop + response.output_item.done -> content_block_delta (signature_delta) + content_block_stop response.completed -> message_delta + message_stop """ @@ -94,6 +101,38 @@ class AnthropicResponsesStreamWrapper: ) return block_idx + @staticmethod + def _field(source: object, name: str) -> object: + return source.get(name) if isinstance(source, dict) else getattr(source, name, None) + + def _close_reasoning_item(self, item: object, item_id: str | None) -> None: + block_idx: Final = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + encrypted_content: Final = self._field(item, "encrypted_content") + signature: Final = ( + encrypted_reasoning_signature(encrypted_content) + if isinstance(encrypted_content, str) and encrypted_content + else None + ) + if block_idx < 0 and signature is None: + return + if block_idx < 0: + redacted_idx: Final = self._open_block( + item_id, + {"type": "redacted_thinking", "data": signature}, # mutable-ok: API message payload + ) + stop: Final = {"type": "content_block_stop", "index": redacted_idx} # mutable-ok: API message payload + self._chunk_queue.append(stop) + return + if signature is not None: + self._chunk_queue.append( + { # mutable-ok: API message payload + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "signature_delta", "signature": signature}, # mutable-ok: API message payload + } + ) + self._chunk_queue.append({"type": "content_block_stop", "index": block_idx}) # mutable-ok: API message payload + def _process_event(self, event: object) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) @@ -175,6 +214,26 @@ class AnthropicResponsesStreamWrapper: ) return + if event_type == "response.reasoning_summary_part.added": + part_item_id: Final = self._field(event, "item_id") + summary_index: Final = self._field(event, "summary_index") + part_block_idx: Final = ( + self._item_id_to_block_index.get(part_item_id, -1) if isinstance(part_item_id, str) else -1 + ) + if part_block_idx < 0 or not isinstance(summary_index, int) or summary_index == 0: + return + self._chunk_queue.append( + { # mutable-ok: API message payload + "type": "content_block_delta", + "index": part_block_idx, + "delta": { # mutable-ok: API message payload + "type": "thinking_delta", + "thinking": REASONING_SUMMARY_PART_SEPARATOR, + }, + } + ) + return + # ---- reasoning summary text delta ---- if event_type == "response.reasoning_summary_text.delta": item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) @@ -220,6 +279,9 @@ class AnthropicResponsesStreamWrapper: item_id = ( getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None ) + if self._field(item, "type") == "reasoning": + self._close_reasoning_item(item, item_id) + return block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index if block_idx < 0: return diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index f1daf2be42a..1fdb0318bab 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -13,7 +13,8 @@ from typing import Any, Final, cast from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, - responses_reasoning_item_from_thinking_blocks, + encrypted_reasoning_signature, + responses_reasoning_items_from_thinking_blocks, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.reasoning_effort_utils import ( @@ -33,6 +34,7 @@ from litellm.types.llms.anthropic import ( AnthropicFinishReason, AnthropicMessagesRequest, AnthropicMessagesToolChoice, + AnthropicResponseContentBlockRedactedThinking, AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, AnthropicResponseContentBlockToolUse, @@ -43,11 +45,13 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicUsage, ) from litellm.types.llms.openai import ( - ChatCompletionThinkingBlock, ResponseAPIUsage, ResponsesAPIResponse, ) +REASONING_SUMMARY_PART_SEPARATOR: Final = "\n\n" +RESPONSES_INCLUDE_ENCRYPTED_REASONING: Final = "reasoning.encrypted_content" + class LiteLLMAnthropicToResponsesAPIAdapter: """ @@ -163,49 +167,55 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return str(getattr(part, "text", None) or "") @classmethod - def _thinking_blocks_from_reasoning_item( + def _thinking_block_from_reasoning_item( cls, summary: Iterable[object], - ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload - """Anthropic thinking blocks for one Responses reasoning item. + encrypted_content: object, + ) -> dict[str, Any] | None: # mutable-ok: API message payload + """The one Anthropic block for a Responses reasoning item. - The signature stays empty: only Anthropic can sign a thinking block, and a stand-in - value would be replayed as a real one and rejected by every backend that verifies it. + The item's encrypted reasoning rides the block's opaque field (`signature`, or + `data` when there is no summary text) so the client echoes it back and the next + turn replays the very item OpenAI produced; without it the signature stays empty, + since only Anthropic can sign a thinking block. """ - return tuple( - AnthropicResponseContentBlockThinking( - type="thinking", - thinking=text, - signature=None, - ).model_dump() - for part in summary - if (text := cls._summary_part_text(part)) + text: Final = REASONING_SUMMARY_PART_SEPARATOR.join( + part_text for part in summary if (part_text := cls._summary_part_text(part)) ) + if not isinstance(encrypted_content, str) or not encrypted_content: + if not text: + return None + return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=None).model_dump() + signature: Final = encrypted_reasoning_signature(encrypted_content) + if not text: + return AnthropicResponseContentBlockRedactedThinking(type="redacted_thinking", data=signature).model_dump() + return AnthropicResponseContentBlockThinking(type="thinking", thinking=text, signature=signature).model_dump() @staticmethod def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: """Group a run of consecutive thinking blocks together; keep every other block alone.""" index, block = indexed_block - return "thinking" if block.get("type") == "thinking" else f"block:{index}" + return "thinking" if block.get("type") in ("thinking", "redacted_thinking") else f"block:{index}" @classmethod - def _assistant_group_to_input_item( + def _assistant_group_to_input_items( cls, group: tuple[Mapping[str, object], ...] - ) -> dict[str, Any] | None: # mutable-ok: API message payload + ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload first: Final = group[0] btype: Final = first.get("type") - if btype == "thinking": - blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group) # cast-ok: untrusted client payload - reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks) - return None if reasoning_item is None else dict(reasoning_item) # mutable-ok: API message payload + if btype in ("thinking", "redacted_thinking"): + replayed: Final = responses_reasoning_items_from_thinking_blocks(group) + return tuple(dict(item) for item in replayed) # mutable-ok: API message payload if btype == "tool_use": - return { # mutable-ok: API message payload - "type": "function_call", - "call_id": first.get("id", ""), - "name": first.get("name", ""), - "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload - } - return None + return ( + { # mutable-ok: API message payload + "type": "function_call", + "call_id": first.get("id", ""), + "name": first.get("name", ""), + "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload + }, + ) + return () def translate_messages_to_responses_input( self, @@ -362,7 +372,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: input_items.extend( item for _, group in groupby(enumerate(blocks), key=self._assistant_block_group_key) - if (item := self._assistant_group_to_input_item(tuple(block for _, block in group))) is not None + for item in self._assistant_group_to_input_items(tuple(block for _, block in group)) ) asst_parts: list[dict[str, Any]] = [ # mutable-ok: API message payload {"type": "output_text", "text": block.get("text", "")} # mutable-ok: API message payload @@ -495,10 +505,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_request( self, anthropic_request: AnthropicMessagesRequest, + include_encrypted_reasoning: bool = True, ) -> dict[str, Any]: """ Translate a full Anthropic /v1/messages request dict to litellm.responses() / litellm.aresponses() kwargs. + + ``include_encrypted_reasoning`` asks the provider for ``reasoning.encrypted_content`` + on every call, so a reasoning model's items can be replayed intact next turn even + when the client sent no ``thinking`` block; pass False for a provider whose + Responses API rejects ``include``. """ model: Final[str] = anthropic_request["model"] messages_list: Final = cast( @@ -528,6 +544,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: "model": model, "input": input_items, } + if include_encrypted_reasoning: + responses_kwargs["include"] = [RESPONSES_INCLUDE_ENCRYPTED_REASONING] # mutable-ok: API request payload if system and not developer_parts: if isinstance(system, str): @@ -634,7 +652,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: for item in response.output: if isinstance(item, ResponseReasoningItem): - content.extend(self._thinking_blocks_from_reasoning_item(item.summary)) + reasoning_block = self._thinking_block_from_reasoning_item(item.summary, item.encrypted_content) + if reasoning_block is not None: + content.append(reasoning_block) elif isinstance(item, ResponseOutputMessage): for part in item.content: @@ -684,11 +704,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ).model_dump() ) elif item_type == "reasoning": - content.extend( - self._thinking_blocks_from_reasoning_item( - cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json - ) + reasoning_block = self._thinking_block_from_reasoning_item( + cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json + item.get("encrypted_content"), ) + if reasoning_block is not None: + content.append(reasoning_block) elif item_type == "function_call": try: input_data = json.loads(item.get("arguments", "{}")) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 880a51eb584..ed16d7f3de0 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -7,8 +7,9 @@ from httpx._models import Headers, Response import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, + flatten_combinators_and_drop_non_python_regex_patterns, hoist_images_from_tool_messages, - tool_with_flattened_parameters, + tool_with_sanitized_parameters, ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, @@ -39,14 +40,17 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) -def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: +def sanitized_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: tools: Final = optional_params.get("tools") if not isinstance(tools, list): return _NO_TOOLS_UPDATE - flattened: Final = [ # mutable-ok: request tools are a JSON list - tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + sanitized: Final = [ # mutable-ok: request tools are a JSON list + tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) + if isinstance(tool, dict) + else tool + for tool in tools ] - return MappingProxyType({"tools": flattened}) + return MappingProxyType({"tools": sanitized}) class AzureOpenAIConfig(BaseConfig): @@ -278,7 +282,7 @@ class AzureOpenAIConfig(BaseConfig): "model": model, "messages": azure_messages, **optional_params, - **flattened_tools_update(optional_params), + **sanitized_tools_update(optional_params), } def transform_response( diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index 246bf69cb5f..09d8075e857 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -20,7 +20,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.utils import get_model_info, supports_reasoning from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig -from .gpt_transformation import flattened_tools_update +from .gpt_transformation import sanitized_tools_update class AzureOpenAIO1Config(OpenAIOSeriesConfig): @@ -111,6 +111,6 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict **optional_params, - **flattened_tools_update(optional_params), + **sanitized_tools_update(optional_params), } return super().transform_request(model, messages, flattened_params, litellm_params, headers) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6cb7d09cec4..e6b3eb1f2bb 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -14,6 +14,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.constants import DEFAULT_MAX_RETRIES from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import BaseOpenAILLM from litellm.secret_managers.get_azure_ad_token_provider import ( @@ -582,7 +583,8 @@ class BaseAzureLLM(BaseOpenAILLM): if scope is None: scope = "https://cognitiveservices.azure.com/.default" - max_retries: Final = litellm_params.get("max_retries") + configured_max_retries: Final = litellm_params.get("max_retries") + max_retries: Final = DEFAULT_MAX_RETRIES if configured_max_retries is None else configured_max_retries timeout: Final = litellm_params.get("timeout") if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret: verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth") @@ -642,8 +644,7 @@ class BaseAzureLLM(BaseOpenAILLM): else: azure_client_params["http_client"] = self._get_sync_http_client() - if max_retries is not None: - azure_client_params["max_retries"] = max_retries + azure_client_params["max_retries"] = max_retries if timeout is not None: azure_client_params["timeout"] = timeout diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 898852e645f..1b5a4083ebe 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,24 +1,102 @@ +import re +from collections.abc import Callable, Collection, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional import httpx from httpx import Response +from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + RelayShape, + logged_relay_shape, + replace_path_segment, + strip_leading_model_segment, +) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse if TYPE_CHECKING: from httpx import URL - from litellm.types.utils import CostResponseTypes + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse + + +class RelayedChatRequest(BaseModel): + messages: Sequence[Mapping[str, object]] | None = None + + +class RelayedCallDetails(BaseModel): + request_data: RelayedChatRequest | None = None + + +def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None: + try: + details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details) + except ValidationError: + return None + return details.request_data.messages if details.request_data else None + + +RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate) + +OPENAI_RELAY_SHAPES: Final = ( + RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate), + RESPONSES_RELAY_SHAPE, + RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate), +) + + +def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None: + """A streaming logging object assembles the logged response from the terminal event, not from its body.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks) + if terminal_event is None: + return None + logging_obj.call_type = ( + RESPONSES_RELAY_SHAPE.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return terminal_event + + +AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? str | None: + parts: Final = endpoint.split("/") + if len(parts) < 2: + return None + return next((part for part in parts if part in router_models), None) + + +def foreign_azure_deployment( + endpoint: str, model_group: str, served_models: Callable[[], Collection[str]] +) -> str | None: + match: Final = AZURE_DEPLOYMENT_SEGMENT.search(endpoint) + if match is None: + return None + deployment: Final = match.group(1) + if deployment == model_group: + return None + served: Final = frozenset(name.casefold() for name in served_models()) + return None if deployment.casefold() in served else deployment + + +def without_api_version(api_base: str) -> str: + url: Final = httpx.URL(api_base) + kept_params: Final = tuple((key, value) for key, value in url.params.multi_items() if key != "api-version") + return str(url.copy_with(params=httpx.QueryParams(kept_params))) class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: - return "stream" in request_data + return bool(request_data.get("stream")) def get_complete_url( self, @@ -36,14 +114,17 @@ class AzurePassthroughConfig(BasePassthroughConfig): litellm_metadata: Final = litellm_params.get("litellm_metadata") or {} model_group: Final = litellm_metadata.get("model_group") - if model_group and model_group in endpoint: - endpoint = endpoint.replace(model_group, model) + routed_endpoint: Final = replace_path_segment(endpoint, model_group, model) if model_group else endpoint + native_endpoint: Final = strip_leading_model_segment(routed_endpoint, (model,)) + caller_api_version: Final = request_query_params.get("api-version") if request_query_params else None + relay_base: Final = without_api_version(base_target_url) if caller_api_version else base_target_url complete_url: Final = BaseAzureLLM._get_base_azure_url( - api_base=base_target_url, - litellm_params=litellm_params, - route=endpoint, - default_api_version=litellm_params.get("api_version"), + api_base=relay_base, + litellm_params=MappingProxyType( + {**litellm_params, "api_version": caller_api_version or litellm_params.get("api_version")} + ), + route=native_endpoint, ) return ( httpx.URL(complete_url), @@ -92,13 +173,13 @@ class AzurePassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: Logging, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> Optional["LoggedRelayResponse"]: from litellm import encoding from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse if "chat/completions" not in endpoint: - return None + return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint) openai_chat_config: Final = OpenAIGPTConfig() @@ -116,3 +197,27 @@ class AzurePassthroughConfig(BasePassthroughConfig): ) return litellm_model_response + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> Optional["LoggedRelayResponse"]: + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + + if f"/{endpoint.strip('/')}".endswith(RESPONSES_RELAY_SHAPE.path_suffix): + return logged_responses_stream(all_chunks, litellm_logging_obj) + if "chat/completions" not in endpoint: + return None + + return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + messages=_relayed_messages(litellm_logging_obj), + ) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 039c462b38a..00e1c1e25ba 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -2,7 +2,6 @@ import copy import enum import re from typing import TYPE_CHECKING, Final, cast -from urllib.parse import urlparse import httpx from httpx import Response @@ -15,7 +14,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base +from litellm.llms.azure_ai.common_utils import ( + api_key_header_for_base, + is_foundry_model_inference_base, +) from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error @@ -146,11 +148,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Returns True if the request should use `api-key` header for authentication. """ - parsed_url: Final = urlparse(api_base) - host: Final = parsed_url.hostname - if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): - return True - return False + return api_key_header_for_base(api_base) == "api-key" def get_complete_url( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index aa34bab5b2e..0665b3f64c5 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -19,6 +19,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool: return "/openai/deployments" not in parsed.path +def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader: + host: Final = urlparse(api_base).hostname if api_base else None + if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): + return "api-key" + return "Authorization" + + def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. diff --git a/litellm/llms/azure_ai/image_edit/__init__.py b/litellm/llms/azure_ai/image_edit/__init__.py index 51a23859058..fda9335a5d6 100644 --- a/litellm/llms/azure_ai/image_edit/__init__.py +++ b/litellm/llms/azure_ai/image_edit/__init__.py @@ -23,7 +23,7 @@ def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig: """ Get the appropriate image edit config for an Azure AI model. - - MAI models use /mai/v1/images/edits with multipart form data and size + - MAI models use /mai/v1/images/edits with multipart form data - FLUX 2 models use JSON with base64 image - FLUX 1 models use multipart/form-data """ diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py index e639c20292b..55b179e9591 100644 --- a/litellm/llms/azure_ai/image_edit/mai_transformation.py +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final import httpx from httpx._types import RequestFiles @@ -13,7 +13,6 @@ from litellm.llms.azure_ai.image_generation.mai_transformation import ( from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str -from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.openai import FileTypes from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageResponse @@ -26,65 +25,8 @@ if TYPE_CHECKING: class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): """Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5).""" - DEFAULT_SIZE = "1024x1024" - def get_supported_openai_params(self, model: str) -> list: - return ["prompt", "image", "model", "n", "size"] - - def map_openai_params( - self, - image_edit_optional_params: ImageEditOptionalRequestParams, - model: str, - drop_params: bool, - ) -> dict: - optional_params: Final[dict[str, Any]] = {} - supported_params: Final = self.get_supported_openai_params(model) - - for key, value in dict(image_edit_optional_params).items(): - if value is None or key in optional_params: - continue - - if key in supported_params: - if key == "size" and value: - size_param = cast(str, value) - self._validate_size_param(size_param) - optional_params[key] = size_param - else: - optional_params[key] = value - elif not drop_params: - raise ValueError( - f"Parameter {key} is not supported for model {model}. " - f"Supported parameters are {supported_params}. " - f"Set drop_params=True to drop unsupported parameters." - ) - - if "size" not in optional_params: - optional_params["size"] = self.DEFAULT_SIZE - - return optional_params - - def _validate_size_param(self, size: str) -> None: - known_sizes: Final = { - "1024x1024", - "1792x1024", - "1024x1792", - "512x512", - "256x256", - } - - if size in known_sizes: - return - - if "x" in size: - try: - tuple(map(int, size.lower().split("x", 1))) - return - except ValueError: - raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") - - raise ValueError( - f"Unsupported size value: '{size}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." - ) + return ["prompt", "image", "model", "n"] def validate_environment( self, diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 64f81956ad7..67b1a8bcab3 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.exceptions import UnsupportedParamsError from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -21,6 +22,10 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_WIDTH = 1024 DEFAULT_HEIGHT = 1024 + MAX_IMAGES_PER_REQUEST: Final = 1 + MIN_DIMENSION_PX: Final = 768 + MAX_TOTAL_PX: Final = 1_056_768 + @staticmethod def get_mai_image_generation_url( api_base: str | None, @@ -145,16 +150,27 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: if k == "size" and v: - self._map_size_param(v, optional_params) + self._map_size_param(v, optional_params, model) + elif k == "n" and v is not None and self._image_count(v, model) != self.MAX_IMAGES_PER_REQUEST: + if not drop_params: + raise self._unsupported( + model, + f"n={v} is not supported for model {model}. The Azure AI MAI image " + f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per " + "request and ignores any count, so a larger value would silently " + "return fewer images than requested. Send one request per image, or " + "set drop_params=True to drop n.", + ) else: optional_params[k] = v elif k in ("width", "height"): optional_params[k] = v elif not drop_params: - raise ValueError( + raise self._unsupported( + model, f"Parameter {k} is not supported for model {model}. " f"Supported parameters are {supported_params} and width/height. " - f"Set drop_params=True to drop unsupported parameters." + f"Set drop_params=True to drop unsupported parameters.", ) if "width" not in optional_params: @@ -165,7 +181,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): optional_params.pop("size", None) return optional_params - def _map_size_param(self, size: str, optional_params: dict) -> None: + @staticmethod + def _unsupported(model: str, message: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, llm_provider="azure_ai", model=model) + + def _image_count(self, n: object, model: str) -> int: + if isinstance(n, int): + return n + try: + return int(str(n)) + except ValueError: + raise self._unsupported(model, f"n={n!r} is not a whole number of images for model {model}.") + + def _map_size_param(self, size: str, optional_params: dict, model: str) -> None: size_mapping: Final = { "1024x1024": (1024, 1024), "1792x1024": (1792, 1024), @@ -176,19 +204,36 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if size in size_mapping: width, height = size_mapping[size] - optional_params["width"] = width - optional_params["height"] = height elif "x" in size: try: width, height = map(int, size.lower().split("x")) - optional_params["width"] = width - optional_params["height"] = height except ValueError: - raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") + raise self._unsupported( + model, f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." + ) else: - raise ValueError( + raise self._unsupported( + model, f"Unsupported size value: '{size}'. " - f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string.", + ) + + self._validate_dimensions(model=model, size=size, width=width, height=height) + optional_params["width"] = width + optional_params["height"] = height + + def _validate_dimensions(self, model: str, size: str, width: int, height: int) -> None: + if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX: + raise self._unsupported( + model, + f"Unsupported size value: '{size}'. Azure AI MAI image models require width and " + f"height of at least {self.MIN_DIMENSION_PX} pixels.", + ) + if width * height > self.MAX_TOTAL_PX: + raise self._unsupported( + model, + f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most " + f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height}).", ) def transform_image_generation_response( diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py new file mode 100644 index 00000000000..f2be1d95593 --- /dev/null +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + api_key_header_for_base, + get_azure_ai_auth_headers, +) +from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + RelayShape, + logged_relay_shape, + strip_leading_model_segment, +) +from litellm.types.llms.openai import AllMessageValues +from litellm.types.rerank import RerankResponse +from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject + +if TYPE_CHECKING: + from httpx import URL, Response + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse + from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse + + +EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({}) + + +class PassthroughMetadata(BaseModel): + model_config = ConfigDict(extra="ignore") + + model_group: str = "" + + +def model_group_from(litellm_params: Mapping[str, object]) -> str: + try: + return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group + except ValidationError: + return "" + + +def api_version_from(litellm_params: Mapping[str, object]) -> str | None: + try: + return TypeAdapter(str | None).validate_python(litellm_params.get("api_version")) + except ValidationError: + return None + + +def foundry_root(api_base: str) -> str: + url: Final = httpx.URL(api_base) + segments: Final = tuple(segment for segment in url.path.split("/") if segment) + root_segments: Final = segments[: segments.index("models")] if "models" in segments else segments + return str(url.copy_with(path="/" + "/".join(root_segments), query=None)).rstrip("/") + + +def is_repeated_native_prefix(native_segments: tuple[str, ...], overlap: int) -> bool: + return overlap == len(native_segments) or native_segments[0] == "openai" + + +def without_repeated_native_prefix(root: str, native_endpoint: str) -> str: + url: Final = httpx.URL(root) + root_segments: Final = tuple(segment for segment in url.path.split("/") if segment) + native_segments: Final = tuple(segment.casefold() for segment in native_endpoint.split("/") if segment) + overlap: Final = next( + ( + length + for length in range(min(len(root_segments), len(native_segments)), 0, -1) + if tuple(segment.casefold() for segment in root_segments[-length:]) == native_segments[:length] + and is_repeated_native_prefix(native_segments, length) + ), + 0, + ) + kept_segments: Final = root_segments[: len(root_segments) - overlap] + return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/") + + +def relay_query_params( + request_query_params: Mapping[str, object] | None, + deployment_api_version: str | None, + api_base: str, +) -> Mapping[str, object] | None: + if request_query_params and "api-version" in request_query_params: + return request_query_params + api_version: Final = deployment_api_version or httpx.URL(api_base).params.get("api-version") + if api_version is None: + return request_query_params + return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version}) + + +def relayed_body(httpx_response: Response) -> str | dict: + try: + body: Final[object] = httpx_response.json() + except ValueError: + return httpx_response.text + return body if isinstance(body, dict) else httpx_response.text + + +FOUNDRY_RELAY_SHAPES: Final = ( + RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate), + RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate), +) + + +class AzureAIPassthroughConfig(AzureFoundryModelInfo, BasePassthroughConfig): + def __init__(self, ocr_config_for: Callable[[str], BaseOCRConfig | None] = get_azure_ai_ocr_config) -> None: + super().__init__() + self.ocr_config_for: Final = ocr_config_for + + def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: + return bool(request_data.get("stream")) + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + endpoint: str, + request_query_params: Mapping[str, object] | None, + litellm_params: Mapping[str, object], + ) -> tuple[URL, str]: + base_target_url: Final = self.get_api_base(api_base) + if base_target_url is None: + raise ValueError("Azure AI api base not found: set `api_base` on the deployment or AZURE_AI_API_BASE") + + native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params))) + root: Final = without_repeated_native_prefix(foundry_root(base_target_url), native_endpoint) + query_params: Final = relay_query_params( + request_query_params, api_version_from(litellm_params), base_target_url + ) + return (self.format_url(native_endpoint, root, query_params), root) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx + auth_headers: Final = get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params, + api_key_header=api_key_header_for_base(api_base), + ) + return {**headers, **auth_headers} # mutable-ok: base class contract returns dict for httpx + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: Mapping[str, object], + logging_obj: Logging, + endpoint: str, + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + chat_result: Final = AzurePassthroughConfig().logging_non_streaming_response( # pyright: ignore[reportUnknownMemberType] # the Azure config still types request_data as a bare dict + model=model, + custom_llm_provider=custom_llm_provider, + httpx_response=httpx_response, + request_data=dict(request_data), # mutable-ok: AzurePassthroughConfig wants a dict + logging_obj=logging_obj, + endpoint=endpoint, + ) + if chat_result is not None: + return chat_result + ocr_result: Final = self.logged_ocr_response(model, httpx_response, logging_obj, endpoint) + if ocr_result is not None: + return ocr_result + foundry_result: Final = logged_relay_shape(FOUNDRY_RELAY_SHAPES, httpx_response, logging_obj, endpoint) + if foundry_result is not None: + return foundry_result + return StandardPassThroughResponseObject(response=relayed_body(httpx_response)) + + def logged_ocr_response( + self, model: str, httpx_response: Response, logging_obj: Logging, endpoint: str + ) -> OCRResponse | None: + ocr_config: Final = self.ocr_config_for(model) + if ocr_config is None or httpx_response.status_code != 200: + return None + relayed_url: Final = httpx_response.request.url + relayed_origin: Final = str(relayed_url.copy_with(path="/", query=None, fragment=None)).rstrip("/") + ocr_url: Final = httpx.URL( + ocr_config.get_complete_url( + api_base=relayed_origin, + model=model, + optional_params={}, # mutable-ok: BaseOCRConfig wants a dict + ) + ) + known_prefixes: Final = (model, model_group_from(logging_obj.litellm_params)) + native_endpoint: Final = strip_leading_model_segment(endpoint, known_prefixes) + if f"/{native_endpoint.strip('/')}" != ocr_url.path: + return None + try: + ocr_response: Final = ocr_config.transform_ocr_response( + model=model, raw_response=httpx_response, logging_obj=logging_obj + ) + except (ValueError, AttributeError) as error: + verbose_logger.warning("azure_ai passthrough: OCR body from %s is not costable: %s", ocr_url, error) + return None + logging_obj.call_type = CallTypes.aocr.value # rebind-ok: routes cost calculation to the per-page OCR path + return ocr_response + + def handle_logging_collected_chunks( + self, + all_chunks: Sequence[str], + litellm_logging_obj: Logging, + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> LoggedRelayResponse | None: + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig + + return AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + custom_llm_provider=custom_llm_provider, + endpoint=endpoint, + ) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index afd8e0f67f7..f1143425ced 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -52,13 +52,28 @@ class StreamingScanKey: class BaseTranslation(ABC): - delivers_ended_stream_text_rewrites: ClassVar[bool] = False + delivers_ended_stream_rewrites: ClassVar[bool] = False """Whether ``process_output_streaming_response`` accepts ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered) - stream, writes guardrail text rewrites back across ``responses_so_far`` so - a buffered pipeline can release rewritten chunks. Tool-call rewrites, and - text rewrites on every other translation, are undeliverable: the pipeline - executor discards them and releases the original chunks.""" + stream, writes guardrail text and tool-call rewrites back across + ``responses_so_far`` so a buffered pipeline can release rewritten chunks, + raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites + on every other translation are undeliverable: the pipeline executor + discards them and releases the original chunks.""" + + assembles_streamed_response: ClassVar[bool] = False + """Whether ``process_output_streaming_response`` stores the assembled response of an + ended stream under ``request_data["response"]`` before scanning it, the way the chat, + Responses, and Messages translations do. A streaming pipeline runs a guardrail that only + has the legacy post-call hook against that response, so on a translation without it such + a guardrail keeps running on its own.""" + + def post_call_hook_response(self, response: object) -> object: + """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from + the object the translation stores under ``request_data["response"]`` while scanning an + ended stream. Chat and Responses scan that shape already; a translation that scans a + different one (Messages scans an OpenAI-shaped ModelResponse) overrides this.""" + return response @staticmethod def transform_user_api_key_dict_to_metadata( @@ -175,9 +190,9 @@ class BaseTranslation(ABC): transformations (see ``StreamTransformSink``); base handlers ignore it. ``deliver_ended_stream_rewrites`` is passed True only when the caller holds the whole buffered stream and the subclass declares - ``delivers_ended_stream_text_rewrites``: the handler then writes - guardrail text rewrites back across ``responses_so_far`` instead of - discarding them. + ``delivers_ended_stream_rewrites``: the handler then writes + guardrail text and tool-call rewrites back across ``responses_so_far`` + instead of discarding them. """ return responses_so_far diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index adbf2e126fb..20180c5cfa2 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,5 +1,14 @@ +from __future__ import annotations + +import re from abc import abstractmethod -from typing import TYPE_CHECKING, Final, Optional, Union +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm.types.utils import CallTypes from ..base_utils import BaseLLMModelInfo @@ -7,9 +16,68 @@ if TYPE_CHECKING: from httpx import URL, Headers, Response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import CostResponseTypes + from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent + from litellm.types.rerank import RerankResponse + from litellm.types.utils import CostResponseTypes, StandardPassThroughResponseObject from ..chat.transformation import BaseLLMException + from ..ocr.transformation import OCRResponse + + LoggedRelayResponse: TypeAlias = CostResponseTypes | RerankResponse | ResponsesAPIResponse | ResponsesTerminalEvent + + +RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) + + +def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str: + path: Final = endpoint.lstrip("/") + for model_name in model_names: + if not model_name: + continue + if path == model_name: + return "" + if path.startswith(f"{model_name}/"): + return path[len(model_name) + 1 :] + return path + + +def replace_path_segment(endpoint: str, segment: str, replacement: str) -> str: + bounded_segment: Final = re.compile(rf"(? Mapping[str, object] | None: + if httpx_response.status_code != 200: + return None + try: + return RELAYED_JSON_OBJECT.validate_python(httpx_response.json()) + except (ValueError, ValidationError): + return None + + +@dataclass(frozen=True, slots=True) +class RelayShape: + path_suffix: str + call_type: CallTypes + parse: Callable[[Mapping[str, object]], LoggedRelayResponse] + + +def logged_relay_shape( + shapes: Sequence[RelayShape], httpx_response: Response, logging_obj: LiteLLMLoggingObj, endpoint: str +) -> LoggedRelayResponse | None: + relayed_path: Final = f"/{endpoint.strip('/')}" + shape: Final = next((candidate for candidate in shapes if relayed_path.endswith(candidate.path_suffix)), None) + body: Final = relayed_json_object(httpx_response) if shape else None + if shape is None or body is None: + return None + try: + parsed: Final = shape.parse(body) + except ValidationError: + return None + logging_obj.call_type = ( + shape.call_type.value + ) # rebind-ok: routes cost calculation to the relayed shape's pricing path + return parsed class BasePassthroughConfig(BaseLLMModelInfo): @@ -23,8 +91,8 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, endpoint: str, base_target_url: str, - request_query_params: dict | None, - ) -> "URL": + request_query_params: Mapping[str, object] | None, + ) -> URL: """ Helper function to add query params to the url Args: @@ -58,7 +126,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): endpoint: str, request_query_params: dict | None, litellm_params: dict, - ) -> tuple["URL", str]: + ) -> tuple[URL, str]: """ Get the complete url for the request Returns: @@ -88,9 +156,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): """ return headers, None - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, "Headers"] - ) -> "BaseLLMException": + def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: from litellm.llms.base_llm.chat.transformation import BaseLLMException return BaseLLMException(status_code=status_code, message=error_message, headers=headers) @@ -99,21 +165,21 @@ class BasePassthroughConfig(BaseLLMModelInfo): self, model: str, custom_llm_provider: str, - httpx_response: "Response", + httpx_response: Response, request_data: dict, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: pass def handle_logging_collected_chunks( self, all_chunks: list[str], - litellm_logging_obj: "LiteLLMLoggingObj", + litellm_logging_obj: LiteLLMLoggingObj, model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> LoggedRelayResponse | None: return None def _convert_raw_bytes_to_str_lines(self, raw_bytes: list[bytes]) -> list[str]: diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index a197702921f..96804a1fa62 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,13 +1,17 @@ +import asyncio import base64 +import contextvars import hashlib import json import os import re import urllib.parse from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor from datetime import datetime +from functools import partial from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload import httpx from pydantic import BaseModel, ValidationError @@ -16,6 +20,7 @@ from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + AWS_SIGNING_MAX_THREADS, BEDROCK_EMBEDDING_PROVIDERS_LITERAL, BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES, BEDROCK_IAM_CACHE_MAX_ENTRIES, @@ -80,7 +85,11 @@ class AwsAuthError(Exception): super().__init__(self.message) # Call the base class constructor with the parameters it needs -class BaseAWSLLM: +class SignsRequestsWithAWS: + pass + + +class BaseAWSLLM(SignsRequestsWithAWS): # Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request). # Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static # access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient @@ -1668,3 +1677,52 @@ class BaseAWSLLM: request_headers_dict["Authorization"] = incoming_authorization return request_headers_dict, request.body + + +def sign_aws_json_post( + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") + + aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + +_SignParams = ParamSpec("_SignParams") +_SignedRequest = TypeVar("_SignedRequest") + +AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing") + + +async def run_aws_signing( + sign: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature +) -> _SignedRequest: + context: Final = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor( + AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs) + ) + + +async def sign_request_off_loop_if_aws( + provider_config: object, + sign_request: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature +) -> _SignedRequest: + if isinstance(provider_config, SignsRequestsWithAWS): + return await run_aws_signing(sign_request, *args, **kwargs) + return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 984ba371898..6d48ff3f07c 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -136,7 +136,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await run_aws_signing( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, @@ -206,7 +207,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await run_aws_signing( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index fa24f8be893..fa18361e44c 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -4,6 +4,7 @@ Translating between OpenAI's `/chat/completion` format and Amazon's `/converse` import copy import json +import re import time import types from collections.abc import Mapping @@ -293,6 +294,10 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock", ) + @staticmethod + def _is_openai_gpt_reasoning_model(model: str) -> bool: + return re.search(r"openai\.gpt-\d", model) is not None + def _is_nova_2_model(self, model: str) -> bool: """ Check if the model is a Nova 2 model that supports reasoningConfig. @@ -422,15 +427,15 @@ class AmazonConverseConfig(BaseConfig): """ Handle the reasoning_effort parameter based on the model type. - - GPT-OSS models: passed through unchanged via additionalModelRequestFields. - - OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields. + - GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields. + - OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields. - Nova 2 models: transformed to reasoningConfig. - Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on adaptive Claude 4.6 / 4.7). """ - if "gpt-oss" in model: + if "gpt-oss" in model or "deepseek" in model: optional_params["reasoning_effort"] = reasoning_effort - elif "openai.gpt-5" in model: + elif self._is_openai_gpt_reasoning_model(model): reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort} optional_params["reasoning"] = reasoning elif self._is_nova_2_model(model): @@ -509,6 +514,36 @@ class AmazonConverseConfig(BaseConfig): ) thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS + def _is_deepseek_model(self, model: str, base_model: str) -> bool: + return "deepseek" in model or "deepseek" in base_model + + def _is_deepseek_r1_model(self, model: str, base_model: str) -> bool: + return "deepseek.r1" in model or "deepseek.r1" in base_model + + def _model_accepts_anthropic_thinking_param(self, model: str, base_model: str) -> bool: + """Whether the model accepts the Anthropic-shaped ``thinking`` request field. + + Only Claude reasoning models accept it. DeepSeek advertises ``supports_reasoning`` but reasons + natively: R1 returns a 400 when the field is sent and V3 silently ignores it. + """ + if self._is_deepseek_model(model=model, base_model=base_model): + return False + return ( + "claude-3-7" in model + or "claude-sonnet-4" in model + or "claude-opus-4" in model + or supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider) + or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) + ) + + def _model_rejects_reasoning_effort_param(self, model: str, base_model: str) -> bool: + """Whether the model returns a 400 for every ``reasoning_effort`` shape on Converse. + + DeepSeek R1 always reasons and rejects any reasoning request field. DeepSeek V3 accepts a raw + ``reasoning_effort`` like gpt-oss does, and every other model maps it to a shape it accepts. + """ + return self._is_deepseek_r1_model(model=model, base_model=base_model) + def get_supported_openai_params(self, model: str) -> list[str]: from litellm.utils import supports_function_calling @@ -564,23 +599,20 @@ class AmazonConverseConfig(BaseConfig): # only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html supported_params.append("tool_choice") - if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model: + if ( + "gpt-oss" in model + or self._is_openai_gpt_reasoning_model(model) + or self._is_openai_gpt_reasoning_model(base_model) + ): supported_params.append("reasoning_effort") + elif self._is_deepseek_model(model=model, base_model=base_model): + if not self._is_deepseek_r1_model(model=model, base_model=base_model): + supported_params.append("reasoning_effort") elif self._is_nova_2_model(model): # Nova 2 models support reasoning_effort (transformed to reasoningConfig) # These models use a different reasoning structure than Anthropic's thinking parameter supported_params.append("reasoning_effort") - elif ( - "claude-3-7" in model - or "claude-sonnet-4" in model - or "claude-opus-4" in model - or "deepseek.r1" in model - or supports_reasoning( - model=model, - custom_llm_provider=self.custom_llm_provider, - ) - or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) - ): + elif self._model_accepts_anthropic_thinking_param(model=model, base_model=base_model): supported_params.append("thinking") supported_params.append("reasoning_effort") supported_params.append("output_config") @@ -872,6 +904,11 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, ) -> dict: is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) + base_model: Final = BedrockModelInfo.get_base_model(model) + drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model) + drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param( + model=model, base_model=base_model + ) for param, value in non_default_params.items(): if param == "response_format" and isinstance(value, dict): @@ -920,7 +957,12 @@ class AmazonConverseConfig(BaseConfig): optional_params["_parallel_tool_use_config"] = { "tool_choice": {"type": "auto", "disable_parallel_tool_use": not value} } - if param == "thinking" and "openai.gpt-5" not in model: + if param == "thinking" and drop_thinking_param: + verbose_logger.debug( + "Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.", + model, + ) + elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model): if ( isinstance(value, dict) and value.get("type") == "adaptive" @@ -946,6 +988,11 @@ class AmazonConverseConfig(BaseConfig): AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=optional_params, custom_llm_provider="bedrock" ) + elif param == "reasoning_effort" and isinstance(value, str) and drop_reasoning_effort_param: + verbose_logger.debug( + "Dropping unsupported `reasoning_effort` param for Bedrock model=%s; it always reasons and rejects it.", + model, + ) elif param == "reasoning_effort" and isinstance(value, str): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params @@ -1805,6 +1852,7 @@ class AmazonConverseConfig(BaseConfig): data=request_data, messages=messages, encoding=encoding, + json_mode=json_mode, ) def _transform_reasoning_content(self, reasoning_content_blocks: list[BedrockConverseReasoningContentBlock]) -> str: @@ -2237,6 +2285,7 @@ class AmazonConverseConfig(BaseConfig): data: dict | str, messages: list, encoding, + json_mode: bool | None = None, ) -> ModelResponse: ## LOGGING if logging_obj is not None: @@ -2247,7 +2296,9 @@ class AmazonConverseConfig(BaseConfig): additional_args={"complete_input_dict": data}, ) - json_mode: Final[bool | None] = optional_params.get("json_mode", None) + resolved_json_mode: Final[bool | None] = ( + json_mode if json_mode is not None else optional_params.get("json_mode", None) + ) ## RESPONSE OBJECT try: completion_response: Final = ConverseResponseBlock(**response.json()) @@ -2339,7 +2390,7 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks) chat_completion_message["content"] = content_str filtered_tools: Final = self._filter_json_mode_tools( - json_mode=json_mode, + json_mode=resolved_json_mode, tools=tools, chat_completion_message=chat_completion_message, ) @@ -2363,7 +2414,7 @@ class AmazonConverseConfig(BaseConfig): # When json_mode filtered out all synthetic tool calls the response # is plain content, not a pending tool invocation. Fix finish_reason # so callers (e.g. OpenAI SDK) don't misinterpret it. - if json_mode and not filtered_tools and tools: + if resolved_json_mode and not filtered_tools and tools: initial_finish_reason = "stop" ( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index a0e32c8aa22..90a2692f68a 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -340,6 +340,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): optional_params=optional_params, litellm_params=litellm_params, encoding=encoding, + json_mode=json_mode, ) elif provider == "twelvelabs": return litellm.AmazonTwelveLabsPegasusConfig().transform_response( diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 1fb53f6ff0a..d7fb510f057 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -10,9 +10,10 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client class BedrockCountTokensHandler(BedrockCountTokensConfig): @@ -27,6 +28,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): request_data: dict[str, Any], litellm_params: dict[str, Any], resolved_model: str, + client: AsyncHTTPHandler | None = None, ) -> dict[str, Any]: """ Handle a CountTokens request using existing LiteLLM patterns. @@ -75,7 +77,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): # Extract api_key for bearer token auth if provided api_key: Final = litellm_params.get("api_key", None) headers: Final = {"Content-Type": "application/json"} - signed_headers, signed_body = self._sign_request( + signed_headers, signed_body = await run_aws_signing( + self._sign_request, service_name="bedrock", headers=headers, optional_params=litellm_params, @@ -85,7 +88,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): api_key=api_key, ) - async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) + async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) response: Final = await async_client.post( endpoint_url, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index e249feb9ff3..be766eaedd0 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -5,7 +5,7 @@ Handles embedding calls to Bedrock's `/invoke` endpoint import copy import json import urllib.parse -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final, get_args, overload import httpx @@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -41,6 +41,20 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +def _sign_get_request( + credentials: Credentials, url: str, headers: Mapping[str, str], aws_region_name: str +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + request: Final = AWSRequest(method="GET", url=url, data=None, headers=headers) + SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request) + return request.prepare() + + class BedrockEmbedding(BaseAWSLLM): @overload def _load_credentials( @@ -342,7 +356,8 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = self.get_request_headers( + prepped = await run_aws_signing( + self.get_request_headers, credentials=credentials, aws_region_name=aws_region_name, extra_headers=extra_headers, @@ -600,9 +615,6 @@ class BedrockEmbedding(BaseAWSLLM): dict: Status response from AWS Bedrock """ - # Get AWS credentials using the same method as other Bedrock methods - credentials, _ = self._load_credentials(kwargs) - # Get the runtime endpoint endpoint_url, _ = self.get_runtime_endpoint( api_base=None, @@ -619,27 +631,13 @@ class BedrockEmbedding(BaseAWSLLM): # Prepare headers for GET request headers: Final = {"Content-Type": "application/json"} - # Use AWSRequest directly for GET requests (get_request_headers hardcodes POST) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + def sign_status_request() -> AWSPreparedRequest: + credentials, _ = self._load_credentials(kwargs) + return _sign_get_request( + credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name + ) - # Create AWSRequest with GET method and encoded URL - request: Final = AWSRequest( - method="GET", - url=status_url, - data=None, # GET request, no body - headers=headers, - ) - - # Sign the request - SigV4Auth will create canonical string from request URL - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - sigv4.add_auth(request) - - # Prepare the request - prepped: Final = request.prepare() + prepped: Final = await run_aws_signing(sign_status_request) # LOGGING if logging_obj is not None: diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 42fe8941443..ca2370303f2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig @@ -149,7 +149,8 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) - credentials: Final = self.get_credentials( + credentials: Final = await run_aws_signing( + self.get_credentials, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, @@ -169,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM): "or configure credentials in the environment" ), ) - frozen_credentials: Final = credentials.get_frozen_credentials() + frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials) # Initialize Bedrock client with aws_sdk_bedrock_runtime config: Final = Config( diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 850738bc320..ac94d9cb922 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -23,7 +23,7 @@ from botocore.exceptions import ( ProfileNotFound, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS from litellm.secret_managers.main import get_secret_str BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" @@ -55,7 +55,7 @@ def resolve_mantle_region(params: Mapping[str, object]) -> str: ) -class BedrockMantleAuthMixin: +class BedrockMantleAuthMixin(SignsRequestsWithAWS): _aws_signer: BaseAWSLLM @staticmethod diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7587b963a38..e720428847d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -579,7 +580,7 @@ class BaseLLMHTTPHandler: data: dict[str, object], # mutable-ok: async_completion takes dict signed_headers: dict[str, object], # mutable-ok: async_completion takes dict signed_json_body: bytes | None, - ): + ) -> Coroutine[object, object, ModelResponse | CustomStreamWrapper]: async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None if stream is True: return self.acompletion_stream_function( @@ -626,7 +627,7 @@ class BaseLLMHTTPHandler: if acompletion is True and provider_config.uses_async_transform_request: - async def transform_then_dispatch(): + async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper: transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict "dict[str, object]", await provider_config.async_transform_request( @@ -637,7 +638,12 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed)) + signed_request: Final = await ( + run_aws_signing(sign_and_log, transformed) + if isinstance(provider_config, SignsRequestsWithAWS) + else asyncio.to_thread(sign_and_log, transformed) + ) + return await dispatch_async(*signed_request) return transform_then_dispatch() @@ -1973,7 +1979,9 @@ class BaseLLMHTTPHandler: api_key=api_key, ) - signed_headers, signed_json_body = provider_config.sign_request( + signed_headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params, request_data=data, @@ -2074,7 +2082,9 @@ class BaseLLMHTTPHandler: max_attempts, ) provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body) - headers, signed_json_body = provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params_dict, request_data=request_body, @@ -2234,7 +2244,9 @@ class BaseLLMHTTPHandler: stream=stream, ) - headers, signed_json_body = anthropic_messages_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + anthropic_messages_provider_config, + anthropic_messages_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), # dynamic aws_* params are passed under litellm_params request_data=request_body, @@ -2910,7 +2922,9 @@ class BaseLLMHTTPHandler: fake_stream=fake_stream, ) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -4618,7 +4632,9 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -9845,7 +9861,9 @@ class BaseLLMHTTPHandler: ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) - headers, signed_json_body = vector_store_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + vector_store_provider_config, + vector_store_provider_config.sign_request, headers=headers, optional_params=all_optional_params, request_data=request_body, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index e2e2ea5b553..09aaf970dc5 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_reasoning_content, # pyright: ignore[reportPrivateUsage] # same import as the OpenAI transformation strip_litellm_internal_message_fields, strip_name_from_message, ) @@ -23,7 +24,9 @@ from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.databricks import ( AllDatabricksContentValues, DatabricksChoice, + DatabricksDelta, DatabricksFunction, + DatabricksMessage, DatabricksResponse, DatabricksTool, ) @@ -247,8 +250,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = self._get_api_base(api_base) - complete_url: Final = f"{api_base}/chat/completions" + use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2 + api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway) + url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base + complete_url: Final = f"{url_base}/chat/completions" return complete_url def get_supported_openai_params(self, model: str | None = None) -> list: @@ -534,6 +539,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): thinking_blocks.append(thinking_block) return reasoning_content, thinking_blocks + @staticmethod + def extract_top_level_reasoning_content(delta: DatabricksDelta) -> str | None: + return delta.get("reasoning_content") + + @staticmethod + def resolve_reasoning_and_content( + message: DatabricksMessage, block_reasoning_content: str | None + ) -> tuple[str | None, str | None]: + content_str: Final = DatabricksConfig.extract_content_str(message["content"]) + if block_reasoning_content is not None: + return block_reasoning_content, content_str + return _extract_reasoning_content({**message, "content": content_str}) + @staticmethod def extract_citations( content: AllDatabricksContentValues | None, @@ -577,14 +595,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): finish_reason = "stop" if translated_message is None: - ## get the content str - content_str = DatabricksConfig.extract_content_str(choice["message"]["content"]) - - ## get the reasoning content ( - reasoning_content, + block_reasoning_content, thinking_blocks, ) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content")) + reasoning_content, content_str = DatabricksConfig.resolve_reasoning_and_content( + choice["message"], block_reasoning_content + ) citations = DatabricksConfig.extract_citations(choice["message"].get("content")) @@ -738,12 +755,16 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator): # extract the reasoning content ( - reasoning_content, + block_reasoning_content, thinking_blocks, ) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content")) choice["delta"]["content"] = content_str - choice["delta"]["reasoning_content"] = reasoning_content + choice["delta"]["reasoning_content"] = ( + block_reasoning_content + if block_reasoning_content is not None + else DatabricksConfig.extract_top_level_reasoning_content(choice["delta"]) + ) choice["delta"]["thinking_blocks"] = thinking_blocks translated_choices.append(choice) return ModelResponseStream( diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 7695b1cb35e..a4ec2c5378b 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -177,19 +177,13 @@ class DatabricksBase: # Default: just litellm return f"litellm/{version}" - def _get_api_base(self, api_base: str | None) -> str: - """ - Get the Databricks API base URL. - - If not provided, attempts to get it from the Databricks SDK. - """ + def _get_api_base(self, api_base: str | None, use_ai_gateway: bool = False) -> str: if api_base is None: try: from databricks.sdk import WorkspaceClient databricks_client: Final = WorkspaceClient() api_base = f"{databricks_client.config.host}/serving-endpoints" - return api_base except ImportError: raise DatabricksException( status_code=400, @@ -198,6 +192,18 @@ class DatabricksBase: "or install the databricks-sdk Python library." ), ) + + if not use_ai_gateway: + return api_base + + normalized_api_base: Final = api_base.rstrip("/") + if normalized_api_base.endswith("/ai-gateway/mlflow/v1"): + return normalized_api_base + if normalized_api_base.endswith("/serving-endpoints"): + return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1" + api_base_parts: Final = urlsplit(normalized_api_base) + if api_base_parts.path in ("", "/"): + return f"{normalized_api_base}/ai-gateway/mlflow/v1" return api_base def _get_oauth_m2m_token( diff --git a/litellm/llms/hosted_vllm/image_edit/__init__.py b/litellm/llms/hosted_vllm/image_edit/__init__.py new file mode 100644 index 00000000000..27e005e8a0d --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/__init__.py @@ -0,0 +1,9 @@ +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + +from .transformation import HostedVLLMImageEditConfig + +__all__ = ("HostedVLLMImageEditConfig",) + + +def get_hosted_vllm_image_edit_config(model: str) -> BaseImageEditConfig: + return HostedVLLMImageEditConfig() diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py new file mode 100644 index 00000000000..3b8cc437168 --- /dev/null +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -0,0 +1,43 @@ +from typing import Final + +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig +from litellm.secret_managers.main import get_secret_str + +PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT: Final = frozenset({"mask", "quality", "input_fidelity"}) + + +class HostedVLLMImageEditConfig(OpenAIImageEditConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseImageEditConfig contract + return [ # mutable-ok: BaseImageEditConfig returns list + param + for param in super().get_supported_openai_params(model) + if param not in PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT + ] + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseImageEditConfig contract + model: str, + api_key: str | None = None, + litellm_params: dict | None = None, # mutable-ok: BaseImageEditConfig contract + api_base: str | None = None, + ) -> dict: # mutable-ok: BaseImageEditConfig contract + resolved_key: Final = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" + return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseImageEditConfig contract + ) -> str: + resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE") + if resolved_api_base is None: + raise ValueError( + "api_base not set for Hosted VLLM images edits API. " + "Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable" + ) + trimmed: Final = resolved_api_base.rstrip("/") + if trimmed.endswith("/v1"): + return f"{trimmed}/images/edits" + return f"{trimmed}/v1/images/edits" diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9afc6331d96..9b410cf073e 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -19,10 +19,12 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, drop_tool_reference_parts_from_tool_messages, + flatten_combinators_and_drop_non_python_regex_patterns, get_tool_call_names, hoist_images_from_tool_messages, - tool_with_flattened_parameters, + tool_with_sanitized_parameters, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, @@ -432,7 +434,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): custom_llm_provider, api_base ) - def _flattened_tools_update_for_openai( + def _sanitized_tools_update_for_openai( self, optional_params: Mapping[str, object], litellm_params: Mapping[str, object], @@ -440,22 +442,26 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ OpenAI's chat completions validator rejects tool `parameters` carrying 'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every - model family, unlike the Responses API, where GPT-5+ accepts them. + model family, unlike the Responses API, where GPT-5+ accepts them, and + a `pattern` Python's `re` cannot compile for every model family on both. + A custom api_base on the `openai` provider is usually a proxy in front of + the same validator, so regexes are dropped there too, while the lossier + combinator flattening stays limited to api.openai.com hosts. """ tools: Final = optional_params.get("tools") - if not isinstance(tools, list): - return _NO_TOOLS_UPDATE provider: Final = litellm_params.get("custom_llm_provider") - raw_api_base: Final = litellm_params.get("api_base") - if not self._targets_openai_hosted_endpoint( - provider if isinstance(provider, str) else None, - raw_api_base if isinstance(raw_api_base, str) else None, - ): + if not isinstance(tools, list) or provider != "openai": return _NO_TOOLS_UPDATE - flattened: Final = [ # mutable-ok: request tools are a JSON list - tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + raw_api_base: Final = litellm_params.get("api_base") + sanitize: Final = ( + flatten_combinators_and_drop_non_python_regex_patterns + if self._targets_openai_hosted_endpoint(provider, raw_api_base if isinstance(raw_api_base, str) else None) + else drop_non_python_regex_patterns + ) + sanitized: Final = [ # mutable-ok: request tools are a JSON list + tool_with_sanitized_parameters(tool, sanitize) if isinstance(tool, dict) else tool for tool in tools ] - return MappingProxyType({"tools": flattened}) + return MappingProxyType({"tools": sanitized}) def transform_request( self, @@ -489,7 +495,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": messages, **optional_params, - **self._flattened_tools_update_for_openai(optional_params, litellm_params), + **self._sanitized_tools_update_for_openai(optional_params, litellm_params), } async def async_transform_request( @@ -521,7 +527,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "model": model, "messages": transformed_messages, **optional_params, - **self._flattened_tools_update_for_openai(optional_params, litellm_params), + **self._sanitized_tools_update_for_openai(optional_params, litellm_params), } else: ## allow for any object specific behaviour to be handled diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 80292aef2cf..58ff03e6a0d 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import coerce_stream_holdback_value, ) from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + ChatCompletionMessageToolCall, Choices, GenericGuardrailAPIInputs, ModelResponse, @@ -78,7 +80,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True + assembles_streamed_response = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ @@ -610,13 +613,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): deliver_ended_stream_rewrites: bool, ) -> None: """Ended-stream path: rebuild the full response, run the non-streaming - output guardrail against it, and (when opted in) write any text rewrite - back across the buffered chunks.""" + output guardrail against it, and (when opted in) write any text or + tool-call rewrite back across the buffered chunks.""" model_response: Final = cast( ModelResponse, stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), ) pre_guardrail_texts: Final = self._string_choice_contents(model_response) + pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) await self.process_output_response( response=model_response, guardrail_to_apply=guardrail_to_apply, @@ -624,13 +628,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation): user_api_key_dict=user_api_key_dict, request_data=request_data, ) - if deliver_ended_stream_rewrites: - await self._write_ended_stream_text_rewrites( - responses_so_far=responses_so_far, - guardrailed_response=model_response, - pre_guardrail_texts=pre_guardrail_texts, - guardrail_name=guardrail_to_apply.guardrail_name or "unknown", - ) + if not deliver_ended_stream_rewrites: + return + guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown" + await self._write_ended_stream_text_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_texts=pre_guardrail_texts, + guardrail_name=guardrail_name, + ) + self._write_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_name, + ) def build_stream_error_items( self, @@ -1043,6 +1055,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists ) + @staticmethod + def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]: + return tuple( + (tool_call.function.name, tool_call.function.arguments) + for choice in response.choices + for tool_call in choice.message.tool_calls or () + if isinstance(tool_call, ChatCompletionMessageToolCall) + ) + + @staticmethod + def _function_tool_call_fragments( + responses_so_far: Sequence["ModelResponseStream"], + ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: + """Group the stream's function tool-call fragments by their tool-call index, in + the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping + only the indices the builder keeps (an id and a name somewhere in the stream).""" + fragments: Final = tuple( + tool_call + for response in responses_so_far + for choice in response.choices + for tool_call in choice.delta.tool_calls or () + if isinstance(tool_call, ChatCompletionDeltaToolCall) + ) + identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) + named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) + return tuple( + tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) + ) + + def _write_ended_stream_tool_call_rewrites( + self, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrailed_response: "ModelResponse", + pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites back across the buffered + chunks: the rewritten name and full arguments land in the tool call's first + fragment and the arguments of its later fragments are blanked, mirroring the + text write-back. A rewrite on a stream carrying more than one distinct choice + index, or whose fragments do not line up with the rebuilt tool calls, is + reported as undeliverable, so the pipeline executor discards it and releases + the original chunks.""" + post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) + fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) + if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + for before, (name, arguments), fragments in zip( + pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call + ): + if (name, arguments) == before: + continue + head, *tail = fragments + head.function.name = name + head.function.arguments = arguments + for fragment in tail: + fragment.function.arguments = "" + async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b0f79552bc5..33e0a0a923f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -37,14 +37,14 @@ from itertools import accumulate, chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast -from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, OpenAiResponsesToChatCompletionStreamIterator, + tool_call_dict_from_output_item, ) from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -84,7 +84,6 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import ( GenericResponseOutputItem, - OutputFunctionToolCall, OutputText, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -101,6 +100,72 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ResponseInputParam +class _ToolCallShape(NamedTuple): + name: str | None + arguments: str + + +class _ToolCallFunctionFields(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str | None = None + arguments: str = "" + + +class _ToolCallFields(BaseModel): + model_config = ConfigDict(frozen=True) + + function: _ToolCallFunctionFields + + +def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]: + return tuple( + _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", "")) + for tool_call in tool_calls + ) + + +def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None: + payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call + try: + fields: Final = _ToolCallFields.model_validate(payload) + except ValidationError: + return None + return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments) + + +def _post_guardrail_tool_call_shapes( + returned_tool_calls: Sequence[object] | None, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str | None, +) -> tuple[_ToolCallShape, ...]: + if not pre_guardrail_tool_calls: + return pre_guardrail_tool_calls + if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, " + "leaving the tool call output items unchanged", + guardrail_name, + "no" if returned_tool_calls is None else len(returned_tool_calls), + len(pre_guardrail_tool_calls), + ) + return pre_guardrail_tool_calls + returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls) + validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None) + if len(validated_shapes) != len(returned_shapes): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, " + "leaving the tool call output items unchanged", + guardrail_name, + ) + return pre_guardrail_tool_calls + return validated_shapes + + +def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape: + return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments) + + class ResponseOutputEnvelope(TypedDict, total=False): """Dict form of a Responses API response, as far as guardrail write-back reads it.""" @@ -128,6 +193,20 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( ) +_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"}) +_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"function_call": "arguments", "custom_tool_call": "input"} +) +_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset( + {"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"} +) +_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"} +) +_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset( + _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS +) +_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( {"function_call_output": "output", "message": "content"} ) @@ -164,8 +243,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts -def _is_function_call_item(item: object) -> bool: - return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call") +def _is_tool_call_item(item: object) -> bool: + return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES + + +def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None: + if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES: + return None + if isinstance(item, Mapping): + return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects + return item.model_dump() if isinstance(item, BaseModel) else None + + +def _is_tool_call_output_item(item: object) -> bool: + return _tool_call_output_item_mapping(item) is not None def _last_message_role(messages: Sequence[object]) -> str | None: @@ -189,7 +280,7 @@ def _provenance_unit_bounds( start_indexes: Final = tuple( index for index in range(len(raw_input)) - if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") ) return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) @@ -340,7 +431,8 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True + assembles_streamed_response = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ @@ -587,7 +679,7 @@ class OpenAIResponsesHandler(BaseTranslation): - response.output is a list of output items - Each output item can be: * GenericResponseOutputItem with a content list of OutputText objects - * ResponseFunctionToolCall with tool call data + * ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data - Each OutputText object has a text field """ @@ -652,6 +744,7 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -660,6 +753,11 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, + ) # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( @@ -667,6 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation): responses=guardrailed_texts, task_mappings=task_mappings, ) + self._write_tool_call_rewrites_to_output( + tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=post_guardrail_tool_calls, + ) verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response) @@ -754,6 +857,7 @@ class OpenAIResponsesHandler(BaseTranslation): if response_model: inputs["model"] = response_model + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, @@ -762,6 +866,11 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, + ) # Write guardrailed texts back into the output items in-place. # final_chunk is a reference into responses_so_far so this @@ -784,6 +893,13 @@ class OpenAIResponsesHandler(BaseTranslation): stream_events=responses_so_far[:-1], rewrites_by_position=rewrites_by_position, ) + self._deliver_ended_stream_tool_call_rewrites( + responses_so_far=responses_so_far, + outputs=outputs, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + post_guardrail_tool_calls=post_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", + ) return responses_so_far # ------------------------------------------------------------------ # @@ -894,6 +1010,148 @@ class OpenAIResponsesHandler(BaseTranslation): continue OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) + def _deliver_ended_stream_tool_call_rewrites( + self, + responses_so_far: Sequence[object], + outputs: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str, + ) -> None: + """Write ended-stream guardrail tool-call rewrites into the completed + envelope's ``function_call`` and ``custom_tool_call`` items and sync the + earlier stream events, keyed by ``call_id``. The guardrail sees the + envelope's tool calls in output order, which is how a rewritten call + finds its ``call_id``; the stream events find their call through the + ``call_id`` on ``output_item`` events and the ``item_id`` on argument + and custom-input events, since an + event's ``output_index`` need not match the envelope's (the chat bridge + numbers tool calls from 1 while the envelope lists them after the + message). A rewrite whose calls do not line up with the envelope, or + whose events cannot be found, is reported as undeliverable, so the + pipeline executor discards it and releases the original events.""" + if post_guardrail_tool_calls == pre_guardrail_tool_calls: + return + tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item)) + call_ids: Final = tuple( + call_id + for output_item in tool_call_items + if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id + ) + stream_events: Final = responses_so_far[:-1] + call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events) + event_call_ids: Final = tuple( + self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events + ) + rewrites_by_call_id: Final = MappingProxyType( + { + call_id: _tool_call_rewrite(before, after) + for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + } + ) + unresolved_argument_event: Final = any( + call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES + for event, call_id in zip(stream_events, event_call_ids) + ) + if ( + len(call_ids) != len(tool_call_items) + or len(frozenset(call_ids)) != len(call_ids) + or len(call_ids) != len(post_guardrail_tool_calls) + or unresolved_argument_event + or not rewrites_by_call_id.keys() <= frozenset(event_call_ids) + ): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + for output_item, rewrite in ( + (output_item, rewrites_by_call_id[call_id]) + for output_item, call_id in zip(tool_call_items, call_ids) + if call_id in rewrites_by_call_id + ): + self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments) + delta_replacements: Final = MappingProxyType( + {call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()} + ) + for event, call_id in zip(stream_events, event_call_ids): + if call_id not in rewrites_by_call_id: + continue + match stream_item_field(event, "type"): + case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: + self._write_event_field(event, "delta", next(delta_replacements[call_id])) + case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: + self._write_event_field( + event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments + ) + case "response.output_item.added": + self._write_tool_call_item( + stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None + ) + case "response.output_item.done": + self._write_tool_call_item( + stream_item_field(event, "item"), + rewrites_by_call_id[call_id].name, + rewrites_by_call_id[call_id].arguments, + ) + case _: + pass + + def _write_tool_call_rewrites_to_output( + self, + tool_call_items: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + post_guardrail_tool_calls: tuple[_ToolCallShape, ...], + ) -> None: + if len(tool_call_items) != len(post_guardrail_tool_calls): + return + for output_item, rewrite in ( + (output_item, _tool_call_rewrite(before, after)) + for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if after != before + ): + self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments) + + @staticmethod + def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: + items: Final = tuple( + stream_item_field(event, "item") + for event in stream_events + if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES + ) + return MappingProxyType( + { + item_id: call_id + for item in items + if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES + and isinstance(item_id := stream_item_field(item, "id"), str) + and isinstance(call_id := stream_item_field(item, "call_id"), str) + } + ) + + @staticmethod + def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None: + event_type: Final = stream_item_field(event, "type") + if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES: + item_id: Final = stream_item_field(event, "item_id") + return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None + if event_type not in _OUTPUT_ITEM_EVENT_TYPES: + return None + item: Final = stream_item_field(event, "item") + call_id: Final = stream_item_field(item, "call_id") + return ( + call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None + ) + + @staticmethod + def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None: + if item is None: + return + if name is not None: + OpenAIResponsesHandler._write_event_field(item, "name", name) + item_type: Final = stream_item_field(item, "type") + if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS: + OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload) + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. @@ -920,7 +1178,7 @@ class OpenAIResponsesHandler(BaseTranslation): def _completed_response_scan_key(response: object) -> StreamingScanKey: output_items: Final = stream_item_items(response, "output") message_items: Final = tuple( - item for item in output_items if stream_item_field(item, "type") != "function_call" + item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES ) return StreamingScanKey( texts=tuple( @@ -932,7 +1190,7 @@ class OpenAIResponsesHandler(BaseTranslation): tool_calls=tuple( stream_item_fingerprint(item) for item in output_items - if stream_item_field(item, "type") == "function_call" + if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES ), stream_ended=True, ) @@ -1043,34 +1301,10 @@ class OpenAIResponsesHandler(BaseTranslation): Override this method to customize text/image/tool extraction logic. """ - # Check if this is a tool call (OutputFunctionToolCall) - if isinstance(output_item, OutputFunctionToolCall) or ( - isinstance(output_item, BaseModel) - and hasattr(output_item, "type") - and getattr(output_item, "type") == "function_call" - ): + tool_call_item: Final = _tool_call_output_item_mapping(output_item) + if tool_call_item is not None: if tool_calls_to_check is not None: - tool_call_dict = ( - LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=output_item, - index=output_idx, - ) - ) - tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) - return - elif isinstance(output_item, dict) and output_item.get("type") == "function_call": - # Handle dict representation of tool call - if tool_calls_to_check is not None: - # Convert dict to ResponseFunctionToolCall for processing - try: - tool_call_obj: Final = ResponseFunctionToolCall(**output_item) - tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=tool_call_obj, - index=output_idx, - ) - tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) - except Exception: - pass + tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx)) return # Handle both GenericResponseOutputItem and dict diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 926de3e8854..666e9b32011 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints @@ -15,6 +15,10 @@ from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + flatten_combinators_and_drop_non_python_regex_patterns, +) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -40,7 +44,7 @@ else: _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4") -_PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) +_PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) @@ -293,7 +297,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): model=model, input=validated_input, tools=tools ) object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools) - sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( + sanitized_tools: Final = self._sanitized_tool_schemas_for_openai( model=model, tools=object_schema_tools, litellm_params=litellm_params ) return self._drop_foreign_tool_call_item_ids(stripped_input), sanitized_tools @@ -378,35 +382,35 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return item return {key: value for key, value in item.items() if key != "id"} # mutable-ok: outgoing JSON request item - def _flatten_tool_schema_combinators_for_openai( + def _sanitized_tool_schemas_for_openai( self, model: str, tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, litellm_params: GenericLiteLLMParams, ) -> Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None: - """Flatten top-level schema combinators only where OpenAI's validator rejects them. + """Rewrite tool schemas only where OpenAI's validator rejects them. - OpenAI-compatible backends reusing this config (and the ChatGPT backend - Codex talks to natively) accept them, and so do GPT-5 and later models, - which also call tools better with the union intact. Codex wraps MCP tools - inside namespace entries, so nested ``tools`` arrays are walked too. - Azure OpenAI shares the validator but names deployments arbitrarily, so - the router's declared ``model_info.base_model`` wins over the deployment - name and an unrecognized name without one is left untouched. + Every model family refuses a ``pattern`` Python's ``re`` cannot compile, + while top-level schema combinators are flattened only for the families + whose validator rejects them: OpenAI-compatible backends reusing this + config (and the ChatGPT backend Codex talks to natively) accept them, + and so do GPT-5 and later models, which also call tools better with the + union intact. Codex wraps MCP tools inside namespace entries, so nested + ``tools`` arrays are walked too. Azure OpenAI shares the validator but + names deployments arbitrarily, so the router's declared + ``model_info.base_model`` wins over the deployment name and an + unrecognized name without one keeps its combinators. """ - if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: + if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: return tools gate_model: Final = self._combinator_gate_model(model=model, litellm_params=litellm_params) - if not self._rejects_top_level_schema_combinators(gate_model): - return tools - flattened: Final = [ # mutable-ok: request tools are a JSON list - self._flattened_tool_or_passthrough(tool) for tool in tools - ] - return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", flattened) # cast-ok: spread keeps each tool's shape - - @staticmethod - def _flattened_tool_or_passthrough(tool: object) -> object: - return OpenAIResponsesAPIConfig._flattened_tool_entry(tool) if isinstance(tool, dict) else tool + sanitize: Final = ( + flatten_combinators_and_drop_non_python_regex_patterns + if self._rejects_top_level_schema_combinators(gate_model) + else drop_non_python_regex_patterns + ) + sanitized: Final = self._sanitized_tools(tools, sanitize) + return cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", sanitized) # cast-ok: spread keeps each tool's shape @staticmethod def _rejects_top_level_schema_combinators(model: str) -> bool: @@ -421,35 +425,42 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return base_model if isinstance(base_model, str) and base_model else model @staticmethod - def _flattened_tool_entry( + def _sanitized_tool_entry( entry: Mapping[str, object], - ) -> dict[str, object]: # mutable-ok: request tools are JSON dicts - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - flatten_top_level_schema_combinators, - ) - + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], + ) -> Mapping[str, object]: parameters: Final = entry.get("parameters") nested_tools: Final = entry.get("tools") + sanitized_parameters: Final = sanitize(parameters) if isinstance(parameters, dict) else parameters + sanitized_nested_tools: Final = ( + OpenAIResponsesAPIConfig._sanitized_tools(nested_tools, sanitize) + if isinstance(nested_tools, list) + else nested_tools + ) parameters_update: Final = ( - MappingProxyType({"parameters": flatten_top_level_schema_combinators(parameters)}) - if isinstance(parameters, dict) + MappingProxyType({"parameters": sanitized_parameters}) + if sanitized_parameters is not parameters else _NO_TOOL_UPDATE ) tools_update: Final = ( - MappingProxyType({"tools": OpenAIResponsesAPIConfig._flattened_nested_tools(nested_tools)}) - if isinstance(nested_tools, list) + MappingProxyType({"tools": sanitized_nested_tools}) + if sanitized_nested_tools is not nested_tools else _NO_TOOL_UPDATE ) + if not parameters_update and not tools_update: + return entry return {**entry, **parameters_update, **tools_update} # mutable-ok: request tools are JSON dicts @staticmethod - def _flattened_nested_tools( - nested_tools: Sequence[object], - ) -> list[object]: # mutable-ok: namespace tools are a JSON list - return [ # mutable-ok: namespace tools are a JSON list - OpenAIResponsesAPIConfig._flattened_tool_entry(item) if isinstance(item, dict) else item - for item in nested_tools + def _sanitized_tools( + tools: Sequence[object], + sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], + ) -> Sequence[object]: + sanitized: Final = [ # mutable-ok: request tools are a JSON list + OpenAIResponsesAPIConfig._sanitized_tool_entry(item, sanitize) if isinstance(item, dict) else item + for item in tools ] + return tools if all(new is old for new, old in zip(sanitized, tools, strict=True)) else sanitized def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam: """ @@ -620,15 +631,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model.model_construct(**parsed_chunk) @staticmethod - def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + def parse_terminal_event_from_stream_chunks(all_chunks: Sequence[str]) -> ResponsesTerminalEvent | None: for chunk_str in reversed(all_chunks): for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent): try: - return event_model.model_validate_json(chunk_str.removeprefix("data: ")).response + return event_model.model_validate_json(chunk_str.removeprefix("data: ")) except ValueError: continue return None + @staticmethod + def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks) + return None if terminal_event is None else terminal_event.response + @staticmethod def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]: """ diff --git a/litellm/main.py b/litellm/main.py index 04b0963851a..4f1c57adb53 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -19,7 +19,7 @@ import random import sys import time import traceback -from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy @@ -5398,6 +5398,14 @@ def completion( if dynamic_api_key is not None: api_key = dynamic_api_key # check if user passed in any of the OpenAI optional params + bridges_to_responses_api: Final = ( + responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge + ) + allowed_openai_params: Final[list[str] | None] = ( + [*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"] + if bridges_to_responses_api + else kwargs.get("allowed_openai_params") + ) optional_param_args: Final = { "functions": functions, "function_call": function_call, @@ -5442,7 +5450,7 @@ def completion( "service_tier": service_tier, "store": store, "prompt_cache_key": prompt_cache_key, - "allowed_openai_params": kwargs.get("allowed_openai_params"), + "allowed_openai_params": allowed_openai_params, "base_model": base_model, } optional_params = get_optional_params(**optional_param_args, **non_default_params) @@ -8587,7 +8595,7 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion(chunks: list, messages: list | None = None) -> TextCompletionResponse: +def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse: id: Final = chunks[0]["id"] object: Final = chunks[0]["object"] created: Final = chunks[0]["created"] @@ -8704,10 +8712,11 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o def stream_chunk_builder( chunks: list, - messages: list | None = None, + messages: Sequence | None = None, start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, + count_prompt_tokens: Callable[[], int] | None = None, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -8781,6 +8790,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=0, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) @@ -8958,6 +8968,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, + count_prompt_tokens=count_prompt_tokens, ) setattr(response, "usage", usage) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 54ebdc85be9..0d2eda93323 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -364,7 +364,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -380,6 +381,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -399,6 +401,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -416,6 +419,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -435,6 +439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -452,6 +457,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -471,6 +477,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -488,6 +495,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -507,6 +515,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -537,7 +546,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -551,7 +561,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -756,6 +767,14 @@ "mode": "chat", "supports_video_input": true }, + "global.twelvelabs.pegasus-1-2-v1:0": { + "input_cost_per_video_per_second": 0.00049, + "output_cost_per_token": 7.5e-06, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_video_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "amazon.titan-text-express-v1": { "input_cost_per_token": 1.3e-06, "litellm_provider": "bedrock", @@ -2876,7 +2895,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.7e-08, @@ -2888,7 +2908,8 @@ "output_cost_per_token": 1.48e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "apac.amazon.nova-pro-v1:0": { "input_cost_per_token": 8.4e-07, @@ -2902,7 +2923,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.anthropic.claude-3-5-sonnet-20240620-v1:0": { "deprecation_date": "2026-07-30", @@ -8064,7 +8086,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8108,7 +8131,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8152,7 +8176,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -12317,7 +12342,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12496,7 +12522,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "input_cost_per_token": 4.2e-08, @@ -12508,7 +12535,8 @@ "output_cost_per_token": 1.68e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, @@ -12522,7 +12550,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -13973,7 +14002,8 @@ "max_output_tokens": 3072, "max_tokens": 3072, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/meta/llama-2-7b-chat-int8": { "input_cost_per_token": 1.923e-06, @@ -13982,7 +14012,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.1": { "input_cost_per_token": 1.923e-06, @@ -13991,7 +14022,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@hf/thebloke/codellama-7b-instruct-awq": { "input_cost_per_token": 1.923e-06, @@ -14000,7 +14032,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/openai/gpt-oss-120b": { "input_cost_per_token": 3.5e-07, @@ -14010,6 +14043,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14020,7 +14054,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-3b-instruct": { "input_cost_per_token": 5.09e-08, @@ -14029,7 +14064,8 @@ "max_output_tokens": 80000, "max_tokens": 80000, "mode": "chat", - "output_cost_per_token": 3.35e-07 + "output_cost_per_token": 3.35e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-guard-3-8b": { "input_cost_per_token": 4.84e-07, @@ -14038,7 +14074,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-08 + "output_cost_per_token": 3e-08, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { "input_cost_per_token": 0.0, @@ -14047,7 +14084,8 @@ "max_output_tokens": 15000, "max_tokens": 15000, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 1.9e-07, @@ -14058,6 +14096,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -14069,6 +14108,7 @@ "max_tokens": 80000, "mode": "chat", "output_cost_per_token": 4.881e-06, + "rpm": 300, "supports_reasoning": true }, "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { @@ -14078,7 +14118,8 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 2.87e-07 + "output_cost_per_token": 2.87e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-1b-instruct": { "input_cost_per_token": 2.7e-08, @@ -14087,7 +14128,8 @@ "max_output_tokens": 60000, "max_tokens": 60000, "mode": "chat", - "output_cost_per_token": 2.01e-07 + "output_cost_per_token": 2.01e-07, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.6": { "cache_read_input_token_cost": 1.6e-07, @@ -14098,6 +14140,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -14109,6 +14152,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14119,7 +14163,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { "input_cost_per_token": 2.93e-07, @@ -14129,6 +14174,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 2.253e-06, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { @@ -14139,6 +14185,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.12e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { @@ -14148,7 +14195,8 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-06 + "output_cost_per_token": 1e-06, + "rpm": 300 }, "cloudflare/@cf/zai-org/glm-5.2": { "cache_read_input_token_cost": 2.6e-07, @@ -14159,6 +14207,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4.4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -14170,6 +14219,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14180,7 +14230,8 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.55e-07 + "output_cost_per_token": 5.55e-07, + "rpm": 300 }, "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { "input_cost_per_token": 5.09e-08, @@ -14190,6 +14241,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3.35e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14200,7 +14252,8 @@ "max_output_tokens": 3500, "max_tokens": 3500, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/google/gemma-4-26b-a4b-it": { "input_cost_per_token": 1e-07, @@ -14210,6 +14263,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14221,6 +14275,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5.55e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { @@ -14231,6 +14286,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.76e-07, + "rpm": 300, "supports_vision": true }, "cloudflare/@cf/openai/gpt-oss-20b": { @@ -14241,6 +14297,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14252,6 +14309,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 8.5e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwq-32b": { @@ -14262,6 +14320,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 1e-06, + "rpm": 300, "supports_reasoning": true }, "codestral/codestral-2405": { @@ -14388,6 +14447,28 @@ "output_vector_size": 1536, "supports_embedding_image_input": true }, + "us.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, + "global.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "cohere", @@ -20959,7 +21040,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.amazon.nova-micro-v1:0": { "input_cost_per_token": 4.6e-08, @@ -20971,7 +21053,8 @@ "output_cost_per_token": 1.84e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "eu.amazon.nova-pro-v1:0": { "input_cost_per_token": 1.05e-06, @@ -20986,24 +21069,25 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.anthropic.claude-3-5-haiku-20241022-v1:0": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 8e-07, "litellm_provider": "bedrock", "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 4e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 2.5e-08, - "cache_creation_input_token_cost": 3.125e-07, + "cache_read_input_token_cost": 8e-08, + "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -24010,9 +24094,9 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, @@ -24035,12 +24119,12 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_url_context": false, "supports_vision": true, "supports_web_search": true, "search_context_cost_per_query": { @@ -27832,6 +27916,70 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.5": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.5", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-luna": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-luna", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-sol": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-sol", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-terra": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-terra", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", "max_input_tokens": 1050000, @@ -28438,6 +28586,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -29262,7 +29411,12 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + } }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29279,9 +29433,9 @@ "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29380,9 +29534,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29407,9 +29561,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29520,9 +29674,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29547,9 +29701,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29631,6 +29785,66 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2.5-flare": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-flare-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, "low/1024-x-1024/gpt-image-1.5": { "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, @@ -41711,6 +41925,28 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "rerank-v4.0-fast": { + "input_cost_per_query": 0.002, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, + "rerank-v4.0-pro": { + "input_cost_per_query": 0.0025, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -43671,7 +43907,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -43683,7 +43920,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -43712,7 +43950,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, @@ -45891,8 +46130,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -45916,8 +46155,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -47897,9 +48136,9 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47913,9 +48152,9 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47928,14 +48167,17 @@ }, "vertex_ai/xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -47944,14 +48186,17 @@ }, "vertex_ai/xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -47959,6 +48204,44 @@ "supports_vision": true, "supports_web_search": true }, + "vertex_ai/xai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "vertex_ai/xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -48217,6 +48500,16 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-multilingual-2": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-3-large": { "input_cost_per_token": 1.8e-07, "litellm_provider": "voyage", @@ -52290,7 +52583,8 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-07, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -52329,7 +52623,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-01" }, "scaleway/hcompany/holo2-30b-a3b": { "input_cost_per_token": 3e-07, @@ -52340,7 +52635,8 @@ "mode": "chat", "output_cost_per_token": 7e-07, "supports_reasoning": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-09" }, "scaleway/mistralai/mistral-medium-3.5-128b": { "input_cost_per_token": 1.5e-06, @@ -52363,7 +52659,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2e-06, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/voxtral-small-24b-2507": { "input_cost_per_audio_token": 1.5e-07, @@ -52374,7 +52671,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.5e-07, - "supports_audio_input": true + "supports_audio_input": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": { "input_cost_per_token": 1.5e-07, @@ -52396,7 +52694,8 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_vision": true, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/BAAI/bge-multilingual-gemma2": { "input_cost_per_token": 1e-07, @@ -54846,7 +55145,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "output_cost_per_token": 3e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, @@ -54880,8 +55179,8 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 65536, @@ -55001,6 +55300,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, @@ -55196,6 +55528,96 @@ "supports_reasoning": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "us.openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "global.openai.gpt-6-astra": { + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, @@ -56861,8 +57283,8 @@ "rpm": 10 }, "vertex_ai/gemini-3.5-transcribe-preview": { - "input_cost_per_audio_token": 2.5e-06, - "input_cost_per_token": 2.5e-06, + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai", "mode": "audio_transcription", "output_cost_per_token": 1.2e-05, @@ -56897,6 +57319,27 @@ ], "supports_audio_input": true }, + "vertex_ai/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", @@ -59588,6 +60031,77 @@ "image" ] }, + "xai/grok-imagine-video": { + "input_cost_per_image": 0.002, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_720p": 0.07, + "source": "https://docs.x.ai/docs/models/grok-imagine-video", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-2026-05-30": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-preview": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", @@ -60855,6 +61369,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60865,6 +61380,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60920,6 +61436,31 @@ "supports_web_search": false, "output_cost_per_image": 0.08 }, + "gemini/lyria-3.5": { + "input_cost_per_token": 0, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 0, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": false, + "supports_web_search": false, + "output_cost_per_image": 0.08 + }, "perplexity/anthropic/claude-fable-5": { "litellm_provider": "perplexity", "mode": "responses", diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 4bac65125b4..7e6de474b0b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -179,16 +179,7 @@ def _gateway_dcr_challenge_target( mcp_servers: list[str] | None, client_ip: str | None, ) -> str | None: - """The single path-named server this request targets, iff it resolves to a - gateway-managed oauth2 server — the one per-server shape the gateway's own keyless - DCR flow serves end to end, so the 401 challenge may advertise the per-server - protected-resource metadata (whose ``authorization_servers`` names the gateway). - - Multi-server CSV paths, header/path mismatches, unknown names, and every - client-forwarded or delegated mode return ``None``: those cells keep their existing - challenge (or absence of one), and a challenge is never emitted for a name the - public discovery routes would 404, so this reveals exactly the server set the - per-server protected-resource metadata already reveals.""" + """Resolve a single path target whose sign-in metadata advertises the gateway.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -217,7 +208,7 @@ def _is_gateway_dcr_challenge_scope( the caller is not a cold-start DCR client), on the scopes the gateway's keyless flow serves: the aggregate ``/mcp`` endpoint, an ``x-mcp-servers``-scoped request (the resource the client configured is still ``/mcp``), or a per-server path whose - single target is a gateway-managed oauth2 server. Every other named target keeps + single target advertises gateway-owned sign-in. Every other named target keeps its existing behavior, failing closed to the original admission error.""" if not _is_litellm_auth_admission_error(exc): return False @@ -236,7 +227,7 @@ def _gateway_dcr_challenge( ) -> HTTPException: """The RFC 9728 challenge pointing the client at the protected-resource metadata matching the scope it requested: the per-server document (same URL spelling the - request arrived on) when the single target is a gateway-managed oauth2 server, + request arrived on) when the single target advertises gateway-owned sign-in, else the gateway's aggregate document. Either way the client discovers the gateway as its authorization server and starts the same sign-in flow. diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index cab4b6c161a..f18acb7d88d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2310,8 +2310,7 @@ async def _build_oauth_protected_resource_response( it. Only the legacy ``is_oauth_passthrough`` opt-in rewrites ``resource`` to the gateway's own URL so clients present the bearer token back to the gateway. - An explicitly named gateway-managed oauth2 server (interactive with - gateway-vaulted per-user tokens, or M2M) advertises the gateway's own + An explicitly named server with gateway-owned sign-in advertises the gateway's own authorization server (``{base}/mcp``): a keyless DCR client that configured the per-server URL completes the same sign-in flow the aggregate ``/mcp`` endpoint supports and is admitted with a gateway session bearer. The per-server relay @@ -2401,11 +2400,6 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response - # An OBO server with no configured issuer falls through to the gateway default so discovery still - # returns metadata; every other non-oauth2 named server 404s to avoid enumeration. - if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: - _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource") - if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], @@ -2413,6 +2407,9 @@ async def _build_oauth_protected_resource_response( "scopes_supported": (mcp_server.scopes if mcp_server.scopes else []), } + if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: + _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource") + return { "authorization_servers": [ (f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}") diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 3d94fa345d0..f4889008e94 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -411,7 +411,7 @@ def relative_request_url(request: Request) -> str: def resolve_scoped_resource_server(request: Request, resource: str | None) -> MCPServer | None: - """Resolve an RFC 8707 ``resource`` value to the single gateway-managed oauth2 server it + """Resolve an RFC 8707 ``resource`` value to the single gateway-owned server it names, or ``None`` for every other shape: absent, the aggregate resource, a foreign host, an unparseable value, a multi-server path, an unknown name, or any server mode the keyless gateway flow does not serve (whose protected-resource metadata never directs a @@ -443,7 +443,7 @@ def resolve_scoped_resource_server(request: Request, resource: str | None) -> MC if len(names) != 1: return None server: Final = global_mcp_server_manager.get_mcp_server_by_name(names[0]) - if server is None or not server.is_gateway_managed_oauth2: + if server is None or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server): return None return server @@ -729,11 +729,15 @@ async def _flow_target( server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) if ( server is None - or not server.is_gateway_managed_oauth2 + or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server) or not await lookup_server_reachability(flow.user_id, server.server_id) ): return "stale", None - state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive" + state: Final = ( + "interactive" + if server.is_gateway_managed_oauth2 and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" + else "m2m" + ) return state, server diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 6aaa00ae415..ff30ef99ebd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -22,7 +22,16 @@ Response headers returned (all values are masked for safety): x-mcp-debug-auth-resolution Which auth priority was used for the outbound MCP call: ``per-request-header``, ``m2m-client-credentials``, ``static-token``, - ``oauth2-passthrough``, or ``no-auth``. + ``oauth2-passthrough``, ``stored-user-token``, ``token-exchange``, + ``id-jag``, ``aws-sigv4``, ``extra-headers``, or ``no-auth``. + ``unresolved`` means no outcome was available before the first response + frame; ``multiple`` means several servers resolved credentials; + ``not-applicable`` covers stdio; ``resolution-failed`` is a resolver error. + + x-mcp-debug-auth-resolutions + For multiple servers, a JSON map of server IDs to resolution labels. + At most 32 entries are included; x-mcp-debug-auth-resolutions-truncated + is true when additional servers were omitted. No credentials are included. x-mcp-debug-outbound-url The upstream MCP server URL that will receive the request. @@ -58,10 +67,16 @@ header is free for OAuth2 discovery:: Symptom: ``x-mcp-debug-oauth2-token`` shows ``(none)`` and ``x-mcp-debug-auth-resolution`` shows ``no-auth``. -This means the client didn't go through the OAuth2 flow. Check that: -1. The ``Authorization`` header is NOT set as a static header in the client config. -2. The ``.well-known/oauth-protected-resource`` endpoint returns valid metadata. -3. The MCP server in LiteLLM config has ``auth_type: oauth2``. +``no-auth`` means the resolved upstream client carries no authentication. +An absent inbound OAuth2 token does not imply the user skipped OAuth: the gateway +can retrieve a stored per-user token, reported as ``stored-user-token``. +``unresolved`` is used when a stream starts before credential resolution, or a +request (such as initialization or a cached tool listing) resolves no credential. +Debug reporting does not fetch credentials or delay a streaming frame to resolve them. +``extra-headers`` identifies supplied headers that won over the resolver or were +the only headers supplied; their values are never inspected to guess a scheme. +``per-request-header`` denotes a legacy credential override, including a BYOK +credential supplied by the gateway; it does not imply a caller-supplied token. **Common issue: M2M token used instead of user token** @@ -69,8 +84,8 @@ Symptom: ``x-mcp-debug-auth-resolution`` shows ``m2m-client-credentials``. This means the server has ``client_id``/``client_secret``/``token_url`` configured and LiteLLM is fetching a machine-to-machine token instead of -using the per-user OAuth2 token. If you want per-user tokens, remove the -client credentials from the server config. +using the per-user OAuth2 token. For gateway-stored per-user tokens, +configure ``oauth2_flow: authorization_code``. Usage from Claude Code:: @@ -85,14 +100,16 @@ Usage with curl:: http://localhost:4000/mcp/atlassian_mcp """ -from typing import TYPE_CHECKING, Final +import json +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final +from starlette.requests import HTTPConnection from starlette.types import Message, Send from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker - -if TYPE_CHECKING: - from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution # Header the client sends to opt into debug mode MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug" @@ -101,6 +118,83 @@ MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug" _RESPONSE_HEADER_PREFIX: Final = "x-mcp-debug" +MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: Final = "litellm.mcp.auth_diagnostics" + + +def record_auth_resolution(server_id: str, source: AuthResolution) -> None: + from mcp.server.lowlevel.server import request_ctx + + context: Final[object] = request_ctx.get(None) + request: Final[object] = getattr(context, "request", None) + if isinstance(request, HTTPConnection): + diagnostics: Final[object] = request.scope.get(MCP_AUTH_DIAGNOSTICS_SCOPE_KEY) + if isinstance(diagnostics, MCPAuthDiagnostics): + diagnostics.record(server_id, source) + + +class MCPAuthDiagnostics: + def __init__(self) -> None: + self._outcomes: tuple[tuple[str, AuthResolution], ...] = () + + def record(self, server_id: str, resolution: AuthResolution) -> None: + self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),) + + def resolution(self) -> str: + match self._outcomes: + case (): + return AuthResolution.unresolved.value + case ((_, source),): + return source.value + case _: + return AuthResolution.multiple.value + + def headers(self) -> Mapping[str, str]: + if len(self._outcomes) <= 1: + return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()}) + return MappingProxyType( + { + "x-mcp-debug-auth-resolution": AuthResolution.multiple.value, + "x-mcp-debug-auth-resolutions": json.dumps( + { + server_id: source.value for server_id, source in self._outcomes[:32] + }, # mutable-ok: JSON encoder requires a concrete dict + separators=(",", ":"), + ensure_ascii=True, + ), + **( + MappingProxyType({"x-mcp-debug-auth-resolutions-truncated": "true"}) + if len(self._outcomes) > 32 + else MappingProxyType({}) + ), + } + ) + + +class _DiagnosticSend: + def __init__(self, send: Send, headers: Mapping[str, str], resolution: Callable[[], Mapping[str, str]]) -> None: + self._send = send + self._headers = headers + self._resolution = resolution + self._start: Message | None = None + + async def __call__(self, message: Message) -> None: + if message["type"] == "http.response.start": + self._start = message + return + if self._start is not None: + start: Final = self._start + self._start = None + headers: Final = MappingProxyType({**self._headers, **self._resolution()}) + await self._send( + { # mutable-ok: ASGI send consumes a mutable message mapping + **start, + "headers": tuple(start.get("headers", ())) + + tuple((key.encode(), value.encode()) for key, value in headers.items()), + } + ) + await self._send(message) + + class MCPDebug: """ Static helper class for MCP OAuth2 debug headers. @@ -144,37 +238,6 @@ class MCPDebug: return val.strip().lower() in ("true", "1", "yes") return False - @staticmethod - def resolve_auth_resolution( - server: "MCPServer", - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - oauth2_headers: dict[str, str] | None, - ) -> str: - """ - Determine which auth priority will be used for the outbound MCP call. - - Returns one of: ``per-request-header``, ``m2m-client-credentials``, - ``static-token``, ``oauth2-passthrough``, or ``no-auth``. - """ - from litellm.types.mcp import MCPAuth - - has_server_specific: Final = bool( - mcp_server_auth_headers - and ( - mcp_server_auth_headers.get(server.alias or "") or mcp_server_auth_headers.get(server.server_name or "") - ) - ) - if has_server_specific or mcp_auth_header: - return "per-request-header" - if server.has_client_credentials: - return "m2m-client-credentials" - if server.authentication_token: - return "static-token" - if oauth2_headers and server.auth_type == MCPAuth.oauth2: - return "oauth2-passthrough" - return "no-auth" - @staticmethod def build_debug_headers( *, @@ -244,12 +307,21 @@ class MCPDebug: return debug @staticmethod - def wrap_send_with_debug_headers(send: Send, debug_headers: dict[str, str]) -> Send: + def wrap_send_with_debug_headers( + send: Send, + debug_headers: Mapping[str, str], + resolution: Callable[[], Mapping[str, str]] | None = None, + *, + request_method: str | None = None, + ) -> Send: """ Return a new ASGI ``send`` callable that injects *debug_headers* into the ``http.response.start`` message. """ + if resolution is not None and request_method == "POST": + return _DiagnosticSend(send, debug_headers, resolution) + async def _send_with_debug(message: Message) -> None: if message["type"] == "http.response.start": headers: Final = list(message.get("headers", [])) @@ -266,8 +338,6 @@ class MCPDebug: raw_headers: dict[str, str] | None, scope: dict, mcp_servers: list[str] | None, - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, oauth2_headers: dict[str, str] | None, client_ip: str | None, ) -> dict[str, str]: @@ -288,16 +358,13 @@ class MCPDebug: server_url: str | None = None server_auth_type: str | None = None - auth_resolution = "no-auth" + auth_resolution: Final = AuthResolution.unresolved.value for server_name in mcp_servers or []: server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) if server: server_url = server.url server_auth_type = server.auth_type - auth_resolution = MCPDebug.resolve_auth_resolution( - server, mcp_auth_header, mcp_server_auth_headers, oauth2_headers - ) break scope_headers: Final = MCPRequestHandler._safe_get_headers_from_scope(scope) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d7f238f142c..116bcc715d4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -80,6 +80,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( raise_classified_list_failure, upstream_auth_challenge, ) +from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, mcp_per_user_token_cache, @@ -108,12 +109,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_provider import ( build_token_exchanger, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, AuthorizationCodeConfig, + AuthResolution, ClientCredentialsConfig, CredError, IdJagConfig, @@ -3832,13 +3835,21 @@ class MCPServerManager: (authorization_code's browser-OAuth 401, token_exchange's RFC 9728 challenge) or maps any other ``CredError`` onto its public HTTP status; it never returns an error as a value. """ - match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): - case Ok(auth): + match await resolve_credentials_with_source(provider, to_subject(user_api_key_auth, subject_token), spec): + case Ok(credential): + auth: Final = credential.auth # NoOpAuth has no header_name and so never conflicts. header_name: Final[str | None] = getattr(auth, "header_name", None) if header_name is None or not extra_headers: + source: Final = ( + AuthResolution.extra_headers + if credential.source == AuthResolution.no_auth and extra_headers + else credential.source + ) + record_auth_resolution(server.server_id, source) return auth, extra_headers if not has_header(extra_headers, header_name): + record_auth_resolution(server.server_id, credential.source) return auth, extra_headers if isinstance( spec.config, @@ -3853,11 +3864,14 @@ class MCPServerManager: # one-shot 401 refetch is lost with it). Drop only the header the resolved # credential is about to occupy, so a static credential the operator aimed at a # DIFFERENT header still reaches upstream. + record_auth_resolution(server.server_id, credential.source) return auth, without_header(extra_headers, header_name) # Other modes: an Authorization already supplied via extra_headers (a forwarded caller # header or static_headers) is intentional and wins; v1 applies those last. + record_auth_resolution(server.server_id, AuthResolution.extra_headers) return None, extra_headers case Error(err): + record_auth_resolution(server.server_id, AuthResolution.failed) if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig): # authorization_code's missing per-user token -> the per-server browser-OAuth # challenge, built here where the full MCPServer is in hand. @@ -3960,6 +3974,7 @@ class MCPServerManager: Returns: Configured MCP client instance. """ + record_auth_resolution(server.server_id, AuthResolution.unresolved) resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) transport: Final = resolved_server.transport or MCPTransport.sse spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server) @@ -4032,6 +4047,7 @@ class MCPServerManager: env=resolved_env, ) + record_auth_resolution(server.server_id, AuthResolution.not_applicable) return MCPClient( server_url="", # Not used for stdio transport_type=transport, @@ -4086,6 +4102,20 @@ class MCPServerManager: aws_session_name=resolved_server.aws_session_name, ) + legacy_source: Final = ( + AuthResolution.aws_sigv4 + if aws_auth is not None + else AuthResolution.extra_headers + if extra_headers and has_header(extra_headers, auth_header_name or "Authorization") + else AuthResolution.per_request_header + if mcp_auth_header + else AuthResolution.static_token + if auth_value + else AuthResolution.extra_headers + if extra_headers + else AuthResolution.no_auth + ) + record_auth_resolution(server.server_id, legacy_source) return MCPClient( server_url=server_url, transport_type=transport, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 8328aae01ab..85c7f68719d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -65,6 +65,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + AuthResolution, AuthSpecKind, AwsSigV4Config, Byok, @@ -76,6 +77,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, + ResolvedCredential, ServerSpec, SharedKey, Subject, @@ -448,3 +450,32 @@ def _client_auth_fingerprint(client_auth: ClientAuth) -> str: def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet")) + + +async def resolve_credentials_with_source( + provider: UpstreamCredentialProvider, subject: Subject, server: ServerSpec +) -> Result[ResolvedCredential, CredError]: + match await provider.resolve_credentials(subject, server): + case Error(err): + return Error(err) + case Ok(auth): + if isinstance(auth, NoOpAuth): + return Ok(ResolvedCredential(auth, AuthResolution.no_auth)) + match server.config: + case NoneConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.no_auth)) + case ApiKeyConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.static_token)) + case PassthroughConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.oauth2_passthrough)) + case ClientCredentialsConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.client_credentials)) + case TokenExchangeConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.token_exchange)) + case IdJagConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.id_jag)) + case AuthorizationCodeConfig(): + return Ok(ResolvedCredential(auth, AuthResolution.stored_user_token)) + case AwsSigV4Config(): + return Ok(ResolvedCredential(auth, AuthResolution.aws_sigv4)) + assert_never(server.config) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 632dc57dcf6..d186724fd9f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -26,10 +26,11 @@ union (see `result.py`), not `expression.Result`. from __future__ import annotations from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import Annotated, Final, Literal +import httpx from expression import case, tag, tagged_union from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator from typing_extensions import assert_never @@ -46,6 +47,29 @@ from litellm.types.mcp import ( ) +class AuthResolution(str, Enum): + no_auth = "no-auth" + stored_user_token = "stored-user-token" + static_token = "static-token" + per_request_header = "per-request-header" + oauth2_passthrough = "oauth2-passthrough" + client_credentials = "m2m-client-credentials" + token_exchange = "token-exchange" + id_jag = "id-jag" + aws_sigv4 = "aws-sigv4" + extra_headers = "extra-headers" + not_applicable = "not-applicable" + unresolved = "unresolved" + failed = "resolution-failed" + multiple = "multiple" + + +@dataclass(frozen=True, slots=True) +class ResolvedCredential: + auth: httpx.Auth = field(repr=False) + source: AuthResolution + + class AuthSpecKind(str, Enum): """The server's statically-declared upstream-auth mode — derived from its `config`. diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 329dddbdf05..27cb632c843 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -3,12 +3,15 @@ import importlib from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime +from traceback import walk_tb from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +from uuid import uuid4 import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status +from pydantic import ValidationError from starlette.datastructures import Headers from litellm._logging import verbose_logger @@ -30,6 +33,8 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( list_fault_http_status, outcome_wire_value, ) +from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree +from litellm.proxy._experimental.mcp_server.oauth_utils import _redact_mcp_resource_url from litellm.proxy._experimental.mcp_server.ui_session_utils import ( acting_user_auth, build_effective_auth_contexts, @@ -78,11 +83,39 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + reference: Final = uuid4().hex + verbose_logger.error( + "MCP connection test failed (reference=%s): %s", + reference, + tuple( + ( + type(cause).__name__, + tuple( + (frame.f_code.co_filename, lineno, frame.f_code.co_name) + for frame, lineno in walk_tb(cause.__traceback__) + ), + ) + for cause in iter_exception_tree(exc) + ), + ) + return next( + ( + message + for cause in iter_exception_tree(exc) + if (message := _known_connection_error_message(cause, url, timeout_seconds)) is not None + ), + "An unexpected error occurred while testing the MCP connection. " + f"Retry; if it persists, share reference {reference} with your gateway administrator.", + ) + + +def _known_connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str | None: if isinstance(exc, MCPServerURLCredentialsError): return str(exc.detail) if isinstance(exc, TimeoutError): return ( - f"Failed to connect to MCP server: no response from {url or 'the server'} " + "Failed to connect to MCP server: no valid MCP response received from " + f"{_redact_mcp_resource_url(url) or 'the server'} " f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." ) @@ -99,13 +132,45 @@ def _connection_error_message(exc: BaseException, url: str | None, timeout_secon return "Failed to connect to MCP server: the connection timed out." if isinstance(exc, httpx.HTTPStatusError): return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}." - return "Failed to connect to MCP server. Check proxy logs for details." + if isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)): + return ( + "Failed to connect to MCP server: the connection was interrupted. " + "Check the server and network connection, then retry." + ) + if isinstance(exc, ValueError) and str(exc).startswith("Unexpected content type:"): + return ( + "Failed to connect to MCP server: the endpoint returned an unsupported content type. " + "Check that the URL is an MCP endpoint, not a web page, and matches the selected transport." + ) + if isinstance(exc, ValidationError) and exc.title in ("JSONRPCMessage", "InitializeResult", "ListToolsResult"): + return ( + "Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. " + "Check the MCP endpoint URL and the server's protocol implementation." + ) + if MCP_AVAILABLE and isinstance(exc, McpError): + if exc.error.code == -32000 and exc.error.message == "Connection closed": + return ( + "Failed to connect to MCP server: the connection was closed before the request completed. " + "Check that the server stays running and returns a complete MCP response, then retry." + ) + if exc.error.code == 32600 and exc.error.message == "Session terminated": + return ( + "Failed to connect to MCP server: the MCP session was terminated. " + "Check that the URL points to an MCP endpoint and matches the selected transport, " + "then retry to start a new session." + ) + return ( + f"Failed to connect to MCP server: the MCP request failed (JSON-RPC code {exc.error.code}). " + "Check that the endpoint supports MCP initialization and tool listing, and check the upstream server logs." + ) + return None if MCP_AVAILABLE: + from mcp.shared.exceptions import McpError from mcp.types import Tool as MCPTool - from litellm.experimental_mcp_client.client import MCPClient + from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout from litellm.llms.litellm_proxy.skills.skill_search import ( DEFAULT_SKILL_SEARCH_TOP_K, ) @@ -1169,7 +1234,7 @@ if MCP_AVAILABLE: return client_id, client_secret, scopes _STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset( - (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization) + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization, MCPAuth.token) ) @dataclass(frozen=True, slots=True) @@ -1178,6 +1243,17 @@ if MCP_AVAILABLE: mcp_auth_header: str | None oauth2_headers: dict[str, str] | None + def _preview_origin(url: str | None) -> tuple[str, str, int | None] | None: + if not url: + return None + try: + parsed: Final = httpx.URL(url) + except httpx.InvalidURL: + return None + if parsed.scheme not in ("http", "https") or not parsed.host: + return None + return parsed.scheme, parsed.host, parsed.port + def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest: """ Resolve the credentials a not-yet-saved server config carries for a preview call. @@ -1190,7 +1266,19 @@ if MCP_AVAILABLE: MCPRequestHandler, ) - request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request) + saved_server: Final = ( + global_mcp_server_manager.get_mcp_server_by_id(new_mcp_server_request.server_id) + if new_mcp_server_request.server_id + else None + ) + saved_origin: Final = _preview_origin(saved_server.url) if saved_server else None + preview_origin: Final = _preview_origin(new_mcp_server_request.url) + may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or ( + saved_origin is not None and saved_origin == preview_origin + ) + request: Final = ( + _inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request + ) mcp_auth_header: Final = ( request.credentials.get("auth_value") if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict) @@ -1253,8 +1341,15 @@ if MCP_AVAILABLE: if _oauth2_flow == "client_credentials" and not request.token_url: _oauth2_flow = None + # Static previews inherit credentials before this step, but must not resolve back to + # the saved record during client creation and discard the edited connection settings. + preview_server_id: Final = ( + "" + if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES or request.auth_type in (None, MCPAuth.none) + else request.server_id or "" + ) server_model: Final = MCPServer( - server_id=request.server_id or "", + server_id=preview_server_id, name=request.alias or request.server_name or "", url=request.url, transport=request.transport, @@ -1342,11 +1437,18 @@ if MCP_AVAILABLE: except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise except BaseException as e: - verbose_logger.error("Error in MCP operation: %s", e, exc_info=True) + effective_timeout: Final = ( + min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds) + if any( + isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None + for cause in iter_exception_tree(e) + ) + else timeout_seconds + ) return { "status": "error", "error": True, - "message": _connection_error_message(e, request.url, timeout_seconds), + "message": _connection_error_message(e, request.url, effective_timeout), } async def _preview_openapi_tools(spec_path: str) -> dict: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 9a52da0cab1..d259fa32cd9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -49,7 +49,11 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_server_name, _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode ) -from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug +from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + MCPAuthDiagnostics, + MCPDebug, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, get_passthrough_www_authenticate, @@ -1053,7 +1057,7 @@ if MCP_AVAILABLE: route="/mcp/call_tool", traceback_str=failure_traceback, ) - except Exception: + except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error verbose_logger.exception("Error logging failed MCP proxy tool call") raise if proxy_logging_obj is not None: @@ -4472,13 +4476,15 @@ if MCP_AVAILABLE: raw_headers=raw_headers, scope=dict(scope), mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, client_ip=_client_ip, ) - if _debug_headers: - send = MCPDebug.wrap_send_with_debug_headers(send, _debug_headers) + diagnostics: Final = MCPAuthDiagnostics() if _debug_headers else None + if diagnostics is not None: + scope[MCP_AUTH_DIAGNOSTICS_SCOPE_KEY] = diagnostics + send = MCPDebug.wrap_send_with_debug_headers( + send, _debug_headers, diagnostics.headers, request_method=scope.get("method") + ) # Ensure session managers are initialized if not _SESSION_MANAGERS_INITIALIZED: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d746cfd38d8..c22bb76629d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3,6 +3,7 @@ import json import os from collections.abc import Callable, Mapping from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple import httpx @@ -1294,6 +1295,13 @@ class UpdateKeyRequest(KeyRequestBase): rotation_interval: str | None = None organization_id: str | None = None + @model_validator(mode="before") + @classmethod + def drop_blank_team_id(cls, values: object) -> object: + if isinstance(values, Mapping) and values.get("team_id") == "": + return MappingProxyType({k: v for k, v in values.items() if k != "team_id"}) + return values + @field_validator("organization_id", mode="before") @classmethod def treat_cleared_organization_id_as_unset(cls, v: object) -> object: @@ -2828,6 +2836,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "UI username/password login. Default is False." ), ) + disable_env_credential_login: bool | None = Field( + None, + description=( + "If True, disables signing in to the Admin UI with the environment credentials: " + "UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback " + "means env-credential login is always live by default). Database users with passwords " + "are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password " + "before enabling, or nobody can sign in to the UI. A locked-out admin can still " + "administer the proxy over the API with the master key, and can unset this setting " + "and restart the proxy to restore env-credential login. Default is False." + ), + ) disable_budget_reservation: bool | None = Field( None, description=( @@ -5148,9 +5168,26 @@ class CostEstimateRequest(LiteLLMPydanticObjectBase): model: str = Field(description="Model name (from /model_group/info)") input_tokens: int = Field(description="Expected input tokens per request", ge=0) output_tokens: int = Field(description="Expected output tokens per request", ge=0) + cache_read_input_tokens: int = Field( + default=0, description="Input tokens read from the prompt cache; counted within input_tokens", ge=0 + ) + cache_creation_input_tokens: int = Field( + default=0, description="Input tokens written to the prompt cache; counted within input_tokens", ge=0 + ) + reasoning_tokens: int = Field( + default=0, description="Reasoning tokens the model emits; counted within output_tokens", ge=0 + ) num_requests_per_day: int | None = Field(default=None, description="Number of requests per day", ge=0) num_requests_per_month: int | None = Field(default=None, description="Number of requests per month", ge=0) + @model_validator(mode="after") + def validate_token_subsets(self) -> "CostEstimateRequest": + if self.cache_read_input_tokens + self.cache_creation_input_tokens > self.input_tokens: + raise ValueError("cache_read_input_tokens plus cache_creation_input_tokens cannot exceed input_tokens") + if self.reasoning_tokens > self.output_tokens: + raise ValueError("reasoning_tokens cannot exceed output_tokens") + return self + class CostEstimateResponse(LiteLLMPydanticObjectBase): """Response body for /cost/estimate endpoint.""" @@ -5158,6 +5195,9 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): model: str input_tokens: int output_tokens: int + cache_read_input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + reasoning_tokens: int = 0 num_requests_per_day: int | None = None num_requests_per_month: int | None = None # Per-request costs @@ -5165,17 +5205,33 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): input_cost_per_request: float = Field(description="Input token cost per request (before margin)") output_cost_per_request: float = Field(description="Output token cost per request (before margin)") margin_cost_per_request: float = Field(default=0.0, description="Margin/fee added per request") + cache_read_cost_per_request: float = Field(default=0.0, description="Cache-read share of input_cost_per_request") + cache_creation_cost_per_request: float = Field( + default=0.0, description="Cache-write share of input_cost_per_request" + ) + reasoning_cost_per_request: float = Field(default=0.0, description="Reasoning share of output_cost_per_request") # Daily costs (if num_requests_per_day provided) daily_cost: float | None = Field(default=None, description="Total daily cost (includes margin)") daily_input_cost: float | None = Field(default=None, description="Daily input token cost") daily_output_cost: float | None = Field(default=None, description="Daily output token cost") daily_margin_cost: float | None = Field(default=None, description="Daily margin/fee") + daily_cache_read_cost: float | None = Field(default=None, description="Cache-read share of daily_input_cost") + daily_cache_creation_cost: float | None = Field(default=None, description="Cache-write share of daily_input_cost") + daily_reasoning_cost: float | None = Field(default=None, description="Reasoning share of daily_output_cost") # Monthly costs (if num_requests_per_month provided) monthly_cost: float | None = Field(default=None, description="Total monthly cost (includes margin)") monthly_input_cost: float | None = Field(default=None, description="Monthly input token cost") monthly_output_cost: float | None = Field(default=None, description="Monthly output token cost") monthly_margin_cost: float | None = Field(default=None, description="Monthly margin/fee") - # Pricing info - input_cost_per_token: float | None = None - output_cost_per_token: float | None = None + monthly_cache_read_cost: float | None = Field(default=None, description="Cache-read share of monthly_input_cost") + monthly_cache_creation_cost: float | None = Field( + default=None, description="Cache-write share of monthly_input_cost" + ) + monthly_reasoning_cost: float | None = Field(default=None, description="Reasoning share of monthly_output_cost") + # Pricing info: the rates this request's usage bills at, after token tiers and regional multipliers + input_cost_per_token: float | None = Field(default=None, description="Rate billed per input token") + output_cost_per_token: float | None = Field(default=None, description="Rate billed per output token") + cache_read_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-read token") + cache_creation_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-write token") + output_cost_per_reasoning_token: float | None = Field(default=None, description="Rate billed per reasoning token") provider: str | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dc693317de0..1efc9611fe6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -94,6 +94,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_auth_cache_key, team_membership_reservation_cache_key, ) +from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, @@ -475,6 +476,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) +_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool: @@ -2858,7 +2860,9 @@ class TeamNotFoundError(HTTPException): async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None ) -> "_PrismaTeamRow | None": - response = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) + response = await _team_table(TeamRepository(prisma_client)).find_unique( + where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS + ) if response is None and team_id_upsert: from litellm.proxy.management_endpoints.team_endpoints import new_team @@ -3158,7 +3162,9 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(where={"team_alias": team_alias}) + teams: Final = await _team_table(TeamRepository(prisma_client)).find_many( + where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS + ) if not teams: raise HTTPException( @@ -3445,36 +3451,37 @@ async def _fetch_key_object_from_db_with_reconnect( """ Fetch key object from DB and retry once if a DB connection error can be healed. """ - try: - return await prisma_client.get_data( - token=hashed_token, - table_name="combined_view", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - if PrismaDBExceptionHandler.is_database_transport_error(e): - did_reconnect = False - if hasattr(prisma_client, "attempt_db_reconnect"): - auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0) - if not isinstance(auth_reconnect_timeout, (int, float)): - auth_reconnect_timeout = 2.0 - auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1) - if not isinstance(auth_reconnect_lock_timeout, (int, float)): - auth_reconnect_lock_timeout = 0.1 - did_reconnect = await prisma_client.attempt_db_reconnect( - reason="auth_get_key_object_lookup_failure", - timeout_seconds=auth_reconnect_timeout, - lock_timeout_seconds=auth_reconnect_lock_timeout, - ) - if did_reconnect: - return await prisma_client.get_data( - token=hashed_token, - table_name="combined_view", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - raise + async with db_lookup_gate.current(): + try: + return await prisma_client.get_data( + token=hashed_token, + table_name="combined_view", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + if PrismaDBExceptionHandler.is_database_transport_error(e): + did_reconnect = False + if hasattr(prisma_client, "attempt_db_reconnect"): + auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0) + if not isinstance(auth_reconnect_timeout, (int, float)): + auth_reconnect_timeout = 2.0 + auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1) + if not isinstance(auth_reconnect_lock_timeout, (int, float)): + auth_reconnect_lock_timeout = 0.1 + did_reconnect = await prisma_client.attempt_db_reconnect( + reason="auth_get_key_object_lookup_failure", + timeout_seconds=auth_reconnect_timeout, + lock_timeout_seconds=auth_reconnect_lock_timeout, + ) + if did_reconnect: + return await prisma_client.get_data( + token=hashed_token, + table_name="combined_view", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + raise def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index f78c4221f5a..aa77835e858 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.url_utils import ( provider_url_destination_candidates, validate_url, ) +from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint from litellm.proxy._types import * from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2003,9 +2004,20 @@ def get_model_from_request( bedrock_model: Final = _model_from_bedrock_route(route) return model if bedrock_model is None else bedrock_model + if route.lower().startswith(("/azure/", "/azure_ai/")): + azure_model: Final = _router_model_from_azure_route(route, llm_router) + return model if azure_model is None else azure_model + return model +def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: + if llm_router is None: + return None + endpoint: Final = re.sub(r"^/azure(?:_ai)?/", "", route, flags=re.IGNORECASE) + return azure_router_model_in_endpoint(endpoint, frozenset(llm_router.get_model_names())) + + def _model_from_bedrock_route(route: str) -> str | None: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _extract_model_from_bedrock_endpoint, diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 0795cee7409..69091ee8344 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -53,6 +53,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.team_grants import team_model_aliases from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, @@ -1595,7 +1596,7 @@ class JWTAuthManager: model=requested_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=None, + team_model_aliases=team_model_aliases(team_object), ) ): is_allowed = allowed_routes_check( @@ -2132,7 +2133,7 @@ class JWTAuthManager: model=requested_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=None, + team_model_aliases=team_model_aliases(team_object), ) except ProxyException: continue diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 8d4f6f81363..c0a76a4fc20 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -85,6 +85,29 @@ def get_ui_credentials(master_key: str | None) -> tuple[str, str]: return ui_username, ui_password +def _matches_env_credentials(username: str, password: str, master_key: str | None) -> bool: + ui_username, ui_password = get_ui_credentials(master_key) + return secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( + password.encode("utf-8"), ui_password.encode("utf-8") + ) + + +def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool: + """Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed. + + Two settings can turn it off: `disable_env_credential_login` unconditionally, and + `disable_password_login_when_sso_enabled` as a side effect, since its gate rejects + every username/password login before the env comparison runs. Feeds both the + `authenticate_user` gate and the Admin UI warning banner, so the banner never nags + about a login path that is already unreachable. + """ + if general_settings.get("disable_env_credential_login") is True: + return False + if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured(): + return False + return True + + class LoginResult: """Result object containing authentication data from login.""" @@ -129,7 +152,8 @@ async def authenticate_user( master_key: Master key for the proxy (required) prisma_client: Prisma database client (optional) general_settings: Proxy general_settings, checked for - `disable_password_login_when_sso_enabled` + `disable_password_login_when_sso_enabled` and + `disable_env_credential_login` Returns: LoginResult: Object containing authentication data @@ -170,8 +194,6 @@ async def authenticate_user( code=500, ) - ui_username, ui_password = get_ui_credentials(master_key) - # Check if we can find the `username` in the db. On the UI, users can enter username=their email _user_row: LiteLLM_UserTable | None = None user_role: ( @@ -197,8 +219,8 @@ async def authenticate_user( - Login with UI_USERNAME and UI_PASSWORD - Login with Invite Link `user_email` and `password` combination """ - if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( - password.encode("utf-8"), ui_password.encode("utf-8") + if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials( + username, password, master_key ): # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin user_role = LitellmUserRoles.PROXY_ADMIN @@ -340,8 +362,13 @@ async def authenticate_user( code=401, ) else: + env_credentials_hint: Final = ( + "\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file" + if is_env_credential_login_enabled(general_settings) + else "" + ) raise ProxyException( - message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file", + message=f"Invalid credentials used to access UI.{env_credentials_hint}", type=ProxyErrorTypes.auth_error, param="invalid_credentials", code=401, diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py new file mode 100644 index 00000000000..1196011dcdd --- /dev/null +++ b/litellm/proxy/auth/team_grants.py @@ -0,0 +1,122 @@ +"""Project a team row (plus the caller's membership in it) onto the ``team_*`` fields of ``UserAPIKeyAuth``. + +The virtual-key path gets these fields for free from the combined-view SQL join. Every other auth path +starts from a ``LiteLLM_TeamTable`` object instead and has to copy them over by hand, which is how JWT +callers kept losing grants (aliases, permissions, limits) one field at a time. Build the badge through +``team_grants`` and the two paths cannot drift. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Annotated, Final + +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError +from pydantic.main import IncEx +from typing_extensions import ReadOnly, TypedDict + +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + Member, +) + +_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) +_JSON_COLUMNS: Final[Mapping[str, IncEx | bool]] = MappingProxyType( + {"metadata": True, "litellm_model_table": MappingProxyType({"model_aliases": True})} +) + + +def _decode_model_aliases(value: object) -> object: + """``LiteLLM_ModelTable.model_aliases`` is typed ``str | dict``; writers hand Prisma ``json.dumps(...)``, so take both.""" + if not isinstance(value, str): + return value + try: + return _MODEL_ALIASES_ADAPTER.validate_json(value) + except ValidationError: + return None + + +class TeamModelAliasTable(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + + model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None + + +class _TeamJsonColumns(BaseModel): + """The two loosely typed columns on ``LiteLLM_TeamTable``, re-read with the shape the badge needs.""" + + metadata: Mapping[str, object] | None = None + litellm_model_table: TeamModelAliasTable | None = None + + +class TeamGrants(TypedDict, total=False): + """Keyword arguments for ``UserAPIKeyAuth``. Empty when the caller has no team, so the model's own defaults apply.""" + + team_alias: ReadOnly[str | None] + team_tpm_limit: ReadOnly[int | None] + team_rpm_limit: ReadOnly[int | None] + team_max_budget: ReadOnly[float | None] + team_soft_budget: ReadOnly[float | None] + team_spend: ReadOnly[float | None] + team_models: ReadOnly[Sequence[str]] + team_blocked: ReadOnly[bool] + team_metadata: ReadOnly[Mapping[str, object] | None] + team_model_aliases: ReadOnly[Mapping[str, str] | None] + team_object_permission_id: ReadOnly[str | None] + team_object_permission: ReadOnly[LiteLLM_ObjectPermissionTable | None] + team_member: ReadOnly[Member | None] + team_member_spend: ReadOnly[float | None] + team_member_tpm_limit: ReadOnly[int | None] + team_member_rpm_limit: ReadOnly[int | None] + + +def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns: + try: + return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS)) + except ValidationError: + return _TeamJsonColumns() + + +def team_model_aliases(team_object: LiteLLM_TeamTable | None) -> Mapping[str, str] | None: + if team_object is None: + return None + alias_table: Final = _json_columns(team_object).litellm_model_table + return alias_table.model_aliases if alias_table is not None else None + + +def team_grants( + team_object: LiteLLM_TeamTable | None, + team_membership: LiteLLM_TeamMembership | None, + user_id: str | None, +) -> TeamGrants: + if team_object is None: + return TeamGrants() + json_columns: Final = _json_columns(team_object) + return TeamGrants( + team_alias=team_object.team_alias, + team_tpm_limit=team_object.tpm_limit, + team_rpm_limit=team_object.rpm_limit, + team_max_budget=team_object.max_budget, + team_soft_budget=team_object.soft_budget, + team_spend=team_object.spend, + team_models=tuple(team_object.models), + team_blocked=team_object.blocked, + team_metadata=json_columns.metadata, + team_model_aliases=( + json_columns.litellm_model_table.model_aliases if json_columns.litellm_model_table is not None else None + ), + team_object_permission_id=team_object.object_permission_id, + team_object_permission=team_object.object_permission, + team_member=next( + (m for m in team_object.members_with_roles if user_id is not None and m.user_id == user_id), + None, + ), + team_member_spend=team_membership.spend if team_membership is not None else None, + team_member_tpm_limit=( + team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None + ), + team_member_rpm_limit=( + team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None + ), + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6b000489d5a..64c0da1c28f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -82,6 +82,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.resolvers import CredentialRef, Principal from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.team_grants import team_grants from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( @@ -91,6 +92,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_set_request_parsed_body, populate_request_with_path_params, ) +from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, @@ -182,6 +184,44 @@ def _get_model_from_request_context( ) +_CLAUDE_MODEL_ROUTES: Final = frozenset( + f"/{prefix}{endpoint}" for prefix in ("", "v1/") for endpoint in ("messages", "chat/completions", "responses") +) +_CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized" + + +async def _normalize_claude_model( + request_data: dict, valid_token: UserAPIKeyAuth, request: Request | None, route: str +) -> None: + from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + + if route not in _CLAUDE_MODEL_ROUTES or llm_router is None: + return + if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True: + return + requested: Final = _get_model_from_request_context(request_data, route, request, llm_router) + if not isinstance(requested, str) or requested != request_data.get("model"): + return + if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"): + return + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + aliases: Final = settings.get("model_group_alias") if isinstance(settings, Mapping) else None + source: Final = claude_code_requested_group( + requested, llm_router, valid_token.team_id, (valid_token.aliases, valid_token.team_model_aliases, aliases) + ) + if request is not None: + request.scope[_CLAUDE_MODEL_NORMALIZED] = True + if source is None: + return + request_data["model"] = source + _safe_set_request_parsed_body(request=request, parsed_body=request_data) + if request is not None: + request._json = request_data + request._body = orjson.dumps(request_data) + + def _get_model_names_for_budget_checks( model: str | list[str] | None, ) -> list[str]: @@ -1476,24 +1516,16 @@ async def _user_api_key_auth_builder( user_id=user_id, user_email=user_email, team_id=team_id, - team_alias=(team_object.team_alias if team_object is not None else None), - team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), - team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), - team_models=(team_object.models if team_object is not None else []), - team_metadata=(team_object.metadata if team_object is not None else None), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, jwt_claims=jwt_claims, + **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) valid_token = UserAPIKeyAuth( api_key=None, team_id=team_id, - team_alias=(team_object.team_alias if team_object is not None else None), - team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), - team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), - team_models=(team_object.models if team_object is not None else []), user_role=( LitellmUserRoles(user_object.user_role) if user_object is not None and user_object.user_role is not None @@ -1507,17 +1539,8 @@ async def _user_api_key_auth_builder( user_tpm_limit=(user_object.tpm_limit if user_object is not None else None), user_rpm_limit=(user_object.rpm_limit if user_object is not None else None), user_model_max_budget=(user_object.model_max_budget if user_object is not None else None), - team_member_rpm_limit=( - team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None - ), - team_member_tpm_limit=( - team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None - ), - team_metadata=(team_object.metadata if team_object is not None else None), jwt_claims=jwt_claims, - ) - valid_token.team_object_permission = ( - team_object.object_permission if team_object is not None else None + **team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id), ) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. @@ -2784,6 +2807,7 @@ async def _authorize_authenticated_request( """ ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) + await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route) # Single authorization point. Builder paths MUST NOT call common_checks. # Route through the same exception handler the builder uses so @@ -3147,6 +3171,7 @@ async def _enforce_key_and_fallback_model_access( Key-level model allowlist and client fallbacks (same as standard auth). Not included in common_checks — common_checks enforces team/user/project model access only. """ + await _normalize_claude_model(request_data, valid_token, request, route) config: Final = valid_token.config if config != {}: diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index f7d9eb7da9a..4045857a237 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -490,7 +490,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. @@ -508,7 +508,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. @@ -532,12 +532,28 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi lite --base-url https://your-proxy.example.com login --config-claude ``` -It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. +It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops. +#### Configuring Claude Code Once, With a Virtual Key or Your Login + +`lite configure claude` wires Claude Code up persistently and `lite unconfigure claude` puts things back. It is what `lite login --config-claude` does, plus a pinned model and an undo, and it also takes a long-lived virtual key when that is what you have: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh +lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto +claude +``` + +With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control + +Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt + +What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Like `--config-claude`, both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any login prompt or request + ### QA Complexity-Based Auto-Routing Against Your Real Proxy `lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session. @@ -584,7 +600,7 @@ An interactive wizard. It runs the same model-group discovery as above, splits t The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key. -You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) +You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets the top-level `model` and `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_OPUS_MODEL` and `ANTHROPIC_DEFAULT_FABLE_MODEL` to `autorouter` in `~/.claude/settings.json` (and `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when missing, like every other wiring), so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 7a0ae9dc955..bf2a784f590 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -4,6 +4,7 @@ import subprocess import sys from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias @@ -12,10 +13,12 @@ import requests from pydantic import BaseModel, TypeAdapter, ValidationError from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login +from .claude_settings import claude_settings_path, lite_api_key_helper_configured from .cmd_quoting import quote_for_cmd from .pi import ( LITELLM_PROXY_API_KEY_ENV, PI_PROVIDER_NAME, + ListingFailure, PiSyncError, fetch_model_ids, fetch_model_limits, @@ -83,6 +86,8 @@ def build_agent_env( base_url: str, api_key: str, profiles: frozenset[str], + *, + export_anthropic_token: bool = True, ) -> dict[str, str]: """Return a copy of base_env wired to route the agent through the proxy. @@ -97,12 +102,19 @@ def build_agent_env( proxy's /v1/models; likewise left alone when already set. pi ignores both base URL variables and instead resolves $LITELLM_PROXY_API_KEY from its synced models.json provider entry. + + With export_anthropic_token=False the bearer is left out (and any inherited + one dropped) so Claude Code asks its configured apiKeyHelper instead; Claude + Code prefers ANTHROPIC_AUTH_TOKEN over the helper and warns when both are set. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") if PROFILE_ANTHROPIC in profiles: env[ANTHROPIC_BASE_URL_ENV] = root - env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key + if export_anthropic_token: + env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key + else: + env.pop(ANTHROPIC_AUTH_TOKEN_ENV, None) env.pop(ANTHROPIC_API_KEY_ENV, None) if ENABLE_TOOL_SEARCH_ENV not in env: env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE @@ -165,7 +177,9 @@ def prepare_pi( """ ids: Final = fetch_model_ids(base_url, api_key, get=get) if isinstance(ids, PiSyncError): - raise AgentRunError(ids.message) + raise AgentRunError( + f"{ids.message} pi would have nothing to run." if ids.kind is ListingFailure.EMPTY else ids.message + ) limits: Final = fetch_model_limits(base_url, api_key, get=get) path: Final = models_json_path(base_env) error: Final = sync_models_json(path, base_url, ids, limits) @@ -460,6 +474,7 @@ def run_agent( launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off, reattach_terminal: Callable[[], None] | None = None, preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS), + export_anthropic_token: bool = True, ) -> None: """Validate, wire the environment, and hand off to the agent. @@ -491,7 +506,9 @@ def run_agent( env: Final = MappingProxyType( { - **build_agent_env(env_before_sync, base_url, api_key, profiles), + **build_agent_env( + env_before_sync, base_url, api_key, profiles, export_anthropic_token=export_anthropic_token + ), **(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced), } ) @@ -529,14 +546,26 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." +def _helper_supplies_token( + ctx_obj: CliContextObj, base_url: str, profiles: frozenset[str], settings_path: Path +) -> bool: + if PROFILE_ANTHROPIC not in profiles or not ctx_obj.get("api_key_from_token_file"): + return False + return lite_api_key_helper_configured(base_url, settings_path) + + def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] started_interactive: Final = _is_interactive() api_key: Final = resolve_api_key(ctx) - display_name, _ = agent_profile(binary) + display_name, profiles = agent_profile(binary) + settings_path: Final = claude_settings_path(os.environ) + helper_supplies_token: Final = _helper_supplies_token(ctx_obj, base_url, profiles, settings_path) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") + if helper_supplies_token: + click.echo(f"litellm: {display_name} reads its key from the apiKeyHelper in {settings_path}") try: run_agent( @@ -545,6 +574,7 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify [binary, *args], skip_verify=skip_verify, reattach_terminal=(_restore_controlling_terminal if started_interactive else None), + export_anthropic_token=not helper_supplies_token, ) except AgentRunError as e: raise click.ClickException(str(e)) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 12a288202b6..4f704afe9d6 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -1,3 +1,4 @@ +import os import sys import time import webbrowser @@ -40,10 +41,16 @@ from litellm.litellm_core_utils.cli_token_utils import ( ) from .claude_settings import ( - CLAUDE_SETTINGS_PATH, - SETTINGS_FILE_OWNERS, + STARTING_MODEL_ROLE, + ApiKeyHelper, ClaudeSettingsError, - write_claude_settings, + KeepModel, + claude_settings_path, + configure_claude_settings, + configure_state_path, + refuse_while_owned, + resolve_api_key_helper, + settings_file_owners, ) from .pkce_login import ( Http, @@ -778,13 +785,24 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None: def _configure_claude_code(base_url: str) -> None: - """Point Claude Code at base_url by patching ~/.claude/settings.json.""" + """Point Claude Code at base_url by patching the settings.json it reads, undoable with `lite unconfigure claude`.""" + settings_path: Final = claude_settings_path(os.environ) try: - write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS) + configure_claude_settings( + base_url, + ApiKeyHelper(resolve_api_key_helper(base_url)), + KeepModel(), + settings_path, + configure_state_path(settings_path), + settings_file_owners(settings_path), + ) except ClaudeSettingsError as e: raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}") - click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.") - click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.") + click.echo(f"\nConfigured Claude Code: {settings_path} now routes through {base_url.rstrip('/')}.") + click.echo( + "Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. " + f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}." + ) def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None: @@ -853,6 +871,12 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] + if config_claude: + settings_path: Final = claude_settings_path(os.environ) + try: + refuse_while_owned(settings_path, settings_file_owners(settings_path)) + except ClaudeSettingsError as e: + raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}") try: if pkce: diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py index 26d45138a27..86701f186fd 100644 --- a/litellm/proxy/client/cli/commands/autoroute/commands.py +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -14,11 +14,13 @@ from ..claude_settings import ( AUTOROUTE_BACKUP_PATH, CLAUDE_SETTINGS_PATH, ClaudeSettingsError, + StaticToken, load_json_or_empty, + merge_claude_settings, ) from ..up import BackupRecord as ClaudeBackupRecord from ..up import restore_claude_settings, write_backup -from .config import master_key_from_config +from .config import AUTOROUTER_MODEL_NAME, master_key_from_config from .process import ( CONFIG_PATH, DEFAULT_AUTOROUTE_PORT, @@ -37,7 +39,6 @@ from .process import ( terminate, write_pid_record, ) -from .settings import merge_claude_settings_static_token from .wizard import run_configure_wizard _GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) @@ -156,7 +157,9 @@ def up(port: int) -> None: ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None), AUTOROUTE_BACKUP_PATH, ) - merged: Final = merge_claude_settings_static_token(original_settings, base_url, master_key) + merged: Final = merge_claude_settings( + original_settings, base_url, StaticToken(master_key), AUTOROUTER_MODEL_NAME, AUTOROUTER_MODEL_NAME + ) CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) with secure_create(CLAUDE_SETTINGS_PATH) as f: json.dump(merged, f, indent=2) diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py deleted file mode 100644 index 60729b5410d..00000000000 --- a/litellm/proxy/client/cli/commands/autoroute/settings.py +++ /dev/null @@ -1,51 +0,0 @@ -from typing import Final - -from pydantic import JsonValue - -from .config import AUTOROUTER_MODEL_NAME - -ENV_KEY: Final = "env" -API_KEY_HELPER_KEY: Final = "apiKeyHelper" -ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" -ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" -ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" -ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" -ENABLE_TOOL_SEARCH_VALUE: Final = "true" -# Force every one of Claude Code's own model tiers to request the auto-router by name. -# Router's auto-router registry is keyed by the literal requested model string -# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*" -# model_name can never work as a catch-all -- these overrides are what actually makes -# Claude Code send "autorouter" regardless of /model or its own version-specific defaults. -ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = ( - "ANTHROPIC_DEFAULT_SONNET_MODEL", - "ANTHROPIC_DEFAULT_HAIKU_MODEL", - "ANTHROPIC_DEFAULT_OPUS_MODEL", -) - - -def merge_claude_settings_static_token( - settings: dict[str, JsonValue], base_url: str, auth_token: str -) -> dict[str, JsonValue]: - """Return a new settings dict wired to a local ephemeral proxy with a static token. - - Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real - remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key is the - locally persisted autoroute master key, so a plain env var is simpler and correct. Any - existing apiKeyHelper is cleared so it can't fight with the static token. - """ - raw_env: Final = settings.get(ENV_KEY, {}) - base_env: Final = raw_env if isinstance(raw_env, dict) else {} - env: Final[dict[str, JsonValue]] = { - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, - **base_env, - ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - ANTHROPIC_AUTH_TOKEN_KEY: auth_token, - **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, - } - env.pop(ANTHROPIC_API_KEY_KEY, None) - merged: Final[dict[str, JsonValue]] = {**settings, ENV_KEY: env} - merged.pop(API_KEY_HELPER_KEY, None) - return merged - - -__all__ = ["merge_claude_settings_static_token"] diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 5e3ce95f088..d0ee507f0b2 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -1,37 +1,71 @@ """Shared handling of Claude Code's ~/.claude/settings.json. -`lite up` patches this file temporarily and restores it on exit; `lite login ---config-claude` patches it persistently. Both need the same merge and the same -apiKeyHelper command, and `up` already imports from `auth`, so the shared parts -live here rather than in either command module. +`lite up` and `lite autoroute up` patch this file temporarily and restore it on +exit; `lite login --config-claude` and `lite configure claude` patch it +persistently and record how to undo it. All of them need the same merge and the +same apiKeyHelper command, and `up` already imports from `auth`, so the shared +parts live here rather than in any one command module. """ +import hashlib +import json import shlex import shutil import sys -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from functools import reduce +from itertools import chain from pathlib import Path -from typing import Final +from types import MappingProxyType +from typing import Final, TypeAlias -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError -from litellm.litellm_core_utils.private_json import write_private_json +from litellm.litellm_core_utils.private_json import ( + commit_staged_json, + discard_staged_json, + ensure_private_dir, + stage_private_json, +) from .cmd_quoting import quote_for_cmd ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" +MODEL_KEY: Final = "model" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" +ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" ENABLE_TOOL_SEARCH_VALUE: Final = "true" ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1" +ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = ( + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", + "ANTHROPIC_DEFAULT_FABLE_MODEL", +) +OWNED_ENV_KEYS: Final = ( + ENABLE_TOOL_SEARCH_KEY, + ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, + ANTHROPIC_BASE_URL_KEY, + ANTHROPIC_AUTH_TOKEN_KEY, + ANTHROPIC_API_KEY_KEY, +) +OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY) +OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS) +_CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY)) +_CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY) +_BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}" +STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts on" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" +CLAUDE_CONFIG_DIR_ENV: Final = "CLAUDE_CONFIG_DIR" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" +CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json" @dataclass(frozen=True, slots=True) @@ -55,6 +89,129 @@ class ClaudeSettingsError(Exception): """Raised for any user-actionable failure while reading or writing Claude Code settings.""" +def claude_settings_path(environ: Mapping[str, str]) -> Path: + """The settings.json Claude Code reads: under CLAUDE_CONFIG_DIR when set, else ~/.claude/settings.json.""" + config_dir: Final = environ.get(CLAUDE_CONFIG_DIR_ENV, "") + if not config_dir: + return CLAUDE_SETTINGS_PATH + return Path(config_dir).expanduser() / "settings.json" + + +def _is_default_settings_file(settings_path: Path) -> bool: + return settings_path.resolve() == CLAUDE_SETTINGS_PATH.resolve() + + +def settings_file_owners(settings_path: Path) -> tuple[SettingsFileOwner, ...]: + """The commands whose backups guard settings_path: `lite up` and `lite autoroute up` only ever manage the default file.""" + return SETTINGS_FILE_OWNERS if _is_default_settings_file(settings_path) else () + + +def configure_state_path(settings_path: Path) -> Path: + """The receipt describing settings_path: the default file keeps CONFIGURE_STATE_PATH, and any other file + (a CLAUDE_CONFIG_DIR) gets its own beside it, keyed by its resolved path, so two settings files never + share one undo record.""" + if _is_default_settings_file(settings_path): + return CONFIGURE_STATE_PATH + digest: Final = hashlib.sha256(str(settings_path.resolve()).encode()).hexdigest() + return CONFIGURE_STATE_PATH.parent / CONFIGURE_STATE_PATH.stem / f"{digest}.json" + + +@dataclass(frozen=True, slots=True) +class StaticToken: + """A long-lived virtual key, written into env.ANTHROPIC_AUTH_TOKEN.""" + + token: str + + +@dataclass(frozen=True, slots=True) +class ApiKeyHelper: + """A `lite auth print-token` command Claude Code runs per request, so a login renews in place.""" + + command: str + + +ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper + + +@dataclass(frozen=True, slots=True) +class KeepModel: + """Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login).""" + + +@dataclass(frozen=True, slots=True) +class UnpinModel: + """Let go of a `model` an earlier configure pinned; one the user set themselves stays.""" + + +@dataclass(frozen=True, slots=True) +class StartOn: + """Pin the top-level `model`, the row Claude Code starts on.""" + + model: str + + +ModelChoice: TypeAlias = KeepModel | UnpinModel | StartOn + + +class OwnedValue(BaseModel): + """What one key held at a moment in time; `present=False` is an absent key, not a null one.""" + + model_config = ConfigDict(frozen=True) + + present: bool + value: JsonValue = None + + +class ConfigureReceipt(BaseModel): + """What `lite configure claude` found and what it owns, keyed by dotted path (`env.X` or a top-level key). + + Ownership moves only by a write: `written` fingerprints the keys some configure changed, at the + value it wrote; a repeat configure refreshes a fingerprint only for a key its merge changed and + carries the earlier one otherwise, so a key the user edited in between stops matching and is left + alone. `previous` is what each key held before configure took it over; a repeat keeps the earlier + snapshot while the key still holds our value and snapshots afresh otherwise, so whatever the + repeat displaces is what comes back. `endpoints` is the ANTHROPIC_BASE_URL each credential slot + was captured beside, so a credential is only ever put back next to the server it was issued for. + No fingerprint is a second copy of a token. + """ + + model_config = ConfigDict(frozen=True) + + file_existed: bool + env_present: bool + env_was_object: bool + previous: Mapping[str, OwnedValue] + written: Mapping[str, str] + endpoints: Mapping[str, OwnedValue] + + +@dataclass(frozen=True, slots=True) +class WithheldCredential: + """A credential left removed: captured beside `endpoint`, while the restored file points elsewhere.""" + + key: str + endpoint: str + + +@dataclass(frozen=True, slots=True) +class UnconfigureOutcome: + """Keys whose value unconfigure changed back, keys the user changed since and so were left as they + are, credentials withheld (the receipt is kept for them, so a later unconfigure can finish once the + URL points back), and whether no settings file remains.""" + + restored: tuple[str, ...] + kept: tuple[str, ...] + withheld: tuple[WithheldCredential, ...] = () + file_removed: bool = False + + +@dataclass(frozen=True, slots=True) +class _Claim: + previous: OwnedValue + written: str | None + endpoint: OwnedValue | None + + def load_json_or_empty(path: Path) -> dict[str, JsonValue]: try: content: Final = path.read_bytes() if path.exists() else b"" @@ -70,29 +227,104 @@ def load_json_or_empty(path: Path) -> dict[str, JsonValue]: ) -def merge_claude_settings( - settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str -) -> dict[str, JsonValue]: - """Return a new settings dict wired to route Claude Code through the proxy. +def _env_object(settings: Mapping[str, JsonValue], path: Path) -> Mapping[str, JsonValue]: + raw_env: Final = settings.get(ENV_KEY) + if raw_env is None: + return MappingProxyType({}) + if not isinstance(raw_env, dict): + raise ClaudeSettingsError( + f'{path} has a non-object "{ENV_KEY}" value, which this would discard. Fix or remove it, then retry.' + ) + return raw_env - Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a - stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued - token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH - defaults to true because Claude Code turns tool search off when - ANTHROPIC_BASE_URL is not a first-party Anthropic host, and - CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker - is filled from the proxy's /v1/models; existing values of both are left - alone. Every other key is preserved untouched. + +def refuse_while_owned(settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: + """Refuse while `lite up` or `lite autoroute up` holds a backup it will restore over any write; a + purely local check, so commands run it before any login prompt or request.""" + for owner in owners: + if owner.backup_path.exists(): + raise ClaudeSettingsError( + f"`{owner.start_command}` is currently managing {settings_path} (backup at " + f"{owner.backup_path}) and will restore it when it stops. " + f"Run `{owner.stop_command}` first, then retry." + ) + + +def _write_target(settings_path: Path) -> Path: + """Write through a symlinked settings.json rather than replacing the link, which would silently + detach a file symlinked into a dotfiles repo.""" + try: + return settings_path.resolve() if settings_path.is_symlink() else settings_path + except OSError as e: + raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e + + +def _stage(path: Path, document: Mapping[str, object]) -> str: + try: + return stage_private_json(str(path), document) + except OSError as e: + raise ClaudeSettingsError(f"Could not write {path}: {e}") from e + + +def _land( + path: Path, + staged: str | None, + also_discard: Sequence[str | None] = (), + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + """Commit a staged file to `path`, or remove `path` when nothing is staged for it. The one place a + filesystem error becomes a ClaudeSettingsError; on failure the operation's other staged files are + discarded, so no temp file holding a token is left behind.""" + try: + if staged is None: + path.unlink(missing_ok=True) + else: + commit(staged, str(path)) + except OSError as e: + for other in also_discard: + if other is not None: + discard_staged_json(other) + raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e + + +def merge_claude_settings( + settings: Mapping[str, JsonValue], + base_url: str, + credential: ClaudeCredential, + default_model: str | None = None, + tier_model: str | None = None, +) -> Mapping[str, JsonValue]: + """Return a new settings mapping wired to route Claude Code through the proxy. + + A StaticToken lands in env.ANTHROPIC_AUTH_TOKEN, an ApiKeyHelper in the top-level apiKeyHelper; + the other credential slots are removed either way, since Claude Code given two credentials may + send the wrong one. ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their + defaults only when missing. `default_model` is the top-level `model`, the row Claude Code starts + on; `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one + group. Apart from those tier keys, exactly OWNED_PATHS are touched. """ raw_env: Final = settings.get(ENV_KEY, {}) - base_env: Final = raw_env if isinstance(raw_env, dict) else {} - env: Final = { - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, - ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE, - **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, - ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - } - return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} + current_env: Final = raw_env if isinstance(raw_env, dict) else {} + env: Final = dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain( + ( + (ENABLE_TOOL_SEARCH_KEY, ENABLE_TOOL_SEARCH_VALUE), + (ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE), + ), + ((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS), + ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),), + ((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (), + ((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None), + ) + ) + return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain( + ((key, value) for key, value in settings.items() if key not in (API_KEY_HELPER_KEY, ENV_KEY)), + ((ENV_KEY, env),), + ((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (), + ((MODEL_KEY, default_model),) if default_model is not None else (), + ) + ) def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: @@ -121,56 +353,273 @@ def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token")) -def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: - """Persistently point Claude Code at base_url, preserving every unrelated setting. +def lite_api_key_helper_configured(base_url: str, settings_path: Path) -> bool: + """Whether settings_path already carries the apiKeyHelper `lite login --config-claude` writes for base_url. - Refuses while any owner holds a backup: each restores its backup when it - stops, which would silently undo this write. + Only an exact match counts: a helper for another proxy, a hand-written one, or + settings that cannot be read leave the caller on the env-token path. """ - for owner in owners: - if owner.backup_path.exists(): - raise ClaudeSettingsError( - f"`{owner.start_command}` is currently managing {settings_path} (backup at " - f"{owner.backup_path}) and will restore it when it stops. " - f"Run `{owner.stop_command}` first, then retry." - ) - normalized_base_url: Final = base_url.rstrip("/") - api_key_helper: Final = resolve_api_key_helper(normalized_base_url) - existing: Final = load_json_or_empty(settings_path) - raw_env: Final = existing.get(ENV_KEY) - if raw_env is not None and not isinstance(raw_env, dict): - raise ClaudeSettingsError( - f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. ' - "Fix or remove it, then retry." - ) - merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper) - # os.replace() swaps the symlink itself for a regular file, silently detaching a - # settings.json that is symlinked into a dotfiles repo. There is no backup to undo - # that here, unlike `lite up`, so write through to the link's target instead. - target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path try: - write_private_json(str(target), merged) + configured_helper: Final = load_json_or_empty(settings_path).get(API_KEY_HELPER_KEY) + return configured_helper == resolve_api_key_helper(base_url.rstrip("/")) + except ClaudeSettingsError: + return False + + +def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue: + return OwnedValue(present=key in container, value=container.get(key)) + + +def _fingerprint(owned: OwnedValue) -> str: + return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest() + + +def _env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + raw_env: Final = settings.get(ENV_KEY) + return raw_env if isinstance(raw_env, dict) else MappingProxyType({}) + + +def _lookup(settings: Mapping[str, JsonValue], path: str) -> OwnedValue: + section, _, key = path.rpartition(".") + return _owned(_env(settings) if section else settings, key) + + +def _with_key(container: Mapping[str, JsonValue], key: str, owned: OwnedValue) -> Mapping[str, JsonValue]: + return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping + chain(((k, v) for k, v in container.items() if k != key), ((key, owned.value),) if owned.present else ()) + ) + + +def _with(settings: Mapping[str, JsonValue], path: str, owned: OwnedValue) -> Mapping[str, JsonValue]: + """`settings` with the key at `path` set (or removed when `owned` is absent); nothing else changes.""" + section, _, key = path.rpartition(".") + if not section: + return _with_key(settings, key, owned) + return _with_key(settings, section, OwnedValue(present=True, value=_with_key(_env(settings), key, owned))) + + +def _with_all(settings: Mapping[str, JsonValue], updates: Mapping[str, OwnedValue]) -> Mapping[str, JsonValue]: + return reduce(lambda acc, item: _with(acc, *item), updates.items(), settings) + + +def _ours(settings: Mapping[str, JsonValue], path: str, receipt: ConfigureReceipt) -> bool: + """Whether the key still holds what a configure wrote (a key no configure ever changed is never ours).""" + return receipt.written.get(path) == _fingerprint(_lookup(settings, path)) + + +def _claim( + path: str, + current: Mapping[str, JsonValue], + merged: Mapping[str, JsonValue], + earlier: ConfigureReceipt | None, + url_now: OwnedValue, +) -> _Claim: + """What this configure records for one key; see ConfigureReceipt for the rules.""" + before, after = _lookup(current, path), _lookup(merged, path) + carried: Final = earlier if earlier is not None and _ours(current, path, earlier) else None + return _Claim( + previous=before if carried is None else carried.previous.get(path, before), + written=_fingerprint(after) if before != after else (None if earlier is None else earlier.written.get(path)), + endpoint=None + if path not in _CREDENTIAL_PATHS + else (url_now if carried is None else carried.endpoints.get(path, url_now)), + ) + + +def _receipt( + current: Mapping[str, JsonValue], + merged: Mapping[str, JsonValue], + earlier: ConfigureReceipt | None, + file_exists: bool, +) -> ConfigureReceipt: + url_now: Final = _lookup(current, _BASE_URL_PATH) + claims: Final = MappingProxyType({path: _claim(path, current, merged, earlier, url_now) for path in OWNED_PATHS}) + return ConfigureReceipt( + file_existed=file_exists if earlier is None else earlier.file_existed, + env_present=ENV_KEY in current if earlier is None else earlier.env_present, + env_was_object=isinstance(current.get(ENV_KEY), dict) if earlier is None else earlier.env_was_object, + previous=MappingProxyType({path: claim.previous for path, claim in claims.items()}), + written=MappingProxyType({path: claim.written for path, claim in claims.items() if claim.written is not None}), + endpoints=MappingProxyType( + {path: claim.endpoint for path, claim in claims.items() if claim.endpoint is not None} + ), + ) + + +def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None: + if not state_path.exists(): + return None + try: + return ConfigureReceipt.model_validate_json(state_path.read_bytes()) + except (OSError, ValidationError) as e: + raise ClaudeSettingsError( + f"{state_path} is not a readable `lite configure claude` receipt ({e}). " + "Remove it and edit Claude Code's settings by hand if they still point at the proxy." + ) from e + + +def configure_claude_settings( + base_url: str, + credential: ClaudeCredential, + model: ModelChoice, + settings_path: Path, + state_path: Path, + owners: Sequence[SettingsFileOwner], + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + """Persistently route Claude Code through base_url, recording how to undo it. + + Both files are staged before either is committed, so a full disk or a read-only directory fails + before anything changes. The two commits are still two renames: a receipt rename that fails + discards the staged settings, and a settings rename that fails after the receipt landed puts the + earlier receipt back (or removes the new one), so the receipt on disk never describes settings + that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an + earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login). + """ + refuse_while_owned(settings_path, owners) + current: Final = load_json_or_empty(settings_path) + _env_object(current, settings_path) + earlier: Final = read_configure_receipt(state_path) + existing: Final = ( + _with(current, MODEL_KEY, earlier.previous[MODEL_KEY]) + if isinstance(model, UnpinModel) and earlier is not None and _ours(current, MODEL_KEY, earlier) + else current + ) + merged: Final = merge_claude_settings( + existing, base_url, credential, model.model if isinstance(model, StartOn) else None + ) + receipt: Final = _receipt(current, merged, earlier, settings_path.exists()) + target: Final = _write_target(settings_path) + try: + ensure_private_dir(state_path.parent) except OSError as e: - raise ClaudeSettingsError(f"Could not write {target}: {e}") from e + raise ClaudeSettingsError(f"Could not write {state_path}: {e}") from e + staged_receipt: Final = _stage(state_path, receipt.model_dump(mode="json")) + try: + staged_settings: Final = _stage(target, merged) + except ClaudeSettingsError: + discard_staged_json(staged_receipt) + raise + _land(state_path, staged_receipt, (staged_settings,), commit) + try: + _land(target, staged_settings, commit=commit) + except ClaudeSettingsError as settings_error: + try: + _land(state_path, None if earlier is None else _stage(state_path, earlier.model_dump(mode="json"))) + except ClaudeSettingsError as receipt_error: + raise ClaudeSettingsError( + f"{settings_error} The receipt at {state_path} now describes settings that were not written and " + f"could not be put back either ({receipt_error}); remove it before retrying." + ) from settings_error + raise + + +def _endpoint_text(endpoint: OwnedValue) -> str: + if not endpoint.present: + return f"no {ANTHROPIC_BASE_URL_KEY} (Anthropic's default endpoint)" + return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value) + + +def unconfigure_claude_settings( + settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner] +) -> UnconfigureOutcome: + """Undo `lite configure claude`: put back every key still holding what configure wrote, leave the + rest alone, and withhold a credential the restored file would send to a different server than it + was issued for (the receipt stays, owning only those slots, so a later unconfigure can finish).""" + refuse_while_owned(settings_path, owners) + receipt: Final = read_configure_receipt(state_path) + if receipt is None: + raise ClaudeSettingsError( + f"Claude Code is not configured by `lite configure claude` (no receipt at {state_path}); nothing to undo." + ) + current: Final = load_json_or_empty(settings_path) + _env_object(current, settings_path) + ours: Final = tuple(path for path in receipt.written if _ours(current, path, receipt)) + kept: Final = tuple(path for path in receipt.written if path not in ours and _lookup(current, path).present) + put_back: Final = _with_all(current, MappingProxyType({path: receipt.previous[path] for path in ours})) + url_after: Final = _lookup(put_back, _BASE_URL_PATH) + withheld: Final = tuple( + WithheldCredential(path, _endpoint_text(receipt.endpoints[path])) + for path in _CREDENTIAL_PATHS + if path in ours and receipt.previous[path].present and receipt.endpoints[path] != url_after + ) + absent: Final = OwnedValue(present=False) + trimmed: Final = _with_all(put_back, MappingProxyType({item.key: absent for item in withheld})) + settings: Final = ( + trimmed + if _env(trimmed) or receipt.env_was_object + else _with_key(trimmed, ENV_KEY, OwnedValue(present=receipt.env_present, value=None)) + ) + target: Final = _write_target(settings_path) + file_removed: Final = not settings and not (receipt.file_existed and target.exists()) + kept_receipt: Final = ( # mutable-ok: pydantic serializes the update as given and rejects a mappingproxy + receipt.model_copy(update={"written": {item.key: _fingerprint(absent) for item in withheld}}) + if withheld + else None + ) + staged_settings: Final = None if file_removed else _stage(target, settings) + try: + staged_receipt: Final = ( + None if kept_receipt is None else _stage(state_path, kept_receipt.model_dump(mode="json")) + ) + except ClaudeSettingsError: + if staged_settings is not None: + discard_staged_json(staged_settings) + raise + _land(target, staged_settings, (staged_receipt,)) + _land(state_path, staged_receipt) + return UnconfigureOutcome( + restored=tuple(path for path in ours if _lookup(current, path) != _lookup(settings, path)), + kept=kept, + withheld=withheld, + file_removed=file_removed, + ) __all__ = ( "ANTHROPIC_API_KEY_KEY", + "ANTHROPIC_AUTH_TOKEN_KEY", "ANTHROPIC_BASE_URL_KEY", + "ANTHROPIC_DEFAULT_MODEL_ENV_KEYS", "API_KEY_HELPER_KEY", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", + "CLAUDE_CONFIG_DIR_ENV", "CLAUDE_SETTINGS_PATH", + "CONFIGURE_STATE_PATH", "ENABLE_GATEWAY_MODEL_DISCOVERY_KEY", "ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE", "ENABLE_TOOL_SEARCH_KEY", "ENABLE_TOOL_SEARCH_VALUE", "ENV_KEY", + "MODEL_KEY", + "OWNED_ENV_KEYS", + "OWNED_PATHS", + "OWNED_TOP_LEVEL_KEYS", "SETTINGS_FILE_OWNERS", + "STARTING_MODEL_ROLE", + "ApiKeyHelper", + "ClaudeCredential", "ClaudeSettingsError", + "ConfigureReceipt", + "KeepModel", + "ModelChoice", + "OwnedValue", "SettingsFileOwner", + "StartOn", + "StaticToken", + "UnconfigureOutcome", + "UnpinModel", + "WithheldCredential", + "claude_settings_path", + "configure_claude_settings", + "configure_state_path", + "lite_api_key_helper_configured", "load_json_or_empty", "merge_claude_settings", + "read_configure_receipt", + "refuse_while_owned", "resolve_api_key_helper", - "write_claude_settings", + "settings_file_owners", + "unconfigure_claude_settings", ) diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py new file mode 100644 index 00000000000..9d329f8d8f2 --- /dev/null +++ b/litellm/proxy/client/cli/commands/configure.py @@ -0,0 +1,292 @@ +"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable.""" + +import os +import sys +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import click +from InquirerPy import inquirer +from InquirerPy.base.control import Choice + +from litellm.proxy.common_utils.model_listing_utils import ( + CLAUDE_CODE_CLIENT, + CLAUDE_CODE_PICKER_PATTERN, + GATEWAY_CLIENT_HEADER, +) + +from .auth import CliContextObj, context_secret_vault, get_stored_api_key +from .claude_settings import ( + STARTING_MODEL_ROLE, + ApiKeyHelper, + ClaudeCredential, + ClaudeSettingsError, + ModelChoice, + StartOn, + StaticToken, + UnconfigureOutcome, + UnpinModel, + claude_settings_path, + configure_claude_settings, + configure_state_path, + refuse_while_owned, + resolve_api_key_helper, + settings_file_owners, + unconfigure_claude_settings, +) +from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing +from .up import ensure_fresh_login + +_LISTED_MODELS_SHOWN: Final = 20 +_CLAUDE_TARGET: Final = "claude" +_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),) +_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default" +_CLAUDE_CODE_VIEW: Final = MappingProxyType( + {"anthropic-version": "2023-06-01", GATEWAY_CLIENT_HEADER: CLAUDE_CODE_CLIENT} +) +_MODEL_OPTION_HELP: Final = ( + f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, " + "Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude " + "Code's sub-agent or background tiers; `lite autoroute up` is the mode that does." +) + + +def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]: + """The credential to write and the key to check the proxy with. + + An explicit key (--api-key, `lite --api-key`, LITELLM_PROXY_API_KEY) is long-lived and goes + into settings.json as a static token. Without one, the stored `lite login` credential is used + the way `lite login --config-claude` uses it, through apiKeyHelper, since it expires within a + day and renews in place there; a missing or stale login is refreshed first, as `lite up` does. + """ + ctx_obj: Final[CliContextObj] = ctx.obj + explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key")) + if explicit: + return StaticToken(explicit), explicit + base_url: Final = ctx_obj["base_url"] + ensure_fresh_login(ctx) + stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) + if not stored: + raise ClaudeSettingsError("Login did not produce a usable token.") + return ApiKeyHelper(resolve_api_key_helper(base_url)), stored + + +@dataclass(frozen=True, slots=True) +class _Listing: + models: tuple[ListedModel, ...] + + @property + def ids(self) -> tuple[str, ...]: + return tuple(model.id for model in self.models) + + +def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, _Listing]: + """Every configure path begins the same way: the local ownership check first, so a `lite up` + session is refused before any login prompt or request, then the credential, then the listing.""" + settings_path: Final = claude_settings_path(os.environ) + try: + refuse_while_owned(settings_path, settings_file_owners(settings_path)) + credential, key = resolve_credential(ctx, api_key) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + return credential, _listed_models(ctx.obj["base_url"], key) + + +def _listing_error(base_url: str, error: PiSyncError) -> str: + """The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question.""" + if error.kind is ListingFailure.REJECTED: + return f"LiteLLM rejected your key (HTTP {error.status}). Run `lite login` to refresh it, or pass a valid --api-key." + if error.kind is ListingFailure.UNREACHABLE: + return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?" + if error.kind is ListingFailure.EMPTY: + return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model." + return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy." + + +def _listed_models(base_url: str, key: str) -> _Listing: + listed: Final = fetch_model_listing(base_url, key, headers=_CLAUDE_CODE_VIEW) + if isinstance(listed, PiSyncError): + raise click.ClickException(_listing_error(base_url, listed)) + return _Listing(listed) + + +def _starting_model(model: str, listing: _Listing) -> str | None: + source: Final = next((listed.id for listed in listing.models if listed.source_model == model), None) + return source or next((listed.id for listed in listing.models if listed.id == model), None) + + +def _model_choice(model: str | None) -> ModelChoice: + return StartOn(model) if model is not None else UnpinModel() + + +def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listing: _Listing, model: str | None) -> None: + ctx_obj: Final[CliContextObj] = ctx.obj + base_url: Final = ctx_obj["base_url"] + listed: Final = listing.ids + starting: Final = _starting_model(model, listing) if model is not None else None + if model is not None and starting is None: + shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN]) + more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else "" + raise click.ClickException( + f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}." + ) + settings_path: Final = claude_settings_path(os.environ) + try: + configure_claude_settings( + base_url, + credential, + _model_choice(starting), + settings_path, + configure_state_path(settings_path), + settings_file_owners(settings_path), + ) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + in_picker: Final = sum(1 for listed_model in listed if CLAUDE_CODE_PICKER_PATTERN.search(listed_model)) + click.echo(f"Configured Claude Code: {settings_path} now routes through {base_url}.") + + click.echo( + "Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN." + if isinstance(credential, StaticToken) + else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it." + ) + click.echo( + f"Starting model: {starting} ({STARTING_MODEL_ROLE}); switch any time with /model." + if starting is not None + else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or " + "pass --model to start on a proxy model." + ) + click.echo( + f"/model will list all {len(listed)} of the proxy's models." + if in_picker == len(listed) + else f"/model will list {in_picker} of the proxy's {len(listed)} models: Claude Code shows only ids containing " + "'claude' or 'anthropic', and this proxy does not list the rest under such names." + ) + click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.") + if isinstance(credential, StaticToken) and settings_path.is_symlink(): + click.echo( + f"Note: {settings_path} is a symlink to {settings_path.resolve()}, so your key now lives in " + "that file; keep it out of version control.", + err=True, + ) + + +def _pick_targets() -> tuple[str, ...]: + picked: Final = inquirer.checkbox( + message="Which agents should route through LiteLLM?", + choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS], + validate=lambda chosen: len(chosen) > 0, + invalid_message="Pick at least one.", + ).execute() + return tuple(str(value) for value in picked) + + +def _pick_model(listed: Sequence[str]) -> str | None: + picked: Final = inquirer.fuzzy( + message="Model Claude Code starts on (type to filter; /model switches any time):", + choices=[_KEEP_DEFAULT_MODEL, *listed], + ).execute() + return None if picked == _KEEP_DEFAULT_MODEL else str(picked) + + +def interactive_configure( + ctx: click.Context, + pick_targets: Callable[[], tuple[str, ...]] = _pick_targets, + pick_model: Callable[[Sequence[str]], str | None] = _pick_model, +) -> None: + """`lite configure` with no agent named: ask which agents to wire and which model to pin.""" + targets: Final = pick_targets() + if _CLAUDE_TARGET not in targets: + return + credential, listing = _start(ctx, None) + _apply_claude( + ctx, credential, listing, pick_model(tuple(model.source_model or model.id for model in listing.models)) + ) + + +@click.group(name="configure", invoke_without_command=True) +@click.pass_context +def configure_group(ctx: click.Context) -> None: + """Persistently route a coding agent through your LiteLLM proxy. + + With no agent named, asks which agents to wire and which proxy model to pin. + """ + if ctx.invoked_subcommand is not None: + return + if not sys.stdin.isatty(): + raise click.ClickException( + "`lite configure` asks questions, so it needs a terminal. Non-interactively, run " + "`lite configure claude --api-key --model `." + ) + interactive_configure(ctx) + + +@click.group(name="unconfigure") +def unconfigure_group() -> None: + """Undo `lite configure` for a coding agent.""" + + +@configure_group.command(name="claude") +@click.option( + "--api-key", + "api_key", + default=None, + help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / " + "LITELLM_PROXY_API_KEY value; with neither, your `lite login` credential is used through apiKeyHelper.", +) +@click.option("--model", default=None, help=_MODEL_OPTION_HELP) +@click.pass_context +def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None: + """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`. + + Patches ~/.claude/settings.json in place: the proxy URL, your credential (a virtual key as a + static token, or your `lite login` through apiKeyHelper), and gateway model discovery so + /model lists the proxy's models; --model picks the one Claude Code starts on. Every other + setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. + Assumes the proxy is already running. + """ + credential, listing = _start(ctx, api_key) + _apply_claude(ctx, credential, listing, model) + + +@unconfigure_group.command(name="claude") +def unconfigure_claude() -> None: + """Return Claude Code's settings to what they were before `lite configure claude`. + + Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are + put back; anything you changed since is left as it is and named in the output. + """ + settings_path: Final = claude_settings_path(os.environ) + state_path: Final = configure_state_path(settings_path) + try: + outcome: Final = unconfigure_claude_settings(settings_path, state_path, settings_file_owners(settings_path)) + except ClaudeSettingsError as e: + raise click.ClickException(str(e)) + _report_unconfigure(settings_path, state_path, outcome) + + +def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None: + """Say what unconfigure did, naming only keys whose value it changed.""" + if outcome.file_removed: + click.echo( + f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys." + ) + elif outcome.restored: + click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.") + else: + click.echo(f"Nothing in {settings_path} was still ours to restore.") + if outcome.kept: + click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.") + if outcome.withheld: + click.echo( + "Left removed, since the file now points at a different server than they were issued for: " + + "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld) + + f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` " + "again to put them back, or delete that file to drop them." + ) + + +__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group") diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 7b0c1970c4e..9810e81ae36 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -10,21 +10,40 @@ import os import tempfile from collections.abc import Callable, Mapping from dataclasses import dataclass +from enum import StrEnum from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Annotated, Final import requests -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic.types import StringConstraints PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR" PI_PROVIDER_NAME: Final = "litellm" LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY" +_REJECTED_STATUSES: Final = frozenset((401, 403)) + + +class ListingFailure(StrEnum): + """Why a proxy could not be listed, decided once where the HTTP outcome is classified. + + `unreachable` means no response at all; the other kinds prove the proxy answered, so callers + must not suggest checking whether it is running. + """ + + UNREACHABLE = "unreachable" + REJECTED = "rejected" + BAD_BODY = "bad_body" + EMPTY = "empty" + OTHER = "other" @dataclass(frozen=True, slots=True) class PiSyncError: message: str + status: int | None = None + kind: ListingFailure | None = None @dataclass(frozen=True, slots=True) @@ -33,12 +52,25 @@ class ModelLimits: max_tokens: int | None -class _Model(BaseModel): - id: str +_NonEmptyString = Annotated[str, StringConstraints(min_length=1)] + + +class ListedModel(BaseModel): + model_config = ConfigDict(frozen=True) + + id: _NonEmptyString + source_model: _NonEmptyString | None = None class _ModelList(BaseModel): - data: tuple[_Model, ...] + data: tuple[ListedModel, ...] + + @model_validator(mode="after") + def unique_id_mappings(self) -> "_ModelList": + mappings: Final = frozenset((model.id, model.source_model or model.id) for model in self.data) + if len(frozenset(model.id for model in self.data)) != len(mappings): + raise ValueError("model ids must not map to multiple source models") + return self class _ModelGroup(BaseModel): @@ -51,31 +83,47 @@ class _ModelGroupList(BaseModel): data: tuple[_ModelGroup, ...] +def fetch_model_listing( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, + headers: Mapping[str, str] = MappingProxyType({}), +) -> tuple[ListedModel, ...] | PiSyncError: + url: Final = base_url.rstrip("/") + "/v1/models" + try: + resp: Final = get( + url, + headers={"Authorization": f"Bearer {api_key}", **headers}, # mutable-ok: requests headers require a dict + timeout=10, + ) + except requests.RequestException as e: + return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE) + if resp.status_code != 200: + return PiSyncError( + f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.", + resp.status_code, + ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER, + ) + try: + listing: Final = _ModelList.model_validate(resp.json()) + except (ValueError, ValidationError) as e: + return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY) + models: Final = tuple(dict.fromkeys(listing.data)) + if not models: + return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY) + return models + + def fetch_model_ids( base_url: str, api_key: str, *, get: Callable[..., requests.Response] = requests.get, + headers: Mapping[str, str] = MappingProxyType({}), ) -> tuple[str, ...] | PiSyncError: - url: Final = base_url.rstrip("/") + "/v1/models" - try: - resp: Final = get( - url, - headers={"Authorization": f"Bearer {api_key}"}, # mutable-ok: requests headers require a dict - timeout=10, - ) - except requests.RequestException as e: - return PiSyncError(f"Could not list models from the proxy: {e}") - if resp.status_code != 200: - return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.") - try: - listing: Final = _ModelList.model_validate(resp.json()) - except (ValueError, ValidationError) as e: - return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}") - ids: Final = tuple(dict.fromkeys(model.id for model in listing.data)) - if not ids: - return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.") - return ids + listed: Final = fetch_model_listing(base_url, api_key, get=get, headers=headers) + return listed if isinstance(listed, PiSyncError) else tuple(dict.fromkeys(model.id for model in listed)) _NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({}) @@ -200,10 +248,13 @@ __all__ = ( "LITELLM_PROXY_API_KEY_ENV", "PI_CONFIG_DIR_ENV", "PI_PROVIDER_NAME", + "ListedModel", + "ListingFailure", "ModelLimits", "PiSyncError", "fetch_model_ids", "fetch_model_limits", + "fetch_model_listing", "models_json_path", "provider_block", "sync_models_json", diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index b7c02866d6f..ffece87ab83 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -23,6 +23,7 @@ from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_ from .claude_settings import ( BACKUP_PATH, CLAUDE_SETTINGS_PATH, + ApiKeyHelper, ClaudeSettingsError, load_json_or_empty, merge_claude_settings, @@ -123,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool: return token_data is not None and token_data.get("refresh_token") is not None -def _ensure_fresh_login(ctx: click.Context) -> None: +def ensure_fresh_login(ctx: click.Context) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"].rstrip("/") vault: Final = context_secret_vault(ctx) @@ -141,7 +142,7 @@ def _ensure_fresh_login(ctx: click.Context) -> None: click.echo("No fresh LiteLLM login found for this proxy; starting login...") ctx.invoke(login, pkce=pkce) if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault): - raise UpError("Login did not produce a usable token; cannot start `lite up`.") + raise UpError("Login did not produce a usable token.") def _restore_and_report() -> None: @@ -169,7 +170,7 @@ def up(ctx: click.Context) -> None: base_url: Final = ctx.obj["base_url"] try: - _ensure_fresh_login(ctx) + ensure_fresh_login(ctx) api_key: Final = resolve_api_key(ctx) verify_proxy_key(base_url, api_key) @@ -190,7 +191,7 @@ def up(ctx: click.Context) -> None: ) CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True) - merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper) + merged: Final = merge_claude_settings(original_settings, base_url, ApiKeyHelper(api_key_helper)) with open(CLAUDE_SETTINGS_PATH, "w") as f: json.dump(merged, f, indent=2) except (AgentRunError, ClaudeSettingsError) as e: diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index eae1b0f5bc9..b0e81a222c0 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -13,6 +13,7 @@ from .commands.auth import auth_group, context_secret_vault, get_stored_api_key, from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.config import config_commands, get_config_value, hidden_command_names +from .commands.configure import configure_group, unconfigure_group from .commands.credentials import credentials from .commands.debug import debug from .commands.encryption import encryption @@ -162,6 +163,9 @@ cli.add_command(model_groups) # Add the autoroute command group (QA auto-routing against your real proxy) cli.add_command(autoroute_group, name="autoroute") cli.add_command(config_commands) +# Add configure/unconfigure (persistently wire a coding agent to the proxy with a virtual key) +cli.add_command(configure_group) +cli.add_command(unconfigure_group) if __name__ == "__main__": diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4c9bdf867cd..e3a2b892721 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -190,6 +190,7 @@ from litellm.proxy.litellm_pre_call_utils import ( refresh_proxy_server_request_body_snapshot, reject_url_valued_destination, ) +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -1849,7 +1850,6 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, ) - # Calculate request queue time after add_litellm_data_to_request # which sets arrival_time in proxy_server_request. Ends at start_time # (not a freshly captured time.time() here) so this window is exactly @@ -1997,6 +1997,12 @@ class ProxyBaseLLMRequestProcessing: data=self.data, call_type=route_type, ) + if route_type == "aget_responses": + attach_post_call_pipelines_to_retrieval( + data=self.data, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in @@ -3294,9 +3300,10 @@ class ProxyBaseLLMRequestProcessing: has completed. Guardrails routed through unified_guardrail are skipped, since they already ran - via its streaming iterator. Guardrails that override - async_post_call_success_hook directly run here, including those that implement - apply_guardrail but keep their native lifecycle hooks. + via its streaming iterator, and so are guardrails a post_call policy pipeline + manages, since the pipeline ran them against the buffered stream. Guardrails + that override async_post_call_success_hook directly run here, including those + that implement apply_guardrail but keep their native lifecycle hooks. This is audit-only — content has already been delivered to the client. @@ -3306,12 +3313,18 @@ class ProxyBaseLLMRequestProcessing: _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import _check_and_merge_model_level_guardrails + from litellm.proxy.utils import ( + _check_and_merge_model_level_guardrails, + stream_gated_guardrail_names, + ) guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) + stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict) for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue + if cb.guardrail_name in stream_gated: + continue if not cb.should_run_guardrail( data=guardrail_data, event_type=GuardrailEventHooks.post_call, diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 54a0f18fd63..635cbf3ba9b 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -39,7 +39,7 @@ def _is_form_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) in _FORM_CONTENT_TYPES -def _is_json_content_type(content_type: str) -> bool: +def is_json_content_type(content_type: str) -> bool: """True iff the body should be parsed as JSON.""" return _normalize_media_type(content_type) == "application/json" @@ -406,7 +406,7 @@ async def get_request_body(request: Request) -> dict[str, Any]: """ if request.method == "POST": content_type: Final = request.headers.get("content-type", "") - if _is_json_content_type(content_type): + if is_json_content_type(content_type): return await _read_request_body(request) elif _is_form_content_type(content_type): return await get_form_data(request) diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 213a697b3dd..4b3ff2711b3 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -10,12 +10,24 @@ legacy internal names with `general_settings.use_team_public_model_name: false`. from __future__ import annotations -from collections.abc import Mapping, Sequence +import re +from collections.abc import Container, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +import litellm + if TYPE_CHECKING: from litellm.router import Router + from litellm.types.proxy.model_listing import ModelInfoResponse + +CLAUDE_CODE_PICKER_PATTERN: Final = re.compile(r"claude|anthropic", re.IGNORECASE) +GATEWAY_CLIENT_HEADER: Final = "x-gateway-client" +CLAUDE_CODE_CLIENT: Final = "claude-code" +_CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-" +_ONE_MILLION_SUFFIX: Final = "[1m]" +_ONE_MILLION_TOKENS: Final = 1_000_000 def configured_display_names( @@ -40,6 +52,115 @@ def configured_display_names( ) +def _unmarked(name: str) -> str: + return name[: -len(_ONE_MILLION_SUFFIX)] if name.lower().endswith(_ONE_MILLION_SUFFIX) else name + + +def _compatibility_id(model_id: str) -> str: + return f"{_CLAUDE_CODE_ALIAS_PREFIX}{model_id.encode().hex()}" + + +def _decoded_compatibility_id(view_id: str) -> str | None: + encoded: Final = _unmarked(view_id).removeprefix(_CLAUDE_CODE_ALIAS_PREFIX) + if encoded == _unmarked(view_id): + return None + try: + model_id: Final = bytes.fromhex(encoded).decode() + except (ValueError, UnicodeDecodeError): + return None + return model_id if _compatibility_id(model_id) == _unmarked(view_id) else None + + +def claude_code_model_id( + model_id: str, + max_input_tokens: float | None, + routing_names: Container[str], +) -> str: + """The collision-free id Claude Code's picker lists a model under.""" + if "*" in model_id: + return model_id + shaped: Final = model_id if CLAUDE_CODE_PICKER_PATTERN.search(model_id) else _compatibility_id(model_id) + one_million: Final = max_input_tokens is not None and max_input_tokens >= _ONE_MILLION_TOKENS + marked: Final = ( + f"{shaped}{_ONE_MILLION_SUFFIX}" if one_million and not shaped.lower().endswith(_ONE_MILLION_SUFFIX) else shaped + ) + return next( + ( + name + for name in (marked, shaped) + if name == model_id or claude_code_group_name(name, routing_names) == model_id + ), + model_id, + ) + + +def claude_code_group_name(view_id: str, routing_names: Container[str]) -> str | None: + """Decode a canonical compatibility id only when no configured route claims it.""" + if view_id in routing_names: + return None + unmarked: Final = _unmarked(view_id) + if unmarked != view_id and unmarked in routing_names: + return unmarked + model_id: Final = _decoded_compatibility_id(view_id) + return model_id if model_id and model_id in routing_names else None + + +def is_claude_code_client(headers: Mapping[str, str]) -> bool: + """Claude Code itself, or a client asking for its view of the listing the way Ramp Router's does""" + from litellm.llms.anthropic.common_utils import is_claude_code_user_agent + + return ( + is_claude_code_user_agent(headers.get("user-agent", "")) + or headers.get(GATEWAY_CLIENT_HEADER, "").lower() == CLAUDE_CODE_CLIENT + ) + + +def claude_code_view_ids( + rows: Sequence[ModelInfoResponse], + headers: Mapping[str, str], + routing_names: Container[str], +) -> Mapping[str, str]: + """served id -> Claude Code id for the requested listing view""" + if not is_claude_code_client(headers): + return MappingProxyType({}) + return MappingProxyType( + {row["id"]: claude_code_model_id(row["id"], row.get("max_input_tokens"), routing_names) for row in rows} + ) + + +@dataclass(frozen=True, slots=True) +class ClaudeCodeRoutingNames: + """Existing routes always own their names, including aliases and wildcard routes.""" + + llm_router: Router | None + team_id: str | None = None + alias_maps: tuple[object, ...] = () + + def __contains__(self, name: object) -> bool: + if not isinstance(name, str): + return False + if name in litellm.model_alias_map or any( + isinstance(aliases, Mapping) and name in aliases for aliases in self.alias_maps + ): + return True + if self.llm_router is None: + return False + return ( + name in self.llm_router.model_group_alias + or self.llm_router.has_model_id(name) + or bool(self.llm_router.get_candidate_model_ids_for_route(name, self.team_id)) + ) + + +def claude_code_requested_group( + requested: str, + llm_router: Router, + team_id: str | None, + alias_maps: tuple[object, ...] = (), +) -> str | None: + return claude_code_group_name(requested, ClaudeCodeRoutingNames(llm_router, team_id, alias_maps)) + + class TeamModelNameTranslator: """Translates internal team routing keys to their public names for the model listing/retrieve responses. Stateless; the live router and general_settings diff --git a/litellm/proxy/db/db_lookup_gate.py b/litellm/proxy/db/db_lookup_gate.py new file mode 100644 index 00000000000..2fd427687bd --- /dev/null +++ b/litellm/proxy/db/db_lookup_gate.py @@ -0,0 +1,23 @@ +import asyncio +from typing import Final + +from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY + + +class LoopBoundSemaphore: + __slots__ = ("_loop", "_semaphore", "_value") + + def __init__(self, value: int) -> None: + self._value: Final = value + self._loop: asyncio.AbstractEventLoop | None = None + self._semaphore: asyncio.Semaphore | None = None + + def current(self) -> asyncio.Semaphore: + loop: Final = asyncio.get_running_loop() + if self._semaphore is None or self._loop is not loop: + self._semaphore = asyncio.Semaphore(self._value) + self._loop = loop + return self._semaphore + + +db_lookup_gate: Final = LoopBoundSemaphore(PROXY_DB_LOOKUP_MAX_CONCURRENCY) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 4a0231ad9df..d1e4b3e92b8 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -34,12 +34,20 @@ writer's connection params (pool size, timeouts, pgbouncer mode) for the ones the reader URL does not pin itself. """ +import _ssl +import hashlib import os +import socket +import ssl +import struct +import sys +import tempfile import urllib.parse -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from functools import partial +from pathlib import Path from types import MappingProxyType -from typing import Annotated, Final, cast +from typing import Annotated, Final, Protocol, TypeAlias, cast from pydantic import AliasChoices, BeforeValidator, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -126,21 +134,100 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"}) +PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----" +PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103) +TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0 + +RootCertResolver: TypeAlias = Callable[[str, str, int], str] # mutable-ok: Callable parameter syntax -def translate_libpq_ssl_params(url: str) -> str: +class _VerifiedChainSource(Protocol): + def get_verified_chain(self) -> Sequence[_ssl.Certificate] | None: ... + + +def _verified_chain_der(tls: ssl.SSLSocket) -> tuple[bytes, ...]: + if sys.version_info >= (3, 13): + return tuple(tls.get_verified_chain()) + legacy: Final = cast( # cast-ok: the stub omits _sslobj, the C object has get_verified_chain since 3.10 + "_VerifiedChainSource | None", + tls._sslobj, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # public API only from 3.13 + ) + chain: Final = () if legacy is None else legacy.get_verified_chain() or () + return tuple(cert.public_bytes(_ssl.ENCODING_DER) for cert in chain) + + +def _server_trust_anchor(cafile: str, host: str, port: int) -> bytes | None: + try: + context: Final = ssl.create_default_context(cafile=cafile) + with socket.create_connection((host, port), timeout=TLS_PROBE_TIMEOUT_SECONDS) as raw: + raw.sendall(PG_SSL_REQUEST) + if raw.recv(1) != b"S": + return None + with context.wrap_socket(raw, server_hostname=host) as tls: + chain: Final = _verified_chain_der(tls) + except (OSError, ValueError): + return None + return chain[-1] if chain else None + + +def pin_bundle_root(cert_path: str, host: str, port: int) -> str: + """Reduce a multi-root CA bundle to the one root that verifies ``host``. + + Prisma's ``sslcert`` loads a single PEM certificate (native-tls + ``Certificate::from_pem``), so pointing it at a bundle such as the AWS RDS + global bundle trusts only the first of its 108 regional roots and the + handshake fails with "unable to get local issuer certificate" for every + other region. A single-certificate file is returned as is. For a bundle, + one verifying handshake (chain and hostname, whole bundle as trust store) + identifies the trust anchor the server actually chains to, which is + written to a single-certificate file for Prisma. If the probe fails the + bundle path is returned unchanged, so Prisma fails closed exactly as + before rather than trusting anything the bundle would not. + """ + try: + if Path(cert_path).read_bytes().count(PEM_CERT_HEADER) < 2: + return cert_path + except OSError: + return cert_path + root: Final = _server_trust_anchor(cert_path, host, port) + if root is None: + return cert_path + pinned: Final = Path(tempfile.gettempdir()) / f"litellm-sslcert-{hashlib.sha256(root).hexdigest()[:16]}.pem" + return str(pinned) if _replace_file(pinned, ssl.DER_cert_to_PEM_cert(root)) else cert_path + + +def _replace_file(target: Path, content: str) -> bool: + """Write ``content`` to a private temp file and rename it over ``target``, so + readers never see a partial file and a symlink planted at ``target`` is + replaced rather than followed.""" + try: + fd, staged = tempfile.mkstemp(dir=target.parent, prefix=f"{target.name}.") + except OSError: + return False + try: + with os.fdopen(fd, "w") as handle: + handle.write(content) + os.replace(staged, target) + except OSError: + Path(staged).unlink(missing_ok=True) + return False + return True + + +def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = pin_bundle_root) -> str: """Rewrite libpq's certificate-verification params into Prisma's dialect. Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert`` - (the CA bundle) and ``sslaccept=strict``. It silently discards + (a single CA certificate) and ``sslaccept=strict``. It silently discards ``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to ``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no certificate check at all. ``verify-ca`` and ``verify-full`` both become ``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes - ``sslcert``, and either one turns on ``sslaccept=strict`` (chain and - hostname), matching libpq where a root cert makes ``require`` verify. - Prisma params the operator pinned themselves win; anything else is left - untouched. + ``sslcert`` (run through ``resolve_root_cert``, which pins a multi-root + bundle down to the server's root), and either one turns on + ``sslaccept=strict`` (chain and hostname), matching libpq where a root + cert makes ``require`` verify. Prisma params the operator pinned + themselves win; anything else is left untouched. """ parsed: Final = urllib.parse.urlsplit(url) pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True)) @@ -154,7 +241,9 @@ def translate_libpq_ssl_params(url: str) -> str: if key != "sslrootcert" ) root_cert: Final = tuple( - ("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys + ("sslcert", resolve_root_cert(value, parsed.hostname or "", parsed.port or int(DEFAULT_POSTGRES_PORT))) + for key, value in pairs + if key == "sslrootcert" and "sslcert" not in keys ) strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),) query: Final = urllib.parse.urlencode(translated + root_cert + strict) diff --git a/litellm/proxy/db/gateway_request_tracking.py b/litellm/proxy/db/gateway_request_tracking.py index bebd74e877c..c9ace68db33 100644 --- a/litellm/proxy/db/gateway_request_tracking.py +++ b/litellm/proxy/db/gateway_request_tracking.py @@ -10,13 +10,28 @@ strings rather than passing the raw path through. Nothing a caller sends can add a key, so the fold and the table it commits to are bounded by (days x routes) however much traffic arrives, and the response path carries no unbounded queue that would block once full. + +A flush commits its whole snapshot as one multi-row ``INSERT ... ON CONFLICT DO +UPDATE`` rather than one upsert per key, so a worker costs the primary one +statement per interval however many routes it served. With +``use_redis_transaction_buffer`` on, workers instead push their snapshot to a +Redis list and one lock-holding pod folds every entry and writes the table, so +the deployment as a whole costs the primary one statement per interval. """ -from dataclasses import asdict +import json +from collections.abc import AsyncIterator, Iterable from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.caching import RedisCache +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory from litellm.types.proxy.gateway_requests import ( GatewayRequestCounts, @@ -28,6 +43,15 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient _EMPTY: Final = GatewayRequestCounts(successful_requests=0, failed_requests=0) +_TABLE: Final = '"LiteLLM_DailyGatewayRequests"' +_COLUMNS_PER_ROW: Final = 5 +_UTC_NOW: Final = "(NOW() AT TIME ZONE 'UTC')" +GATEWAY_REQUESTS_JOB_NAME: Final = "update_gateway_requests_job" + +_BufferedRows: TypeAlias = tuple[tuple[str, str, str, int, int], ...] +_BUFFERED_ROWS: Final = TypeAdapter(_BufferedRows) +_BUFFERED_ENTRIES: Final = TypeAdapter(tuple[str | bytes, ...]) +_NO_COUNTS: Final[GatewayRequestSnapshot] = MappingProxyType({}) def _utc_date() -> str: @@ -59,20 +83,54 @@ class GatewayRequestAccumulator: route) however long the database is unreachable. This buys at-least-once, not exactly-once, and the cost is worth stating. - The batch commits inside its context manager's ``__aexit__``, so a failure - raised after the transaction committed (a connection dropped while reading - the acknowledgement) restores counts that are already persisted, and the - next flush increments them a second time. Exactly-once would need a dedup - key the upserts could ignore on replay. For a traffic-volume metric a rare + The statement commits on the server before its acknowledgement is read, so + a failure raised after the commit (a connection dropped while reading the + acknowledgement) restores counts that are already persisted, and the next + flush increments them a second time. Exactly-once would need a dedup key + the upsert could ignore on replay. For a traffic-volume metric a rare overcount on a dropped acknowledgement beats losing a whole interval to every database blip, so the trade is deliberate. """ - for key, counts in snapshot.items(): - existing = self._counts.get(key, _EMPTY) - self._counts[key] = GatewayRequestCounts( - successful_requests=existing.successful_requests + counts.successful_requests, - failed_requests=existing.failed_requests + counts.failed_requests, - ) + self._counts = dict(fold_counts(chain(self._counts.items(), snapshot.items()))) # mutable-ok: fold replaced + + +def fold_counts(items: Iterable[tuple[GatewayRequestKey, GatewayRequestCounts]]) -> GatewayRequestSnapshot: + """Sum counts key-wise; the result stays bounded by (date x category x route).""" + folded: Final[dict[GatewayRequestKey, GatewayRequestCounts]] = {} # mutable-ok: local fold returned once + for key, counts in items: + existing = folded.get(key, _EMPTY) + folded[key] = GatewayRequestCounts( + successful_requests=existing.successful_requests + counts.successful_requests, + failed_requests=existing.failed_requests + counts.failed_requests, + ) + return folded + + +def build_gateway_requests_upsert(snapshot: GatewayRequestSnapshot) -> tuple[str, tuple[str | int, ...]]: + """ + One ``INSERT ... ON CONFLICT DO UPDATE`` that increments every (date, category, + route) in the snapshot. Rows are ordered by the conflict key so concurrent + writers lock rows in the same order and cannot deadlock. + """ + ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + rows: Final = ", ".join( + f"(${base + 1}::text, ${base + 2}::text, ${base + 3}::text, ${base + 4}::bigint, ${base + 5}::bigint, {_UTC_NOW})" + for base in range(0, len(ordered) * _COLUMNS_PER_ROW, _COLUMNS_PER_ROW) + ) + sql: Final = ( + f'INSERT INTO {_TABLE} ("date", "category", "route", "successful_requests", "failed_requests", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("date", "category", "route") DO UPDATE SET\n' + f' "successful_requests" = {_TABLE}."successful_requests" + EXCLUDED."successful_requests",\n' + f' "failed_requests" = {_TABLE}."failed_requests" + EXCLUDED."failed_requests",\n' + f' "updated_at" = {_UTC_NOW}' + ) + params: Final[tuple[str | int, ...]] = tuple( + value + for key, counts in ordered + for value in (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + ) + return sql, params async def commit_gateway_requests_to_db( @@ -80,50 +138,130 @@ async def commit_gateway_requests_to_db( prisma_client: "PrismaClient", snapshot: GatewayRequestSnapshot, ) -> None: - """Upsert one incrementing row per (date, category, route).""" + """Increment every (date, category, route) in the snapshot with a single statement.""" if not snapshot: return - ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + sql, params = build_gateway_requests_upsert(snapshot) + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client - # pyright: ignore[reportAny] on both lines -- prisma's generated client is untyped, - # so .db and every table action off it resolve to Any at this boundary. The dict - # literals below are the shape prisma's generated inputs require. - async with prisma_client.db.batch_() as batcher: # pyright: ignore[reportAny] # untyped prisma client - for key, counts in ordered: - columns = asdict(key) - batcher.litellm_dailygatewayrequests.upsert( # pyright: ignore[reportAny] # untyped prisma client - where={"date_category_route": columns}, # mutable-ok: prisma input is dict-shaped - data={ # mutable-ok: prisma input is dict-shaped - "create": { # mutable-ok: prisma input is dict-shaped - **columns, - "successful_requests": counts.successful_requests, - "failed_requests": counts.failed_requests, - }, - "update": { # mutable-ok: prisma input is dict-shaped - "successful_requests": {"increment": counts.successful_requests}, # mutable-ok: as above - "failed_requests": {"increment": counts.failed_requests}, # mutable-ok: as above - }, - }, + verbose_proxy_logger.debug( + "Gateway request tracking - committed %d aggregated rows in one statement", len(snapshot) + ) + + +class GatewayRequestRedisBuffer: + """ + Folds every worker's snapshot through one Redis list so a single pod per + interval writes the table, mirroring the spend writer's transaction buffer. + + Each entry is one worker's snapshot as JSON rows; the lock holder pops them, + sums them, and commits one statement. A commit failure pushes the summed + rows back so the next holder retries, keeping the at-least-once guarantee. + If that push fails too, the rows go back to the holder's own accumulator so + they ride along with its next flush instead of vanishing with the pop. + """ + + def __init__(self, *, redis_cache: RedisCache, pod_lock_manager: PodLockManager) -> None: + self._redis_cache: Final = redis_cache + self._pod_lock_manager: Final = pod_lock_manager + + async def push(self, snapshot: GatewayRequestSnapshot) -> None: + if not snapshot: + return + rows: Final[_BufferedRows] = tuple( + (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + for key, counts in snapshot.items() + ) + await self._redis_cache.async_rpush(key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, values=(json.dumps(rows),)) + + async def _pop_batch(self) -> tuple[str | bytes, ...]: + popped: Final[object] = await self._redis_cache.async_lpop( # pyright: ignore[reportAny] # redis returns Any + key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT + ) + if not popped: + return () + return _BUFFERED_ENTRIES.validate_python(popped if isinstance(popped, list) else (popped,)) + + async def _pop_all(self) -> AsyncIterator[str | bytes]: + while True: + batch = await self._pop_batch() + for entry in batch: + yield entry + if len(batch) < MAX_REDIS_BUFFER_DEQUEUE_COUNT: + return + + async def pop(self) -> GatewayRequestSnapshot: + entries: Final = tuple([entry async for entry in self._pop_all()]) + return fold_counts( + ( + GatewayRequestKey(date=date, category=category, route=route), + GatewayRequestCounts(successful_requests=succeeded, failed_requests=failed), ) + for entry in entries + for date, category, route, succeeded, failed in _BUFFERED_ROWS.validate_json(entry) + ) - verbose_proxy_logger.debug("Gateway request tracking - committed %d aggregated rows", len(ordered)) + async def commit_if_leader(self, prisma_client: "PrismaClient") -> GatewayRequestSnapshot: + """ + Drain the list and write it as one statement, but only on the pod holding the job lock. + + The lock is a lease, never released: the holder re-enters it on every flush and + keeps committing alone until the TTL lapses, so the primary sees one statement + per flush interval deployment-wide instead of one per worker. + + Returns the popped rows that could be neither committed nor re-queued, for the + caller to keep in memory. Empty on success. + """ + if not await self._pod_lock_manager.acquire_lock(cronjob_id=GATEWAY_REQUESTS_JOB_NAME): + return _NO_COUNTS + buffered: Final = await self.pop() + try: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=buffered) + except Exception: # noqa: BLE001 -- a failed commit must not stop the scheduler + verbose_proxy_logger.warning( + "Gateway request tracking - failed to commit %d buffered rows, re-queuing to Redis for the next flush", + len(buffered), + exc_info=True, + ) + return await self._requeue(buffered) + return _NO_COUNTS + + async def _requeue(self, snapshot: GatewayRequestSnapshot) -> GatewayRequestSnapshot: + try: + await self.push(snapshot) + except Exception: # noqa: BLE001 -- the rows go back to the caller's accumulator instead + verbose_proxy_logger.warning( + "Gateway request tracking - Redis re-queue failed, keeping %d rows in memory for the next flush", + len(snapshot), + exc_info=True, + ) + return snapshot + return _NO_COUNTS async def flush_gateway_requests( prisma_client: "PrismaClient", accumulator: GatewayRequestAccumulator, + redis_buffer: GatewayRequestRedisBuffer | None = None, ) -> None: """ Scheduler entrypoint. Never raises: a metering failure must not kill the job. + With ``redis_buffer`` the snapshot goes to Redis and only the lease holder + writes to Postgres. Shutdown passes no buffer so a departing worker writes its + own counts directly instead of parking them behind a lease it may not hold. + ``CancelledError`` is deliberately not caught, so a flush cancelled during shutdown drops its snapshot rather than restoring counts onto an accumulator the process is about to discard. """ snapshot: Final = accumulator.drain() try: - await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + if redis_buffer is None: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + else: + await redis_buffer.push(snapshot) except Exception: # noqa: BLE001 -- a failed flush must not stop the scheduler accumulator.restore(snapshot) verbose_proxy_logger.warning( @@ -131,3 +269,13 @@ async def flush_gateway_requests( len(snapshot), exc_info=True, ) + return + if redis_buffer is None: + return + try: + accumulator.restore(await redis_buffer.commit_if_leader(prisma_client)) + except Exception: # noqa: BLE001 -- entries still in Redis are drained by the next flush + verbose_proxy_logger.warning( + "Gateway request tracking - leader drain failed, buffered rows stay in Redis for the next flush", + exc_info=True, + ) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index a38b8a47dbd..e35b1c8c82b 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,6 +23,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, @@ -121,30 +122,33 @@ class SpendCounterReseed: if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None try: - if counter_key.startswith("spend:key:"): - token: Final = counter_key[len("spend:key:") :] - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) - elif counter_key.startswith("spend:team_member:"): - suffix: Final = counter_key[len("spend:team_member:") :] - if ":" not in suffix: + async with db_lookup_gate.current(): + if counter_key.startswith("spend:key:"): + token: Final = counter_key[len("spend:key:") :] + row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) + elif counter_key.startswith("spend:team_member:"): + suffix: Final = counter_key[len("spend:team_member:") :] + if ":" not in suffix: + return None + user_id, team_id = suffix.rsplit(":", 1) + row = await TeamMembershipRepository(prisma_client).table.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + ) + elif counter_key.startswith("spend:team:"): + team_id = counter_key[len("spend:team:") :] + row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + elif counter_key.startswith("spend:user:"): + user_id = counter_key[len("spend:user:") :] + row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): + return None + elif counter_key.startswith("spend:org:"): + org_id: Final = counter_key[len("spend:org:") :] + row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": org_id} + ) + else: return None - user_id, team_id = suffix.rsplit(":", 1) - row = await TeamMembershipRepository(prisma_client).table.find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} - ) - elif counter_key.startswith("spend:team:"): - team_id = counter_key[len("spend:team:") :] - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) - elif counter_key.startswith("spend:user:"): - user_id = counter_key[len("spend:user:") :] - row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): - return None - elif counter_key.startswith("spend:org:"): - org_id: Final = counter_key[len("spend:org:") :] - row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id}) - else: - return None except Exception: verbose_proxy_logger.exception("SpendCounterReseed.from_db: failed for %s", counter_key) return None diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 98707e7ddca..c37b9fff1f0 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -80,17 +80,19 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout def policy_for_model( llm_router: "Router | None", model_alias: str, - team_id: str | None, + request_kwargs: Mapping[str, object], request_tags: Sequence[str], ) -> AutoRouterCompressionPolicy | None: - """The compression policy of the auto router marker `model_alias` resolves to. + """The compression policy of the auto router marker `model_alias` resolves to for this caller. - Pre-call arming and the routing hook both resolve through here, so an alias with - several tag-scoped markers cannot suppress under one and then route under another. + Pre-call arming and the routing hook both resolve through here, and here resolves through the + router's own request-scoped deployment lookup, so an alias with several tag-scoped markers + cannot suppress under one and then route under another, and a team router reached by its + public name carries its policy for every principal that can reach it. """ if llm_router is None: return None - deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or () + deployments: Final = llm_router.deployments_for_request(model_alias, request_kwargs) markers: Final = tuple( litellm_params for deployment in deployments @@ -108,17 +110,6 @@ def policy_for_model( return next((policy for policy in candidates if policy is not None), None) -def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: - """The caller's team id, from whichever metadata bucket this surface writes to.""" - for meta_key in ("metadata", "litellm_metadata"): - meta = request_kwargs.get(meta_key) - if isinstance(meta, Mapping): - team_id = meta.get("user_api_key_team_id") - if isinstance(team_id, str): - return team_id - return None - - def _compression_guardrail_classes() -> tuple[type, ...]: """The registered guardrail classes whose provider compresses prompts.""" from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry @@ -172,7 +163,7 @@ async def arm_pre_call( policy: Final = policy_for_model( llm_router=llm_router, model_alias=model_alias, - team_id=team_id_from_request(data), + request_kwargs=data, request_tags=_get_tags_from_request_kwargs(data), ) if policy is None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 7204839f6d3..6a7ac4361b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -44,7 +44,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -917,7 +917,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await run_aws_signing( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) allow_chunking: Final = not self._content_uses_contextual_grounding(content) completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator @@ -1178,7 +1180,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): **base_request_data, "content": content, } # mutable-ok: outbound JSON request body - prepared_request: Final = self._prepare_request( + prepared_request: Final = await run_aws_signing( + self._prepare_request, credentials=credentials, data=bedrock_request_data, optional_params=self.optional_params, @@ -1875,10 +1878,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() api_key: Final[str | None] = request_data.get("api_key") if request_data else None - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await run_aws_signing( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} - prepared_request: Final = self._prepare_request( + prepared_request: Final = await run_aws_signing( + self._prepare_request, credentials=credentials, data=body, optional_params=self.optional_params, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1785a2f0992..747cfa526d9 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1582,6 +1582,13 @@ async def _show_no_redis_warning() -> bool: return await count_live_proxy_workers(prisma_client) != 1 +def _show_env_credential_login_warning() -> bool: + from litellm.proxy.auth.login_utils import is_env_credential_login_enabled + from litellm.proxy.proxy_server import general_settings + + return is_env_credential_login_enabled(general_settings) + + async def _get_health_readiness_details( response: Response | None = None, ) -> dict[str, Any]: @@ -1623,6 +1630,7 @@ async def _get_health_readiness_details( log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel()) is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG) show_no_redis_warning: Final = await _show_no_redis_warning() + show_env_credential_login_warning: Final = _show_env_credential_login_warning() # check DB if prisma_client is not None: # if db passed in, check if it's connected @@ -1650,6 +1658,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } else: return { @@ -1662,6 +1671,7 @@ async def _get_health_readiness_details( "log_level": log_level_name, "is_detailed_debug": is_detailed_debug, "show_no_redis_warning": show_no_redis_warning, + "show_env_credential_login_warning": show_env_credential_login_warning, } except Exception as e: raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})") diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index de8834449de..a074f02f4e8 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import ( resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral @@ -659,22 +663,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Add additional priority-specific headers - if isinstance(response, ModelResponse): + if response_has_hidden_params(response): priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) - - # Get existing additional headers - additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} - - # Add priority information + additional_headers: Final = ensure_response_additional_headers(response) additional_headers["x-litellm-priority"] = priority or "default" additional_headers["x-litellm-rate-limiter-version"] = "v3" - # Update response - if not hasattr(response, "_hidden_params"): - response._hidden_params = {} - response._hidden_params["additional_headers"] = additional_headers - return response except Exception as e: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 31437af7770..c6c3dde4b6e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -31,6 +31,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( ESTIMATED_OUTPUT_TOKENS_FIELD, @@ -52,6 +53,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.router_utils.add_retry_fallback_headers import ( + ensure_response_additional_headers, + response_has_hidden_params, +) from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( @@ -3303,7 +3308,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - raw_estimated_input_tokens: Final = self._estimate_precise_input_tokens( + raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type ) estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) @@ -4677,34 +4682,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Post-call hook to update rate limit headers in the response. """ try: - from pydantic import BaseModel - stash: Final = get_request_stash() litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None - if litellm_proxy_rate_limit_response is not None: - # Update response headers - if hasattr(response, "_hidden_params"): - _hidden_params = getattr(response, "_hidden_params") - else: - _hidden_params = None - - if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) - ): - if isinstance(_hidden_params, BaseModel): - _hidden_params = _hidden_params.model_dump() - - _additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers( - additional_headers=_hidden_params.get("additional_headers", {}) or {}, + if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response): + additional_headers: Final = ensure_response_additional_headers(response) + additional_headers.update( + self._merge_ratelimit_statuses_into_additional_headers( + additional_headers={}, statuses=litellm_proxy_rate_limit_response["statuses"], ) - - setattr( - response, - "_hidden_params", - {**_hidden_params, "additional_headers": _additional_headers}, - ) + ) except Exception as e: verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3250ae5cca9..da033dc2276 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -773,6 +773,16 @@ def apply_missing_session_id_policy( return if policy == "omit": metadata[SESSION_ID_OMITTED_METADATA_KEY] = True + requester_metadata: Final = data.get("metadata") + requester_session_id: Final = ( + requester_metadata.get("session_id") if isinstance(requester_metadata, dict) else None + ) + if ( + (body_session_id := data.get("litellm_session_id")) + and not metadata.get("session_id") + and not requester_session_id + ): + metadata["session_id"] = body_session_id return if data.get("litellm_session_id") or metadata.get("session_id"): return @@ -1750,7 +1760,9 @@ class LiteLLMProxyRequestSetup: callback_vars_dict.pop("success_callback", None) callback_vars_dict.pop("failure_callback", None) callback_vars_dict = { - key: (litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else value) + key: ( + litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else str(value) + ) for key, value in callback_vars_dict.items() } diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 204051c3715..dc0da63555f 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -18,6 +18,7 @@ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel import litellm +from litellm._internal_context import current_billing_time, pinned_billing_time from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import completion_cost from litellm.proxy._types import ( @@ -27,7 +28,15 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo +from litellm.types.utils import ( + CostBreakdown, + CostPerToken, + LlmProvidersSet, + ModelInfo, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) router: Final = APIRouter() @@ -46,13 +55,15 @@ def _configured_price(key: str, sources: tuple[Mapping[str, object], ...]) -> fl def _extract_custom_pricing( - litellm_params: Mapping[str, object], model_info: Mapping[str, object] + litellm_params: Mapping[str, object], model_info: Mapping[str, object], builtin: ModelInfo | None ) -> CostPerToken | None: """ Pull per-token pricing configured on a deployment so on-prem / self-hosted models (absent from the public cost map) still estimate a real cost. Pricing may live on ``litellm_params`` or ``model_info``; ``litellm_params`` - wins, matching the router's cost-map registration precedence. + wins, matching the router's cost-map registration precedence. Cache rates the + deployment leaves unset come from the backend model's built-in entry, then its + own input rate, again matching what the router registers for live billing. """ sources: Final = (litellm_params, model_info) input_price: Final = _configured_price("input_cost_per_token", sources) @@ -61,15 +72,21 @@ def _extract_custom_pricing( if input_price is None and output_price is None: return None + input_rate: Final = input_price or 0.0 + cache_sources: Final = sources if builtin is None else (*sources, builtin) + cache_read_price: Final = _configured_price("cache_read_input_token_cost", cache_sources) + cache_creation_price: Final = _configured_price("cache_creation_input_token_cost", cache_sources) return CostPerToken( - input_cost_per_token=input_price or 0.0, + input_cost_per_token=input_rate, output_cost_per_token=output_price or 0.0, + cache_read_input_token_cost=input_rate if cache_read_price is None else cache_read_price, + cache_creation_input_token_cost=input_rate if cache_creation_price is None else cache_creation_price, ) -def _lookup_model_info(model: str) -> ModelInfo | None: +def _lookup_model_info(model: str, custom_llm_provider: str | None = None) -> ModelInfo | None: try: - return litellm.get_model_info(model=model) + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: return None @@ -98,17 +115,14 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: model_info: Final = first_deployment.get("model_info", {}) custom_llm_provider: Final = litellm_params.get("custom_llm_provider") provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None - custom_cost_per_token: Final = _extract_custom_pricing(litellm_params, model_info) - - # Check base_model first (needed for Azure custom deployment names) + # base_model wins (needed for Azure custom deployment names) base_model: Final = model_info.get("base_model") or litellm_params.get("base_model") - if base_model: - verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model) - return ResolvedCostModel(str(base_model), provider, custom_cost_per_token) - - resolved_model: Final = litellm_params.get("model") + resolved_model: Final = base_model or litellm_params.get("model") if resolved_model: verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model) + custom_cost_per_token: Final = _extract_custom_pricing( + litellm_params, model_info, _lookup_model_info(str(resolved_model), provider) + ) return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token) except Exception as e: verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e) @@ -117,19 +131,59 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel: return ResolvedCostModel(model, None, None) -def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost): - """ - Calculate costs for a given number of requests. +@dataclass(frozen=True, slots=True) +class CostLines: + """Cost of one request split the way the spend logs split it: the cache lines are + shares of input_cost and the reasoning line is a share of output_cost.""" - Returns tuple of (total_cost, input_cost, output_cost, margin_cost) or all None if num_requests is None/0. - """ - if not num_requests: - return None, None, None, None - return ( - cost_per_request * num_requests, - input_cost * num_requests, - output_cost * num_requests, - margin_cost * num_requests, + total_cost: float + input_cost: float + output_cost: float + margin_cost: float + cache_read_cost: float + cache_creation_cost: float + reasoning_cost: float + + def times(self, num_requests: int | None) -> "CostLines | None": + if not num_requests: + return None + return CostLines( + total_cost=self.total_cost * num_requests, + input_cost=self.input_cost * num_requests, + output_cost=self.output_cost * num_requests, + margin_cost=self.margin_cost * num_requests, + cache_read_cost=self.cache_read_cost * num_requests, + cache_creation_cost=self.cache_creation_cost * num_requests, + reasoning_cost=self.reasoning_cost * num_requests, + ) + + +def _cost_lines(cost_per_request: float, cost_breakdown: CostBreakdown | None) -> CostLines: + breakdown: Final = cost_breakdown if cost_breakdown is not None else CostBreakdown() + return CostLines( + total_cost=cost_per_request, + input_cost=breakdown.get("input_cost", 0.0), + output_cost=breakdown.get("output_cost", 0.0), + margin_cost=breakdown.get("margin_total_amount", 0.0), + cache_read_cost=breakdown.get("cache_read_cost", 0.0), + cache_creation_cost=breakdown.get("cache_creation_cost", 0.0), + reasoning_cost=breakdown.get("reasoning_cost", 0.0), + ) + + +def _usage_for_estimate(request: CostEstimateRequest) -> Usage: + cache_tokens: Final = request.cache_read_input_tokens + request.cache_creation_input_tokens + return Usage( + prompt_tokens=request.input_tokens, + completion_tokens=request.output_tokens, + total_tokens=request.input_tokens + request.output_tokens, + reasoning_tokens=request.reasoning_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=request.cache_read_input_tokens, + cache_creation_tokens=request.cache_creation_input_tokens, + ) + if cache_tokens + else None, ) @@ -530,11 +584,14 @@ async def estimate_cost( - model: Model name (e.g., "gpt-4", "claude-3-opus") - input_tokens: Expected input tokens per request - output_tokens: Expected output tokens per request + - cache_read_input_tokens: Cache-read tokens per request, counted within input_tokens (optional) + - cache_creation_input_tokens: Cache-write tokens per request, counted within input_tokens (optional) + - reasoning_tokens: Reasoning tokens per request, counted within output_tokens (optional) - num_requests_per_day: Number of requests per day (optional) - num_requests_per_month: Number of requests per month (optional) Returns cost breakdown including: - - Per-request costs (input, output, margin) + - Per-request costs (input, output, margin, plus the cache-read, cache-write and reasoning shares) - Daily costs (if num_requests_per_day provided) - Monthly costs (if num_requests_per_month provided) @@ -543,14 +600,15 @@ async def estimate_cost( { "model": "gpt-4", "input_tokens": 1000, + "cache_read_input_tokens": 800, "output_tokens": 500, + "reasoning_tokens": 200, "num_requests_per_day": 100, "num_requests_per_month": 3000 } ``` """ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import ModelResponse, Usage # Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4') resolved: Final = _resolve_model_for_cost_lookup(request.model) @@ -559,15 +617,8 @@ async def estimate_cost( verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model) - # Create a mock response with usage for completion_cost - mock_response: Final = ModelResponse( - model=resolved_model, - usage=Usage( - prompt_tokens=request.input_tokens, - completion_tokens=request.output_tokens, - total_tokens=request.input_tokens + request.output_tokens, - ), - ) + usage: Final = _usage_for_estimate(request) + mock_response: Final = ModelResponse(model=resolved_model, usage=usage) # Create a logging object to capture cost breakdown litellm_logging_obj: Final = LiteLLMLoggingObj( @@ -580,92 +631,73 @@ async def estimate_cost( function_id="cost-estimate", ) - # Use completion_cost which handles all the logic including margins/discounts - try: - cost_per_request: Final = completion_cost( - completion_response=mock_response, - model=resolved_model, - custom_llm_provider=resolved_provider, - custom_cost_per_token=resolved.custom_cost_per_token, - litellm_logging_obj=litellm_logging_obj, - ) - except Exception as e: - raise HTTPException( - status_code=404, - detail={ - "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}" - }, - ) + # Pinning one moment keeps an off-peak window that opens mid-quote from pricing the totals on + # one side of it and the reported rates on the other. + with pinned_billing_time(current_billing_time()): + # Use completion_cost which handles all the logic including margins/discounts + try: + cost_per_request: Final = completion_cost( + completion_response=mock_response, + model=resolved_model, + custom_llm_provider=resolved_provider, + custom_cost_per_token=resolved.custom_cost_per_token, + litellm_logging_obj=litellm_logging_obj, + ) + except Exception as e: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriceable model + raise HTTPException( + status_code=404, + detail={ + "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}" + }, + ) - # Get cost breakdown from the logging object - cost_breakdown: Final = litellm_logging_obj.cost_breakdown + # The rates come back from the pricing call itself rather than a second lookup, so they are the + # ones the cost lines above billed at even when completion_cost infers a provider this endpoint + # never resolved (an unrouted "xai/grok-4" prices on xai's inclusive tier thresholds; a lookup + # here without that provider would report the sub-200k rate for a line billed above it). + rates: Final = litellm_logging_obj.billed_token_rates + per_request: Final = _cost_lines(cost_per_request, litellm_logging_obj.cost_breakdown) + daily: Final = per_request.times(request.num_requests_per_day) + monthly: Final = per_request.times(request.num_requests_per_month) - input_cost: Final = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0 - output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0 - margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 - - model_info: Final = _lookup_model_info(resolved_model) - mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None - mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None + model_info: Final = _lookup_model_info(resolved_model, resolved_provider) mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None - - input_cost_per_token: Final = ( - resolved.custom_cost_per_token["input_cost_per_token"] - if resolved.custom_cost_per_token is not None - else mapped_input_price - ) - output_cost_per_token: Final = ( - resolved.custom_cost_per_token["output_cost_per_token"] - if resolved.custom_cost_per_token is not None - else mapped_output_price - ) custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider - # Calculate daily and monthly costs - ( - daily_cost, - daily_input_cost, - daily_output_cost, - daily_margin_cost, - ) = _calculate_period_costs( - num_requests=request.num_requests_per_day, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) - ( - monthly_cost, - monthly_input_cost, - monthly_output_cost, - monthly_margin_cost, - ) = _calculate_period_costs( - num_requests=request.num_requests_per_month, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) - return CostEstimateResponse( model=request.model, input_tokens=request.input_tokens, output_tokens=request.output_tokens, + cache_read_input_tokens=request.cache_read_input_tokens, + cache_creation_input_tokens=request.cache_creation_input_tokens, + reasoning_tokens=request.reasoning_tokens, num_requests_per_day=request.num_requests_per_day, num_requests_per_month=request.num_requests_per_month, - cost_per_request=cost_per_request, - input_cost_per_request=input_cost, - output_cost_per_request=output_cost, - margin_cost_per_request=margin_cost, - daily_cost=daily_cost, - daily_input_cost=daily_input_cost, - daily_output_cost=daily_output_cost, - daily_margin_cost=daily_margin_cost, - monthly_cost=monthly_cost, - monthly_input_cost=monthly_input_cost, - monthly_output_cost=monthly_output_cost, - monthly_margin_cost=monthly_margin_cost, - input_cost_per_token=input_cost_per_token, - output_cost_per_token=output_cost_per_token, + cost_per_request=per_request.total_cost, + input_cost_per_request=per_request.input_cost, + output_cost_per_request=per_request.output_cost, + margin_cost_per_request=per_request.margin_cost, + cache_read_cost_per_request=per_request.cache_read_cost, + cache_creation_cost_per_request=per_request.cache_creation_cost, + reasoning_cost_per_request=per_request.reasoning_cost, + daily_cost=daily.total_cost if daily is not None else None, + daily_input_cost=daily.input_cost if daily is not None else None, + daily_output_cost=daily.output_cost if daily is not None else None, + daily_margin_cost=daily.margin_cost if daily is not None else None, + daily_cache_read_cost=daily.cache_read_cost if daily is not None else None, + daily_cache_creation_cost=daily.cache_creation_cost if daily is not None else None, + daily_reasoning_cost=daily.reasoning_cost if daily is not None else None, + monthly_cost=monthly.total_cost if monthly is not None else None, + monthly_input_cost=monthly.input_cost if monthly is not None else None, + monthly_output_cost=monthly.output_cost if monthly is not None else None, + monthly_margin_cost=monthly.margin_cost if monthly is not None else None, + monthly_cache_read_cost=monthly.cache_read_cost if monthly is not None else None, + monthly_cache_creation_cost=monthly.cache_creation_cost if monthly is not None else None, + monthly_reasoning_cost=monthly.reasoning_cost if monthly is not None else None, + input_cost_per_token=rates.input_cost_per_token if rates is not None else None, + output_cost_per_token=rates.output_cost_per_token if rates is not None else None, + cache_read_input_token_cost=rates.cache_read_input_token_cost if rates is not None else None, + cache_creation_input_token_cost=rates.cache_creation_input_token_cost if rates is not None else None, + output_cost_per_reasoning_token=rates.output_cost_per_reasoning_token if rates is not None else None, provider=custom_llm_provider, ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f46c4170071..749a940de0e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4559,6 +4559,23 @@ async def delete_verification_tokens( litellm_changed_by=litellm_changed_by, ) + # Snapshot before the delete: the FK cascade drops the mapping rows, but their + # cached jwt_key_mapping entries still resolve to the now-dead token (LIT-5380). + jwt_mapping_cache_keys: Final[tuple[str, ...]] = tuple( + cache_key + for keys_for_token in await asyncio.gather( + *( + get_jwt_key_mapping_cache_keys_for_token( + hashed_token=key.token, + prisma_client=prisma_client, + ) + for key in authorized_keys + if key.token is not None + ) + ) + for cache_key in keys_for_token + ) + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: deleted_tokens = await prisma_client.delete_data(tokens=tokens) if deleted_tokens is not None and len(deleted_tokens) != len(tokens): @@ -4571,6 +4588,8 @@ async def delete_verification_tokens( if len(deleted_tokens) != len(tokens): failed_tokens = [token for token in tokens if token not in deleted_tokens] + await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) + else: raise Exception("DB not connected. prisma_client is None") except Exception as e: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 2ec68a10f65..c050368b3fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5601,7 +5601,7 @@ async def team_model_add( updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"updated_at": datetime.now(timezone.utc)}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if updated_team is None: raise HTTPException( @@ -5688,7 +5688,7 @@ async def team_model_delete( updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, - include={"object_permission": True}, + include={"litellm_model_table": True, "object_permission": True}, ) if updated_team is None: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3e6434a5afd..c60888e298f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -22,7 +22,6 @@ from html import escape from types import MappingProxyType from typing import ( TYPE_CHECKING, - Annotated, Any, Final, Literal, @@ -42,7 +41,7 @@ if TYPE_CHECKING: import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse -from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -95,6 +94,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.team_grants import TeamModelAliasTable from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import ( admin_ui_disabled, @@ -209,31 +209,14 @@ def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]": return repo.table -_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) _SSO_TOKEN_CLAIMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -def _decode_model_aliases(value: object) -> object: - """``/team/new`` stores team model aliases as a JSON-encoded string in the Json column.""" - if not isinstance(value, str): - return value - try: - return _MODEL_ALIASES_ADAPTER.validate_json(value) - except ValidationError: - return None - - -class _TeamModelAliasTable(BaseModel): - model_config = ConfigDict(protected_namespaces=()) - - model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None - - class _TeamRowGrants(BaseModel): team_id: str team_alias: str | None = None models: tuple[str, ...] = () - litellm_model_table: _TeamModelAliasTable | None = None + litellm_model_table: TeamModelAliasTable | None = None class CliSsoTeamDetail(BaseModel): diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 2ea46b740a8..face515ef88 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,6 +15,7 @@ import os import re from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast @@ -33,6 +34,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -52,6 +54,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_set_request_parsed_body, get_form_data, get_request_body, + is_json_content_type, ) from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, @@ -77,6 +80,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials +from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager @@ -119,6 +123,24 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li return False +class RelayRejection(TypedDict): + error: ReadOnly[str] + + +def _deployment_model_name(litellm_params: LiteLLMParamsTypedDict) -> str: + model: Final = litellm_params.get("model", "") + try: + return get_llm_provider(model=model, custom_llm_provider=litellm_params.get("custom_llm_provider"))[0] + except litellm.BadRequestError: + return model + + +def _models_served_by_group(llm_router: litellm.Router, model_group: str) -> frozenset[str]: + return frozenset( + _deployment_model_name(row["litellm_params"]) for row in llm_router.get_model_list(model_name=model_group) or () + ) + + def is_passthrough_request_streaming(request_body: object) -> bool: """ Returns True if the request is streaming. @@ -411,7 +433,7 @@ async def vllm_proxy_route( content=None, data=None, files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), params=None, headers=None, cookies=None, @@ -1099,13 +1121,6 @@ async def bedrock_proxy_route( """ create_request_copy(request) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME") if not _is_bedrock_agent_runtime_route(endpoint=endpoint): return await bedrock_llm_proxy_route( @@ -1136,20 +1151,24 @@ async def bedrock_proxy_route( ) # Add or update query parameters + from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() - credentials: Final[Credentials] = bedrock_llm.get_credentials() - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - headers: Final = {"Content-Type": "application/json"} # Assuming the body contains JSON data, parse it try: data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await run_aws_signing( + sign_aws_json_post, + get_credentials=bedrock_llm.get_credentials, + service_name="bedrock", + aws_region_name=aws_region_name, + url=str(updated_url), + body=json.dumps(data), + headers=MappingProxyType({"Content-Type": "application/json"}), + ) ## check for streaming is_streaming_request = False @@ -1207,13 +1226,6 @@ async def comprehend_medical_proxy_route( [Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical) """ - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call comprehendmedical. Run 'pip install boto3'.") - from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import ( COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS, ) @@ -1244,20 +1256,23 @@ async def comprehend_medical_proxy_route( if "stream" in data: raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member") - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post - credentials: Final[Credentials] = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name) - sigv4: Final = SigV4Auth(credentials, "comprehendmedical", aws_region_name) - headers: Final = MappingProxyType( - { - "Content-Type": "application/x-amz-json-1.1", - "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", - } - ) target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" - _request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await run_aws_signing( + sign_aws_json_post, + get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), + service_name="comprehendmedical", + aws_region_name=aws_region_name, + url=target_url, + body=json.dumps(data), + headers=MappingProxyType( + { + "Content-Type": "application/x-amz-json-1.1", + "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", + } + ), + ) endpoint_func: Final = create_pass_through_route( endpoint=operation, @@ -1505,6 +1520,14 @@ async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> Async await upstream.aclose() +async def _relay_upstream_response(upstream: httpx.Response) -> Response: + return Response( + content=await upstream.aread(), + status_code=upstream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), + ) + + async def _relay_azure_router_model( llm_router: litellm.Router, model: str, @@ -1514,30 +1537,37 @@ async def _relay_azure_router_model( is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, ) -> Response: - result: Final = await llm_router.allm_passthrough_route( - model=model, - method=request.method, - endpoint=endpoint, - request_query_params=request.query_params, - request_headers=_safe_get_request_headers(request), - stream=is_streaming_request, - content=None, - data=None, - files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), - params=None, - headers=None, - cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + foreign_deployment: Final = foreign_azure_deployment( + endpoint, model, lambda: _models_served_by_group(llm_router, model) ) + if foreign_deployment is not None: + rejection: Final[RelayRejection] = { + "error": f"deployment '{foreign_deployment}' in the path is not served by model group '{model}'; " + "put the model group name in the deployments segment" + } + raise HTTPException(status_code=400, detail=rejection) + try: + result: Final = await llm_router.allm_passthrough_route( + model=model, + method=request.method, + endpoint=endpoint, + request_query_params=request.query_params, + request_headers=_safe_get_request_headers(request), + stream=is_streaming_request, + content=None, + data=None, + files=None, + json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None), + params=None, + headers=None, + cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + ) + except httpx.HTTPStatusError as upstream_error: + return await _relay_upstream_response(upstream_error.response) if not is_streaming_request: - upstream: Final = cast(httpx.Response, result) - return Response( - content=await upstream.aread(), - status_code=upstream.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), - ) + return await _relay_upstream_response(cast(httpx.Response, result)) if inspect.isasyncgen(result): sse_headers: Final = {"content-type": "text/event-stream"} diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 5f6489a69ca..93fe3c5b31b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -4,6 +4,7 @@ OpenAI Passthrough Logging Handler Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. """ +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final from urllib.parse import urlparse @@ -16,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig @@ -96,6 +98,47 @@ def _is_openai_compatible_url(url_route: str | None) -> bool: return False +def _is_remote_high_detail_image(part: object) -> bool: + if not isinstance(part, Mapping) or part.get("type") != "image_url": + return False + image_url: Final = part.get("image_url") + if not isinstance(image_url, Mapping): + return False + url: Final = image_url.get("url") + return ( + isinstance(url, str) and url.lower().startswith(("http://", "https://")) and image_url.get("detail") == "high" + ) + + +def _content_parts(message: Mapping[str, object]) -> Sequence[object]: + content: Final = message.get("content") + return content if isinstance(content, list) else () + + +def _without_remote_high_detail_images(message: Mapping[str, object]) -> Mapping[str, object]: + if not isinstance(message.get("content"), list): + return message + kept_parts: Final = [ # mutable-ok: token_counter reads message content only when it is a list + part for part in _content_parts(message) if not _is_remote_high_detail_image(part) + ] + return {**message, "content": kept_parts} # mutable-ok: token_counter rejects any message that is not a dict + + +def count_relayed_prompt_tokens(model: str, messages: Sequence[Mapping[str, object]] | None) -> int: + if messages is None: + return 0 + remote_high_detail_images: Final = sum( + 1 for message in messages for part in _content_parts(message) if _is_remote_high_detail_image(part) + ) + local_messages: Final = [ # mutable-ok: token_counter takes a list of messages + _without_remote_high_detail_images(message) for message in messages + ] + return ( + litellm.token_counter(model=model, messages=local_messages) + + high_detail_image_token_upper_bound() * remote_high_detail_images + ) + + class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): """ OpenAI-specific passthrough logging handler that provides cost tracking for /chat/completions endpoints. @@ -512,9 +555,10 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): def _build_complete_streaming_response( self, - all_chunks: list[str], + all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, + messages: Sequence[Mapping[str, object]] | None = None, ) -> ModelResponse | TextCompletionResponse | None: """ Builds complete response from raw chunks for OpenAI streaming responses. @@ -558,7 +602,11 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return None # Build complete response from chunks - complete_streaming_response: Final = litellm.stream_chunk_builder(chunks=all_openai_chunks) + complete_streaming_response: Final = litellm.stream_chunk_builder( + chunks=all_openai_chunks, + messages=messages, + count_prompt_tokens=lambda: count_relayed_prompt_tokens(model, messages), + ) return complete_streaming_response diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 7bcb79cefc9..ed193c7f434 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None: return None if texts is None else tuple(texts) +def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]: + return tuple(texts or ()) + + def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None: return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls) @@ -78,6 +82,10 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non return sent is not None and returned is not None and returned != sent +def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: + return sent is not None and returned is not None and len(returned) != len(sent) + + _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) @@ -89,10 +97,11 @@ def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT: class _StreamRewriteObserver(CustomGuardrail): """Stand-in handed to the endpoint translation in place of a streaming pipeline step's guardrail. It records whether the guardrail returned different output than it was given, - which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text - rewrites are deliverable on translations that write them back across the buffered chunks - (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any - other translation are discarded by the executor, which releases the original chunks. + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and + tool-call rewrites are deliverable on translations that write them back across the + buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation, + and a rewrite that drops or adds a tool call on any translation, are discarded by the + executor, which releases the original chunks. The inner guardrail's ``apply_guardrail`` already records the guardrail information and span, so the observer's stays out of ``log_guardrail_information``.""" @@ -101,6 +110,7 @@ class _StreamRewriteObserver(CustomGuardrail): self.inner: Final = inner self.rewrote_texts = False self.rewrote_tool_calls = False + self.changed_tool_call_count = False def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -118,13 +128,103 @@ class _StreamRewriteObserver(CustomGuardrail): outputs: Final = await self.inner.apply_guardrail( inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj ) + returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls")) self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) - self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote( - sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls")) + self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes) + self.changed_tool_call_count = self.changed_tool_call_count or _changed_count( + sent_tool_shapes, returned_tool_shapes ) return outputs +class _ScannedTextRecorder(CustomGuardrail): + def __init__(self, guardrail_name: str) -> None: + super().__init__(guardrail_name=guardrail_name) + self.inputs: GenericGuardrailAPIInputs | None = None + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + self.inputs = inputs + return inputs + + +class _LegacyHookStreamAdapter(CustomGuardrail): + """Runs a guardrail that only implements the legacy post-call hook (no unified + ``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The + endpoint translation hands it the texts it scanned plus the assembled response under + ``request_data["response"]``; the hook gets that response in the shape its route gives + non-streaming hooks, an exception it raises ends the stream through the executor's + fail/error classification, and the response it hands back, or the one it changed in place + and returned ``None`` for, is re-scanned by the same translation so its texts reach the + client through the translation's ended-stream write-back. A + replacement whose scanned texts do not line up with the originals, or whose tool calls + differ from them, is undeliverable, so the executor releases the original chunks. A stream + that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as + long as the hook left the tool calls alone.""" + + def __init__( + self, + inner: CustomGuardrail, + endpoint_translation: "BaseTranslation", + user_api_key_dict: "UserAPIKeyAuth", + ) -> None: + super().__init__(guardrail_name=inner.guardrail_name) + self.inner: Final = inner + self.endpoint_translation: Final = endpoint_translation + self.user_api_key_dict: Final = user_api_key_dict + + def structured_messages_cover_full_request(self) -> bool: + return self.inner.structured_messages_cover_full_request() + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response")) + replacement: Final = await self.inner.async_post_call_success_hook( + data=request_data, + user_api_key_dict=self.user_api_key_dict, + response=hooked, + ) + rewrite: Final = hooked if replacement is None else replacement + if rewrite is None: + return inputs + rescanned: Final = await self._rescan(rewrite, logging_obj) + if rescanned is None: + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + rewritten: Final = rescanned.get("texts") + if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if not rewritten: + return inputs + rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} + return rewritten_inputs + + async def _rescan( + self, response: object, logging_obj: "LiteLLMLoggingObj | None" + ) -> GenericGuardrailAPIInputs | None: + recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown") + await self.endpoint_translation.process_output_response( + response=response, + guardrail_to_apply=recorder, + litellm_logging_obj=logging_obj, + user_api_key_dict=self.user_api_key_dict, + ) + return recorder.inputs + + def _prepare_hook_input( step: PipelineStep, callback: CustomGuardrail, @@ -292,18 +392,29 @@ class PipelineExecutor: endpoint_translation: "BaseTranslation", streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place hook_input: dict[str, object], # mutable-ok: same request-payload shape as data - user_api_key_dict: "UserAPIKeyAuth | None", + user_api_key_dict: "UserAPIKeyAuth", litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> None: """Run one streaming post_call step through the endpoint translation, delivering - text rewrites on translations that support ended-stream write-back. A rewrite that - cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation - without write-back, or one the translation refused with - ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the - originals and the step passes, so the client gets the stream the merge base sent.""" - observer: Final = _StreamRewriteObserver(callback) - deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites + text and tool-call rewrites on translations that support ended-stream write-back. A + guardrail without the unified interface runs its legacy post-call hook against the + assembled response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the + client yet (one on a translation without write-back, one that drops or adds a tool call, + or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is + discarded: the buffered chunks go back to the originals and the step passes, so the + client gets the stream the merge base sent, and the guardrail stays out of the + applied-guardrails header since its output never reached the client. The response an + earlier step's translation stored under ``request_data["response"]`` is dropped first, + so this step's hook sees the stream as the steps before it left it.""" + scanner: Final = ( + callback + if PipelineExecutor.supports_unified_execution(callback) + else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict) + ) + observer: Final = _StreamRewriteObserver(scanner) + deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites originals: Final = copy.deepcopy(streaming_chunks) + hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored try: if deliver_rewrites: await endpoint_translation.process_output_streaming_response( @@ -324,9 +435,12 @@ class PipelineExecutor: ) except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) - else: - if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): - _release_original_chunks(step.guardrail, streaming_chunks, originals) + return + if observer.changed_tool_call_count or ( + not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls) + ): + _release_original_chunks(step.guardrail, streaming_chunks, originals) + return if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) @@ -386,11 +500,11 @@ class PipelineExecutor: if isinstance(response, dict): callback.mark_pre_call_hook_ran(response) elif mode == "post_call" and streaming_chunks is not None: - if not use_unified or endpoint_translation is None: + if endpoint_translation is None: return ( "error", None, - f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", + f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation", None, ) await PipelineExecutor._run_streaming_step( @@ -446,10 +560,22 @@ class PipelineExecutor: @staticmethod def supports_unified_execution(callback: CustomGuardrail) -> bool: - """Whether this guardrail runs through the unified apply_guardrail path, - the interface streaming pipeline execution requires.""" + """Whether this guardrail runs through the unified apply_guardrail path.""" return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + @staticmethod + def supports_streaming_execution(callback: CustomGuardrail) -> bool: + """Whether a streaming pipeline step can run this guardrail against the buffered + stream: through the unified path, or through its post-call hook on the assembled + response when that hook is its only streaming path. A guardrail with its own + streaming iterator hook, or with neither hook, keeps running on its own.""" + callback_type: Final = type(callback) + return PipelineExecutor.supports_unified_execution(callback) or ( + callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook + and callback_type.async_post_call_streaming_iterator_hook + is CustomLogger.async_post_call_streaming_iterator_hook + ) + @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: """Look up an initialized guardrail callback by name from litellm.callbacks.""" diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py new file mode 100644 index 00000000000..d284c44397e --- /dev/null +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -0,0 +1,152 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + add_policy_sources_to_metadata, + add_policy_to_applied_policies_header, +) +from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import resolve_model_group_alias +from litellm.types.proxy.policy_engine import PolicyMatchContext +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router + +PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...] + +_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines) + + +@dataclass(frozen=True, slots=True) +class UngovernedRetrieval: + reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"] + + +def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval: + if llm_router is None: + return UngovernedRetrieval("no router") + model_id: Final = ( + ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None + ) + if model_id is None: + return UngovernedRetrieval("response id names no deployment") + deployment: Final = llm_router.get_deployment(model_id) + if deployment is None: + return UngovernedRetrieval("deployment no longer in the router") + hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias) + if hidden_by is not None: + verbose_proxy_logger.warning( + "Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), " + "so a policy attached to the model name it was submitted as does not run on it", + response_id, + deployment.model_name, + hidden_by, + ) + return deployment.model_name + + +def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None: + if "*" in model_group: + return "a wildcard deployment" + aliases: Final = tuple( + alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group + ) + if not aliases: + return None + return f"the target of model_group_alias {', '.join(aliases)}" + + +def _retrieval_context( + data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str +) -> PolicyMatchContext: + team_alias: Final = user_api_key_dict.team_alias + key_alias: Final = user_api_key_dict.key_alias + return PolicyMatchContext( + team_alias=team_alias if isinstance(team_alias, str) else None, + key_alias=key_alias if isinstance(key_alias, str) else None, + model=model_group, + tags=get_tags_from_request_body(data) or None, + ) + + +def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: + matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context) + if not matches: + return (), MappingProxyType({}) + applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( + policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list + context=context, + ) + post_call_pipelines: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context( + context=context, policy_names=applied_policy_names + ) + if pipeline.mode == "post_call" + ) + return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches}) + + +def attach_post_call_pipelines_to_retrieval( + data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place + user_api_key_dict: "UserAPIKeyAuth", + llm_router: "Router | None", +) -> None: + if not get_policy_registry().is_initialized(): + return + model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router) + if isinstance(model_group, UngovernedRetrieval): + verbose_proxy_logger.warning( + "Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)", + data.get("response_id"), + model_group.reason, + ) + return + context: Final = _retrieval_context(data, user_api_key_dict, model_group) + post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context) + _, bucket = get_or_create_metadata_bucket(data) + already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ()) + attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached) + added: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in post_call_pipelines + if policy_name not in attached_policy_names + ) + if not added: + return + pipelines: Final = (*already_attached, *added) + bucket["_guardrail_pipelines"] = pipelines + bucket["_pipeline_managed_guardrails"] = frozenset( + step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps + ) + for policy_name, _pipeline in added: + add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) + for _policy_name, pipeline in added: + for step in pipeline.steps: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail) + add_policy_sources_to_metadata( + request_data=data, + policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict + policy_name: policy_sources[policy_name] for policy_name, _pipeline in added + }, + ) + verbose_proxy_logger.debug( + "Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s", + data.get("response_id"), + model_group, + ", ".join(policy_name for policy_name, _pipeline in added), + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7219b373dc3..9269fd48e6c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17,6 +17,7 @@ import time import traceback import warnings from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -62,11 +63,13 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import ( UI_TEAM_ID, CallbackDelete, @@ -131,6 +134,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -138,11 +142,7 @@ from litellm.types.utils import ( TextCompletionResponse, TokenCountResponse, ) -from litellm.utils import ( - _invalidate_model_cost_lowercase_map, - load_credentials_from_list, - reapply_runtime_model_cost_registrations, -) +from litellm.utils import load_credentials_from_list if TYPE_CHECKING: from aiohttp import ClientSession @@ -268,13 +268,12 @@ from litellm.constants import ( WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, @@ -367,8 +366,11 @@ from litellm.proxy.common_utils.load_config_utils import ( ) from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( + ClaudeCodeRoutingNames, TeamModelNameTranslator, + claude_code_view_ids, configured_display_names, + is_claude_code_client, ) from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, @@ -426,6 +428,7 @@ from litellm.proxy.db.exception_handler import ( ) from litellm.proxy.db.gateway_request_tracking import ( GatewayRequestAccumulator, + GatewayRequestRedisBuffer, flush_gateway_requests, ) from litellm.proxy.db.proxy_worker_heartbeat import ( @@ -2357,6 +2360,17 @@ open_telemetry_logger: OpenTelemetry | None = None gateway_request_accumulator: Final = GatewayRequestAccumulator() ### INITIALIZE GLOBAL LOGGING OBJECT ### proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user) + + +def _gateway_request_redis_buffer() -> GatewayRequestRedisBuffer | None: + """Shares the spend writer's transaction-buffer Redis and pod lock when use_redis_transaction_buffer is on.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + redis_cache: Final = writer.redis_update_buffer.redis_cache + if redis_cache is None or not writer.redis_update_buffer._should_commit_spend_updates_to_redis(): + return None + return GatewayRequestRedisBuffer(redis_cache=redis_cache, pod_lock_manager=writer.pod_lock_manager) + + ### REDIS QUEUE ### async_result: Final = None celery_app_conn: Final = None @@ -2707,6 +2721,12 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) return fallback_spend, False +@dataclass(frozen=True, slots=True) +class _PendingSpendIncrement: + counter_key: str + increment: float + + async def increment_spend_counters( token: str | None, team_id: str | None, @@ -2741,7 +2761,7 @@ async def increment_spend_counters( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> None: + async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -2752,30 +2772,29 @@ async def increment_spend_counters( hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token ) key_counter_key: Final = f"spend:key:{hashed_token}" - if key_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=key_counter_key, - source_cache_key=hashed_token, - increment=cost, + key_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if key_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=key_counter_key, + source_cache_key=hashed_token, + increment=cost, + ), ) - - key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) - if key_obj is None: - return - key_budget_limits = getattr(key_obj, "budget_limits", None) or ( - key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(key_budget_limits, str): - key_budget_limits = json.loads(key_budget_limits) - if not isinstance(key_budget_limits, list): - return - for window in key_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + + async def _key_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + key_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}" key_window_start = get_budget_window_start(window) - if key_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, @@ -2783,6 +2802,9 @@ async def increment_spend_counters( window_start=key_window_start, increment=cost, ) + if key_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.KEY, entity_id=hashed_token, @@ -2792,33 +2814,48 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_scope(scope_team_id: str) -> None: - team_counter_key: Final = f"spend:team:{scope_team_id}" - if team_counter_key not in reserved_counter_keys: - await _init_and_increment_spend_counter( - counter_key=team_counter_key, - source_cache_key=f"team_id:{scope_team_id}", - increment=cost, - ) - - team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") - if team_obj is None: - return - team_budget_limits = getattr(team_obj, "budget_limits", None) or ( - team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) + if key_obj is None: + return key_pending + key_budget_limits = getattr(key_obj, "budget_limits", None) or ( + key_obj.get("budget_limits") if isinstance(key_obj, dict) else None ) - if isinstance(team_budget_limits, str): - team_budget_limits = json.loads(team_budget_limits) - if not isinstance(team_budget_limits, list): - return - for window in team_budget_limits: - duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration - team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at - team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" + if isinstance(key_budget_limits, str): + key_budget_limits = json.loads(key_budget_limits) + if not isinstance(key_budget_limits, list): + return key_pending + window_pending: Final = await asyncio.gather( + *(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True + ) + return key_pending + tuple(item for item in window_pending if item is not None) + + async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: + team_counter_key: Final = f"spend:team:{scope_team_id}" + team_pending: Final[tuple[_PendingSpendIncrement, ...]] = ( + () + if team_counter_key in reserved_counter_keys + else ( + await _prepare_spend_counter_increment( + counter_key=team_counter_key, + source_cache_key=f"team_id:{scope_team_id}", + increment=cost, + ), + ) + ) + + async def _team_window_increment(window: object) -> _PendingSpendIncrement | None: + duration = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + team_window_reset_at = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}" team_window_start = get_budget_window_start(window) - if team_window_counter not in reserved_counter_keys: - await _init_and_increment_window_spend_counter( + pending_window: Final = ( + await _prepare_window_spend_counter_increment( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, @@ -2826,6 +2863,9 @@ async def increment_spend_counters( window_start=team_window_start, increment=cost, ) + if team_window_counter not in reserved_counter_keys + else None + ) await _enqueue_window_spend_row_update( entity_type=Litellm_EntityType.TEAM, entity_id=scope_team_id, @@ -2835,25 +2875,47 @@ async def increment_spend_counters( increment=cost, request_started_at=request_started_at, ) + return pending_window - async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: + team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + if team_obj is None: + return team_pending + team_budget_limits = getattr(team_obj, "budget_limits", None) or ( + team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + ) + if isinstance(team_budget_limits, str): + team_budget_limits = json.loads(team_budget_limits) + if not isinstance(team_budget_limits, list): + return team_pending + window_pending: Final = await asyncio.gather( + *(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True + ) + return team_pending + tuple(item for item in window_pending if item is not None) + + async def _team_member_scope( + scope_user_id: str, scope_team_id: str + ) -> tuple[_PendingSpendIncrement | BaseException, ...]: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" if team_member_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=team_member_counter_key, - source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=team_member_counter_key, + source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}", + increment=cost, + ), ) - async def _user_scope(scope_user_id: str) -> None: + async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]: user_counter_key: Final = f"spend:user:{scope_user_id}" if user_counter_key in reserved_counter_keys: - return - await _init_and_increment_spend_counter( - counter_key=user_counter_key, - source_cache_key=scope_user_id, - increment=cost, + return () + return ( + await _prepare_spend_counter_increment( + counter_key=user_counter_key, + source_cache_key=scope_user_id, + increment=cost, + ), ) scope_coros: Final = tuple( @@ -2863,7 +2925,7 @@ async def increment_spend_counters( _team_scope(team_id) if team_id is not None else None, _team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None, _user_scope(user_id) if user_id is not None else None, - _increment_end_user_and_tag_spend_counters( + _prepare_end_user_and_tag_spend_increments( end_user_id=end_user_id, tags=tags, response_cost=cost, @@ -2871,14 +2933,14 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, - _increment_model_access_group_spend_counters( + _prepare_model_access_group_spend_increments( model_access_groups=model_access_groups, response_cost=cost, reserved_counter_keys=reserved_counter_keys, ) if model_access_groups else None, - _increment_org_spend_counter( + _prepare_org_spend_increment( org_id=org_id, response_cost=cost, reserved_counter_keys=reserved_counter_keys, @@ -2893,7 +2955,20 @@ async def increment_spend_counters( # as orphaned tasks that race the caller's reservation-counter invalidation; # all scopes settle, then the first error propagates as before. scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True) - scope_errors: Final = [r for r in scope_results if isinstance(r, BaseException)] + scope_errors: Final = tuple( + item + for scope in scope_results + for item in (scope if isinstance(scope, tuple) else (scope,)) + if isinstance(item, BaseException) + ) + pending: Final = tuple( + item + for scope in scope_results + if not isinstance(scope, BaseException) + for item in scope + if not isinstance(item, BaseException) + ) + await _apply_spend_counter_increments(pending=pending) if scope_errors: raise scope_errors[0] @@ -2936,41 +3011,49 @@ async def _reconcile_budget_reservation_for_counter_update( return reserved_counter_keys -async def _increment_end_user_and_tag_spend_counters( +async def _prepare_end_user_and_tag_spend_increments( end_user_id: str | None, tags: list[str] | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: - if end_user_id is not None: - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:end_user:{end_user_id}", - source_cache_key=end_user_cache_key(end_user_id), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) - - if tags is None: - return - - seen_tags: Final[set[str]] = set() - for tag_name in tags: - if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags: - continue - seen_tags.add(tag_name) - await _init_and_increment_unreserved_spend_counter( - counter_key=f"spend:tag:{tag_name}", - source_cache_key=tag_cache_key(tag_name), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) +) -> tuple[_PendingSpendIncrement | BaseException, ...]: + unique_tags: Final = ( + tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else () + ) + results: Final = await asyncio.gather( + *( + coro + for coro in ( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:end_user:{end_user_id}", + source_cache_key=end_user_cache_key(end_user_id), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + if end_user_id is not None + else None, + *( + _prepare_unreserved_spend_counter_increment( + counter_key=f"spend:tag:{tag_name}", + source_cache_key=tag_cache_key(tag_name), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for tag_name in unique_tags + ), + ) + if coro is not None + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_model_access_group_spend_counters( +async def _prepare_model_access_group_spend_increments( model_access_groups: Sequence[object], response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement | BaseException, ...]: """Charge the model access groups that authorized this request. Without this the counter auth reads is written only by the reservation path, so @@ -2984,55 +3067,63 @@ async def _increment_model_access_group_spend_counters( unique_groups: Final = tuple( dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) ) - for group in unique_groups: - await _init_and_increment_unreserved_spend_counter( - counter_key=model_access_group_spend_counter_key(group), - source_cache_key=model_access_group_cache_key(group), - increment=response_cost, - reserved_counter_keys=reserved_counter_keys, - ) + results: Final = await asyncio.gather( + *( + _prepare_unreserved_spend_counter_increment( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + for group in unique_groups + ), + return_exceptions=True, + ) + return tuple(item for item in results if item is not None) -async def _increment_org_spend_counter( +async def _prepare_org_spend_increment( org_id: str | None, response_cost: float, reserved_counter_keys: set[str], -) -> None: +) -> tuple[_PendingSpendIncrement, ...]: if org_id is None: - return + return () - await _init_and_increment_unreserved_spend_counter( + pending: Final = await _prepare_unreserved_spend_counter_increment( counter_key=f"spend:org:{org_id}", source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"], increment=response_cost, reserved_counter_keys=reserved_counter_keys, ) + return (pending,) if pending is not None else () -async def _init_and_increment_unreserved_spend_counter( +async def _prepare_unreserved_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, reserved_counter_keys: set[str], -) -> None: +) -> _PendingSpendIncrement | None: if counter_key in reserved_counter_keys: - return + return None - await _init_and_increment_spend_counter( + return await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key=source_cache_key, increment=increment, ) -async def _init_and_increment_spend_counter( +async def _prepare_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], increment: float, -): +) -> _PendingSpendIncrement: """ Initialize counter from the authoritative DB spend value if not yet - set, then atomically increment in both in-memory and Redis. + set, then return the pending increment for the caller to apply in one + pipelined Redis call. On first access per pod: 1. Check spend_counter_cache (in-memory -> Redis via DualCache) @@ -3044,13 +3135,13 @@ async def _init_and_increment_spend_counter( the counter as absent and seed it. Using increment means the worst case is over-counting (conservative, blocks slightly early) rather than under-counting (would allow overspend). - 4. Increment atomically (both in-memory + Redis) + 4. Increment is returned for the caller to apply via pipeline """ await _ensure_spend_counter_initialized( counter_key=counter_key, source_cache_key=source_cache_key, ) - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _enqueue_window_spend_row_update( @@ -3102,20 +3193,20 @@ async def _enqueue_window_spend_row_update( ) -async def _init_and_increment_window_spend_counter( +async def _prepare_window_spend_counter_increment( counter_key: str, entity_type: str, entity_id: str, window_duration: str | None, window_start: datetime | None, increment: float, -): +) -> _PendingSpendIncrement | None: if window_start is None: verbose_proxy_logger.warning( "Skipping spend counter increment for invalid budget window %s", counter_key, ) - return + return None initialized: Final = await _ensure_window_spend_counter_initialized( counter_key=counter_key, @@ -3125,8 +3216,8 @@ async def _init_and_increment_window_spend_counter( window_start=window_start, ) if initialized is False: - return - await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + return None + return _PendingSpendIncrement(counter_key=counter_key, increment=increment) async def _ensure_spend_counter_initialized( @@ -3259,6 +3350,32 @@ async def _invalidate_spend_counter(counter_key: str): ) +async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None: + if not pending: + return + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + for item in pending: + await spend_counter_cache.async_increment_cache( + key=item.counter_key, + value=item.increment, + refresh_ttl=True, + ) + return + ttl: Final = redis_cache.get_ttl() + increment_list: Final = [ # mutable-ok: async_increment_pipeline signature requires list[RedisPipelineIncrementOperation] + RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl) + for item in pending + ] + try: + results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) + except Exception: + await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) + raise + for item, current_value in zip(pending, results or ()): + spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value) + + async def update_cache( token: str | None, user_id: str | None, @@ -4436,20 +4553,9 @@ def resolve_classifier_plugin( def _swap_in_model_cost_map(new_model_cost_map: dict) -> int: - """Adopt a freshly fetched cost map into this process's litellm state, return the model count""" - litellm.model_cost = new_model_cost_map - # Invalidate case-insensitive lookup map since model_cost was replaced - _invalidate_model_cost_lowercase_map() - # Repopulate provider model sets (e.g. litellm.anthropic_models) so that - # wildcard patterns like "anthropic/*" include any newly added models. - litellm.add_known_models(model_cost_map=new_model_cost_map) - # Counted before the re-apply below, which writes into this same dict, so the - # number reported describes the fetched price data alone. - fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 - # The swap discards everything registered at runtime (deployment model_info, - # register_model overrides), so put it back on top of the fresh catalog. - reapply_runtime_model_cost_registrations() - return fetched_model_count + from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map + + return adopt_model_cost_map(new_model_cost_map) def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool: @@ -6536,6 +6642,14 @@ class ProxyConfig: return parsed return None + async def get_hierarchical_router_settings( + self, + user_api_key_dict: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + proxy_logging_obj: ProxyLogging | None = None, + ) -> dict | None: + return await self._get_hierarchical_router_settings(user_api_key_dict, prisma_client, proxy_logging_obj) + async def _get_hierarchical_router_settings( self, user_api_key_dict: Optional["UserAPIKeyAuth"], @@ -8529,6 +8643,7 @@ _STREAM_KEEPALIVE: Final = object() _KEEPALIVE_MIN_SECONDS: Final = 1.0 _KEEPALIVE_MAX_SECONDS: Final = 300.0 _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) +_EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) async def _iter_with_keepalive( @@ -9543,7 +9658,7 @@ class ProxyStartupEvent: flush_gateway_requests, "interval", seconds=batch_writing_interval, - args=(prisma_client, gateway_request_accumulator), + args=(prisma_client, gateway_request_accumulator, _gateway_request_redis_buffer()), id="update_gateway_requests_job", replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, @@ -10366,6 +10481,24 @@ async def model_list( wants_anthropic_format: Final = ( http_request is not None and http_request.headers.get("anthropic-version") is not None ) + client_headers: Final[Mapping[str, str]] = http_request.headers if http_request is not None else _EMPTY_HEADERS + view_router_settings: Final = ( + await proxy_config.get_hierarchical_router_settings(user_api_key_dict, prisma_client, proxy_logging_obj) + if wants_anthropic_format and is_claude_code_client(client_headers) + else None + ) + view_aliases: Final = ( + view_router_settings.get("model_group_alias") if isinstance(view_router_settings, Mapping) else None + ) + routing_names: Final = ClaudeCodeRoutingNames( + llm_router, + team_id or user_api_key_dict.team_id, + ( + user_api_key_dict.aliases, + user_api_key_dict.team_model_aliases, + view_aliases, + ), + ) # Validate scope parameter if provided if scope is not None and scope != "expand": @@ -10452,6 +10585,11 @@ async def model_list( return create_anthropic_model_list_response( admin_listing, display_names=configured_display_names(admin_entries, llm_router), + listed_ids=claude_code_view_ids( + admin_listing, + client_headers, + routing_names, + ), ) return dict( @@ -10500,6 +10638,11 @@ async def model_list( return create_anthropic_model_list_response( listing, display_names=configured_display_names(entries, llm_router), + listed_ids=claude_code_view_ids( + listing, + client_headers, + routing_names, + ), ) return dict( @@ -12714,7 +12857,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer) + _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( + model=model_to_use, custom_tokenizer=custom_tokenizer + ) tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) @@ -12727,7 +12872,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await asyncify(litellm.token_counter)( + total_tokens: Final = await offload_token_count(litellm.token_counter)( model=model_to_use, text=prompt, messages=counted_messages, @@ -17550,6 +17695,66 @@ async def delete_callback( ) +def _normalize_callback_alias(callback_name: str) -> str: + callback_aliases: Final = ( + ("opentelemetry", "otel"), + ("s3_v2", "s3"), + ("aws_sqs", "sqs"), + ("custom_callback_api", "generic_api"), + ) + return next( + (canonical_name for alias, canonical_name in callback_aliases if alias == callback_name), + callback_name, + ) + + +def _callback_module_name(callback: CustomLogger | Callable[..., object]) -> str: + if inspect.ismethod(callback): + return callback.__func__.__module__ + if inspect.isfunction(callback): + return callback.__module__ + return type(callback).__module__ + + +def _is_litellm_internal_callback(callback_name: str, callback: CustomLogger | Callable[..., object]) -> bool: + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + module_owner: Final = _callback_module_name(callback).partition(".")[0] + is_registered_integration: Final = callback_name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE + return not is_registered_integration and module_owner in ("litellm", "litellm_enterprise") + + +def _is_instance_of_configured_callback( + callback_name: str, callback: CustomLogger | Callable[..., object], configured_classes: tuple[type, ...] +) -> bool: + """Self-naming OTel-family instances (`arize`, `weave_otel`) match by name, so a configured `logfire` (a bare + `OpenTelemetry`) does not hide YAML-configured siblings of the same class.""" + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + class_derived_name: Final = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback)) + return isinstance(callback, configured_classes) and callback_name in (class_derived_name, type(callback).__name__) + + +def _hidden_runtime_callback_names(configured_callback_names: frozenset[str]) -> frozenset[str]: + from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + configured_classes: Final = tuple( + CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE[name] + for name in configured_callback_names + if name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE + ) + configured_modules: Final = frozenset(name.rsplit(".", 1)[0] for name in configured_callback_names if "." in name) + internal_callback_names: Final = frozenset({"cache", "vector_store_pre_call_hook"}) + return internal_callback_names | frozenset( + callback_name + for callback_name, callback in litellm.logging_callback_manager.get_callback_objects() + if isinstance(callback, CustomGuardrail) + or _is_litellm_internal_callback(callback_name, callback) + or _is_instance_of_configured_callback(callback_name, callback, configured_classes) + or _callback_module_name(callback) in configured_modules + ) + + @router.get( "/get/config/callbacks", tags=["config.yaml"], @@ -17582,10 +17787,10 @@ async def get_config( # Normalize string callbacks to lists def normalize_callback(callback): if isinstance(callback, str): - return [callback] - elif callback is None: - return [] - return callback + return (callback,) + if callback is None: + return () + return tuple(callback) if isinstance(callback, (list, dict)) else () _success_callbacks = normalize_callback(_success_callbacks) _failure_callbacks = normalize_callback(_failure_callbacks) @@ -17616,6 +17821,30 @@ async def get_config( for _callback in _success_and_failure_callbacks: _data_to_return.append(process_callback(_callback, "success_and_failure", environment_variables)) + configured_callback_names: Final = frozenset( + _normalize_callback_alias(callback) + for callback in (_success_callbacks + _failure_callbacks + _success_and_failure_callbacks) + ) + runtime_callbacks_by_type: Final = litellm.logging_callback_manager.get_callbacks_by_type() + hidden_callback_names: Final = _hidden_runtime_callback_names(configured_callback_names) + runtime_callback_rows: Final = tuple( + (_normalize_callback_alias(callback_name), callback_type) + for callback_type, callback_names in ( + ("success", runtime_callbacks_by_type["success"]), + ("failure", runtime_callbacks_by_type["failure"]), + ("success_and_failure", runtime_callbacks_by_type["success_and_failure"]), + ) + for callback_name in callback_names + if callback_name not in hidden_callback_names + ) + runtime_only_rows: Final = sorted( + frozenset(row for row in runtime_callback_rows if row[0] not in configured_callback_names) + ) + _data_to_return.extend( + dict(process_callback(callback_name, callback_type, environment_variables), read_only=True) + for callback_name, callback_type in runtime_only_rows + ) + _data_to_return = _apply_callback_role_gate(_data_to_return, is_full_admin) # Check if slack alerting is on diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 0fd242f2bc1..fac45d4391c 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import Request, Response from fastapi.responses import StreamingResponse +from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -74,6 +75,20 @@ class _StreamEventParser: parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) +async def _never_receive() -> Message: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +def detach_request_from_client(request: Request) -> Request: + """Same scope (headers, parsed body, auth) but a receive() that never yields http.disconnect. + + The polling client closes its connection right after getting the polling id, so the + upstream call must not be cancelled by the client-disconnect guards. + """ + return Request(request.scope, _never_receive) + + async def background_streaming_task( polling_id: str, data: dict[str, object], @@ -123,7 +138,7 @@ async def background_streaming_task( # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. response: Final[StreamingResponse] = await processor.base_process_llm_request( - request=request, + request=detach_request_from_client(request), fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, route_type="aresponses", diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 53aa372d1a1..00ccad33b6d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -94,12 +94,14 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, + get_or_create_metadata_bucket, independent_snapshot, is_expected_client_error, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms import load_guardrail_translation_mappings from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( @@ -157,6 +159,8 @@ from litellm.proxy.hooks.sensitive_data_routing import ( from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( @@ -172,6 +176,7 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import ( MCPDuringCallResponseObject, MCPPreCallRequestObject, @@ -193,6 +198,7 @@ if TYPE_CHECKING: from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation 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 @@ -455,7 +461,7 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps) -def _pipeline_managed_guardrail_names( +def pipeline_managed_guardrail_names( data: Mapping[str, object], mode: Literal["pre_call", "post_call"] ) -> frozenset[str]: return _pipeline_step_guardrail_names( @@ -518,9 +524,17 @@ def _merge_pipeline_metadata_writes( _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: +def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool: callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - return callback is not None and PipelineExecutor.supports_unified_execution(callback) + if callback is None: + return False + if PipelineExecutor.supports_unified_execution(callback): + return True + return ( + translation is not None + and type(translation).assembles_streamed_response + and PipelineExecutor.supports_streaming_execution(callback) + ) def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: @@ -529,50 +543,174 @@ def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardr ) -def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("background") is not True: - return - policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data)) - if not policy_names: - return - verbose_proxy_logger.warning( - "Policies with post_call guardrail pipelines do not run on background responses yet; " - "the response is released ungoverned by them: %s", - ", ".join(policy_names), +_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress")) + + +def _is_pending_background_response(response: LLMResponseTypes) -> bool: + return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES + + +def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]: + resolved: Final = PolicyResolver.resolve_policy_guardrails( + policy_name=policy_name, policies=get_policy_registry().get_all_policies() + ) + return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) + + +def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]: + return frozenset( + callback.guardrail_name + for callback in litellm.callbacks + if isinstance(callback, CustomGuardrail) + and callback.guardrail_name is not None + and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) ) -def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: - unsupported: Final = tuple( +def _without_names( + bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write + slot: str, + names: frozenset[str], +) -> None: + claimed: Final = bucket.get(slot) + if not isinstance(claimed, list): + return + remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to + name for name in claimed if name not in names + ] + if remaining: + bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place + else: + bucket.pop(slot) + + +def _withdraw_deferred_claims( + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data + deferred: Sequence[tuple[str, "GuardrailPipeline"]], +) -> None: + outside_by_policy: Final = MappingProxyType( + {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred} + ) + running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union( + _guardrails_run_standalone_pre_call(data), *outside_by_policy.values() + ) + withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside) + withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere + _, bucket = get_or_create_metadata_bucket(data) + _without_names(bucket, "applied_policies", withdrawn_policies) + _without_names(bucket, "applied_guardrails", withdrawn_guardrails) + sources: Final = bucket.get("policy_sources") + if not isinstance(sources, dict): + return + remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place + name: reason for name, reason in sources.items() if name not in withdrawn_policies + } + if remaining_sources: + bucket["policy_sources"] = remaining_sources + else: + bucket.pop("policy_sources") + + +def _defer_post_call_pipelines( + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data + response: ResponsesAPIResponse, +) -> None: + deferred: Final = _post_call_pipelines(data) + if not deferred: + return + verbose_proxy_logger.debug( + "Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s", + response.id, + response.status, + ", ".join(policy_name for policy_name, _pipeline in deferred), + ) + tag_matched: Final = _tag_matched_deferrals(data, deferred) + if tag_matched: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: %s", + response.id, + ", ".join(tag_matched), + ) + body_selected: Final = _body_selected_deferrals(data, deferred) + if body_selected: + verbose_proxy_logger.warning( + "Policy engine: background response %s matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: %s", + response.id, + ", ".join(body_selected), + ) + _withdraw_deferred_claims(data, deferred) + + +def _tag_matched_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[str, ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + if not isinstance(sources, dict): + return () + return tuple( + policy_name + for policy_name, _pipeline in deferred + if policy_name in sources and "tag:" in str(sources[policy_name]) + ) + + +def _body_selected_deferrals( + data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]] +) -> tuple[str, ...]: + sources: Final = _policy_state_metadata(data).get("policy_sources") + attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset() + return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed) + + +def _pipeline_unsupported_streaming_guardrails( + pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> tuple[str, ...]: + return tuple( dict.fromkeys( - step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail) + step.guardrail + for step in pipeline.steps + if not _pipeline_step_supports_streaming(step.guardrail, translation) ) ) + + +def _pipeline_is_streamable( + policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> bool: + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation) if not unsupported: return True verbose_proxy_logger.warning( - "Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, " - "which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s", + "Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they " + "need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a " + "route whose translation assembles the streamed response. The stream skips the pipeline and its " + "guardrails run on their own: %s", policy_name, ", ".join(unsupported), ) return False -def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool: - return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None +def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None": + resolved: Final = resolve_endpoint_translation(user_api_key_dict, None) + return None if resolved is None else resolved[1] -def _stream_gated_guardrail_names( +def stream_gated_guardrail_names( request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> frozenset[str]: - if not _route_supports_streaming_pipelines(user_api_key_dict): + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if translation is None: return frozenset() return _pipeline_step_guardrail_names( tuple( (policy_name, pipeline) for policy_name, pipeline in _post_call_pipelines(request_data) - if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) ) ) @@ -584,16 +722,19 @@ def _streamable_post_call_pipelines( The post_call pipelines a streaming response can be gated through. Streaming pipelines scan the buffered stream through the endpoint guardrail - translation of the request route, so every step's guardrail needs the - unified apply_guardrail interface and the route needs a translation. A - pipeline that cannot be run that way yet is left out and its guardrails - run on the stream on their own, the way they did before pipelines ran on - streams at all, with a warning naming the pipeline. + translation of the request route, so every step's guardrail needs either the + unified apply_guardrail interface or, on a route whose translation assembles + the streamed response, a post-call hook that is its only streaming path, and + the route needs a translation. A pipeline that + cannot be run that way yet is left out and its guardrails run on the stream + on their own, the way they did before pipelines ran on streams at all, with + a warning naming the pipeline. """ post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: return () - if not _route_supports_streaming_pipelines(user_api_key_dict): + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if translation is None: verbose_proxy_logger.warning( "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " @@ -605,7 +746,7 @@ def _streamable_post_call_pipelines( return tuple( (policy_name, pipeline) for policy_name, pipeline in post_call_pipelines - if _pipeline_is_streamable(policy_name, pipeline) + if _pipeline_is_streamable(policy_name, pipeline, translation) ) @@ -1985,8 +2126,6 @@ class ProxyLogging: ) try: - _warn_background_skips_post_call_pipelines(data) - # Execute guardrail pipelines before the normal callback loop data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below data=data, @@ -1997,7 +2136,7 @@ class ProxyLogging: ) # Get pipeline-managed guardrails to skip in normal loop - pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call") + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call") caps: Final = ProxyLogging._callback_capabilities() # Skip the per-request callback walk entirely when nothing in @@ -2762,7 +2901,7 @@ class ProxyLogging: original_exception=original_exception, ) - request_data.update(_failure_fields_to_lift(request_data)) + request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data)) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) @@ -2956,6 +3095,24 @@ class ProxyLogging: daemon=True, ).start() + async def _run_post_call_pipelines( + self, + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data + user_api_key_dict: UserAPIKeyAuth, + response: LLMResponseTypes, + ) -> LLMResponseTypes | None: + if _is_pending_background_response(response): + _defer_post_call_pipelines(data, response) + return None + _, pipeline_response = await self._maybe_execute_pipelines( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", + event_hook="post_call", + response=response, + ) + return pipeline_response + async def post_call_success_hook( self, data: dict, @@ -2975,17 +3132,15 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks - _, pipeline_response = await self._maybe_execute_pipelines( + pipeline_response: Final = await self._run_post_call_pipelines( data=data, user_api_key_dict=user_api_key_dict, - call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", - event_hook="post_call", response=response, ) if pipeline_response is not None: response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below - pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call") + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run @@ -3301,7 +3456,7 @@ class ProxyLogging: _cached_guardrail_data: dict | None = None _guardrail_data_computed = False pipeline_gated: Final = ( - _stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() + stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() ) for callback in litellm.callbacks: @@ -3436,12 +3591,16 @@ class ProxyLogging: ), ) - if post_call_pipelines: + pipeline_translation: Final = ( + resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None + ) + if pipeline_translation is not None: current_response = self._pipeline_gated_stream( response=current_response, user_api_key_dict=user_api_key_dict, request_data=request_data, pipelines=post_call_pipelines, + translation=pipeline_translation, ) try: @@ -3465,6 +3624,7 @@ class ProxyLogging: user_api_key_dict: UserAPIKeyAuth, request_data: dict, # mutable-ok: same request-payload shape the hooks mutate pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", + translation: "tuple[str, BaseTranslation]", ) -> "AsyncGenerator[Any, None]": """ Execute post_call policy pipelines against a streamed response. @@ -3474,14 +3634,13 @@ class ProxyLogging: assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow releases the buffered chunks: verbatim when no guardrail rewrote the - output, rewritten in place when one rewrote text and the translation - delivers ended-stream rewrites (later steps then re-scan the rewritten - chunks, so rewrites chain). A rewrite the translation cannot deliver - yet (a tool-call rewrite, or a text rewrite on a route without - write-back) is discarded by the executor and the original chunks are - released, as is a buffered shape no translation resolves; a block or - modify_response terminates with the translation's block chunks or the - raised error. + output, rewritten in place when one rewrote text or a tool call and the + translation delivers ended-stream rewrites (later steps then re-scan the + rewritten chunks, so rewrites chain). A rewrite the translation cannot + deliver yet (one on a route without write-back, or a shape the route + refuses) is discarded by the executor and the original chunks are + released; a block or modify_response terminates with the translation's + block chunks or the raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: @@ -3489,17 +3648,7 @@ class ProxyLogging: if not buffered: return - resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0]) - if resolved is None: - verbose_proxy_logger.warning( - "Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; " - "the stream is released ungoverned by them: %s", - ", ".join(policy_name for policy_name, _pipeline in pipelines), - ) - for buffered_item in buffered: - yield buffered_item - return - call_type, endpoint_translation = resolved + call_type, endpoint_translation = translation for policy_name, pipeline in pipelines: result: PipelineExecutionResult = await PipelineExecutor.execute_steps( diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index d7f8cd8f8bd..660dd8f0c92 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -437,14 +437,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): response_created_event_data["temperature"] = self.responses_api_request["temperature"] if "text" in self.responses_api_request: response_created_event_data["text"] = self.responses_api_request["text"] - if "tool_choice" in self.responses_api_request: - # Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format - response_created_event_data["tool_choice"] = ( - LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"]) - or "auto" + response_created_event_data["tool_choice"] = ( + LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + self.responses_api_request.get("tool_choice") ) - else: - response_created_event_data["tool_choice"] = "auto" + ) if "tools" in self.responses_api_request: response_created_event_data["tools"] = self.responses_api_request["tools"] else: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b2d1a69e0d8..fca5b0d11cf 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -27,8 +27,10 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( ) from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam +from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger @@ -68,6 +70,7 @@ from litellm.types.llms.openai import ( ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ResponsesAPIStatus, + ToolChoice, ValidChatCompletionMessageContentTypes, ValidChatCompletionMessageContentTypesLiteral, ) @@ -126,6 +129,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) +_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) @runtime_checkable @@ -267,6 +271,27 @@ class LiteLLMCompletionResponsesConfig: # Return as-is for unknown formats return tool_choice + @staticmethod + def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice: + if tool_choice is None: + return "auto" + try: + return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(tool_choice) + except ValidationError: + return LiteLLMCompletionResponsesConfig._chat_tool_choice_as_responses_api_tool_choice(tool_choice) + + @staticmethod + def _chat_tool_choice_as_responses_api_tool_choice(tool_choice: object) -> ToolChoice: + match tool_choice, LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice): + case {"type": "custom"}, {"function": {"name": str(custom_name)}}: + return ToolChoiceCustomParam(type="custom", name=custom_name) + case _, {"type": "function", "function": {"name": str(function_name)}}: + return ToolChoiceFunctionParam(type="function", name=function_name) + case _, "none" | "auto" | "required" as normalized: + return normalized + case _, _: + return "auto" + @staticmethod def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool: """ @@ -2263,7 +2288,9 @@ class LiteLLMCompletionResponsesConfig: ), parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False), temperature=getattr(chat_completion_response, "temperature", 0), - tool_choice=getattr(chat_completion_response, "tool_choice", "auto"), + tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + responses_api_request.get("tool_choice") + ), tools=getattr(chat_completion_response, "tools", []), top_p=getattr(chat_completion_response, "top_p", None), max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f9016c5a7f..40ff88fc557 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti import httpx from openai._streaming import SSEDecoder +from pydantic import BaseModel, ValidationError from typing_extensions import TypeIs import litellm @@ -438,18 +439,7 @@ class BaseResponsesAPIStreamingIterator: if self._persist_completed_response_before_logging: self._persist_completed_response_to_cache(is_async=is_async) - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr(self.completed_response, "model_dump"): - try: - logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump()) - except Exception: - # Fallback to original if serialization fails - pass + logging_response: Final[object] = _logging_copy(self.completed_response) self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() @@ -488,10 +478,10 @@ class BaseResponsesAPIStreamingIterator: def _restore_provider_response_headers(self, logging_response: object) -> None: """Re-apply the provider's response headers to the copy handed to logging callbacks. - ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the ``_hidden_params`` the provider transform set on the nested response are lost. Returns early - when that copy fell back to the original event, so logging-only state never lands on the - object the caller is iterating. + when the event was not a pydantic model and logging got the original, so logging-only state + never lands on the object the caller is iterating. """ if logging_response is self.completed_response: return @@ -544,7 +534,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return try: @@ -1293,14 +1283,46 @@ def _add_text_like_part_events( ) +def _logging_copy(event: object) -> object: + """Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never + reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the + deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow + copies of the event and its nested response still keep the caller's ``usage`` attribute separate.""" + if not isinstance(event, BaseModel): + return event + try: + return type(event).model_validate(event.model_dump()) + except Exception: + return _detached_shallow_copy(event) + + +def _detached_shallow_copy(event: BaseModel) -> BaseModel: + nested: Final[object] = getattr(event, "response", None) + if isinstance(nested, BaseModel): + return event.model_copy(update={"response": nested.model_copy()}) + return event.model_copy() + + +def _usage_as_model(usage: object) -> ResponseAPIUsage | None: + if isinstance(usage, ResponseAPIUsage): + return usage + if not isinstance(usage, dict): + return None + try: + return ResponseAPIUsage.model_validate(usage) + except ValidationError: + return None + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: if response_obj is None or logging_obj is None: return - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None)) if usage_obj is None: return + response_obj.usage = usage_obj # rebind-ok: the stamped cost has to ride on the response the client receives if isinstance(getattr(usage_obj, "cost", None), (int, float)): return try: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 540d492beec..599e978df6a 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -543,6 +543,49 @@ class ResponsesAPIRequestUtils: return request_input + @staticmethod + def strip_encrypted_reasoning_from_input(request_input: object) -> None: + """Drop reasoning items the routed deployment cannot decrypt, keeping their readable summary. + + Mutates ``request_input`` in place: the router's fallback snapshot shares this + list object, so a rebound list would replay the stripped items on the fallback hop. + """ + if not isinstance(request_input, list): + return + items: Final = cast(list[object], request_input) # cast-ok: untyped client json + stripped: Final = tuple(ResponsesAPIRequestUtils._without_encrypted_reasoning(item) for item in items) + items[:] = (item for item in stripped if item is not None) # rebind-ok: list shared with fallback snapshot + + @staticmethod + def _without_encrypted_reasoning(item: object) -> object | None: + if not isinstance(item, dict): + return item + reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json + if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"): + return reasoning + readable: Final = any( + ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content") + ) + if not readable: + return None + kept: Final[dict[str, object]] = { # mutable-ok: request item rebuilt without the undecryptable keys + key: value for key, value in reasoning.items() if key not in ("encrypted_content", "id") + } + return kept + + @staticmethod + def _has_readable_text(value: object) -> bool: + """A reasoning item's ``summary``/``content`` carries readable text: a non-empty string, or a + list holding at least one block with a non-empty ``text`` field (summary_text / output_text).""" + if isinstance(value, str): + return bool(value.strip()) + if isinstance(value, list): + return any( + isinstance(block, dict) and bool(cast(Mapping[str, object], block).get("text")) # cast-ok: untyped json + for block in value + ) + return False + @staticmethod def _build_responses_api_response_id( custom_llm_provider: str | None, diff --git a/litellm/router.py b/litellm/router.py index 934a4ac86a9..a20f010c608 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -67,7 +67,7 @@ from litellm.constants import ( SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import asyncify, run_async_function +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, coerce_token_limit, @@ -98,6 +98,8 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.base_llm.passthrough.transformation import replace_path_segment from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, @@ -148,6 +150,8 @@ from litellm.router_utils.common_utils import ( _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, + get_request_team_id, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, warn_on_provider_credential_mismatch, @@ -1317,6 +1321,43 @@ class Router: if isinstance(litellm.input_callback, list): litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] + def _apply_updated_routing_strategy_args(self) -> None: + """ + Re-link the default group's selector to the current `routing_strategy_args`. + + Selectors freeze their `RoutingArgs` at construction, so a runtime args + update would otherwise keep serving the boot-time values until restart. + Latency/usage state survives the rebuild: it lives in the shared router + cache, not on the selector. + """ + strategy: Final = self._normalize_strategy(self.routing_strategy) + if strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(self, self.routing_strategy_args) + return + + attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") + current: Final = getattr(self, attr, None) if attr is not None else None + if attr is None or current is None: + return + + try: + rebuilt: Final = self._build_strategy_selector( + strategy=strategy or "", + routing_strategy_args=self.routing_strategy_args, + ) + except (TypeError, ValidationError): + verbose_router_logger.exception( + "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", + self.routing_strategy_args, + strategy, + ) + return + + self._unregister_router_selectors((current,)) + setattr(self, attr, rebuilt) + def routing_strategy_init(self, routing_strategy: RoutingStrategy | str, routing_strategy_args: dict): verbose_router_logger.info("Routing strategy: %s", routing_strategy) self._validate_routing_strategy(routing_strategy) @@ -5160,7 +5201,7 @@ class Router: # If get_llm_provider fails, fall back to using model_name as-is replacement_model_name = model_name - kwargs["endpoint"] = kwargs["endpoint"].replace(model, replacement_model_name) + kwargs["endpoint"] = replace_path_segment(kwargs["endpoint"], model, replacement_model_name) return kwargs async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): @@ -5196,16 +5237,7 @@ class Router: kwargs=kwargs, model=model, model_name=model_name ) - # Get custom_llm_provider from deployment params - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response_kwargs: Final = { **data, @@ -5716,15 +5748,7 @@ class Router: # Perform pre-call checks for routing strategy self.routing_strategy_pre_call_checks(deployment=deployment) - try: - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - except Exception: - custom_llm_provider = None + custom_llm_provider: Final = provider_for_generic_call(data) response: Final = original_function( **{ @@ -11130,6 +11154,43 @@ class Router: return ids + def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: + """ + Deployment ids that could serve ``model`` for ``team_id``, following the same + precedence ``_common_checks_available_deployment`` uses to build a candidate pool: + ``model_group_alias``, then a routing group, then the first matching early-resolve + path for a name that is not a ``model_name`` (team route, wildcard pattern via + ``get_deployments_by_pattern``, team pattern router, default deployment), then the + ``model_name`` and team indexes. Delegating to the router's own resolvers keeps this + aligned with how a route actually resolves rather than re-deriving it, and unlike + ``_common_checks_available_deployment`` it is read-only: it does not apply request + fallbacks and (with ``include_team_models`` left off) does not raise. Lets a pre-call + check tell a genuine cross-group route from same-group unavailability without leaking + deployment ids into request kwargs bound for the provider. + """ + resolved: Final = self._get_model_from_alias(model=model) or model + routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id) + if routing_group_members is not None: + return self._deployment_ids(routing_group_members) + early: Final = self._try_early_resolve_deployments_for_model_not_in_names( + model=resolved, request_team_id=team_id + ) + if early is not None: + early_deployments: Final = early[1] + return self._deployment_ids( + (early_deployments,) if isinstance(early_deployments, Mapping) else early_deployments + ) + return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id)) + + @staticmethod + def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]: + return frozenset( + str(model_info["id"]) + for deployment in deployments + for model_info in (deployment.get("model_info"),) + if isinstance(model_info, Mapping) and model_info.get("id") is not None + ) + def has_model_id(self, candidate_id: str) -> bool: """ O(1) membership check for a deployment ID without allocating large lists. @@ -11847,7 +11908,7 @@ class Router: _existing_router_settings: Final = self.get_settings() rebuild_routing_groups = False - relink_lar1_from_args = False + routing_args_updated = False for var in kwargs: if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS: if var in _int_settings: @@ -11886,15 +11947,13 @@ class Router: ) rebuild_routing_groups = True elif var == "routing_strategy_args": - relink_lar1_from_args = True + routing_args_updated = True setattr(self, var, value) else: verbose_router_logger.debug("Setting %s is not allowed", var) - if relink_lar1_from_args and self._normalize_strategy(self.routing_strategy) == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self, self.routing_strategy_args) + if routing_args_updated: + self._apply_updated_routing_strategy_args() if rebuild_routing_groups: self._init_routing_groups(self._routing_groups_input) @@ -12040,7 +12099,7 @@ class Router: try: if not self._pre_call_checks_need_token_count(model, healthy_deployments): return None - return await asyncify(self._count_pre_call_check_tokens)( + return await offload_token_count(self._count_pre_call_check_tokens)( messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter request_kwargs=request_kwargs, @@ -12268,27 +12327,7 @@ class Router: if team_deployments: return model, team_deployments elif include_team_models: - team_deployments = [ - self.model_list[index] - for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] - team_ids: Final = { - team_id - for deployment in team_deployments - for team_id in [(deployment.get("model_info") or {}).get("team_id")] - if team_id is not None - } - if len(team_ids) > 1: - raise litellm.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) + team_deployments = self._team_deployments_across_teams(model) if team_deployments: return model, team_deployments @@ -12315,6 +12354,45 @@ class Router: return None + def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: + """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" + team_deployments: Final = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] + team_ids: Final = { + team_id + for deployment in team_deployments + for team_id in [(deployment.get("model_info") or {}).get("team_id")] + if team_id is not None + } + if len(team_ids) > 1: + raise litellm.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments + + def deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[DeploymentTypedDict]: + """The deployments `model` names for this caller, through the same alias, then team-first, then + global, then admin-across-teams resolution `_common_checks_available_deployment` applies, so + strategy selection and compression policy can never disagree with deployment selection about + which marker a name means.""" + registered_name: Final = self._get_model_from_alias(model=model) or model + team_id: Final = get_request_team_id(request_kwargs) + deployments: Final = self._get_all_deployments(model_name=registered_name, team_id=team_id) + if deployments or team_id is not None or not _is_proxy_admin_request(request_kwargs): + return deployments + return self._team_deployments_across_teams(registered_name) + @staticmethod def _is_strategy_marker_deployment(deployment: Mapping[str, object]) -> bool: litellm_params: Final = deployment.get("litellm_params") @@ -12342,11 +12420,7 @@ class Router: - Dict, if specific model chosen """ - request_team_id: str | None = None - if request_kwargs is not None: - metadata: Final = request_kwargs.get("metadata") or {} - litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) # check if aliases set on litellm model alias map if specific_deployment is True: return model, self._get_deployment_by_litellm_model(model=model) @@ -12371,7 +12445,9 @@ class Router: include_team_models=_is_proxy_admin_request(request_kwargs), ) if early is not None: - return early + if not isinstance(early[1], list): + return early + return early[0], self._drop_strategy_markers(early[0], early[1]) ## get healthy deployments ### get all deployments @@ -12448,19 +12524,22 @@ class Router: model ] # update the model to the actual value if an alias has been passed in - marker_flags: Final = tuple(self._is_strategy_marker_deployment(d) for d in healthy_deployments) - if not any(marker_flags): - return model, healthy_deployments - selectable: Final = [ # mutable-ok: matches this function's list contract expected by downstream filters - d for d, is_marker in zip(healthy_deployments, marker_flags, strict=True) if not is_marker + return model, self._drop_strategy_markers(model, healthy_deployments) + + def _drop_strategy_markers( + self, model: str, deployments: Sequence[DeploymentTypedDict] + ) -> list[DeploymentTypedDict]: + """A strategy marker is never a callable deployment, whichever resolution arm produced it.""" + selectable: Final = [ # mutable-ok: matches _common_checks_available_deployment's list contract + d for d in deployments if not self._is_strategy_marker_deployment(d) ] - if not selectable: + if deployments and not selectable: raise litellm.BadRequestError( message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}", model=model, llm_provider="", ) - return model, selectable + return selectable def _filter_deployments_by_model_access_groups( self, @@ -13150,12 +13229,8 @@ class Router: return filtered - def _model_name_has_plain_deployments(self, model: str) -> bool: - indices: Final = self.model_name_to_deployment_indices.get(model) or () - return any(not self._is_strategy_marker_deployment(self.model_list[idx]) for idx in indices) - def _select_pre_routing_strategy( - self, model: str, request_kwargs: dict + self, model: str, request_kwargs: Mapping[str, object] ) -> "TaggedPreRoutingStrategy[PreRoutingStrategy] | None": """ Resolve the pre-routing strategy for `model`, disambiguating deployments @@ -13166,6 +13241,12 @@ class Router: deployment the strategy was registered from via its (model_name, tags) pair. + The registries are keyed by the marker deployment's own `model_name`, which + for a team-scoped router is the internal `model_name_{team}_{uuid}` while + the caller sends the team's public name. So the names looked up are the + `model_name`s of whatever deployments this caller's request resolves `model` + to, and `model` itself when it resolves to none. + With tag filtering enabled, router-wide or by the request's enable_tag_filtering (which the proxy sets from key/team router_settings), strategies that all carry real tags matching none of @@ -13173,12 +13254,14 @@ class Router: deployments: returning None hands the request to ordinary tag-aware deployment selection. """ - candidates: Final[list[TaggedPreRoutingStrategy[PreRoutingStrategy]]] = [ - *self.auto_routers.get(model, []), - *self.complexity_routers.get(model, []), - *self.adaptive_routers.get(model, []), - *self.quality_routers.get(model, []), - ] + registries: Final = (self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers) + if not any(registries): + return None + deployments: Final = self.deployments_for_request(model, request_kwargs) + registered_names: Final = tuple(dict.fromkeys(str(d["model_name"]) for d in deployments)) or (model,) + candidates: Final = tuple( + tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) + ) if not candidates: return None @@ -13196,7 +13279,7 @@ class Router: if ( (self.enable_tag_filtering or request_scoped_filtering) and all(tagged.tags for tagged in candidates) - and self._model_name_has_plain_deployments(model) + and any(not self._is_strategy_marker_deployment(d) for d in deployments) ): return None return candidates[0] @@ -13308,11 +13391,12 @@ class Router: Used for the litellm auto-router to modify the request before the routing decision is made. - `model` is whatever the caller asked for, which may be a `model_group_alias` key, while the - strategy registries and the marker deployment are keyed by the marker's own `model_name`, so - every lookup below resolves the alias first. Only the lookups: the caller-facing name stays - the alias, since spend metadata is stamped before routing and the response carries the tier - group the strategy picked. + `model` is whatever the caller asked for, which may be a `model_group_alias` key or a team's + public model name, while the strategy registries and the marker deployment are keyed by the + marker's own `model_name`, so every lookup below resolves the alias first and the team name + through the deployment path. Only the lookups: the caller-facing name stays the alias, since + spend metadata is stamped before routing and the response carries the tier group the + strategy picked. """ requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model registered_model_name: Final = await self._resolve_claude_code_session_router( @@ -13349,7 +13433,6 @@ class Router: messages_for_routing, model_hop_compression_armed, policy_for_model, - team_id_from_request, ) # Same tag-aware lookup the proxy's pre-call arming used, so an alias with @@ -13357,7 +13440,7 @@ class Router: compression_policy: Final = policy_for_model( llm_router=self, model_alias=registered_model_name, - team_id=team_id_from_request(request_kwargs), + request_kwargs=request_kwargs, request_tags=_get_tags_from_request_kwargs(request_kwargs), ) # Shared compression already ran in the pre-call hook, so reuse it rather than @@ -13426,7 +13509,9 @@ class Router: # Per-tier `litellm_params` on the hook response are deliberate overrides # the caller applies on top, so those keys are never forwarded here. marker_params: Final = ( - self._forwardable_alias_marker_params(model=registered_model_name, strategy_tags=selected_strategy.tags) + self._forwardable_alias_marker_params( + model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs + ) if pre_routing_hook_response is not None else () ) @@ -13444,13 +13529,14 @@ class Router: return pre_routing_hook_response def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...] + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] ) -> tuple[tuple[str, object], ...]: marker_params: Final = tuple( litellm_params - for idx in self.model_name_to_deployment_indices.get(model, ()) - if isinstance(litellm_params := self.model_list[idx].get("litellm_params", {}), dict) - and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + for deployment in self.deployments_for_request(model, request_kwargs) + if str((litellm_params := deployment["litellm_params"]).get("model", "")).startswith( + AUTO_ROUTER_MODEL_PREFIX + ) ) tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 62b30365f4a..faafcea404a 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2568,14 +2568,14 @@ class ComplexityRouter(CustomLogger): """Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the event loop; None when counting fails, and the gate then leaves the placement alone.""" import litellm - from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.token_counter import offload_token_count out_of_band: Final = self._out_of_band_request_text(request_kwargs) try: - counted: Final = await asyncify(litellm.token_counter)( + counted: Final = await offload_token_count(litellm.token_counter)( messages=cast(list, resolved_messages) # cast-ok: token_counter only iterates the sequence ) - return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0) + return counted + (await offload_token_count(litellm.token_counter)(text=out_of_band) if out_of_band else 0) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e) return None diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 805d4ff9080..e902192811c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -3,8 +3,11 @@ import random from collections.abc import Sequence from datetime import datetime, timedelta +from math import ceil from typing import TYPE_CHECKING, Any, Final +from pydantic import Field + import litellm from litellm import ModelResponse, token_counter, verbose_logger from litellm.caching.caching import DualCache @@ -24,6 +27,7 @@ class RoutingArgs(LiteLLMPydanticObjectBase): ttl: float = 1 * 60 * 60 # 1 hour lowest_latency_buffer: float = 0 max_latency_list_size: int = 10 + ttft_percentile: float | None = Field(default=None, gt=0, le=1) def _average_latency(samples: Sequence[float]) -> float: @@ -32,6 +36,12 @@ def _average_latency(samples: Sequence[float]) -> float: return sum(samples) / len(samples) +def _percentile_latency(samples: Sequence[float], percentile: float) -> float: + values: Final = sorted(samples) + index: Final = ceil(len(values) * percentile) - 1 + return values[index] + + def _ttft_seconds(elapsed: timedelta | float) -> float: if isinstance(elapsed, timedelta): return elapsed.total_seconds() @@ -427,14 +437,17 @@ class LowestLatencyLoggingHandler(CustomLogger): item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) - # get average latency or average ttft (depending on streaming/non-streaming) use_ttft = ( request_kwargs is not None and request_kwargs.get("stream", None) is not None and request_kwargs["stream"] is True and len(item_ttft_latency) > 0 ) - average_latency = _average_latency(item_ttft_latency if use_ttft else item_latency) + selected_latency = ( + _percentile_latency(item_ttft_latency, self.routing_args.ttft_percentile) + if use_ttft and self.routing_args.ttft_percentile is not None + else _average_latency(item_ttft_latency if use_ttft else item_latency) + ) # -------------- # # Debugging Logic @@ -443,7 +456,7 @@ class LowestLatencyLoggingHandler(CustomLogger): # this helps a user to debug why the router picked a specfic deployment # _deployment_api_base = _deployment.get("litellm_params", {}).get("api_base", "") if _deployment_api_base is not None: - _latency_per_deployment[_deployment_api_base] = average_latency + _latency_per_deployment[_deployment_api_base] = selected_latency # -------------- # # End of Debugging Logic # -------------- # @@ -453,7 +466,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ): # if user passed in tpm / rpm in the model_list continue else: - potential_deployments.append((_deployment, average_latency)) + potential_deployments.append((_deployment, selected_latency)) if len(potential_deployments) == 0: return None diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 3ec92ad226a..3251ea457cf 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None: return response +def response_has_hidden_params(response: object) -> bool: + if isinstance(response, dict): + return "_hidden_params" in response + return hasattr(response, "_hidden_params") + + def ensure_response_additional_headers(response: object) -> dict[str, object]: hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict)) _write_hidden_params(response, hidden_params) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 280a7defcf8..49cca8ee99e 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject +import litellm from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError @@ -26,6 +27,18 @@ def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" +def get_request_team_id(request_kwargs: Mapping[str, object] | None) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + if request_kwargs is None: + return None + for bucket_name in ("metadata", "litellm_metadata"): + bucket = request_kwargs.get(bucket_name) + team_id = bucket.get("user_api_key_team_id") if isinstance(bucket, Mapping) else None + if isinstance(team_id, str) and team_id: + return team_id + return None + + def resolve_model_group_alias(model_group_alias: object, model: str) -> str | None: """ Resolve ``model`` through a ``model_group_alias`` map. @@ -110,7 +123,7 @@ def filter_team_based_models( metadata: Final = request_kwargs.get("metadata") or {} litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id: Final = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list): requested_model: Final = ( request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group") @@ -244,6 +257,32 @@ PROVIDER_SCOPED_CREDENTIAL_PARAMS: Final[Mapping[str, frozenset[str]]] = Mapping ) +def provider_for_generic_call(litellm_params: Mapping[str, object]) -> str | None: + """ + The provider the router hands a deployment's generic SDK call, or None when it cannot be resolved. + + A model that carries its own provider prefix keeps that prefix even where get_llm_provider + would resolve it to a sibling provider (azure_ai/ on an Azure OpenAI host + resolves to azure): the SDK call still receives the prefixed model, and an explicit provider + that contradicts the prefix makes get_llm_provider re-prefix it into a deployment name that + does not exist upstream. + """ + declared: Final = litellm_params.get("custom_llm_provider") + if isinstance(declared, str) and declared: + return declared + model: Final = litellm_params.get("model") + if not isinstance(model, str) or not model: + return None + prefix: Final = model.split("/", 1)[0] + if "/" in model and prefix in litellm.provider_list: + return prefix + try: + _, inferred, _, _ = get_llm_provider(model=model) + except BadRequestError: + return None + return inferred + + def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mapping[str, object]) -> str | None: """ Warn when a deployment carries one provider's credentials but resolves to another. diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index b623e31ce06..cdd70e6baf2 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,17 +37,21 @@ Safe to enable globally: """ import time +from collections.abc import Iterator, Mapping from typing import TYPE_CHECKING, Final, Optional, Protocol, cast import httpx from litellm._logging import verbose_router_logger from litellm.exceptions import ( - BadRequestError, RateLimitError, ServiceUnavailableError, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_content_of_block, + strip_encrypted_reasoning_from_messages, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues @@ -138,15 +142,48 @@ class EncryptedContentAffinityCheck(CustomLogger): # If no encoded ID, check if encrypted_content itself is wrapped encrypted_content = item.get("encrypted_content") if encrypted_content and isinstance(encrypted_content, str): - ( - model_id, - _, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) + model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content) if model_id: return model_id return None + @staticmethod + def _anthropic_content_blocks(messages: object) -> Iterator[Mapping[str, object]]: + if not isinstance(messages, list): + return iter(()) + return ( + cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance + for message in cast(list[object], messages) # cast-ok: narrowed by isinstance + if isinstance(message, Mapping) + for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance + if isinstance(content, list) + for block in cast(list[object], content) # cast-ok: narrowed by isinstance + if isinstance(block, Mapping) + ) + + @staticmethod + def _model_id_from_wrapped_encrypted_content(encrypted_content: str) -> str | None: + model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) + return model_id or None + + @staticmethod + def _extract_model_id_from_anthropic_messages(messages: object) -> str | None: + return next( + ( + model_id + for block in EncryptedContentAffinityCheck._anthropic_content_blocks(messages) + if (encrypted_content := encrypted_content_of_block(block)) is not None + if ( + model_id := EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content( + encrypted_content + ) + ) + is not None + ), + None, + ) + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -158,6 +195,23 @@ class EncryptedContentAffinityCheck(CustomLogger): return deployment return None + @staticmethod + def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None: + containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping)) + return next((tid for tid in team_ids if isinstance(tid, str)), None) + + def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]: + """ + Deployment ids that could serve this turn's routed ``model``, as the router + resolves a route (model_group_alias / routing group / model_name / team / + pattern). Delegates to the router so the full precedence is not re-derived here + and no deployment ids are written into request kwargs bound for the provider. + """ + if self.router is None: + return frozenset() + return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs)) + @staticmethod def _encryption_boundary_key( litellm_params: object, @@ -223,12 +277,17 @@ class EncryptedContentAffinityCheck(CustomLogger): parent_otel_span: Span | None = None, ) -> list[dict]: """ - If the request ``input`` contains litellm-encoded item IDs, decode the - embedded ``model_id`` and pin the request to that deployment. Raises - ``RateLimitError`` / ``ServiceUnavailableError`` / ``BadRequestError`` - when the originating deployment is unavailable and no encryption-boundary - peer exists, rather than dispatching a doomed request to a non-peer - deployment. The 429/503 split mirrors the originating cooldown's status: + If the request ``input`` contains litellm-encoded item IDs, or its Anthropic + ``messages`` replay a bridge-tagged thinking block, decode the embedded + ``model_id`` and pin the request to that deployment. Raises + ``RateLimitError`` / ``ServiceUnavailableError`` when the originating + deployment is a member of the routed model group but currently unavailable + and no encryption-boundary peer exists, rather than dispatching a doomed + request to a non-peer deployment. When the origin is not a member of the + routed group (an auto-router tier change, a model switch with no peer, a + removed deployment, or an unknown/forged marker), the encrypted reasoning is + stripped and the request dispatches with its readable history instead. The + 429/503 split mirrors the originating cooldown's status: a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the remaining cooldown window) so OpenAI-compatible clients back off and retry after the deployment is eligible again. @@ -249,12 +308,15 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True request_input: Final = request_kwargs.get("input") - model_id: Final = self._extract_model_id_from_input(request_input) + anthropic_messages: Final = messages or request_kwargs.get("messages") + model_id: Final = self._extract_model_id_from_input( + request_input + ) or self._extract_model_id_from_anthropic_messages(anthropic_messages) if not model_id: return typed_healthy_deployments verbose_router_logger.debug( - "EncryptedContentAffinityCheck: decoded model_id=%s from input item IDs", + "EncryptedContentAffinityCheck: decoded model_id=%s from the request's encrypted content markers", model_id, ) @@ -285,12 +347,35 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["_encrypted_content_affinity_pinned"] = True return boundary_matches - # Dispatching to a non-peer would guarantee an upstream - # `invalid_encrypted_content` 400, so fail fast with a clearer error. + # The origin cannot serve this turn's routed group and no peer shares the boundary, so its + # encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch + # to the routed group instead of failing. Membership is tested by deployment id against the set + # the router actually resolved for this route, not by model-group name, so an alias, a + # provider-qualified spelling, a team-public name, or a pattern route of the same group is not + # mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is + # treated the same as a cross-group one, which also denies an authenticated caller a + # deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and + # dispatch rather than returning distinguishable responses. Only a genuine same-group member + # that is currently unavailable falls through to the fail-fast, preserving the cooldown contract. + routed_group_model_ids: Final = ( + self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset() + ) + if str(model_id) not in routed_group_model_ids: + verbose_router_logger.debug( + "EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; " + "forwarding without its encrypted reasoning", + model_id, + model, + ) + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + strip_encrypted_reasoning_from_messages(anthropic_messages) + return typed_healthy_deployments + + # The origin is a member of the routed group but currently unavailable (cooled down); fail fast + # rather than dispatching to a non-peer, which would guarantee an upstream 400. raise await self._unavailable_origin_error( model=model, model_id=model_id, - originating=originating, parent_otel_span=parent_otel_span, ) @@ -298,25 +383,11 @@ class EncryptedContentAffinityCheck(CustomLogger): self, model: str, model_id: str, - originating: Deployment | None, parent_otel_span: Span | None, ) -> Exception: # Public error messages intentionally omit the originating ``model_id`` so # an authenticated caller forging encrypted-content markers cannot use the # error surface to enumerate which deployment IDs exist on this router. - if originating is None: - return BadRequestError( - message=( - "The deployment that produced this encrypted_content is no " - "longer configured on this router, and no deployment on the " - "same encryption boundary is available. Re-issue the request " - "without the stale encrypted_content items, or restore the " - "originating deployment." - ), - model=model, - llm_provider="", - ) - cooldown: Final = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span) if cooldown is not None and str(cooldown.get("status_code")) == "429": diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 01d42627001..fbd3e18e357 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -21,6 +21,7 @@ import litellm from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.router import RouterCacheEnum, RouterErrors from litellm.utils import get_utc_datetime @@ -466,7 +467,7 @@ async def async_io_token_pre_call_check( request_kwargs: Final = get_io_token_rate_limit_request_kwargs() _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" - estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) + estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) dt: Final = get_utc_datetime() diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 70362e60495..0589e290b47 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -14,6 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -61,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger): if request_kwargs is not None and request_kwargs.get("_target_order") is not None: return healthy_deployments - if messages is not None and is_prompt_caching_valid_prompt( + if messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( messages=messages, model=model, min_token_count=_get_min_token_count_for_deployments(healthy_deployments), @@ -139,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger): return ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider - if is_prompt_caching_valid_prompt( + if await offload_token_count(is_prompt_caching_valid_prompt)( model=model, messages=cast(list[AllMessageValues], messages), ): diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 8498b6f6d00..a024581f600 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -154,6 +154,22 @@ LATENCY_BUCKETS: Final = ( float("inf"), ) +UNKNOWN_INPUT_SEQUENCE_LENGTH: Final = "unknown" +INPUT_SEQUENCE_LENGTH_BUCKETS: Final = ( + (1_000, "0-1k"), + (4_000, "1k-4k"), + (16_000, "4k-16k"), + (64_000, "16k-64k"), + (float("inf"), "64k+"), +) + + +def get_input_sequence_length_bucket(prompt_tokens: object) -> str: + if not isinstance(prompt_tokens, int) or isinstance(prompt_tokens, bool) or prompt_tokens < 0: + return UNKNOWN_INPUT_SEQUENCE_LENGTH + return next(label for upper, label in INPUT_SEQUENCE_LENGTH_BUCKETS if prompt_tokens < upper) + + # Batch jobs can run for minutes to hours; buckets span 1 min → 24 h. BATCH_DURATION_BUCKETS: Final = ( 60.0, @@ -205,6 +221,7 @@ class UserAPIKeyLabelNames(Enum): MCP_TOOL_NAME = "mcp_tool_name" MCP_SERVER_NAME = "mcp_server_name" SERVICE_TIER = "service_tier" + INPUT_SEQUENCE_LENGTH = "input_sequence_length" DEFINED_PROMETHEUS_METRICS = Literal[ @@ -857,6 +874,13 @@ class PrometheusMetricLabels: "litellm_images_generated_metric", } ) + _input_sequence_length_metrics: ClassVar[frozenset[str]] = frozenset( + { + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", + } + ) # Managed batch metrics _batch_user_labels = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -955,14 +979,23 @@ class PrometheusMetricLabels: custom_labels.append(label) if label_name in PrometheusMetricLabels._org_label_metrics: - for label in [ + for label in ( UserAPIKeyLabelNames.ORG_ID.value, UserAPIKeyLabelNames.ORG_ALIAS.value, - ]: + ): if label not in default_labels and label not in custom_labels: custom_labels.append(label) - return default_labels + custom_labels + input_sequence_length_labels: Final = ( + (UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value,) + if ( + label_name in PrometheusMetricLabels._input_sequence_length_metrics + and litellm.prometheus_emit_input_sequence_length_label is True + and UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in custom_labels + ) + else () + ) + return [*default_labels, *custom_labels, *input_sequence_length_labels] _USER_API_KEY_LABEL_VALUE_INIT_ALIASES: Final[Mapping[str, str]] = MappingProxyType( @@ -1015,6 +1048,7 @@ class UserAPIKeyLabelValues: mcp_tool_name: str | None = None mcp_server_name: str | None = None service_tier: str | None = None + input_sequence_length: str | None = None # Added for test compatibility. def __init__(self, **kwargs: Any) -> None: diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index e87a684aab8..a9c027bd2de 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -2,6 +2,7 @@ from typing import Any, Literal from pydantic import BaseModel from typing_extensions import ( + ReadOnly, Required, TypedDict, ) @@ -57,6 +58,14 @@ class DatabricksMessage(TypedDict, total=False): role: Required[str] content: Required[AllDatabricksContentValues] tool_calls: list[DatabricksTool] | None + reasoning_content: ReadOnly[str | None] + reasoning: ReadOnly[str | None] + + +class DatabricksDelta(TypedDict, total=False): + role: ReadOnly[str] + content: ReadOnly[AllDatabricksContentValues | None] + reasoning_content: ReadOnly[str | None] class DatabricksChoice(TypedDict, total=False): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..b7c4371f32f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1564,6 +1564,9 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): response: ResponsesAPIResponse +ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent + + class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED] item_id: str diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 84ffd50eea1..ab2e8a70754 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -250,7 +250,23 @@ class MCPServer(BaseModel): @property def advertises_gateway_authorization_server(self) -> bool: """Whether named discovery should advertise the aggregate gateway authorization server.""" - return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + if self.auth_type == MCPAuth.oauth2: + return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + if self.auth_type not in ( + None, + MCPAuth.none, + MCPAuth.api_key, + MCPAuth.bearer_token, + MCPAuth.basic, + MCPAuth.authorization, + MCPAuth.token, + MCPAuth.aws_sigv4, + ): + return False + return not any( + header.lower() in ("authorization", "x-api-key", "api-key", "apikey") + for header in (self.extra_headers or ()) + ) @property def is_true_passthrough(self) -> bool: diff --git a/litellm/types/router.py b/litellm/types/router.py index 5c9eab30f3d..6b707a544a2 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -525,6 +525,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): input_cost_per_second: float | None output_cost_per_second: float | None output_cost_per_second_480p: ReadOnly[float | None] + output_cost_per_second_720p: ReadOnly[float | None] output_cost_per_second_1080p: float | None output_cost_per_second_4k: ReadOnly[float | None] num_retries: int | None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d62f00f3676..ab0cc5f959c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -318,6 +318,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): float | None ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) output_cost_per_second_480p: ReadOnly[float | None] + output_cost_per_second_720p: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_credit: float | None # for OCR models priced by credit @@ -3522,6 +3523,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_second: float | None = None output_cost_per_second_1080p: float | None = None output_cost_per_second_480p: float | None = None + output_cost_per_second_720p: float | None = None output_cost_per_second_4k: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 36b48d3b8d8..917af2b89d4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2293,15 +2293,7 @@ def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: st dict: A dictionary with the tokenizer and its type. """ - try: - tokenizer = Tokenizer.from_pretrained( - identifier, - revision=revision, - auth_token=auth_token, - ) - except Exception as e: - verbose_logger.error("Error creating pretrained tokenizer: %s. Defaulting to version without 'auth_token'.", e) - tokenizer = Tokenizer.from_pretrained(identifier, revision=revision) + tokenizer: Final = Tokenizer.from_pretrained(identifier, revision=revision, token=auth_token) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} @@ -3079,7 +3071,7 @@ def register_model( # Convert stringified numbers to appropriate numeric types loaded_model_cost = model_cost elif isinstance(model_cost, str): - loaded_model_cost = litellm.get_model_cost_map(url=model_cost) + loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1) if persist_across_reloads: _registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost @@ -3412,7 +3404,7 @@ def get_optional_params_image_gen( non_default_params=non_default_params, optional_params=optional_params, model=model or "", - drop_params=drop_params if drop_params is not None else False, + drop_params=litellm.drop_params is True or drop_params is True, ) elif ( custom_llm_provider == "openai" @@ -5913,6 +5905,7 @@ def _get_model_info_helper( output_cost_per_second=_model_info.get("output_cost_per_second", None), output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), + output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None), output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), @@ -8857,6 +8850,12 @@ class ProviderConfigManager: ) return AzurePassthroughConfig() + elif LlmProviders.AZURE_AI == provider: + from litellm.llms.azure_ai.passthrough.transformation import ( + AzureAIPassthroughConfig, + ) + + return AzureAIPassthroughConfig() elif LlmProviders.GIGACHAT == provider: from litellm.llms.gigachat.passthrough.transformation import ( GigaChatPassthroughConfig, @@ -9241,6 +9240,10 @@ class ProviderConfigManager: from litellm.llms.openai.image_edit import get_openai_image_edit_config return get_openai_image_edit_config(model=model) + elif LlmProviders.HOSTED_VLLM == provider: + from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config + + return get_hosted_vllm_image_edit_config(model=model) elif LlmProviders.AZURE == provider: from litellm.llms.azure.image_edit.transformation import ( AzureImageEditConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 54ebdc85be9..0d2eda93323 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -364,7 +364,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -380,6 +381,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -399,6 +401,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -416,6 +419,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -435,6 +439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -452,6 +457,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -471,6 +477,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -488,6 +495,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -507,6 +515,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -537,7 +546,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -551,7 +561,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -756,6 +767,14 @@ "mode": "chat", "supports_video_input": true }, + "global.twelvelabs.pegasus-1-2-v1:0": { + "input_cost_per_video_per_second": 0.00049, + "output_cost_per_token": 7.5e-06, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_video_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "amazon.titan-text-express-v1": { "input_cost_per_token": 1.3e-06, "litellm_provider": "bedrock", @@ -2876,7 +2895,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.7e-08, @@ -2888,7 +2908,8 @@ "output_cost_per_token": 1.48e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "apac.amazon.nova-pro-v1:0": { "input_cost_per_token": 8.4e-07, @@ -2902,7 +2923,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "apac.anthropic.claude-3-5-sonnet-20240620-v1:0": { "deprecation_date": "2026-07-30", @@ -8064,7 +8086,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8108,7 +8131,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, @@ -8152,7 +8176,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2027-10-26" }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -12317,7 +12342,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-east-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -12496,7 +12522,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": { "input_cost_per_token": 4.2e-08, @@ -12508,7 +12535,8 @@ "output_cost_per_token": 1.68e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": { "input_cost_per_token": 9.6e-07, @@ -12522,7 +12550,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "bedrock/us-gov-west-1/amazon.titan-embed-text-v1": { "input_cost_per_token": 1e-07, @@ -13973,7 +14002,8 @@ "max_output_tokens": 3072, "max_tokens": 3072, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/meta/llama-2-7b-chat-int8": { "input_cost_per_token": 1.923e-06, @@ -13982,7 +14012,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.1": { "input_cost_per_token": 1.923e-06, @@ -13991,7 +14022,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@hf/thebloke/codellama-7b-instruct-awq": { "input_cost_per_token": 1.923e-06, @@ -14000,7 +14032,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.923e-06 + "output_cost_per_token": 1.923e-06, + "rpm": 300 }, "cloudflare/@cf/openai/gpt-oss-120b": { "input_cost_per_token": 3.5e-07, @@ -14010,6 +14043,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 7.5e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14020,7 +14054,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-3b-instruct": { "input_cost_per_token": 5.09e-08, @@ -14029,7 +14064,8 @@ "max_output_tokens": 80000, "max_tokens": 80000, "mode": "chat", - "output_cost_per_token": 3.35e-07 + "output_cost_per_token": 3.35e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-guard-3-8b": { "input_cost_per_token": 4.84e-07, @@ -14038,7 +14074,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-08 + "output_cost_per_token": 3e-08, + "rpm": 300 }, "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { "input_cost_per_token": 0.0, @@ -14047,7 +14084,8 @@ "max_output_tokens": 15000, "max_tokens": 15000, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.7-code": { "cache_read_input_token_cost": 1.9e-07, @@ -14058,6 +14096,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -14069,6 +14108,7 @@ "max_tokens": 80000, "mode": "chat", "output_cost_per_token": 4.881e-06, + "rpm": 300, "supports_reasoning": true }, "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { @@ -14078,7 +14118,8 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 2.87e-07 + "output_cost_per_token": 2.87e-07, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.2-1b-instruct": { "input_cost_per_token": 2.7e-08, @@ -14087,7 +14128,8 @@ "max_output_tokens": 60000, "max_tokens": 60000, "mode": "chat", - "output_cost_per_token": 2.01e-07 + "output_cost_per_token": 2.01e-07, + "rpm": 300 }, "cloudflare/@cf/moonshotai/kimi-k2.6": { "cache_read_input_token_cost": 1.6e-07, @@ -14098,6 +14140,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -14109,6 +14152,7 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14119,7 +14163,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { "input_cost_per_token": 2.93e-07, @@ -14129,6 +14174,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 2.253e-06, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { @@ -14139,6 +14185,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.12e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { @@ -14148,7 +14195,8 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-06 + "output_cost_per_token": 1e-06, + "rpm": 300 }, "cloudflare/@cf/zai-org/glm-5.2": { "cache_read_input_token_cost": 2.6e-07, @@ -14159,6 +14207,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4.4e-06, + "rpm": 20, "supports_function_calling": true, "supports_reasoning": true }, @@ -14170,6 +14219,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14180,7 +14230,8 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.55e-07 + "output_cost_per_token": 5.55e-07, + "rpm": 300 }, "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { "input_cost_per_token": 5.09e-08, @@ -14190,6 +14241,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 3.35e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14200,7 +14252,8 @@ "max_output_tokens": 3500, "max_tokens": 3500, "mode": "chat", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "rpm": 300 }, "cloudflare/@cf/google/gemma-4-26b-a4b-it": { "input_cost_per_token": 1e-07, @@ -14210,6 +14263,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14221,6 +14275,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5.55e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { @@ -14231,6 +14286,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.76e-07, + "rpm": 300, "supports_vision": true }, "cloudflare/@cf/openai/gpt-oss-20b": { @@ -14241,6 +14297,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-07, + "rpm": 300, "supports_function_calling": true, "supports_reasoning": true }, @@ -14252,6 +14309,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 8.5e-07, + "rpm": 300, "supports_function_calling": true }, "cloudflare/@cf/qwen/qwq-32b": { @@ -14262,6 +14320,7 @@ "max_tokens": 24000, "mode": "chat", "output_cost_per_token": 1e-06, + "rpm": 300, "supports_reasoning": true }, "codestral/codestral-2405": { @@ -14388,6 +14447,28 @@ "output_vector_size": 1536, "supports_embedding_image_input": true }, + "us.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, + "global.cohere.embed-v4:0": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1536, + "supports_embedding_image_input": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, "litellm_provider": "cohere", @@ -20959,7 +21040,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.amazon.nova-micro-v1:0": { "input_cost_per_token": 4.6e-08, @@ -20971,7 +21053,8 @@ "output_cost_per_token": 1.84e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "eu.amazon.nova-pro-v1:0": { "input_cost_per_token": 1.05e-06, @@ -20986,24 +21069,25 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "eu.anthropic.claude-3-5-haiku-20241022-v1:0": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 8e-07, "litellm_provider": "bedrock", "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 4e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 2.5e-08, - "cache_creation_input_token_cost": 3.125e-07, + "cache_read_input_token_cost": 8e-08, + "cache_creation_input_token_cost": 1e-06, "prompt_cache_min_tokens": 2048 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { @@ -24010,9 +24094,9 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "realtime", "output_cost_per_audio_token": 1.2e-05, "output_cost_per_token": 2e-06, @@ -24035,12 +24119,12 @@ "supports_audio_output": true, "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_pdf_input": false, + "supports_prompt_caching": false, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_url_context": true, + "supports_url_context": false, "supports_vision": true, "supports_web_search": true, "search_context_cost_per_query": { @@ -27832,6 +27916,70 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.5": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.5", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-luna": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-luna", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-sol": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-sol", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.6-terra": { + "litellm_provider": "chatgpt", + "source": "https://platform.openai.com/docs/models/gpt-5.6-terra", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "chatgpt/gpt-5.4": { "litellm_provider": "chatgpt", "max_input_tokens": 1050000, @@ -28438,6 +28586,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true }, @@ -29262,7 +29411,12 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + } }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, @@ -29279,9 +29433,9 @@ "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29380,9 +29534,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29407,9 +29561,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.03, + "search_context_size_high": 0.025, "search_context_size_low": 0.025, - "search_context_size_medium": 0.0275 + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29520,9 +29674,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29547,9 +29701,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.05, - "search_context_size_low": 0.03, - "search_context_size_medium": 0.035 + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 }, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -29631,6 +29785,66 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2.5-flare": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-flare-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, + "gpt-image-2.5-sunburst-2026-09-08": { + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true, + "source": "https://developers.openai.com/api/docs/pricing" + }, "low/1024-x-1024/gpt-image-1.5": { "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, @@ -41711,6 +41925,28 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "rerank-v4.0-fast": { + "input_cost_per_query": 0.002, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, + "rerank-v4.0-pro": { + "input_cost_per_query": 0.0025, + "input_cost_per_token": 0.0, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/pricing" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -43671,7 +43907,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -43683,7 +43920,8 @@ "output_cost_per_token": 1.4e-07, "supports_function_calling": true, "supports_prompt_caching": true, - "supports_response_schema": true + "supports_response_schema": true, + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -43712,7 +43950,8 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, @@ -45891,8 +46130,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -45916,8 +46155,8 @@ "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06, "regional_endpoint_uplift_multiplier": 1.1, @@ -47897,9 +48136,9 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47913,9 +48152,9 @@ "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", @@ -47928,14 +48167,17 @@ }, "vertex_ai/xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -47944,14 +48186,17 @@ }, "vertex_ai/xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 2000000, "max_output_tokens": 2000000, "max_tokens": 2000000, "mode": "chat", - "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/developers/models", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -47959,6 +48204,44 @@ "supports_vision": true, "supports_web_search": true }, + "vertex_ai/xai/grok-4.3": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 200000, + "max_output_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "vertex_ai/xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -48217,6 +48500,16 @@ "mode": "embedding", "output_cost_per_token": 0.0 }, + "voyage/voyage-multilingual-2": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-3-large": { "input_cost_per_token": 1.8e-07, "litellm_provider": "voyage", @@ -52290,7 +52583,8 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-07, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -52329,7 +52623,8 @@ "mode": "chat", "output_cost_per_token": 5e-07, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-01" }, "scaleway/hcompany/holo2-30b-a3b": { "input_cost_per_token": 3e-07, @@ -52340,7 +52635,8 @@ "mode": "chat", "output_cost_per_token": 7e-07, "supports_reasoning": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-08-09" }, "scaleway/mistralai/mistral-medium-3.5-128b": { "input_cost_per_token": 1.5e-06, @@ -52363,7 +52659,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2e-06, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/voxtral-small-24b-2507": { "input_cost_per_audio_token": 1.5e-07, @@ -52374,7 +52671,8 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3.5e-07, - "supports_audio_input": true + "supports_audio_input": true, + "deprecation_date": "2026-08-01" }, "scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": { "input_cost_per_token": 1.5e-07, @@ -52396,7 +52694,8 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_vision": true, - "supports_function_calling": true + "supports_function_calling": true, + "deprecation_date": "2026-10-01" }, "scaleway/BAAI/bge-multilingual-gemma2": { "input_cost_per_token": 1e-07, @@ -54846,7 +55145,7 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-20b": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 7e-08, "output_cost_per_token": 3e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, @@ -54880,8 +55179,8 @@ "supports_tool_choice": true }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2e-07, "litellm_provider": "bedrock_mantle", "max_input_tokens": 131072, "max_output_tokens": 65536, @@ -55001,6 +55300,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, @@ -55196,6 +55528,96 @@ "supports_reasoning": true, "supports_vision": true }, + "bedrock_mantle/openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "us.openai.gpt-6-astra": { + "input_cost_per_token": 1.1e-05, + "input_cost_per_token_above_272k_tokens": 2.2e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, + "cache_read_input_token_cost": 1.1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.2e-06, + "output_cost_per_token": 5.5e-05, + "output_cost_per_token_above_272k_tokens": 8.25e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, + "global.openai.gpt-6-astra": { + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, @@ -56861,8 +57283,8 @@ "rpm": 10 }, "vertex_ai/gemini-3.5-transcribe-preview": { - "input_cost_per_audio_token": 2.5e-06, - "input_cost_per_token": 2.5e-06, + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai", "mode": "audio_transcription", "output_cost_per_token": 1.2e-05, @@ -56897,6 +57319,27 @@ ], "supports_audio_input": true }, + "vertex_ai/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "realtime", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio", + "text" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", @@ -59588,6 +60031,77 @@ "image" ] }, + "xai/grok-imagine-video": { + "input_cost_per_image": 0.002, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_480p": 0.05, + "output_cost_per_second_720p": 0.07, + "source": "https://docs.x.ai/docs/models/grok-imagine-video", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-2026-05-30": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, + "xai/grok-imagine-video-1.5-preview": { + "input_cost_per_image": 0.01, + "litellm_provider": "xai", + "mode": "video_generation", + "output_cost_per_second": 0.08, + "output_cost_per_second_1080p": 0.25, + "output_cost_per_second_480p": 0.08, + "output_cost_per_second_720p": 0.14, + "source": "https://docs.x.ai/docs/models/grok-imagine-video-1.5", + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "video" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", @@ -60855,6 +61369,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60865,6 +61380,7 @@ "litellm_provider": "cloudflare", "mode": "audio_transcription", "output_cost_per_second": 0.0, + "rpm": 720, "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", "supported_endpoints": [ "/v1/audio/transcriptions" @@ -60920,6 +61436,31 @@ "supports_web_search": false, "output_cost_per_image": 0.08 }, + "gemini/lyria-3.5": { + "input_cost_per_token": 0, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 0, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": false, + "supports_web_search": false, + "output_cost_per_image": 0.08 + }, "perplexity/anthropic/claude-fable-5": { "litellm_provider": "perplexity", "mode": "responses", diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 47a1934a703..7ed1e7e568b 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -478,6 +478,10 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_720p": { + "type": "number", + "minimum": 0 + }, "output_cost_per_token": { "type": "number", "minimum": 0, diff --git a/schema.prisma b/schema.prisma index 3d254cd2ea2..05c5aad9303 100644 --- a/schema.prisma +++ b/schema.prisma @@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping { updated_at DateTime @default(now()) @updatedAt updated_by String? - litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token]) + litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) @@unique([jwt_claim_name, jwt_claim_value]) @@index([jwt_claim_name, jwt_claim_value, is_active]) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 842bfb4bdb1..e5e0a164a83 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,6 +16,7 @@ longer signal it. ### Added +- **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it - **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users @@ -38,12 +39,14 @@ longer signal it. - **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state - **key**: Read now unwraps the `info` envelope `/key/info` actually returns; previously reads mapped nothing back into state, so drift on a key was never detected +- **key**: Read now picks up `model_rpm_limit`, `model_tpm_limit`, `guardrails`, `tags`, `enforced_params`, `allowed_passthrough_routes`, `rpm_limit_type`, `tpm_limit_type` and `prompts` from `info.metadata`, where the proxy actually stores them; previously they stayed empty in state, so a matching config showed a permanent phantom diff on them and out-of-band changes to them were never detected - **key**: Updates no longer send an empty `budget_duration`, which the proxy rejects with a 400; any update to a key without a configured `budget_duration` previously failed outright - **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead - **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs ### Changed +- **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying ## [0.4.0] - 2026-08-06 diff --git a/terraform/provider/README.md b/terraform/provider/README.md index 0a6d15c7844..b392fd6279d 100644 --- a/terraform/provider/README.md +++ b/terraform/provider/README.md @@ -103,9 +103,12 @@ resource "litellm_key" "example_key" { permissions = { can_create_keys = "true" } - model_max_budget = { - "gpt-4" = 50.0 - } + model_max_budget = jsonencode({ + "gpt-4" = { + budget_limit = 50.0 + time_period = "30d" + } + }) model_rpm_limit = { "claude-3.5-sonnet" = 30 } diff --git a/terraform/provider/docs/resources/key.md b/terraform/provider/docs/resources/key.md index 5094b77cbec..0ef0688830f 100644 --- a/terraform/provider/docs/resources/key.md +++ b/terraform/provider/docs/resources/key.md @@ -30,9 +30,12 @@ resource "litellm_key" "example" { permissions = { "can_create_keys" = "true" } - model_max_budget = { - "gpt-4" = 50.0 - } + model_max_budget = jsonencode({ + "gpt-4" = { + budget_limit = 50.0 + time_period = "30d" + } + }) model_rpm_limit = { "gpt-3.5-turbo" = 30 } @@ -73,7 +76,7 @@ The following arguments are supported: * `key_alias` - (Optional) Alias for this key. This provides a human-readable identifier for the key. -* `duration` - (Optional) Duration for which this key is valid. This sets an expiration time for the key. +* `duration` - (Optional) How long the key stays valid, e.g. "30d" or "12h". The proxy stores this as an absolute `expires` timestamp. Changing the value resets the expiry to the time of the update plus the new duration; removing it from the configuration leaves the current expiry in place. * `aliases` - (Optional) Map of model aliases. This allows you to create custom names for models when using this key. @@ -81,7 +84,7 @@ The following arguments are supported: * `permissions` - (Optional) Permissions associated with this key. This defines what actions are allowed with this key. -* `model_max_budget` - (Optional) Maximum budget per model. This allows setting different budget limits for each model. +* `model_max_budget` - (Optional) JSON string of per-model budget config, e.g. `jsonencode({"gpt-4" = {budget_limit = 50.0, time_period = "30d"}})`. Each model maps to an object with `budget_limit` (or `max_budget`), `time_period` (or `budget_duration`), `tpm_limit` and `rpm_limit`. * `model_rpm_limit` - (Optional) Requests per minute limit per model. This allows setting different RPM limits for each model. diff --git a/terraform/provider/docs/resources/team.md b/terraform/provider/docs/resources/team.md index 821d8c1dee3..19575907269 100644 --- a/terraform/provider/docs/resources/team.md +++ b/terraform/provider/docs/resources/team.md @@ -14,6 +14,16 @@ resource "litellm_team" "engineering" { } ``` +### Team with a Custom ID + +```hcl +resource "litellm_team" "platform" { + team_id = "platform-team" + team_alias = "platform" + models = ["gpt-4-proxy"] +} +``` + ### Team with Comprehensive Configuration ```hcl @@ -92,6 +102,8 @@ resource "litellm_team" "model_dependent_team" { The following arguments are supported: +* `team_id` - (Optional) A stable, human-readable ID for the team (for example `platform-team`). If omitted, the provider generates a random UUID. Changing this forces a new team to be created. + * `team_alias` - (Required) A human-readable identifier for the team. * `organization_id` - (Optional) The ID of the organization this team belongs to. @@ -152,7 +164,7 @@ The following arguments are supported: In addition to the arguments above, the following attributes are exported: -* `id` - The unique identifier for the team. +* `id` - The unique identifier for the team, equal to `team_id`. ## Import @@ -162,7 +174,7 @@ Teams can be imported using the team ID: terraform import litellm_team.engineering ``` -Note: The team ID is generated when the team is created and is different from the `team_alias`. +Note: Unless `team_id` is set, the team ID is generated when the team is created and is different from the `team_alias`. ## Note on Team Members diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index 0f825d85d31..e68b8a3a80b 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -4,6 +4,7 @@ import ( "bytes" "crypto/tls" "encoding/json" + "errors" "fmt" "io" "log" @@ -19,6 +20,20 @@ type Client struct { InsecureSkipVerify bool } +type apiError struct { + StatusCode int + Body string +} + +func (e *apiError) Error() string { + return fmt.Sprintf("API request failed with status code %d: %s", e.StatusCode, e.Body) +} + +func isNotFound(err error) bool { + var apiErr *apiError + return errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound +} + func NewClient(apiBase, apiKey string, insecureSkipVerify bool) *Client { tr := &http.Transport{ TLSClientConfig: &tls.Config{InsecureSkipVerify: insecureSkipVerify}, @@ -57,6 +72,9 @@ func (c *Client) CreateKey(key *Key) (*Key, error) { func (c *Client) GetKey(keyID string) (*Key, error) { resp, err := c.sendRequest("GET", fmt.Sprintf("/key/info?key=%s", keyID), nil) + if isNotFound(err) { + return nil, nil + } if err != nil { return nil, err } @@ -69,32 +87,71 @@ func (c *Client) GetKey(keyID string) (*Key, error) { info["key"] = k } } + hoistKeyFieldsStoredInMetadata(info) return c.parseKeyResponse(info) } return c.parseKeyResponse(resp) } +var keyFieldsStoredInMetadata = []string{ + "model_rpm_limit", + "model_tpm_limit", + "guardrails", + "tags", + "enforced_params", + "allowed_passthrough_routes", + "rpm_limit_type", + "tpm_limit_type", + "prompts", +} + +func hoistKeyFieldsStoredInMetadata(info map[string]interface{}) { + metadata, ok := info["metadata"].(map[string]interface{}) + if !ok { + return + } + for _, field := range keyFieldsStoredInMetadata { + if existing, present := info[field]; present && existing != nil { + continue + } + if v, present := metadata[field]; present { + info[field] = v + } + } +} + func (c *Client) UpdateKey(key *Key) (*Key, error) { // Create a new map with only the fields that can be updated updateData := map[string]interface{}{ "key": key.Key, "team_id": key.TeamID, - "metadata": key.Metadata, "key_alias": key.KeyAlias, "aliases": key.Aliases, "permissions": key.Permissions, "model_max_budget": key.ModelMaxBudget, - "model_rpm_limit": key.ModelRPMLimit, - "model_tpm_limit": key.ModelTPMLimit, "blocked": key.Blocked, } + // The proxy keeps the stored metadata only when the field is absent, so nil means omit. + if key.Metadata != nil { + updateData["metadata"] = key.Metadata + } + if key.ModelRPMLimit != nil { + updateData["model_rpm_limit"] = key.ModelRPMLimit + } + if key.ModelTPMLimit != nil { + updateData["model_tpm_limit"] = key.ModelTPMLimit + } + // The proxy rejects an empty-string budget_duration with a 400, so only // send it when set. if key.BudgetDuration != "" { updateData["budget_duration"] = key.BudgetDuration } + if key.Duration != "" { + updateData["duration"] = key.Duration + } // Only add pointer fields if they are explicitly set if key.MaxBudget != nil { @@ -366,7 +423,7 @@ func (c *Client) sendRequest(method, path string, body interface{}) (map[string] log.Printf("Response body: %s", c.redactSensitiveData(string(bodyBytes))) if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status code %d: %s", resp.StatusCode, string(bodyBytes)) + return nil, &apiError{StatusCode: resp.StatusCode, Body: string(bodyBytes)} } var result map[string]interface{} diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 0d8674f2d4c..ffdc448f5fe 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -2,7 +2,9 @@ package litellm import ( "context" + "encoding/json" "fmt" + "log" "github.com/hashicorp/go-cty/cty" "github.com/hashicorp/terraform-plugin-sdk/v2/diag" @@ -10,7 +12,7 @@ import ( ) func resourceKey() *schema.Resource { - return &schema.Resource{ + r := &schema.Resource{ CreateContext: resourceKeyCreate, ReadContext: resourceKeyRead, UpdateContext: resourceKeyUpdate, @@ -18,6 +20,7 @@ func resourceKey() *schema.Resource { Importer: &schema.ResourceImporter{ StateContext: schema.ImportStatePassthroughContext, }, + SchemaVersion: 1, Schema: map[string]*schema.Schema{ "key": { Type: schema.TypeString, @@ -86,8 +89,9 @@ func resourceKey() *schema.Resource { Optional: true, }, "duration": { - Type: schema.TypeString, - Optional: true, + Type: schema.TypeString, + Optional: true, + Description: "How long the key stays valid, e.g. \"30d\" or \"12h\". Changing it resets the expiry to the time of the update plus the new duration; removing it leaves the current expiry in place", }, "aliases": { Type: schema.TypeMap, @@ -105,9 +109,11 @@ func resourceKey() *schema.Resource { Elem: &schema.Schema{Type: schema.TypeString}, }, "model_max_budget": { - Type: schema.TypeMap, - Optional: true, - Elem: &schema.Schema{Type: schema.TypeFloat, Computed: true}, + Type: schema.TypeString, + Optional: true, + ValidateFunc: validateKeyModelMaxBudget, + DiffSuppressFunc: budgetSuppressEquivalentJSON, + Description: "JSON string of per-model budget config (e.g. '{\"gpt-4o-mini\": {\"budget_limit\": 50, \"time_period\": \"30d\"}}')", }, "model_rpm_limit": { Type: schema.TypeMap, @@ -182,6 +188,79 @@ func resourceKey() *schema.Resource { }, }, } + r.StateUpgraders = []schema.StateUpgrader{{ + Version: 0, + Type: resourceKeyV0Type(r.Schema), + Upgrade: resourceKeyStateUpgradeV0, + }} + return r +} + +// Schema version 0 typed model_max_budget as map(number), which the proxy +// rejects; version 1 stores the per-model BudgetConfig objects as a JSON string. +func resourceKeyV0Type(current map[string]*schema.Schema) cty.Type { + v0 := make(map[string]*schema.Schema, len(current)) + for k, v := range current { + v0[k] = v + } + v0["model_max_budget"] = &schema.Schema{ + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeFloat}, + } + return (&schema.Resource{Schema: v0}).CoreConfigSchema().ImpliedType() +} + +func resourceKeyStateUpgradeV0(_ context.Context, rawState map[string]interface{}, _ interface{}) (map[string]interface{}, error) { + delete(rawState, "model_max_budget") + return rawState, nil +} + +var keyModelBudgetFields = map[string]bool{ + "budget_limit": true, + "max_budget": true, + "time_period": true, + "budget_duration": true, + "tpm_limit": true, + "rpm_limit": true, +} + +func validateKeyModelMaxBudget(v interface{}, k string) ([]string, []error) { + var parsed map[string]json.RawMessage + if err := json.Unmarshal([]byte(v.(string)), &parsed); err != nil || parsed == nil { + return nil, []error{fmt.Errorf("%q must be a JSON object keyed by model name, got %s", k, v)} + } + for model, cfg := range parsed { + var budget map[string]json.RawMessage + if err := json.Unmarshal(cfg, &budget); err != nil || len(budget) == 0 { + return nil, []error{fmt.Errorf("%q[%q] must be a budget object such as {\"budget_limit\": 50, \"time_period\": \"30d\"}, got %s", k, model, cfg)} + } + for field := range budget { + if !keyModelBudgetFields[field] { + return nil, []error{fmt.Errorf("%q[%q] has unknown budget field %q; supported fields are budget_limit, max_budget, time_period, budget_duration, tpm_limit, rpm_limit", k, model, field)} + } + } + } + return nil, nil +} + +func parseKeyModelMaxBudget(raw string) map[string]interface{} { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(raw), &parsed); err != nil || parsed == nil { + return map[string]interface{}{} + } + return parsed +} + +func keyModelMaxBudgetJSON(modelMaxBudget map[string]interface{}) string { + if len(modelMaxBudget) == 0 { + return "" + } + encoded, err := json.Marshal(modelMaxBudget) + if err != nil { + return "" + } + return string(encoded) } func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics { @@ -219,10 +298,12 @@ func resourceKeyRead(ctx context.Context, d *schema.ResourceData, m interface{}) } if key == nil { + log.Printf("[WARN] Key %s not found, removing from state", d.Id()) d.SetId("") return nil } + key.Metadata = declaredKeyMetadata(key.Metadata, d.Get("metadata").(map[string]interface{})) mapKeyToResourceData(d, key) return nil } @@ -232,15 +313,75 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ key := &Key{Key: d.Id()} mapResourceDataToKey(d, key) + if !d.HasChange("duration") { + key.Duration = "" + } + key.ModelRPMLimit = changedMap(d, "model_rpm_limit") + key.ModelTPMLimit = changedMap(d, "model_tpm_limit") - _, err := c.UpdateKey(key) + metadata, err := plannedKeyMetadata(c, d) if err != nil { + d.Partial(true) + return diag.FromErr(fmt.Errorf("error updating key: %s", err)) + } + key.Metadata = metadata + + if _, err := c.UpdateKey(key); err != nil { return diag.FromErr(fmt.Errorf("error updating key: %s", err)) } return resourceKeyRead(ctx, d, m) } +func changedMap(d *schema.ResourceData, name string) map[string]interface{} { + if !d.HasChange(name) { + return nil + } + return d.Get(name).(map[string]interface{}) +} + +func plannedKeyMetadata(c *Client, d *schema.ResourceData) (map[string]interface{}, error) { + if !d.HasChange("metadata") { + return nil, nil + } + current, err := c.GetKey(d.Id()) + if err != nil { + return nil, err + } + if current == nil { + return nil, fmt.Errorf("key %s no longer exists", d.Id()) + } + oldDeclared, newDeclared := d.GetChange("metadata") + return mergeKeyMetadata(current.Metadata, oldDeclared.(map[string]interface{}), newDeclared.(map[string]interface{})), nil +} + +func declaredKeyMetadata(server, declared map[string]interface{}) map[string]interface{} { + if server == nil { + return nil + } + result := make(map[string]interface{}, len(declared)) + for k := range declared { + if v, ok := server[k]; ok { + result[k] = v + } + } + return result +} + +func mergeKeyMetadata(server, oldDeclared, newDeclared map[string]interface{}) map[string]interface{} { + result := make(map[string]interface{}, len(server)+len(newDeclared)) + for k, v := range server { + result[k] = v + } + for k := range oldDeclared { + delete(result, k) + } + for k, v := range newDeclared { + result[k] = v + } + return result +} + func resourceKeyDelete(ctx context.Context, d *schema.ResourceData, m interface{}) diag.Diagnostics { c := m.(*Client) @@ -285,7 +426,7 @@ func mapResourceDataToKey(d *schema.ResourceData, key *Key) { key.Aliases = d.Get("aliases").(map[string]interface{}) key.Config = d.Get("config").(map[string]interface{}) key.Permissions = d.Get("permissions").(map[string]interface{}) - key.ModelMaxBudget = d.Get("model_max_budget").(map[string]interface{}) + key.ModelMaxBudget = parseKeyModelMaxBudget(d.Get("model_max_budget").(string)) key.ModelRPMLimit = d.Get("model_rpm_limit").(map[string]interface{}) key.ModelTPMLimit = d.Get("model_tpm_limit").(map[string]interface{}) key.Guardrails = expandStringList(d.Get("guardrails").([]interface{})) @@ -358,9 +499,7 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) { if key.Permissions != nil { d.Set("permissions", key.Permissions) } - if key.ModelMaxBudget != nil { - d.Set("model_max_budget", key.ModelMaxBudget) - } + d.Set("model_max_budget", keyModelMaxBudgetJSON(key.ModelMaxBudget)) if key.ModelRPMLimit != nil { d.Set("model_rpm_limit", key.ModelRPMLimit) } diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go index 91f0061a9ef..66291eadcc5 100644 --- a/terraform/provider/litellm/resource_key_test.go +++ b/terraform/provider/litellm/resource_key_test.go @@ -6,9 +6,11 @@ import ( "io" "net/http" "net/http/httptest" + "reflect" "testing" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { @@ -193,6 +195,102 @@ func TestCreateKeySendsConfigSuppliedKey(t *testing.T) { } } +// The proxy validates each model_max_budget entry as a BudgetConfig object and +// 500s on a bare number, so the JSON string must reach /key/generate as nested +// objects and the proxy's response must map back to equivalent JSON in state. +func TestCreateKeySendsModelMaxBudgetAsBudgetObjects(t *testing.T) { + var captured map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/key/generate" { + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &captured) + w.Write([]byte(`{"key": "sk-test", "token_id": "hash-1"}`)) + return + } + w.Write([]byte(`{"key": "hash-1", "info": {"model_max_budget": {"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d", "rpm_limit": 60}}}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyResourceData(t, map[string]interface{}{ + "model_max_budget": `{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`, + }) + + if diags := resourceKeyCreate(context.Background(), d, client); diags.HasError() { + t.Fatalf("create returned error: %v", diags) + } + + budgets, ok := captured["model_max_budget"].(map[string]interface{}) + if !ok { + t.Fatalf("create payload model_max_budget = %v, want object", captured["model_max_budget"]) + } + cfg, ok := budgets["gpt-4o-mini"].(map[string]interface{}) + if !ok { + t.Fatalf("model_max_budget[gpt-4o-mini] = %v, want BudgetConfig object", budgets["gpt-4o-mini"]) + } + if cfg["budget_limit"] != float64(50) || cfg["time_period"] != "30d" { + t.Errorf("BudgetConfig = %v, want budget_limit 50 and time_period 30d", cfg) + } + + var state map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("model_max_budget").(string)), &state); err != nil { + t.Fatalf("state model_max_budget %q is not JSON: %v", d.Get("model_max_budget"), err) + } + if got, _ := state["gpt-4o-mini"].(map[string]interface{}); got["budget_limit"] != float64(50) || got["rpm_limit"] != float64(60) { + t.Errorf("state model_max_budget = %v, want the BudgetConfig read back from /key/info", state) + } +} + +// Schema version 0 stored model_max_budget as map(number); that state cannot +// decode into the version 1 string attribute, so the upgrader must drop it. +func TestKeyStateUpgradeV0DropsMapModelMaxBudget(t *testing.T) { + upgraded, err := resourceKey().StateUpgraders[0].Upgrade(context.Background(), map[string]interface{}{ + "id": "hash-1", + "key_alias": "legacy", + "model_max_budget": map[string]interface{}{"gpt-4o-mini": 50.0}, + }, nil) + if err != nil { + t.Fatalf("upgrade returned error: %v", err) + } + if _, present := upgraded["model_max_budget"]; present { + t.Errorf("upgraded state still carries map model_max_budget: %v", upgraded["model_max_budget"]) + } + if upgraded["key_alias"] != "legacy" { + t.Errorf("upgrade dropped unrelated attribute: %v", upgraded) + } +} + +func TestKeyModelMaxBudgetValidationRequiresBudgetObjects(t *testing.T) { + validate := resourceKey().Schema["model_max_budget"].ValidateFunc + for _, valid := range []string{ + `{}`, + `{"gpt-4o-mini": {"budget_limit": 50, "time_period": "30d"}}`, + `{"gpt-4o-mini": {"max_budget": 50, "rpm_limit": 60}, "gpt-4o": {"budget_duration": "1d", "tpm_limit": 1000}}`, + } { + if _, errs := validate(valid, "model_max_budget"); len(errs) != 0 { + t.Errorf("validate(%s) = %v, want accepted", valid, errs) + } + } + for _, invalid := range []string{ + `null`, + `[]`, + `"gpt-4o-mini"`, + `50`, + `{"gpt-4o-mini": 50}`, + `{"gpt-4o-mini": null}`, + `{"gpt-4o-mini": [50]}`, + `{"gpt-4o-mini": {}}`, + `{"gpt-4o-mini": {"budget_limt": 50}}`, + `{"gpt-4o-mini": {"budget_limit": 50, "max_tokens": 100}}`, + `not json`, + } { + if _, errs := validate(invalid, "model_max_budget"); len(errs) == 0 { + t.Errorf("validate(%s) accepted a value that would send no per-model budget", invalid) + } + } +} + // The proxy 400s on budget_duration: "", so an unset duration must be // omitted from the update payload entirely. func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) { @@ -221,6 +319,47 @@ func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) { } } +func TestResourceKeyUpdateFailureKeepsPriorState(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/key/update" { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"error":{"message":"Invalid budget_duration 'bad'"}}`)) + return + } + w.Write([]byte(`{"key":"hash-1","info":{"key_alias":"demo","models":["fake-model"]}}`)) + })) + defer srv.Close() + + res := resourceKey() + priorData := newKeyResourceData(t, map[string]interface{}{ + "key_alias": "demo", + "models": []interface{}{"fake-model"}, + }) + priorData.SetId("hash-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "key_alias": "demo", + "models": []interface{}{"fake-model"}, + "budget_duration": "bad", + }) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + + newState, diags := res.Apply(context.Background(), prior, diff, NewClient(srv.URL, "test-key", true)) + if !diags.HasError() { + t.Fatal("apply succeeded, want the proxy's 400 surfaced as an error") + } + if got, ok := newState.Attributes["budget_duration"]; ok { + t.Errorf("failed update persisted budget_duration=%q into state, want it absent", got) + } + if newState.Attributes["key_alias"] != "demo" { + t.Errorf("prior key_alias lost from state: %v", newState.Attributes) + } +} + // /key/info nests the key's fields under "info"; GetKey must unwrap that // envelope or reads map nothing back into state. func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { @@ -254,3 +393,296 @@ func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { t.Errorf("RPMLimit not parsed: %+v", key.RPMLimit) } } + +func TestGetKeyReadsFieldsStoredInMetadata(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "hash-1", + "info": { + "models": ["gpt-4o-mini"], + "metadata": { + "team": "core-infra", + "model_rpm_limit": {"gpt-4o-mini": 7}, + "model_tpm_limit": {"gpt-4o-mini": 10000}, + "guardrails": ["pii-guard"], + "tags": ["prod"], + "enforced_params": ["user"], + "allowed_passthrough_routes": ["/v1/foo"], + "rpm_limit_type": "guaranteed_throughput", + "tpm_limit_type": "dynamic", + "prompts": ["p1"] + } + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + key, err := client.GetKey("hash-1") + if err != nil { + t.Fatalf("GetKey returned error: %v", err) + } + if got, ok := key.ModelRPMLimit["gpt-4o-mini"].(float64); !ok || got != 7 { + t.Errorf("ModelRPMLimit = %v, want gpt-4o-mini=7 read from metadata", key.ModelRPMLimit) + } + if got, ok := key.ModelTPMLimit["gpt-4o-mini"].(float64); !ok || got != 10000 { + t.Errorf("ModelTPMLimit = %v, want gpt-4o-mini=10000 read from metadata", key.ModelTPMLimit) + } + if len(key.Guardrails) != 1 || key.Guardrails[0] != "pii-guard" { + t.Errorf("Guardrails = %v, want [pii-guard]", key.Guardrails) + } + if len(key.Tags) != 1 || key.Tags[0] != "prod" { + t.Errorf("Tags = %v, want [prod]", key.Tags) + } + if len(key.EnforcedParams) != 1 || key.EnforcedParams[0] != "user" { + t.Errorf("EnforcedParams = %v, want [user]", key.EnforcedParams) + } + if len(key.AllowedPassthroughRoutes) != 1 || key.AllowedPassthroughRoutes[0] != "/v1/foo" { + t.Errorf("AllowedPassthroughRoutes = %v, want [/v1/foo]", key.AllowedPassthroughRoutes) + } + if key.RPMLimitType != "guaranteed_throughput" || key.TPMLimitType != "dynamic" { + t.Errorf("limit types = %q/%q, want guaranteed_throughput/dynamic", key.RPMLimitType, key.TPMLimitType) + } + if len(key.Prompts) != 1 || key.Prompts[0] != "p1" { + t.Errorf("Prompts = %v, want [p1]", key.Prompts) + } + if key.Metadata["team"] != "core-infra" { + t.Errorf("Metadata = %v, want team=core-infra preserved", key.Metadata) + } +} + +func TestGetKeyPrefersTopLevelOverMetadataCopy(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "hash-1", + "info": { + "tags": ["top-level"], + "guardrails": null, + "metadata": { + "tags": ["from-metadata"], + "guardrails": ["from-metadata"] + } + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + key, err := client.GetKey("hash-1") + if err != nil { + t.Fatalf("GetKey returned error: %v", err) + } + if len(key.Tags) != 1 || key.Tags[0] != "top-level" { + t.Errorf("Tags = %v, want [top-level]", key.Tags) + } + if len(key.Guardrails) != 1 || key.Guardrails[0] != "from-metadata" { + t.Errorf("Guardrails = %v, want [from-metadata] (null top-level must not shadow)", key.Guardrails) + } +} + +func TestResourceKeyReadDropsMissingKeyFromState(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + w.Write([]byte(`{"error":{"message":"Key not found in database","type":"not_found_error","param":"key","code":"404"}}`)) + })) + defer srv.Close() + + d := newKeyResourceData(t, map[string]interface{}{"key_alias": "stale"}) + d.SetId("deleted-out-of-band") + + diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true)) + if diags.HasError() { + t.Fatalf("read of a missing key must not error, got: %v", diags) + } + if d.Id() != "" { + t.Errorf("Id = %q, want empty so Terraform plans a recreate", d.Id()) + } +} + +func TestResourceKeyReadStillFailsOnNon404Errors(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":{"message":"db down"}}`)) + })) + defer srv.Close() + + d := newKeyResourceData(t, map[string]interface{}{"key_alias": "live"}) + d.SetId("still-exists") + + diags := resourceKeyRead(context.Background(), d, NewClient(srv.URL, "test-key", true)) + if !diags.HasError() { + t.Fatal("a 500 from /key/info must surface as an error, not be treated as a deleted key") + } + if d.Id() != "still-exists" { + t.Errorf("Id = %q, want unchanged on a transient error", d.Id()) + } +} + +// fakeKeyProxy serves /key/info from stored metadata and applies /key/update +// the way the proxy does: an absent "metadata" keeps the stored map, a +// present one replaces it wholesale. +type fakeKeyProxy struct { + metadata map[string]interface{} + updates []map[string]interface{} +} + +func (p *fakeKeyProxy) handler() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/key/info": + json.NewEncoder(w).Encode(map[string]interface{}{ + "key": "hash-1", + "info": map[string]interface{}{"key_alias": "alias-1", "models": []string{"gpt-4o-mini"}, "metadata": p.metadata}, + }) + case "/key/update": + var body map[string]interface{} + json.NewDecoder(r.Body).Decode(&body) + p.updates = append(p.updates, body) + if m, ok := body["metadata"].(map[string]interface{}); ok { + p.metadata = m + } + json.NewEncoder(w).Encode(map[string]interface{}{"key": "hash-1", "metadata": p.metadata}) + default: + http.NotFound(w, r) + } + } +} + +func applyKeyUpdate(t *testing.T, client *Client, stateAttrs map[string]string, config map[string]interface{}) *terraform.InstanceState { + t.Helper() + r := resourceKey() + state := &terraform.InstanceState{ID: "hash-1", Attributes: stateAttrs} + diff, err := r.Diff(context.Background(), state, terraform.NewResourceConfigRaw(config), client) + if err != nil { + t.Fatalf("Diff returned error: %v", err) + } + if diff == nil { + t.Fatalf("expected a non-empty diff between %v and %v", stateAttrs, config) + } + newState, diags := r.Apply(context.Background(), state, diff, client) + if diags.HasError() { + t.Fatalf("Apply returned error: %v", diags) + } + return newState +} + +func TestKeyUpdateWithoutMetadataChangePreservesServerMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": float64(5)}}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + newState := applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "max_budget": "10", "metadata.%": "1", "metadata.a": "1"}, + map[string]interface{}{"key_alias": "alias-1", "max_budget": 20, "metadata": map[string]interface{}{"a": "1"}}, + ) + + if len(proxy.updates) != 1 { + t.Fatalf("expected one /key/update call, got %d", len(proxy.updates)) + } + for _, field := range []string{"metadata", "model_rpm_limit", "model_tpm_limit"} { + if _, present := proxy.updates[0][field]; present { + t.Errorf("unchanged %q was sent on /key/update: %v", field, proxy.updates[0][field]) + } + } + if proxy.metadata["server_side"] != "x" { + t.Errorf("server-side metadata lost: %v", proxy.metadata) + } + if got := newState.Attributes["metadata.%"]; got != "1" { + t.Errorf("state metadata should hold only the declared entry, got %v", newState.Attributes) + } + if got := newState.Attributes["metadata.a"]; got != "1" { + t.Errorf("metadata.a = %q, want 1", got) + } +} + +func TestKeyUpdateWithMetadataChangeMergesOverServerMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "b": "2", "server_side": "x"}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "metadata.%": "2", "metadata.a": "1", "metadata.b": "2"}, + map[string]interface{}{"key_alias": "alias-1", "metadata": map[string]interface{}{"a": "2", "c": "3"}}, + ) + + want := map[string]interface{}{"a": "2", "c": "3", "server_side": "x"} + if !reflect.DeepEqual(proxy.metadata, want) { + t.Errorf("metadata after update = %v, want %v", proxy.metadata, want) + } +} + +func TestKeyUpdateSendsChangedModelLimits(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "model_rpm_limit.%": "1", "model_rpm_limit.gpt-4o-mini": "5"}, + map[string]interface{}{"key_alias": "alias-1", "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 7}}, + ) + + got, ok := proxy.updates[0]["model_rpm_limit"].(map[string]interface{}) + if !ok || got["gpt-4o-mini"] != float64(7) { + t.Errorf("changed model_rpm_limit not sent: %v", proxy.updates[0]) + } +} + +func TestKeyReadKeepsOnlyDeclaredMetadata(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{"a": "1", "server_side": "x"}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + d := newKeyResourceData(t, map[string]interface{}{"metadata": map[string]interface{}{"a": "1"}}) + d.SetId("hash-1") + if diags := resourceKeyRead(context.Background(), d, client); diags.HasError() { + t.Fatalf("Read returned error: %v", diags) + } + + want := map[string]interface{}{"a": "1"} + if got := d.Get("metadata"); !reflect.DeepEqual(got, want) { + t.Errorf("metadata in state = %v, want %v", got, want) + } +} + +func TestKeyUpdateSendsChangedDuration(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "duration": "30d"}, + map[string]interface{}{"key_alias": "alias-1", "duration": "90d"}, + ) + + if got := proxy.updates[0]["duration"]; got != "90d" { + t.Errorf("update payload duration = %v, want 90d", got) + } +} + +func TestKeyUpdateOmitsUnchangedDuration(t *testing.T) { + proxy := &fakeKeyProxy{metadata: map[string]interface{}{}} + srv := httptest.NewServer(proxy.handler()) + defer srv.Close() + client := NewClient(srv.URL, "test-key", true) + + applyKeyUpdate(t, client, + map[string]string{"key_alias": "alias-1", "duration": "30d"}, + map[string]interface{}{"key_alias": "alias-2", "duration": "30d"}, + ) + + if got := proxy.updates[0]["key_alias"]; got != "alias-2" { + t.Fatalf("update payload key_alias = %v, want alias-2", got) + } + if v, present := proxy.updates[0]["duration"]; present { + t.Errorf("update payload unexpectedly contains duration = %v", v) + } +} diff --git a/terraform/provider/litellm/resource_key_utils.go b/terraform/provider/litellm/resource_key_utils.go index d426fec05b2..7046ef7097a 100644 --- a/terraform/provider/litellm/resource_key_utils.go +++ b/terraform/provider/litellm/resource_key_utils.go @@ -65,7 +65,7 @@ func buildKeyData(d *schema.ResourceData) map[string]interface{} { keyData["permissions"] = v.(map[string]interface{}) } if v, ok := d.GetOkExists("model_max_budget"); ok { - keyData["model_max_budget"] = v.(map[string]interface{}) + keyData["model_max_budget"] = parseKeyModelMaxBudget(v.(string)) } if v, ok := d.GetOkExists("model_rpm_limit"); ok { keyData["model_rpm_limit"] = v.(map[string]interface{}) @@ -107,7 +107,7 @@ func setKeyResourceData(d *schema.ResourceData, key *Key) error { "aliases": key.Aliases, "config": key.Config, "permissions": key.Permissions, - "model_max_budget": key.ModelMaxBudget, + "model_max_budget": keyModelMaxBudgetJSON(key.ModelMaxBudget), "model_rpm_limit": key.ModelRPMLimit, "model_tpm_limit": key.ModelTPMLimit, "guardrails": key.Guardrails, diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index 24c47843cd1..bf7d2508077 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -31,6 +31,13 @@ func ResourceLiteLLMTeam() *schema.Resource { }, Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Optional: true, + Computed: true, + ForceNew: true, + Description: "Unique ID for the team. Generated by the provider if not provided", + }, "team_alias": { Type: schema.TypeString, Required: true, @@ -162,7 +169,7 @@ func ResourceLiteLLMTeam() *schema.Resource { func resourceLiteLLMTeamCreate(d *schema.ResourceData, m interface{}) error { client := m.(*Client) - teamID := uuid.New().String() + teamID := resolveTeamID(d) teamData := buildTeamData(d, teamID) // Throughput limit types are only accepted by /team/new, not /team/update. @@ -214,6 +221,7 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { teamResp := infoResp.TeamInfo // Update the state with values from the response or fall back to the data passed in during creation + d.Set("team_id", d.Id()) d.Set("team_alias", GetStringValue(teamResp.TeamAlias, d.Get("team_alias").(string))) d.Set("organization_id", GetStringValue(teamResp.OrganizationID, d.Get("organization_id").(string))) @@ -263,11 +271,11 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { d.Set("team_member_tpm_limit", *teamResp.TeamMemberTPMLimit) } d.Set("team_member_key_duration", GetStringValue(teamResp.TeamMemberKeyDuration, d.Get("team_member_key_duration").(string))) - if teamResp.ModelRPMLimit != nil { - d.Set("model_rpm_limit", teamResp.ModelRPMLimit) + if v := teamModelLimit(teamResp.ModelRPMLimit, teamResp.Metadata, "model_rpm_limit"); v != nil { + d.Set("model_rpm_limit", v) } - if teamResp.ModelTPMLimit != nil { - d.Set("model_tpm_limit", teamResp.ModelTPMLimit) + if v := teamModelLimit(teamResp.ModelTPMLimit, teamResp.Metadata, "model_tpm_limit"); v != nil { + d.Set("model_tpm_limit", v) } if teamResp.AllowedPassthroughRoutes != nil { d.Set("allowed_passthrough_routes", teamResp.AllowedPassthroughRoutes) @@ -354,6 +362,13 @@ func resourceLiteLLMTeamDelete(d *schema.ResourceData, m interface{}) error { return nil } +func resolveTeamID(d *schema.ResourceData) string { + if v, ok := d.GetOk("team_id"); ok { + return v.(string) + } + return uuid.New().String() +} + func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} { teamData := map[string]interface{}{ "team_id": teamID, @@ -364,14 +379,19 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} "organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions", "model_aliases", "guardrails", "prompts", "team_member_budget", "team_member_budget_duration", "team_member_rpm_limit", - "team_member_tpm_limit", "team_member_key_duration", "model_rpm_limit", - "model_tpm_limit", "allowed_passthrough_routes", + "team_member_tpm_limit", "team_member_key_duration", "allowed_passthrough_routes", } { if v, ok := d.GetOk(key); ok { teamData[key] = v } } + for _, key := range []string{"model_rpm_limit", "model_tpm_limit"} { + if v, ok := d.GetOk(key); ok || d.HasChange(key) { + teamData[key] = v + } + } + if v, ok := d.GetOk("soft_budget"); ok { teamData["soft_budget"] = v } else if d.HasChange("soft_budget") { @@ -404,6 +424,14 @@ func buildTeamMetadata(d *schema.ResourceData) map[string]interface{} { return metadata } +func teamModelLimit(topLevel, metadata map[string]interface{}, key string) map[string]interface{} { + if topLevel != nil { + return topLevel + } + nested, _ := metadata[key].(map[string]interface{}) + return nested +} + func splitTeamMetadata(raw map[string]interface{}) (map[string]string, []string, []string) { metadata := map[string]string{} var tags, alertEmails []string diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go index 9638378cdfe..35d60401d30 100644 --- a/terraform/provider/litellm/resource_team_test.go +++ b/terraform/provider/litellm/resource_team_test.go @@ -9,6 +9,7 @@ import ( "reflect" "testing" + "github.com/google/uuid" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) @@ -85,6 +86,87 @@ func TestTeamCreateSendsSoftBudgetTagsAndAlertEmails(t *testing.T) { } } +func TestTeamCreateSendsConfiguredTeamID(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"platform-team","team_info":{"team_id":"platform-team","team_alias":"platform"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_id": "platform-team", + "team_alias": "platform", + }) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if got := captured["team_id"]; got != "platform-team" { + t.Fatalf("payload team_id = %v, want platform-team", got) + } + if got := d.Id(); got != "platform-team" { + t.Fatalf("resource id = %q, want platform-team", got) + } + if got := d.Get("team_id"); got != "platform-team" { + t.Fatalf("state team_id = %v, want platform-team", got) + } +} + +func TestTeamCreateGeneratesTeamIDWhenUnset(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"x","team_info":{"team_alias":"eng"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"}) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + sent, _ := captured["team_id"].(string) + if _, err := uuid.Parse(sent); err != nil { + t.Fatalf("payload team_id = %q, want a generated UUID: %v", sent, err) + } + if d.Id() != sent || d.Get("team_id") != sent { + t.Fatalf("id = %q, state team_id = %v, want both to equal the sent id %q", d.Id(), d.Get("team_id"), sent) + } +} + +func TestTeamReadSetsTeamIDFromResourceID(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"imported-team","team_info":{"team_id":"imported-team","team_alias":"imported"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{}) + d.SetId("imported-team") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + if got := d.Get("team_id"); got != "imported-team" { + t.Fatalf("team_id = %v, want imported-team", got) + } +} + +func TestTeamIDChangeForcesReplacement(t *testing.T) { + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_id": "old-team", + "team_alias": "eng", + }) + priorData.SetId("old-team") + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "team_id": "new-team", + "team_alias": "eng", + }) + diff, err := res.Diff(context.Background(), priorData.State(), config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + if diff == nil || !diff.RequiresNew() { + t.Fatalf("changing team_id must force replacement, diff = %+v", diff) + } +} + func TestTeamReadMapsTeamInfoEnvelope(t *testing.T) { var captured map[string]interface{} srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget) @@ -250,6 +332,77 @@ func TestTeamReadMapsNewFields(t *testing.T) { } } +func TestTeamReadMapsPerModelLimitsFromMetadata(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{ + "team_id": "team-1", + "team_info": { + "team_id": "team-1", + "team_alias": "eng", + "model_rpm_limit": null, + "model_tpm_limit": null, + "metadata": { + "department": "eng", + "model_rpm_limit": {"gpt-4o-mini": 250}, + "model_tpm_limit": {"gpt-4o-mini": 5000} + } + } + }`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100}, + }) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read returned error: %v", err) + } + if got := d.Get("model_rpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 250}) { + t.Errorf("model_rpm_limit = %v, want server value 250", got) + } + if got := d.Get("model_tpm_limit"); !reflect.DeepEqual(got, map[string]interface{}{"gpt-4o-mini": 5000}) { + t.Errorf("model_tpm_limit = %v, want server value 5000", got) + } + if got := d.Get("metadata"); !reflect.DeepEqual(got, map[string]interface{}{"department": "eng"}) { + t.Errorf("metadata = %v, want per-model limits kept out of the string map", got) + } +} + +func TestTeamUpdateClearsRemovedPerModelLimits(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_alias": "eng", + "model_rpm_limit": map[string]interface{}{"gpt-4o-mini": 100}, + "model_tpm_limit": map[string]interface{}{"gpt-4o-mini": 5000}, + }) + priorData.SetId("team-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{"team_alias": "eng"}) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(prior, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + for _, k := range []string{"model_rpm_limit", "model_tpm_limit"} { + if got, ok := captured[k]; !ok || !reflect.DeepEqual(got, map[string]interface{}{}) { + t.Errorf("payload %s = %v (present=%v), want explicit empty map", k, got, ok) + } + } +} + // rpm_limit_type / tpm_limit_type are accepted by /team/new but not // /team/update, so create must send them and update must not. func TestTeamLimitTypesSentOnCreateOnly(t *testing.T) { diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 0af29f069c6..a11f015743b 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -87,6 +87,7 @@ ignored_function_names = [ "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py + "_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name) ] diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..34fbe9d9247 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -181,13 +181,15 @@ quota_management... | team_multi_window | fallback | spend_counter chat_completions | stream | messages_bridge | embeddings | cache_hit | key_rollup | concurrent_burst | tags | end_user - | per_model | failure | spend_calculate | pagination + | per_model | failure | spend_calculate | pagination | key_attribution assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows - | writes_failure_row | returns_cost | keeps_total + | writes_failure_row | returns_cost | keeps_total | joins_key | reports_alias_and_email + | health_rows_keep_service_account | retrieve_batch_cost_joins_retrieving_key + | poller_batch_cost_joins_creating_key e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] quota_management.budget.key.blocks_over_limit exercised_on=[chat_completions] ``` diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index d0afcaca848..ad0914d455b 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -58,3 +58,8 @@ - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} - {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"} +- {id: quota_management.spend_tracking.key_attribution.joins_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [joins_key], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Every spend row a virtual key writes across chat, queued chat, messages, responses, embeddings, the Gemini passthrough, file upload, batch create, and a replayed callback log carries api_key equal to the key's token hash and the key alias, the join the usage APIs depend on; a re-hashed token shows up as an unattributed key-hash-* row (#39568, #39572)"} +- {id: quota_management.spend_tracking.key_attribution.reports_alias_and_email, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [reports_alias_and_email], exercised_on: [chat_completions, messages, responses, embeddings, batches, files, google_native, rust_control_plane], source: "proxy/management_endpoints/internal_user_endpoints.py", rationale: "/spend/logs?api_key= returns every one of the key's rows with its alias and /user/daily/activity aggregates them under the key's token with key_alias and user_email; /spend/logs carries no email field, so the email is asserted on daily activity only"} +- {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} +- {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"} +- {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..faf8557498b 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -84,6 +84,7 @@ class KeyGenerateBody(BaseModel): class KeyGenerateResponse(BaseModel): key: str + token: str | None = None key_alias: str | None = None models: list[str] = [] max_budget: float | None = None @@ -672,6 +673,7 @@ class GuardrailRunRecord(BaseModel): class SpendLogMetadata(BaseModel): + user_api_key_alias: str | None = None applied_guardrails: list[str] | None = None guardrail_information: list[GuardrailRunRecord] | None = None diff --git a/tests/e2e/quota_management/spend_tracking/conftest.py b/tests/e2e/quota_management/spend_tracking/conftest.py index 0597c9af400..9c8ffd18144 100644 --- a/tests/e2e/quota_management/spend_tracking/conftest.py +++ b/tests/e2e/quota_management/spend_tracking/conftest.py @@ -36,6 +36,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = ( ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"), ("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"), ("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"), + ("openai-gpt-4o-mini", "openai/gpt-4o-mini", "OPENAI_API_KEY"), ) diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 056799b8499..9ac97f57f47 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -15,9 +15,12 @@ import time from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from typing import Final from e2e_config import unique_marker from e2e_http import ( + FileUploadForm, + Headers, NoBody, ProbeResult, Result, @@ -35,6 +38,8 @@ from models import ( DateRangeParams, EmbedBody, EmbedResponse, + KeyGenerateBody, + KeyGenerateResponse, OpenAPISchema, SpendCalculateBody, SpendCalculateResponse, @@ -43,13 +48,27 @@ from models import ( SpendLogsPageParams, SpendTagsResponse, TagSpend, + UserDeleteBody, + UserDeleteResponse, + UserNewBody, + UserNewResponse, + UserRole, ) -from proxy_client import ProxyClient +from proxy_client import Converged, ProxyClient, await_converged +from pydantic import BaseModel, Field __all__ = [ + "BatchCreateBody", + "CallbackLogMetadata", + "CallbackLogPayload", + "BatchObject", + "DailyActivityKeyBreakdown", + "FileObject", "ProbeResult", + "ResponseIdentity", "SpendClient", "SpendLogRow", + "StreamingResponse", "build_client", "is_ok", "unique_marker", @@ -57,6 +76,139 @@ __all__ = [ ] +class GeminiApiKeyHeaders(Headers): + x_goog_api_key: str = Field(serialization_alias="x-goog-api-key") + content_type: str = Field(default="application/json", serialization_alias="Content-Type") + + +class GeminiPart(BaseModel): + text: str + + +class GeminiContent(BaseModel): + parts: list[GeminiPart] + + +class GeminiGenerationConfig(BaseModel): + maxOutputTokens: int + + +class GeminiGenerateBody(BaseModel): + contents: list[GeminiContent] + generationConfig: GeminiGenerationConfig + + +class ResponsesBody(BaseModel): + model: str + input: str + cache: dict[str, bool] | None = {"no-cache": True} + + +class QueuedChatBody(ChatBody): + priority: int = 0 + + +class ResponseIdentity(BaseModel): + id: str | None = None + + +class HealthParams(BaseModel): + model: str + + +class ModelQuery(BaseModel): + model: str + + +class FileObject(BaseModel): + id: str + + +class BatchCreateBody(BaseModel): + input_file_id: str + endpoint: str = "/v1/chat/completions" + completion_window: str = "24h" + model: str + metadata: dict[str, str] + + +class BatchObject(BaseModel): + id: str + status: str + + +class ProviderQuery(BaseModel): + provider: str + + +class CallbackLogMetadata(BaseModel): + user_api_key_hash: str + user_api_key_alias: str + user_api_key_user_id: str + + +class CallbackLogPayload(BaseModel): + id: str + litellm_call_id: str + model: str + call_type: str = "acompletion" + start_time: float = Field(serialization_alias="startTime") + end_time: float = Field(serialization_alias="endTime") + response_cost: float + prompt_tokens: int + completion_tokens: int + total_tokens: int + metadata: CallbackLogMetadata + + +class CallbackLogRecord(BaseModel): + status: str = "success" + standard_logging_payload: CallbackLogPayload + + +class CallbackLogsRequest(BaseModel): + records: list[CallbackLogRecord] + + +class CallbackLogsResponse(BaseModel): + processed: int + failed: int + + +class DailyActivityParams(BaseModel): + start_date: str + end_date: str + api_key: str + + +class DailyActivityKeyMetadata(BaseModel): + key_alias: str | None = None + team_id: str | None = None + user_email: str | None = None + + +class DailyActivityKeyMetrics(BaseModel): + api_requests: int = 0 + + +class DailyActivityKeyBreakdown(BaseModel): + metrics: DailyActivityKeyMetrics + metadata: DailyActivityKeyMetadata + + +class DailyActivityBreakdown(BaseModel): + api_keys: dict[str, DailyActivityKeyBreakdown] = {} + + +class DailyActivityRow(BaseModel): + date: str + breakdown: DailyActivityBreakdown + + +class DailyActivityResponse(BaseModel): + results: list[DailyActivityRow] = [] + + def _chat_body( model: str, content: str, @@ -207,6 +359,166 @@ class SpendClient: def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: return self.proxy.transport.probe(path, params=params) + def create_user(self, *, email: str, role: UserRole, user_id: str) -> str: + return unwrap( + self.proxy.transport.post( + "/user/new", + headers=self.proxy.transport.master, + json=UserNewBody(user_email=email, user_role=role, user_id=user_id), + response_type=UserNewResponse, + ) + ).user_id + + def delete_user(self, user_id: str) -> None: + _ = unwrap( + self.proxy.transport.post( + "/user/delete", + headers=self.proxy.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) + ) + + def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse: + return unwrap( + self.proxy.transport.post( + "/key/generate", + headers=self.proxy.transport.master, + json=body, + response_type=KeyGenerateResponse, + ) + ) + + def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/chat/completions", + headers=self.proxy.transport.bearer(key), + json=_chat_body(model, content, max_tokens=max_tokens), + ) + + def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/queue/chat/completions", + headers=self.proxy.transport.bearer(key), + json=QueuedChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + "/v1/messages", + headers=self.proxy.transport.bearer(key), + json=AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + ), + ) + + def send_responses(self, key: str, model: str, content: str) -> StreamingResponse: + return self.proxy.transport.send( + "/v1/responses", + headers=self.proxy.transport.bearer(key), + json=ResponsesBody(model=model, input=content), + ) + + def send_embed(self, key: str, model: str, content: str) -> StreamingResponse: + return self.proxy.transport.send( + "/embeddings", + headers=self.proxy.transport.bearer(key), + json=EmbedBody(model=model, input=content), + ) + + def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse: + return self.proxy.transport.send( + f"/gemini/v1beta/models/{model}:generateContent", + headers=GeminiApiKeyHeaders(x_goog_api_key=key), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=content)])], + generationConfig=GeminiGenerationConfig(maxOutputTokens=max_tokens), + ), + ) + + def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject: + return unwrap( + self.proxy.transport.upload( + "/v1/files", + headers=self.proxy.transport.bearer(key), + form=FileUploadForm(purpose="batch"), + filename="key_attribution.jsonl", + content=content, + params=ModelQuery(model=model), + response_type=FileObject, + ) + ) + + def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject: + return unwrap( + self.proxy.transport.post( + "/v1/batches", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=BatchObject, + ) + ) + + def retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject: + return unwrap( + self.proxy.transport.get( + f"/v1/batches/{batch_id}", + headers=self.proxy.transport.bearer(key), + params=ProviderQuery(provider=provider), + response_type=BatchObject, + ) + ) + + def replay_callback_log(self, key: str, payload: CallbackLogPayload) -> CallbackLogsResponse: + return unwrap( + self.proxy.transport.post( + "/v1/rust_control_plane/logs", + headers=self.proxy.transport.bearer(key), + json=CallbackLogsRequest(records=[CallbackLogRecord(standard_logging_payload=payload)]), + response_type=CallbackLogsResponse, + ) + ) + + def health(self, model: str) -> ProbeResult: + return self.proxy.transport.probe("/health", params=HealthParams(model=model)) + + def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None: + response: Final = unwrap( + self.proxy.transport.get( + "/user/daily/activity", + headers=self.proxy.transport.master, + params=DailyActivityParams( + start_date=start.strftime("%Y-%m-%d"), + end_date=end.strftime("%Y-%m-%d"), + api_key=token, + ), + response_type=DailyActivityResponse, + ) + ) + return next( + (row.breakdown.api_keys[token] for row in response.results if token in row.breakdown.api_keys), + None, + ) + + def poll_daily_activity_for_key( + self, token: str, *, start: datetime, end: datetime, min_requests: int + ) -> DailyActivityKeyBreakdown | None: + outcome: Final = await_converged( + lambda: self.daily_activity_for_key(token, start=start, end=end), + converged=lambda found: found is not None and found.metrics.api_requests >= min_requests, + timeout=self.proxy.poll_timeout, + interval=self.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + def openapi(self) -> OpenAPISchema: return unwrap( self.proxy.transport.get( diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py new file mode 100644 index 00000000000..4a2c23927c6 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -0,0 +1,405 @@ +"""Every spend row a live proxy writes joins its virtual key (MAT-180). + +One virtual key with an alias, owned by a user with an email, drives every spend +write path a key can reach: /chat/completions, /queue/chat/completions, +/v1/messages, /v1/responses, /embeddings, the Gemini native passthrough, a batch +input file upload, a batch create, and a replayed callback log (POST +/v1/rust_control_plane/logs, the writer an external gateway feeds). Each row those calls write must carry +`api_key` equal to the key's LiteLLM_VerificationToken.token (the sha256 hash +/key/generate returns as `token`), which is the join /spend/logs?api_key= and +/user/daily/activity rely on to report key_alias and user_email. A row keyed by a +re-hashed token (v1.99.0's regression, #39568 and #39572) shows up as a +key-hash-* row with no alias and no email in the customer's usage exports. + +The health-check service account writes rows too; those must stay keyed by the +literal service-account name, never by a hash of it. A batch's cost row is +written by the retrieve that first sees the batch in a terminal state, so the +batch the run creates is one OpenAI fails at validation within seconds (its one +line targets /v1/embeddings under a /v1/chat/completions batch), and the test +retrieves it by its raw provider id with the same key until it is failed. A raw +id is never owned by the CheckBatchCost poller, so that retrieve prices the batch +inline against the retrieving key and its {provider_batch_id}_batch_cost row +must join the key's token with its alias. A completed batch with a positive +cost is out of a single run's reach (OpenAI's completion window is 24h, and a +stack booted fresh per run lists no earlier run's batches), so the poller's own +row is not asserted here. + +/spend/logs carries no email field, so the email assertion lives on +/user/daily/activity alone; /spend/logs is held to the alias in metadata. +""" + +import base64 +import time +from collections.abc import Iterator +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from models import KeyGenerateBody +from proxy_client import Converged, await_converged +from pydantic import BaseModel +from spend_e2e_client import ( + BatchCreateBody, + BatchObject, + CallbackLogMetadata, + CallbackLogPayload, + DailyActivityKeyBreakdown, + ResponseIdentity, + SpendClient, + SpendLogRow, + StreamingResponse, + unique_marker, +) + +pytestmark = pytest.mark.e2e + +CHAT_MODEL: Final = "gemini-2.5-flash" +MESSAGES_MODEL: Final = "claude-haiku-4-5" +RESPONSES_MODEL: Final = "openai-responses-codex" +EMBED_MODEL: Final = "openai-text-embedding-3-small" +BATCH_MODEL: Final = "openai-gpt-4o-mini" +BATCH_BACKEND_MODEL: Final = "gpt-4o-mini" +BATCH_PROVIDER: Final = "openai" +HEALTH_SERVICE_ACCOUNT: Final = "litellm-internal-health-check" +BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) +FAILED_BATCH_POLL_SECONDS: Final = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS: Final = 5.0 +MAX_TOKENS: Final = 8 +REPLAY_RESPONSE_COST: Final = 0.0001 +REPLAY_PROMPT_TOKENS: Final = 5 +REPLAY_COMPLETION_TOKENS: Final = 1 +WRITE_PATHS: Final = ( + "chat_completions", + "queue_chat_completions", + "messages", + "responses", + "embeddings", + "gemini_passthrough", + "batch_file_upload", + "batch_create", + "callback_replay", +) + + +class EmbeddingLineBody(BaseModel): + model: str + input: str + + +class EmbeddingLine(BaseModel): + custom_id: str + method: str = "POST" + url: str = "/v1/embeddings" + body: EmbeddingLineBody + + +@dataclass(frozen=True, slots=True) +class AttributedKey: + key: str + token: str + alias: str + email: str + user_id: str + + +@dataclass(frozen=True, slots=True) +class WritePath: + name: str + request_id: str + + +@dataclass(frozen=True, slots=True) +class DrivenKey: + identity: AttributedKey + paths: tuple[WritePath, ...] + started_at: datetime + + +def _body_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + response_id: Final = ResponseIdentity.model_validate_json(sent.body).id + assert response_id, f"{name} answered without a response id: {sent.body[:300]}" + return WritePath(name=name, request_id=response_id) + + +def _call_id(name: str, sent: StreamingResponse) -> WritePath: + assert sent.ok, f"{name} failed with {sent.status_code}: {sent.body[:300]}" + assert sent.call_id, f"{name} answered without an x-litellm-call-id header" + return WritePath(name=name, request_id=sent.call_id) + + +def _endpoint_mismatched_jsonl(marker: str) -> bytes: + line: Final = EmbeddingLine(custom_id=marker, body=EmbeddingLineBody(model=BATCH_BACKEND_MODEL, input=marker)) + return f"{line.model_dump_json()}\n".encode() + + +def _drive_batch(client: SpendClient, identity: AttributedKey, marker: str) -> tuple[WritePath, WritePath]: + uploaded: Final = client.upload_batch_file(identity.key, BATCH_MODEL, _endpoint_mismatched_jsonl(marker)) + created: Final = client.create_batch( + identity.key, + BatchCreateBody( + input_file_id=uploaded.id, + model=BATCH_MODEL, + metadata={"run": marker}, + ), + ) + return ( + WritePath(name="batch_file_upload", request_id=uploaded.id), + WritePath(name="batch_create", request_id=created.id), + ) + + +def _drive_callback_replay(client: SpendClient, identity: AttributedKey, marker: str) -> WritePath: + request_id: Final = f"callback-replay-{marker}" + finished_at: Final = time.time() + replayed: Final = client.replay_callback_log( + identity.key, + CallbackLogPayload( + id=request_id, + litellm_call_id=request_id, + model=CHAT_MODEL, + start_time=finished_at - 1, + end_time=finished_at, + response_cost=REPLAY_RESPONSE_COST, + prompt_tokens=REPLAY_PROMPT_TOKENS, + completion_tokens=REPLAY_COMPLETION_TOKENS, + total_tokens=REPLAY_PROMPT_TOKENS + REPLAY_COMPLETION_TOKENS, + metadata=CallbackLogMetadata( + user_api_key_hash=identity.token, + user_api_key_alias=identity.alias, + user_api_key_user_id=identity.user_id, + ), + ), + ) + assert replayed.processed == 1 and replayed.failed == 0, f"callback replay rejected the payload: {replayed}" + return WritePath(name="callback_replay", request_id=request_id) + + +def _drive_every_write_path(client: SpendClient, identity: AttributedKey) -> tuple[WritePath, ...]: + marker: Final = unique_marker() + prompt: Final = f"Reply with the word ok. {marker}" + key: Final = identity.key + return ( + _body_id("chat_completions", client.send_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("queue_chat_completions", client.send_queued_chat(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("messages", client.send_messages(key, MESSAGES_MODEL, prompt, max_tokens=MAX_TOKENS)), + _body_id("responses", client.send_responses(key, RESPONSES_MODEL, prompt)), + _call_id("embeddings", client.send_embed(key, EMBED_MODEL, prompt)), + _call_id("gemini_passthrough", client.send_gemini_generate(key, CHAT_MODEL, prompt, max_tokens=MAX_TOKENS)), + *_drive_batch(client, identity, marker), + _drive_callback_replay(client, identity, marker), + ) + + +def _provider_batch_id(unified_batch_id: str) -> str: + encoded: Final = unified_batch_id.removeprefix("batch_") + decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() + return decoded.removeprefix("litellm:").split(";", 1)[0] + + +def _driven_batch_id(driven: DrivenKey) -> str: + return next(path.request_id for path in driven.paths if path.name == "batch_create") + + +def _await_terminal_batch(client: SpendClient, key: str, provider_batch_id: str) -> BatchObject: + outcome: Final = await_converged( + lambda: client.retrieve_batch(key, provider_batch_id, provider=BATCH_PROVIDER), + converged=lambda batch: batch.status in BATCH_TERMINAL_STATUSES, + timeout=FAILED_BATCH_POLL_SECONDS, + interval=FAILED_BATCH_POLL_INTERVAL_SECONDS, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + +def _health_rows_between(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + return [ + row + for row in client.proxy.spend_logs_window( + start=started_at - timedelta(minutes=1), end=datetime.now(timezone.utc) + timedelta(minutes=1) + ) + if HEALTH_SERVICE_ACCOUNT in (row.request_tags or []) + ] + + +def _health_rows_since(client: SpendClient, started_at: datetime) -> list[SpendLogRow]: + outcome: Final = await_converged( + lambda: _health_rows_between(client, started_at), + converged=lambda rows: bool(rows), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + +class TestKeyAttribution: + @pytest.fixture(scope="class") + def driven(self, client: SpendClient) -> Iterator[DrivenKey]: + marker: Final = unique_marker() + user_id: Final = client.create_user( + email=f"key-attribution-{marker}@example.com", + role="proxy_admin", + user_id=f"key-attribution-{marker}", + ) + record: Final = client.generate_key_record( + KeyGenerateBody(models=[], user_id=user_id, key_alias=f"key-attribution-{marker}") + ) + assert record.token, "/key/generate answered without the key's token hash" + assert record.key_alias, "/key/generate dropped the key alias" + identity: Final = AttributedKey( + key=record.key, + token=record.token, + alias=record.key_alias, + email=f"key-attribution-{marker}@example.com", + user_id=user_id, + ) + started_at: Final = datetime.now(timezone.utc) + try: + yield DrivenKey( + identity=identity, + paths=_drive_every_write_path(client, identity), + started_at=started_at, + ) + finally: + client.proxy.delete_key(identity.key) + client.delete_user(identity.user_id) + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.joins_key", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_every_write_path_row_joins_the_key(self, client: SpendClient, driven: DrivenKey) -> None: + assert tuple(path.name for path in driven.paths) == WRITE_PATHS + found: Final = tuple((path, client.proxy.poll_logs_for_request_id(path.request_id)) for path in driven.paths) + unwritten: Final = [path.name for path, rows in found if not rows] + assert not unwritten, f"write paths that produced no spend row within the poll window: {unwritten}" + unjoined: Final = [ + (path.name, row.call_type, row.api_key) + for path, rows in found + for row in rows + if row.api_key != driven.identity.token + ] + assert not unjoined, ( + "spend rows whose api_key does not join LiteLLM_VerificationToken.token " + f"{driven.identity.token}: {unjoined}" + ) + unaliased: Final = [ + (path.name, row.call_type, row.metadata.user_api_key_alias if row.metadata else None) + for path, rows in found + for row in rows + if row.metadata is None or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unaliased, f"spend rows written without key alias {driven.identity.alias!r}: {unaliased}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_spend_logs_by_key_return_every_row_with_the_alias(self, client: SpendClient, driven: DrivenKey) -> None: + expected_ids: Final = frozenset(path.request_id for path in driven.paths) + rows: Final = client.poll_logs_for_key( + driven.identity.key, + min_rows=len(driven.paths), + predicate=lambda found: expected_ids <= frozenset(row.request_id or "" for row in found), + ) + missing: Final = expected_ids - frozenset(row.request_id or "" for row in rows) + assert not missing, ( + f"/spend/logs?api_key= does not return {len(missing)} of {len(expected_ids)} rows for the key: " + f"{sorted(path.name for path in driven.paths if path.request_id in missing)}" + ) + aliases: Final = frozenset(row.metadata.user_api_key_alias if row.metadata else None for row in rows) + assert aliases == {driven.identity.alias}, f"/spend/logs rows carry aliases {sorted(map(str, aliases))}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.reports_alias_and_email", + exercised_on=[ + "chat_completions", + "messages", + "responses", + "embeddings", + "batches", + "files", + "google_native", + "rust_control_plane", + ], + ) + def test_user_daily_activity_reports_alias_and_email(self, client: SpendClient, driven: DrivenKey) -> None: + breakdown: Final[DailyActivityKeyBreakdown | None] = client.poll_daily_activity_for_key( + driven.identity.token, + start=driven.started_at - timedelta(days=1), + end=datetime.now(timezone.utc) + timedelta(days=1), + min_requests=len(driven.paths), + ) + assert breakdown is not None, ( + f"/user/daily/activity?api_key={driven.identity.token} has no api_keys breakdown: " + "the key's rows did not aggregate under its token" + ) + assert breakdown.metrics.api_requests >= len(driven.paths), ( + f"/user/daily/activity counts {breakdown.metrics.api_requests} requests for the key, " + f"expected at least {len(driven.paths)}" + ) + assert breakdown.metadata.key_alias == driven.identity.alias, f"key_alias={breakdown.metadata.key_alias!r}" + assert breakdown.metadata.user_email == driven.identity.email, f"user_email={breakdown.metadata.user_email!r}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.health_rows_keep_service_account", + exercised_on=["chat_completions"], + ) + def test_health_check_rows_keep_the_service_account_key(self, client: SpendClient) -> None: + started_at: Final = datetime.now(timezone.utc) + probe: Final = client.health(CHAT_MODEL) + assert probe.healthy, f"/health?model={CHAT_MODEL} answered {probe.status_code}: {probe.body[:300]}" + rows: Final = _health_rows_since(client, started_at) + assert rows, f"/health?model={CHAT_MODEL} wrote no {HEALTH_SERVICE_ACCOUNT}-tagged spend row" + rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT] + assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}" + + @pytest.mark.covers( + "quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key", + exercised_on=["batches"], + ) + def test_terminal_batch_cost_row_joins_the_retrieving_key(self, client: SpendClient, driven: DrivenKey) -> None: + provider_batch_id: Final = _provider_batch_id(_driven_batch_id(driven)) + fetched: Final = _await_terminal_batch(client, driven.identity.key, provider_batch_id) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch {provider_batch_id} is {fetched.status!r} after " + f"{FAILED_BATCH_POLL_SECONDS:.0f}s, so its terminal cost row cannot be asserted" + ) + cost_request_id: Final = f"{provider_batch_id}_batch_cost" + rows: Final = client.proxy.poll_logs_for_request_id(cost_request_id) + assert rows, f"retrieving failed batch {provider_batch_id} wrote no cost row under {cost_request_id}" + call_types: Final = tuple(sorted({row.call_type or "" for row in rows})) + assert call_types == ("aretrieve_batch",), f"cost rows under {cost_request_id} carry call types {call_types}" + unjoined: Final = [ + (row.call_type, row.api_key, row.metadata.user_api_key_alias if row.metadata else None) + for row in rows + if row.api_key != driven.identity.token + or row.metadata is None + or row.metadata.user_api_key_alias != driven.identity.alias + ] + assert not unjoined, ( + f"batch cost rows that do not join the retrieving key's token {driven.identity.token} " + f"with alias {driven.identity.alias!r}: {unjoined}" + ) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 3fab20a28ad..e5826b18668 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -2,6 +2,7 @@ import glob import os import re import sys +from pathlib import Path import pytest @@ -870,3 +871,69 @@ class TestMigrateDeployAttemptAccounting: harness.run() assert len(harness.deploy_calls) == 1 assert harness.resolved == [] + + +class TestJWTKeyMappingCascade: + """Regression tests for issue #33702. + + A virtual key referenced by a LiteLLM_JWTKeyMapping row could not be deleted + because LiteLLM_JWTKeyMapping_token_fkey was created ON DELETE RESTRICT, so + deleting the key (Admin UI, /key/delete, team delete, ...) raised a foreign + key violation. The mapping must be removed automatically when its key is + deleted, which the FK now enforces via ON DELETE CASCADE. + """ + + _FK_NAME = "LiteLLM_JWTKeyMapping_token_fkey" + + def _effective_on_delete(self): + """Replay every migration in order and return the last ON DELETE action + declared for the JWT key mapping FK.""" + action = None + for _migration_name, sql in _get_all_migrations(): + for match in re.finditer( + rf'ADD\s+CONSTRAINT\s+"{re.escape(self._FK_NAME)}".*?' + r"ON\s+DELETE\s+(CASCADE|RESTRICT|SET\s+NULL|NO\s+ACTION|SET\s+DEFAULT)", + sql, + re.IGNORECASE | re.DOTALL, + ): + action = re.sub(r"\s+", " ", match.group(1).upper()) + return action + + def test_fk_effective_on_delete_is_cascade(self): + """The final FK definition across all migrations must cascade deletes.""" + assert self._effective_on_delete() == "CASCADE", ( + f"{self._FK_NAME} must end up ON DELETE CASCADE so deleting a " + "virtual key removes its JWT key mapping (issue #33702)" + ) + + def test_schema_declares_cascade_on_relation(self): + """schema.prisma must declare onDelete: Cascade on the mapping relation + so the generated client and DB agree.""" + schema_paths = glob.glob( + os.path.abspath( + os.path.join( + os.path.dirname(__file__), "../../**/schema.prisma" + ) + ), + recursive=True, + ) + declaring = tuple( + (path, schema) + for path, schema in ((p, Path(p).read_text()) for p in schema_paths) + if "model LiteLLM_JWTKeyMapping" in schema + ) + assert declaring, "No schema.prisma declaring LiteLLM_JWTKeyMapping found" + for path, schema in declaring: + match = re.search( + r"litellm_verification_token\s+LiteLLM_VerificationToken\s+@relation\(([^)]*)\)", + schema, + ) + assert match is not None, ( + f"{path} declares LiteLLM_JWTKeyMapping but its verification token " + "relation could not be parsed, so this test cannot vouch for it " + "(issue #33702)" + ) + assert "onDelete: Cascade" in match.group(1), ( + f"{path} must declare onDelete: Cascade on the JWT key mapping " + "relation (issue #33702)" + ) diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index b6e30ddc711..31c554985a7 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1623,15 +1623,11 @@ class TestMissingChoicesGuard: assert "no 'choices'" in exc_info.value.message - def test_convert_to_model_response_object_empty_choices_raises_api_error(self): - """Empty choices list raises APIError, same as missing/null choices. + def test_convert_to_model_response_object_empty_choices_returns_empty_list(self): + """An empty choices list is a real provider answer, so it converts to choices=[] instead of raising. - Provider-specific repair (e.g. github_copilot synthesizing choices for - Anthropic-native responses) happens before this guard, in the provider - config; the core utility keeps treating empty choices as an error. + See: https://github.com/BerriAI/litellm/issues/40276 """ - from litellm.exceptions import APIError - response_object = { "id": "msg_123", "model": "some-model", @@ -1639,16 +1635,17 @@ class TestMissingChoicesGuard: "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, } - with pytest.raises(APIError) as exc_info: - convert_to_model_response_object( - response_object=response_object, - model_response_object=ModelResponse(), - ) + result = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) - assert "no 'choices'" in exc_info.value.message + assert isinstance(result, ModelResponse) + assert result.choices == [] + assert result.usage.prompt_tokens == 10 def test_convert_to_model_response_object_null_choices_raises_api_error(self): - """choices=None raises APIError.""" + """choices=None raises APIError that names the type instead of claiming the key is missing.""" from litellm.exceptions import APIError response_object = { @@ -1664,7 +1661,7 @@ class TestMissingChoicesGuard: model_response_object=ModelResponse(), ) - assert "no 'choices'" in exc_info.value.message + assert "'choices' that is not a list (NoneType)" in exc_info.value.message def test_convert_to_streaming_response_no_choices_raises_api_error(self): """Missing choices in streaming cache-hit path raises APIError.""" diff --git a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py index 89899d3e762..c4b1f4f3afd 100644 --- a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py +++ b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py @@ -23,9 +23,9 @@ from litellm.proxy.proxy_server import token_counter def _fake_hf_tokenizer(num_tokens: int) -> MagicMock: encoding = MagicMock() - encoding.ids = list(range(num_tokens)) + encoding.__len__.return_value = num_tokens tokenizer = MagicMock() - tokenizer.encode.return_value = encoding + tokenizer.encode_batch_fast.return_value = [encoding] return tokenizer @@ -68,13 +68,11 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch): ) ) - mock_tokenizer_cls.from_pretrained.assert_called_once_with( - "my-org/custom-tokenizer", revision="v2", auth_token=None - ) + mock_tokenizer_cls.from_pretrained.assert_called_once_with("my-org/custom-tokenizer", revision="v2", token=None) assert response.tokenizer_type == "huggingface_tokenizer" assert response.request_model == "my-embedding-model" assert response.model_used == "self-hosted-embedder" - assert response.total_tokens > 0 + assert response.total_tokens >= 7 @pytest.mark.asyncio diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 54cce9cdd78..d06eb0426c9 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2920,7 +2920,7 @@ async def test_get_config_callbacks_with_all_types(client_no_auth): assert result["status"] == "success" assert "callbacks" in result - callbacks = result["callbacks"] + callbacks = [cb for cb in result["callbacks"] if not cb.get("read_only", False)] # Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure) assert len(callbacks) == 5 diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 772d3622745..467c1332325 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -14,13 +14,14 @@ These tests ensure the polling handler correctly manages response state following the OpenAI Response API format. """ +import asyncio import json from datetime import datetime, timezone from typing import Any, Dict, Optional from unittest.mock import AsyncMock, Mock, patch import pytest - +from fastapi import Request from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -1414,7 +1415,7 @@ def _make_background_streaming_kwargs( polling_id=polling_id, data={"model": "gpt-4o", "stream": False, "background": True}, polling_handler=polling_handler, - request=Mock(), + request=Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}), fastapi_response=Mock(), user_api_key_dict=Mock(), general_settings={}, @@ -1663,6 +1664,63 @@ class TestBackgroundStreamingTerminalEvents: final_call = handler.update_state.call_args_list[-1] assert final_call.kwargs["status"] == "completed" + @pytest.mark.asyncio + async def test_polling_client_disconnect_does_not_cancel_upstream_call(self): + """The polling client hangs up right after getting its polling id. The detached task + must still stream the upstream response through the client-disconnect guards.""" + from litellm.proxy.common_request_processing import create_response + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + async def client_already_left(): + return {"type": "http.disconnect"} + + async def slow_upstream_stream(): + await asyncio.sleep(0.05) + for event in ( + {"type": "response.in_progress"}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "status": "completed", + "usage": {"input_tokens": 13, "output_tokens": 10}, + "model": "gpt-4o", + "output": [{"id": "item_1", "type": "message"}], + }, + }, + ): + yield f"data: {json.dumps(event)}\n\n" + + async def upstream_call_behind_disconnect_guard(**kwargs): + return await create_response( + slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"] + ) + + handler = AsyncMock(spec=ResponsePollingHandler) + kwargs = _make_background_streaming_kwargs("poll_7", handler) + kwargs["request"] = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"x-litellm-call-id", b"call-123")], + "query_string": b"", + }, + client_already_left, + ) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = upstream_call_behind_disconnect_guard + await background_streaming_task(**kwargs) + + final_call = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "completed" + assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10} + class TestEdgeCases: """Test edge cases and error scenarios""" diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 4db131da62c..713330a8280 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1,25 +1,33 @@ import asyncio import base64 +import json import os import sys +from collections.abc import AsyncIterator from importlib import metadata from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import anyio import httpx import pytest +import respx from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth from mcp import McpError +from mcp.client.streamable_http import streamable_http_client +from pydantic import ValidationError from mcp.shared.message import SessionMessage from mcp.types import ( LATEST_PROTOCOL_VERSION, + CallToolResult, ErrorData, Implementation, InitializeResult, JSONRPCError, JSONRPCMessage, JSONRPCResponse, + LoggingMessageNotificationParams, ServerCapabilities, ) @@ -29,8 +37,9 @@ import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import ( MCP_STREAMABLE_HTTP_REQUIREMENT, MCPClient, - _as_read_timeout, _first_non_cancelled_cause, + _TransportContext, + as_mcp_read_timeout, missing_streamable_http_client_error, strip_auth_scheme, ) @@ -859,25 +868,25 @@ def _raise_mcp_error_while_handling_a_timeout(code: int, message: str) -> McpErr return raised -def test_as_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): +def test_as_mcp_read_timeout_separates_the_sdk_timeout_from_a_relayed_upstream_error(): """Neither signal alone is enough. The code alone cannot separate the SDK's own timeout from an upstream JSON-RPC error that happens to use 408, and the context chain alone cannot separate it from any other relayed error that surfaces while a timeout is being handled, so both must hold. """ timeout_code = int(httpx.codes.REQUEST_TIMEOUT) - translated = _as_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) + translated = as_mcp_read_timeout(_raise_mcp_error_while_handling_a_timeout(timeout_code, "Timed out while waiting")) assert isinstance(translated, TimeoutError) assert str(translated) == "Timed out while waiting" relayed_408 = McpError(ErrorData(code=timeout_code, message="upstream said 408")) - assert _as_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" + assert as_mcp_read_timeout(relayed_408) is None, "an upstream 408 with no elapsed timeout is not our timeout" relayed_other = _raise_mcp_error_while_handling_a_timeout(-32603, "upstream internal error") - assert _as_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" + assert as_mcp_read_timeout(relayed_other) is None, "a non-timeout code is not our timeout, whatever the chain" - assert _as_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None - assert _as_read_timeout(RuntimeError("not an McpError")) is None + assert as_mcp_read_timeout(McpError(ErrorData(code=-32603, message="boom"))) is None + assert as_mcp_read_timeout(RuntimeError("not an McpError")) is None @pytest.mark.asyncio @@ -1224,14 +1233,14 @@ def test_without_a_configured_slot_the_existing_precedence_is_unchanged(): _REDIRECT_CASES = [ - ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin + ("https://upstream.example.com/mcp", "https://upstream.example.com/other"), # same origin ("https://upstream.example.com/mcp", "https://upstream.example.com:443/other"), # explicit default port - ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host - ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade - ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port - ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host - ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade - ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http + ("https://upstream.example.com/mcp", "https://attacker.example.com/collect"), # different host + ("https://upstream.example.com/mcp", "http://upstream.example.com/collect"), # scheme downgrade + ("https://upstream.example.com/mcp", "https://upstream.example.com:8443/other"), # different port + ("https://upstream.example.com/mcp", "https://sub.upstream.example.com/x"), # different host + ("http://upstream.example.com/mcp", "https://upstream.example.com/other"), # http -> https upgrade + ("http://upstream.example.com/mcp", "http://upstream.example.com/other"), # same origin, plain http ] @@ -1283,3 +1292,606 @@ def test_a_differently_cased_injected_header_cannot_shadow_the_slot() -> None: headers = client._get_auth_headers() assert [v for k, v in headers.items() if k.lower() == "esb-oauth"] == ["Bearer minted-token"] assert headers["X-Trace"] == "keep" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("content_type", "body", "expected_type"), + [ + ("text/html", b"secret-page", ValueError), + ("application/json", b"secret-invalid-json", ValidationError), + ("application/json", b"", ValidationError), + ("application/json", b'{"secret":"invalid-rpc"}', ValidationError), + ("application/json", b'{"jsonrpc":"2.0","id":0,"result":{"secret":"invalid-schema"}}', ValidationError), + ], +) +async def test_invalid_http_response_surfaces_without_waiting_for_timeout( + content_type: str, body: bytes, expected_type: type[Exception] +) -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": content_type}, content=body) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(expected_type) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "unsupported content type" in message or "invalid MCP response" in message + assert "secret" not in message + assert "timed out" not in message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [200, 401, 503]) +async def test_http_response_handler_preserves_success_and_http_errors(status_code: int) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": []} + ) + return httpx.Response(status_code, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + operation: Final = client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ) + if status_code == 200: + result: Final = await asyncio.wait_for(operation, timeout=3) + assert result.tools == [] + else: + with pytest.raises(httpx.HTTPStatusError) as caught: + await asyncio.wait_for(operation, timeout=3) + assert caught.value.response.status_code == status_code + + +@pytest.mark.asyncio +async def test_http_response_handler_preserves_notifications_and_tool_listing() -> None: + notification: Final = { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + } + logging_callback: Final = AsyncMock() + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + if payload["method"] == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload["id"], + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"logging": {}, "tools": {}}, + "serverInfo": {"name": "test", "version": "1"}, + }, + }, + ) + response: Final = { + "jsonrpc": "2.0", + "id": payload["id"], + "result": {"tools": [{"name": "search", "inputSchema": {"type": "object"}}]}, + } + return httpx.Response( + 200, + headers={"Content-Type": "text/event-stream"}, + content="".join(f"event: message\ndata: {json.dumps(message)}\n\n" for message in (notification, response)), + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30, logging_callback=logging_callback) + result: Final = await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), lambda session: session.list_tools() + ), + timeout=3, + ) + + assert [tool.name for tool in result.tools] == ["search"] + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + + +@pytest.mark.asyncio +async def test_invalid_tool_list_schema_is_identified_as_an_upstream_response() -> None: + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": "secret-invalid-tools"} + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(ValidationError) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + message: Final = _connection_error_message(caught.value, client.server_url, 30) + assert "invalid MCP response" in message + assert "secret" not in message + + +class _DiagnosticSSEStream(httpx.AsyncByteStream): + def __init__(self, messages: asyncio.Queue[bytes | Exception | None]) -> None: + self.messages = messages + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"event: endpoint\ndata: /messages\n\n" + while True: + message: Final = await self.messages.get() + if message is None: + return + if isinstance(message, Exception): + raise message + yield b"event: message\ndata: " + message + b"\n\n" + + +_DIAGNOSTIC_STDIO_SERVER: Final = """ +import json, sys +mode, failure_method = sys.argv[1:] +for line in sys.stdin: + request = json.loads(line) + if "method" not in request or "id" not in request: + continue + if request["method"] == failure_method: + if mode == "bad-json": + print("secret-invalid-json", flush=True) + continue + if mode == "closed": + sys.exit(0) + if mode == "silent": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Waiting"}}), flush=True) + continue + if request["method"] == "initialize": + result = {"protocolVersion": request["params"]["protocolVersion"], "capabilities": {"tools": {}, "logging": {}}, "serverInfo": {"name": "diagnostic", "version": "1"}} + elif request["method"] == "tools/list": + print(json.dumps({"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info", "data": "Listing tools"}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "unmatched", "result": {}}), flush=True) + print(json.dumps({"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}), flush=True) + result = {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + else: + result = {"content": [{"type": "text", "text": "pong"}], "isError": False} + print(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}), flush=True) +""" + + +def _diagnostic_transport(transport: MCPTransport, mode: str, failure_method: str) -> _TransportContext: + from mcp import StdioServerParameters + from mcp.client.sse import sse_client + from mcp.client.stdio import stdio_client + + if transport == MCPTransport.stdio: + return stdio_client( + StdioServerParameters( + command=sys.executable, args=["-u", "-c", _DIAGNOSTIC_STDIO_SERVER, mode, failure_method] + ) + ) + messages: Final[asyncio.Queue[bytes | Exception | None]] = asyncio.Queue() + + async def respond(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, headers={"Content-Type": "text/event-stream"}, stream=_DiagnosticSSEStream(messages) + ) + payload: Final = json.loads(request.content) + if "method" not in payload or "id" not in payload: + return httpx.Response(202) + if payload["method"] == failure_method and mode != "ok": + if mode == "bad-json": + await messages.put(b"secret-invalid-json") + elif mode == "io-error": + await messages.put(httpx.ReadError("secret-read-error")) + elif mode == "closed": + await messages.put(None) + elif mode == "silent": + await messages.put( + b'{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Waiting"}}' + ) + return httpx.Response(202) + if payload["method"] == "tools/list": + for message in ( + { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": {"level": "info", "data": "Listing tools"}, + }, + {"jsonrpc": "2.0", "id": "unmatched", "result": {}}, + {"jsonrpc": "2.0", "id": "server-ping", "method": "ping"}, + ): + await messages.put(json.dumps(message).encode()) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"tools": {}, "logging": {}}, + "serverInfo": {"name": "diagnostic", "version": "1"}, + } + if payload["method"] == "initialize" + else {"tools": [{"name": "ping", "inputSchema": {"type": "object"}}]} + if payload["method"] == "tools/list" + else {"content": [{"type": "text", "text": "pong"}], "isError": False} + ) + await messages.put(json.dumps({"jsonrpc": "2.0", "id": payload["id"], "result": result}).encode()) + return httpx.Response(202) + + def factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond), headers=headers, timeout=timeout, auth=auth) + + return sse_client("https://example.com/sse", httpx_client_factory=factory) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("failure_method", ["initialize", "tools/list"]) +async def test_transport_parsing_failure_is_preserved(transport: MCPTransport, failure_method: str) -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=transport, timeout=0.2) + with pytest.raises(ValidationError): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(transport, "bad-json", failure_method), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +async def test_sse_read_failure_is_preserved() -> None: + client: Final = MCPClient(server_url="https://example.com/sse", transport_type=MCPTransport.sse, timeout=0.2) + with pytest.raises(httpx.ReadError, match="secret-read-error"): + await asyncio.wait_for( + client._execute_session_operation( + _diagnostic_transport(MCPTransport.sse, "io-error", "tools/list"), lambda session: session.list_tools() + ), + timeout=3, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +@pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) +async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str) -> None: + from mcp import ClientSession + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message + + logging_callback: Final = AsyncMock() + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=0.2, logging_callback=logging_callback + ) + + async def operation(session: ClientSession) -> CallToolResult: + tools: Final = await session.list_tools() + assert [tool.name for tool in tools.tools] == ["ping"] + return await session.call_tool("ping", {}) + + pending: Final = client._execute_session_operation(_diagnostic_transport(transport, mode, "tools/list"), operation) + if mode == "ok": + result: Final = await asyncio.wait_for(pending, timeout=3) + assert result.isError is False + assert result.content[0].text == "pong" + logging_callback.assert_awaited_once_with(LoggingMessageNotificationParams(level="info", data="Listing tools")) + else: + with pytest.raises(McpError) as caught: + await asyncio.wait_for(pending, timeout=3) + if mode == "closed": + assert "connection was closed" in _connection_error_message(caught.value, client.server_url, 0.2) + else: + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) +async def test_transport_cancellation_cleans_up_a_pending_request(transport: MCPTransport) -> None: + ready: Final = asyncio.Event() + + async def on_log(message: LoggingMessageNotificationParams) -> None: + if message.data == "Waiting": + ready.set() + + client: Final = MCPClient( + server_url="https://example.com/sse", transport_type=transport, timeout=30, logging_callback=on_log + ) + task: Final = asyncio.create_task( + client._execute_session_operation( + _diagnostic_transport(transport, "silent", "tools/list"), lambda session: session.list_tools() + ) + ) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + + +class _InterruptedHTTPBody(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b'{"jsonrpc":' + raise httpx.RemoteProtocolError("secret-incomplete-response") + + +@pytest.mark.asyncio +async def test_interrupted_http_response_preserves_the_transport_failure() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "application/json"}, stream=_InterruptedHTTPBody()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=30) + with pytest.raises(httpx.RemoteProtocolError, match="secret-incomplete-response"): + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + + +@pytest.mark.asyncio +async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=b"") + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + client: Final = MCPClient(server_url="https://example.com/mcp", timeout=0.2) + with pytest.raises(McpError) as caught: + await asyncio.wait_for( + client._execute_session_operation( + streamable_http_client(client.server_url, http_client=http_client), + lambda session: session.list_tools(), + ), + timeout=3, + ) + assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize( + "outcome", + ( + "absent", + "other_capability", + "supported", + "method_not_found", + "internal_error", + "unauthorized", + "timeout", + "initialize_not_found", + ), +) +async def test_optional_discovery_capabilities_and_errors( + method: str, outcome: str, caplog: pytest.LogCaptureFixture +) -> None: + import logging + from unittest.mock import Mock + + from mcp.types import JSONRPCRequest + + capability: Final = "prompts" if method == "prompts/list" else "resources" + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + advertised: Final = "resources" if capability == "prompts" else "prompts" + entry: Final = { + "prompts/list": {"name": "example"}, + "resources/list": {"name": "example", "uri": "test://example"}, + "resources/templates/list": {"name": "example", "uriTemplate": "test://{name}"}, + }[method] + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + if outcome == "initialize_not_found": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": {"code": -32601, "message": "Initialization rejected"}, + }, + ) + if payload.method == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {} + if outcome == "absent" + else {advertised if outcome == "other_capability" else capability: {}}, + "serverInfo": {"name": "discovery", "version": "1"}, + }, + }, + ) + if outcome == "timeout": + raise httpx.ReadTimeout("Optional list timed out", request=request) + if outcome == "unauthorized": + return httpx.Response(401) + if outcome in ("method_not_found", "internal_error", "absent", "other_capability"): + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": { + "code": -32603 if outcome == "internal_error" else -32601, + "message": "Optional list rejected", + }, + }, + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry]}}) + + responder: Final = Mock(side_effect=respond) + caplog.set_level(logging.DEBUG, logger="LiteLLM") + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=responder) + client: Final = MCPClient(server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + result: Final = await operation() + + requests: Final = tuple( + JSONRPCMessage.model_validate_json(call.args[0].content).root + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == method for request in requests) == ( + 0 if outcome in ("absent", "other_capability", "initialize_not_found") else 1 + ) + assert [item.name for item in result] == (["example"] if outcome == "supported" else []) + failures: Final = tuple( + record for record in caplog.records if record.name == "LiteLLM" and record.levelno >= logging.WARNING + ) + if outcome in ("internal_error", "unauthorized", "timeout", "initialize_not_found"): + assert any(record.levelno == logging.ERROR and "failed" in record.message for record in failures) + else: + assert failures == () + if outcome == "method_not_found": + assert any( + record.levelno == logging.DEBUG and "Optional list rejected" in record.message for record in caplog.records + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("supports_first", (True, False)) +async def test_optional_discovery_uses_each_sessions_capabilities(supports_first: bool) -> None: + from unittest.mock import Mock + from mcp.types import JSONRPCRequest + + capabilities: Final = iter(({"resources": {}}, {}) if supports_first else ({}, {"resources": {}})) + + def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + result: Final = ( + { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": next(capabilities), + "serverInfo": {"name": "changing", "version": "1"}, + } + if payload.method == "initialize" + else {"resources": [{"name": "example", "uri": "test://example"}]} + ) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + responder: Final = Mock(side_effect=respond) + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=responder) + client: Final = MCPClient(server_url="https://example.com/mcp") + first: Final = await client.list_resources() + second: Final = await client.list_resources() + + assert [item.name for item in first] == (["example"] if supports_first else []) + assert [item.name for item in second] == ([] if supports_first else ["example"]) + requests: Final = tuple( + JSONRPCMessage.model_validate_json(call.args[0].content).root + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == "resources/list" for request in requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +async def test_optional_discovery_preserves_cancellation(method: str) -> None: + from mcp.types import JSONRPCRequest + + ready: Final = asyncio.Event() + pending: Final = asyncio.Event() + + async def respond(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + payload: Final = JSONRPCMessage.model_validate_json(request.content).root + if not isinstance(payload, JSONRPCRequest): + return httpx.Response(202) + if payload.method == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_PROTOCOL_VERSION, + "capabilities": {"resources": {}, "prompts": {}}, + "serverInfo": {"name": "pending", "version": "1"}, + }, + }, + ) + ready.set() + await pending.wait() + return httpx.Response(202) + + with respx.mock(base_url="https://example.com") as router: + router.route().mock(side_effect=respond) + client: Final = MCPClient(server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + task: Final = asyncio.create_task(operation()) + try: + await asyncio.wait_for(ready.wait(), timeout=3) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index cd8d609cf71..ddc8439a83a 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1842,12 +1842,14 @@ class _ApplyStyleGuardrail(CustomGuardrail): self.block = block self.apply_called = False self.seen_texts = None + self.seen_request_data = None async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): from fastapi import HTTPException self.apply_called = True self.seen_texts = inputs.get("texts") + self.seen_request_data = request_data if self.block: raise HTTPException(status_code=400, detail={"error": "Violated moderation policy"}) return inputs @@ -2646,6 +2648,91 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: which starved every later callback in litellm.callbacks (notably the lazily-appended VectorStorePreCallHook that attaches provider_specific_fields["search_results"]).""" + @pytest.mark.asyncio + async def test_apply_guardrail_retains_request_identity(self) -> None: + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail: Final = _ApplyStyleGuardrail(block=False) + guardrail.event_hook = GuardrailEventHooks.post_call + request_data: Final = {"guardrails": ["apply-style-guardrail"]} + response: Final = ModelResponse(choices=[Choices(message=Message(content="review me"))]) + + await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=response, call_type=CallTypes.acompletion + ) + + assert guardrail.seen_request_data is request_data + assert guardrail.seen_texts == ["review me"] + assert "guardrail_to_apply" not in request_data + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", (None, CallTypes.acompletion)) + async def test_apply_guardrail_masks_response_and_records_metadata(self, call_type: CallTypes | None) -> None: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks + from litellm.types.utils import Choices, Message, ModelResponse + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="response-filter", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="secret", action=ContentFilterAction.MASK)], + ) + request_data: Final = {"guardrails": ["response-filter"]} + response: Final = ModelResponse(choices=[Choices(message=Message(content="a secret"))]) + + result: Final = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=response, call_type=call_type + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == f"a {guardrail.keyword_redaction_tag}" + entries: Final = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "response-filter" + assert entries[0]["guardrail_mode"] == "post_call" + assert "guardrail_to_apply" not in request_data + + @pytest.mark.asyncio + @pytest.mark.parametrize("error_type", (None, RuntimeError, asyncio.CancelledError)) + async def test_dispatch_cleans_up_request_on_every_exit(self, error_type: type[BaseException] | None) -> None: + from contextlib import nullcontext + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import LLMResponseTypes, ModelResponse + + error: Final = error_type("dispatch interrupted") if error_type is not None else None + + class Dispatch(CustomLogger): + request_data: dict[str, object] | None = None + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> LLMResponseTypes: + self.request_data = data + if error is not None: + raise error + return response + + dispatch: Final = Dispatch() + + class Guardrail(_ApplyStyleGuardrail): + def _deployment_hook_target(self) -> CustomLogger: + return dispatch + + guardrail: Final = Guardrail(block=False) + guardrail.event_hook = GuardrailEventHooks.post_call + request_data: Final = {"guardrails": ["apply-style-guardrail"]} + with pytest.raises(error_type) if error_type is not None else nullcontext(): + await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, response=ModelResponse(), call_type=CallTypes.acompletion + ) + + assert dispatch.request_data is request_data + assert "guardrail_to_apply" not in request_data + @pytest.mark.asyncio async def test_returns_none_when_request_has_no_guardrails(self): from litellm.types.utils import ModelResponse @@ -2740,4 +2827,5 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: assert result is response assert response.choices[0].message.content == "filtered response" - assert request_data == {"guardrails": ["test-guardrail"]} + assert "guardrail_to_apply" not in request_data + assert len(_guardrail_entries(request_data)) == 1 diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index d36878e455f..87e76499b84 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1341,6 +1341,257 @@ def _emit(logger: LangFuseLogger, *, metadata=None, headers=None): ) +@pytest.mark.parametrize("level", ["DEFAULT", "ERROR"]) +@pytest.mark.parametrize( + "headers,metadata,expected_id", + [ + ({"x-litellm-session-id": "session-7125"}, {}, "call"), + ({"X-Claude-Code-Session-Id": "session-7125"}, {}, "call"), + ({"x-session-id": "session-7125"}, {}, "call"), + ({"session-id": "session-7125", "user-agent": "codex_cli_rs/1.0"}, {}, "call"), + ({"thread-id": "session-7125", "user-agent": "codex-tui"}, {}, "call"), + ({"session_id": "session-7125", "user-agent": "Codex 1.0"}, {}, "call"), + ({"conversation_id": "session-7125", "user-agent": "codex_vscode/1.0"}, {}, "call"), + ({"x-litellm-session-id": "short"}, {}, "call"), + ({"x-litellm-trace-id": "session-7125"}, {}, "session-7125"), + ( + {"X-LiteLLM-Trace-Id": "session-7125", "x-litellm-session-id": "session-7125"}, + {}, + "session-7125", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_trace_id": "session-7125"}, + {}, + "session-7125", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_trace_id": "explicit-trace"}, + {}, + "explicit-trace", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_existing_trace_id": "existing-trace"}, + {}, + "existing-trace", + ), + ( + {"x-litellm-session-id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"x-litellm-session-id": "short", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"X-Claude-Code-Session-Id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + {"x-session-id": "session-7125", "langfuse_session_id": "custom-session"}, + {}, + "call", + ), + ( + { + "session-id": "session-7125", + "user-agent": "codex_cli_rs/1.0", + "langfuse_session_id": "custom-session", + }, + {}, + "call", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "x-litellm-trace-id": "explicit-trace", + }, + {}, + "explicit-trace", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "langfuse_trace_id": "explicit-trace", + }, + {}, + "explicit-trace", + ), + ( + { + "x-litellm-session-id": "session-7125", + "langfuse_session_id": "custom-session", + "langfuse_existing_trace_id": "existing-trace", + }, + {}, + "existing-trace", + ), + ({}, {"trace_id": "session-7125", "session_id": "session-7125"}, "session-7125"), + ({}, {"trace_id": "explicit-trace", "session_id": "session-7125"}, "explicit-trace"), + ( + {"x-vendor-session-id": "short"}, + {"trace_id": "short", "session_id": "short"}, + "short", + ), + ( + {"x-session-id": "invalid value"}, + {"trace_id": "invalid value", "session_id": "invalid value"}, + "invalid value", + ), + ( + {"session-id": "session-7125", "user-agent": "codexfoo/1.0"}, + {"trace_id": "session-7125", "session_id": "session-7125"}, + "session-7125", + ), + ( + {"x-vendor-session-id": "short"}, + {"trace_id": "session-7125", "session_id": "session-7125"}, + "session-7125", + ), + ({}, {}, "call"), + ], +) +def test_session_header_trace_provenance(headers, metadata, expected_id, level): + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + clean_headers, + redact_credential_headers, + ) + + logger: Final = _steering_logger() + for turn in range(2): + call_id = f"call-{turn}" + request_headers = Headers(headers) + data = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=request_headers, data={"metadata": dict(metadata)}, _metadata_variable_name="metadata" + ) + original_metadata = dict(data["metadata"]) + now = datetime.datetime.now() + result = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": call_id, + "litellm_trace_id": data.get("litellm_trace_id"), + "litellm_params": { + "metadata": data["metadata"], + "proxy_server_request": {"headers": redact_credential_headers(clean_headers(request_headers))}, + }, + "messages": [{"role": "user", "content": f"turn {turn}"}], + "optional_params": {}, + }, + response_obj=( + None + if level == "ERROR" + else litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]) + ), + start_time=now, + end_time=now, + level=level, + status_message="provider error" if level == "ERROR" else None, + ) + trace_params = logger.Langfuse.trace.call_args.kwargs + assert trace_params["id"] == (call_id if expected_id == "call" else expected_id) + assert result["trace_id"] == trace_params["id"] + if expected_id != "existing-trace": + assert trace_params["session_id"] == headers.get("langfuse_session_id", original_metadata.get("session_id")) + steering = {key[len("langfuse_") :]: value for key, value in headers.items() if key.startswith("langfuse_")} + assert data["metadata"] == {**original_metadata, **steering} + + +def test_session_header_trace_without_call_id_keeps_session_alias(): + logger: Final = _steering_logger() + now: Final = datetime.datetime.now() + + result: Final = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": "", + "litellm_params": { + "metadata": {"trace_id": "session-7125", "session_id": "session-7125"}, + "proxy_server_request": {"headers": {"x-litellm-session-id": "session-7125"}}, + }, + "messages": [{"role": "user", "content": "no call id"}], + "optional_params": {}, + }, + response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]), + start_time=now, + end_time=now, + ) + + assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125" + assert result["trace_id"] == "session-7125" + + +def test_every_proxy_session_header_shape_is_classified_as_a_session_alias(): + """The classifier must cover every header shape the proxy turns into a chain id.""" + from litellm.integrations.langfuse.langfuse import _is_session_header_trace + from litellm.proxy.litellm_pre_call_utils import ( + _CODEX_SESSION_ID_HEADERS, + get_chain_id_from_headers, + ) + + session: Final = "session-7125-abcdef" + session_shapes: Final = ( + {"x-litellm-session-id": session}, + {"X-Claude-Code-Session-Id": session}, + {"x-session-id": session}, + *({header: session, "user-agent": "codex_cli_rs/1.0"} for header in _CODEX_SESSION_ID_HEADERS), + ) + for headers in session_shapes: + assert get_chain_id_from_headers(dict(headers)) == session, headers + assert _is_session_header_trace(session, session, {"headers": headers}) is True, headers + + explicit_trace: Final = {"x-litellm-trace-id": session, "x-litellm-session-id": session} + assert get_chain_id_from_headers(dict(explicit_trace)) == session + assert _is_session_header_trace(session, session, {"headers": explicit_trace}) is False + + +@pytest.mark.parametrize( + "proxy_server_request", + [None, {}, {"headers": None}], + ids=["no-proxy-request", "no-headers-key", "null-headers"], +) +def test_sdk_caller_without_request_headers_keeps_its_trace(proxy_server_request): + """A direct SDK caller has no request headers, so a session-shaped trace id stays the caller's.""" + logger: Final = _steering_logger() + now: Final = datetime.datetime.now() + + result: Final = logger.log_event_on_langfuse( + kwargs={ + "call_type": "completion", + "litellm_call_id": "call-0", + "litellm_params": { + "metadata": {"trace_id": "session-7125", "session_id": "session-7125"}, + "proxy_server_request": proxy_server_request, + }, + "messages": [{"role": "user", "content": "sdk turn"}], + "optional_params": {}, + }, + response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]), + start_time=now, + end_time=now, + ) + + assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125" + assert result["trace_id"] == "session-7125" + + +def test_session_header_classifier_survives_non_string_header_keys(): + """A non-string header key must not cost the caller its whole trace.""" + from litellm.integrations.langfuse.langfuse import _is_session_header_trace + + session: Final = "session-7125-abcdef" + headers: Final = {7: "numeric key", "x-litellm-session-id": session} + assert _is_session_header_trace(session, session, {"headers": headers}) is True + assert _is_session_header_trace(session, session, {"headers": {7: "numeric key"}}) is False + + def test_mask_input_header_false_keeps_the_prompt(): logger = _steering_logger() diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py index ea661d2ea78..004ac4dbffb 100644 --- a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -84,6 +84,7 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent(): "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None ): logger = PrometheusLogger() + logger._emit_input_sequence_length_label = False logger.litellm_proxy_total_requests_metric = MagicMock() logger.get_labels_for_metric = MagicMock( return_value=["client_ip", "user_agent"] diff --git a/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py new file mode 100644 index 00000000000..bc922061544 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py @@ -0,0 +1,428 @@ +import asyncio +import datetime +from collections.abc import Mapping +from copy import deepcopy +from typing import Final, cast + +import pytest +from prometheus_client import REGISTRY +from prometheus_client.samples import Sample + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, + UserAPIKeyLabelValues, + get_input_sequence_length_bucket, +) +from litellm.types.utils import StandardLoggingPayload + +LATENCY_METRICS: Final = ( + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", +) +FLAG: Final = "prometheus_emit_input_sequence_length_label" + + +def _clear_prometheus_registry() -> None: + for collector in tuple(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] # test registry reset + REGISTRY.unregister(collector) + + +@pytest.fixture(autouse=True) +def isolated_registry(monkeypatch: pytest.MonkeyPatch): + _clear_prometheus_registry() + monkeypatch.setattr(litellm, FLAG, False) + yield + _clear_prometheus_registry() + + +@pytest.mark.parametrize("metric", LATENCY_METRICS) +def test_input_sequence_length_label_is_opt_in(monkeypatch: pytest.MonkeyPatch, metric: str): + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in PrometheusMetricLabels.get_labels(metric) + + monkeypatch.setattr(litellm, FLAG, True) + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value in PrometheusMetricLabels.get_labels(metric) + + +def test_input_sequence_length_label_stays_off_non_latency_metrics(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, FLAG, True) + assert UserAPIKeyLabelNames.INPUT_SEQUENCE_LENGTH.value not in PrometheusMetricLabels.get_labels( + "litellm_proxy_total_requests_metric" + ) + + +@pytest.mark.parametrize( + "prompt_tokens, expected", + [ + (None, "unknown"), + (0, "0-1k"), + (999, "0-1k"), + (1_000, "1k-4k"), + (3_999, "1k-4k"), + (4_000, "4k-16k"), + (15_999, "4k-16k"), + (16_000, "16k-64k"), + (63_999, "16k-64k"), + (64_000, "64k+"), + (10_000_000, "64k+"), + (-1, "unknown"), + ], +) +def test_input_sequence_length_bucket_boundaries(prompt_tokens: int | None, expected: str): + assert get_input_sequence_length_bucket(prompt_tokens) == expected + + +def test_user_api_key_label_values_carries_input_sequence_length(): + values: Final = UserAPIKeyLabelValues(input_sequence_length="4k-16k") + + assert values.input_sequence_length == "4k-16k" + assert values.model_dump()["input_sequence_length"] == "4k-16k" + + +def _assert_latency_metrics(expected: str | None, stream: bool = True, queue_time: float = 0) -> None: + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + for metric, duration in zip(LATENCY_METRICS, (2, 1, 3 + queue_time)): + counts: Final = tuple(sample for sample in samples if sample.name == f"{metric}_count") + sums: Final = tuple(sample for sample in samples if sample.name == f"{metric}_sum") + buckets: Final = tuple(sample for sample in samples if sample.name == f"{metric}_bucket") + if not stream and metric == "litellm_llm_api_time_to_first_token_metric": + assert not counts and not sums and not buckets + continue + assert len(counts) == len(sums) == 1 + assert counts[0].value == 1 + assert sums[0].value == pytest.approx(duration) + assert buckets and any(sample.labels["le"] == "+Inf" for sample in buckets) + assert all(sample.value == int(float(sample.labels["le"]) >= duration) for sample in buckets) + assert all(sample.labels.get("input_sequence_length") == expected for sample in (*counts, *sums, *buckets)) + + +def _non_target_samples() -> tuple[Sample, ...]: + return tuple( + sample + for metric in REGISTRY.collect() + if metric.name not in LATENCY_METRICS + for sample in metric.samples + if "input_sequence_length" in sample.labels and not sample.name.endswith("_created") + ) + + +def _standard_logging_payload(now: datetime.datetime, prompt_tokens: int) -> StandardLoggingPayload: + return cast( + StandardLoggingPayload, + { + "id": "t", + "call_type": "completion", + "response_cost": 0.001, + "status": "success", + "total_tokens": prompt_tokens + 20, + "prompt_tokens": prompt_tokens, + "completion_tokens": 20, + "startTime": now - datetime.timedelta(seconds=3), + "endTime": now, + "completionStartTime": now - datetime.timedelta(seconds=1), + "model": "gpt-4o-mini", + "model_id": "model-123", + "model_group": "gpt-4o-mini", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "request_tags": [], + "stream": True, + "metadata": { + "user_api_key_hash": "h", + "user_api_key_alias": "a", + "user_api_key_team_id": "t", + "user_api_key_team_alias": "ta", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": None, + }, + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + }, + ) + + +def _success_kwargs( + now: datetime.datetime, prompt_tokens: int, requester_metadata: Mapping[str, object] | None = None +) -> Mapping[str, object]: + payload: Final = _standard_logging_payload(now, prompt_tokens) + return { + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + **payload, + "metadata": {**payload["metadata"], "requester_metadata": requester_metadata}, + }, + "stream": True, + "start_time": now - datetime.timedelta(seconds=3), + "api_call_start_time": now - datetime.timedelta(seconds=2), + "completion_start_time": now - datetime.timedelta(seconds=1), + "end_time": now, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flag_at_request_time", (True, False)) +async def test_logger_emits_bucket_from_its_startup_label_set( + monkeypatch: pytest.MonkeyPatch, flag_at_request_time: bool +): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, flag_at_request_time) + + await logger.async_log_success_event(dict(_success_kwargs(now, prompt_tokens=4_000)), None, now, now) + + _assert_latency_metrics("4k-16k") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "combined_usage", "expected"), + ( + ({"id": "moderation", "results": []}, None, "unknown"), + ({"usage": None}, None, "unknown"), + ({"usage": {}}, None, "unknown"), + ({"usage": {"completion_tokens": 3}}, None, "unknown"), + ({"usage": {"total_tokens": 5}}, None, "unknown"), + ({"usage": {"prompt_tokens": 0}}, None, "0-1k"), + ({"usage": {"prompt_tokens": 4_000}}, None, "4k-16k"), + ({"usage": {"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}}, None, "0-1k"), + ({"usage": {"input_tokens": 4_000, "output_tokens": 3, "total_tokens": 4_003}}, None, "4k-16k"), + (litellm.ModelResponse(usage=litellm.Usage(prompt_tokens=0)), None, "0-1k"), + (None, litellm.Usage(prompt_tokens=0), "0-1k"), + ), +) +@pytest.mark.parametrize("include_usage_metadata", (True, False)) +async def test_logger_distinguishes_missing_usage_from_reported_zero( + monkeypatch: pytest.MonkeyPatch, + response: object, + combined_usage: object, + expected: str, + include_usage_metadata: bool, +): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + usage: Final = StandardLoggingPayloadSetup.get_usage_as_dict( + response_obj=response if isinstance(response, dict) else None + ) + + payload: Final = _standard_logging_payload(now, usage.get("prompt_tokens", 0)) + await logger.async_log_success_event( + { + **_success_kwargs(now, prompt_tokens=usage.get("prompt_tokens", 0)), + "combined_usage_object": combined_usage, + "standard_logging_object": { + **payload, + "metadata": {**payload["metadata"], "usage_object": usage if include_usage_metadata else None}, + }, + }, + response, + now, + now, + ) + + _assert_latency_metrics(expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("total_tokens", (None, 0, 5_000)) +@pytest.mark.parametrize("prompt_tokens", (0, 4_000)) +async def test_upstream_total_only_usage_has_unknown_input_length( + monkeypatch: pytest.MonkeyPatch, total_tokens: int | None, prompt_tokens: int +): + import httpx + + from litellm.litellm_core_utils.litellm_logging import Logging, StandardLoggingPayloadSetup + from litellm.proxy.pass_through_endpoints.upstream_usage_headers import apply_upstream_reported_usage + + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, True) + logger: Final = PrometheusLogger() + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=now, + litellm_call_id="test-call-id", + function_id="1", + ) + headers: Final = httpx.Headers( + { + "x-litellm-response-cost": "0.001", + **({"x-litellm-total-tokens": str(total_tokens)} if total_tokens is not None else {}), + } + ) + reported: Final = apply_upstream_reported_usage(logging_obj=logging_obj, headers=headers) + assert reported is not None + combined_usage: Final = logging_obj.model_call_details.get("combined_usage_object") + response: Final = {"usage": {"prompt_tokens": prompt_tokens}} + usage: Final = StandardLoggingPayloadSetup.get_usage_as_dict(response, combined_usage) + payload: Final = _standard_logging_payload(now, usage.get("prompt_tokens", 0)) + + await logger.async_log_success_event( + { + **logging_obj.model_call_details, + **_success_kwargs(now, usage.get("prompt_tokens", 0)), + "standard_logging_object": {**payload, "metadata": {**payload["metadata"], "usage_object": usage}}, + }, + response, + now, + now, + ) + + _assert_latency_metrics("unknown" if total_tokens is not None else get_input_sequence_length_bucket(prompt_tokens)) + + +@pytest.mark.asyncio +async def test_logger_built_with_flag_off_emits_no_bucket_label(monkeypatch: pytest.MonkeyPatch): + now: Final = datetime.datetime.now() + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, True) + + await logger.async_log_success_event(dict(_success_kwargs(now, prompt_tokens=4_000)), None, now, now) + + _assert_latency_metrics(None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flag_at_startup", (True, False)) +@pytest.mark.parametrize("stream", (True, False)) +@pytest.mark.parametrize( + "metadata", + ( + None, + {}, + {"input_sequence_length": None}, + {"input_sequence_length": False}, + {"input_sequence_length": True}, + {"input_sequence_length": 0}, + {"input_sequence_length": []}, + {"input_sequence_length": {}}, + {"input_sequence_length": ""}, + {"input_sequence_length": "from-metadata"}, + ), +) +async def test_custom_input_length_label_is_scoped_to_target_histograms( + monkeypatch: pytest.MonkeyPatch, flag_at_startup: bool, stream: bool, metadata: Mapping[str, object] | None +): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", ["input_sequence_length"]) + kwargs: Final = { + **_success_kwargs(now, prompt_tokens=4_000, requester_metadata=metadata), + "stream": stream, + "litellm_params": {"metadata": {"queue_time_seconds": 0.25}}, + } + original_kwargs: Final = deepcopy(kwargs) + baseline_logger: Final = PrometheusLogger() + await baseline_logger.async_log_success_event(kwargs, None, now, now) + baseline_samples: Final = _non_target_samples() + _clear_prometheus_registry() + monkeypatch.setattr(litellm, FLAG, flag_at_startup) + logger: Final = PrometheusLogger() + monkeypatch.setattr(litellm, FLAG, not flag_at_startup) + + await logger.async_log_success_event(kwargs, None, now, now) + + assert kwargs == original_kwargs + custom_value: Final = (metadata or {}).get("input_sequence_length") + expected: Final = custom_value if isinstance(custom_value, str) else ("4k-16k" if flag_at_startup else "None") + _assert_latency_metrics(expected, stream=stream, queue_time=0.25) + assert all(logger.get_labels_for_metric(metric).count("input_sequence_length") == 1 for metric in LATENCY_METRICS) + non_target_samples: Final = _non_target_samples() + assert { + "litellm_requests_metric_total", + "litellm_spend_metric_total", + "litellm_total_tokens_metric_total", + "litellm_request_queue_time_seconds_count", + "litellm_deployment_success_responses_total", + }.issubset({sample.name for sample in non_target_samples}) + assert non_target_samples == baseline_samples + queue_sum: Final = tuple( + sample for sample in non_target_samples if sample.name == "litellm_request_queue_time_seconds_sum" + ) + assert len(queue_sum) == 1 and queue_sum[0].value == 0.25 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", (True, False)) +async def test_concurrent_requests_keep_independent_buckets(monkeypatch: pytest.MonkeyPatch, enabled: bool): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, enabled) + logger: Final = PrometheusLogger() + cases: Final = ( + (None, "unknown"), + (0, "0-1k"), + (1_000, "1k-4k"), + (4_000, "4k-16k"), + (16_000, "16k-64k"), + (64_000, "64k+"), + ) + calls: Final = tuple( + ( + {**_success_kwargs(now, prompt_tokens=tokens or 0), "stream": stream}, + {"usage": {"prompt_tokens": tokens}} if tokens is not None else None, + ) + for tokens, _ in cases + for stream in (True, False) + for _ in range(2) + ) + original_calls: Final = deepcopy(calls) + + await asyncio.gather(*(logger.async_log_success_event(kwargs, response, now, now) for kwargs, response in calls)) + + assert calls == original_calls + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + for metric, duration in zip(LATENCY_METRICS, (2, 1, 3)): + expected_count: Final = 2 if metric == "litellm_llm_api_time_to_first_token_metric" else 4 + counts: Final = tuple(sample for sample in samples if sample.name == f"{metric}_count") + sums: Final = tuple(sample for sample in samples if sample.name == f"{metric}_sum") + buckets: Final = tuple(sample for sample in samples if sample.name == f"{metric}_bucket") + expected: Final = ( + {bucket: expected_count for _, bucket in cases} if enabled else {None: expected_count * len(cases)} + ) + assert len(counts) == len(sums) == len(expected) + assert {sample.labels.get("input_sequence_length"): sample.value for sample in counts} == expected + assert {sample.labels.get("input_sequence_length"): sample.value for sample in sums} == { + bucket: count * duration for bucket, count in expected.items() + } + assert sum(sample.value for sample in buckets if sample.labels["le"] == "+Inf") == expected_count * len(cases) + assert all( + sample.value + == expected[sample.labels.get("input_sequence_length")] * int(float(sample.labels["le"]) >= duration) + for sample in buckets + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", (True, False)) +async def test_failed_request_does_not_observe_latency(monkeypatch: pytest.MonkeyPatch, enabled: bool): + now: Final = datetime.datetime.now() + monkeypatch.setattr(litellm, FLAG, enabled) + monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", ["input_sequence_length"]) + logger: Final = PrometheusLogger() + kwargs: Final = { + **_success_kwargs(now, prompt_tokens=4_000), + "standard_logging_object": {**_standard_logging_payload(now, 4_000), "status": "failure"}, + "exception": RuntimeError("upstream request failed"), + } + + await logger.async_log_failure_event(kwargs, None, now, now) + + samples: Final = tuple(sample for metric in REGISTRY.collect() for sample in metric.samples) + assert not any(sample.name.startswith(LATENCY_METRICS) for sample in samples) + for metric in ("litellm_llm_api_failed_requests_metric_total", "litellm_deployment_failure_responses_total"): + counts: Final = tuple(sample for sample in samples if sample.name == metric) + assert len(counts) == 1 + assert counts[0].value == 1 + assert counts[0].labels["input_sequence_length"] == "None" diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index a037284d7c1..08d37297ab1 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1,10 +1,19 @@ import asyncio +import re +import sys +import textwrap +import uuid +from contextlib import asynccontextmanager from datetime import datetime -from unittest.mock import MagicMock, patch +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest from litellm.integrations.s3_v2 import S3Logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.integrations.s3_v2 import s3BatchLoggingElement from litellm.types.utils import StandardLoggingPayload @@ -21,9 +30,7 @@ class TestS3V2UnitTests: source_code = inspect.getsource(s3_v2) # Verify that json.dumps is not used directly in the code - assert ( - "json.dumps(" not in source_code - ), "S3 v2 should not use json.dumps directly" + assert "json.dumps(" not in source_code, "S3 v2 should not use json.dumps directly" @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") @@ -86,12 +93,8 @@ class TestS3V2UnitTests: call_args_minio = s3_logger_minio.async_httpx_client.put.call_args assert call_args_minio is not None url_minio = call_args_minio[0][0] - expected_minio_url = ( - "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json" - ) - assert ( - url_minio == expected_minio_url - ), f"Expected MinIO URL {expected_minio_url}, got {url_minio}" + expected_minio_url = "https://minio.example.com:9000/litellm-logs/2025-09-14/test-key.json" + assert url_minio == expected_minio_url, f"Expected MinIO URL {expected_minio_url}, got {url_minio}" # Test 3: Custom endpoint without bucket name (should fall back to default) s3_logger_no_bucket = S3Logger( @@ -136,12 +139,8 @@ class TestS3V2UnitTests: call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] - expected_sync_url = ( - "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json" - ) - assert ( - url_sync == expected_sync_url - ), f"Expected sync URL {expected_sync_url}, got {url_sync}" + expected_sync_url = "https://custom.s3.endpoint.com/sync-bucket/2025-09-14/test-key.json" + assert url_sync == expected_sync_url, f"Expected sync URL {expected_sync_url}, got {url_sync}" # Test 5: Download method with custom endpoint s3_logger_download = S3Logger( @@ -158,19 +157,15 @@ class TestS3V2UnitTests: s3_logger_download.async_httpx_client = AsyncMock() s3_logger_download.async_httpx_client.get.return_value = mock_download_response - result = asyncio.run( - s3_logger_download._download_object_from_s3( - "2025-09-14/download-test-key.json" - ) - ) + result = asyncio.run(s3_logger_download._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download.s3.endpoint.com/download-bucket/2025-09-14/download-test-key.json" - assert ( - url_download == expected_download_url - ), f"Expected download URL {expected_download_url}, got {url_download}" + assert url_download == expected_download_url, ( + f"Expected download URL {expected_download_url}, got {url_download}" + ) assert result == {"downloaded": "data"} @@ -216,12 +211,8 @@ class TestS3V2UnitTests: call_args = s3_logger_virtual.async_httpx_client.put.call_args assert call_args is not None url = call_args[0][0] - expected_url = ( - "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" - ) - assert ( - url == expected_url - ), f"Expected virtual-hosted-style URL {expected_url}, got {url}" + expected_url = "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" + assert url == expected_url, f"Expected virtual-hosted-style URL {expected_url}, got {url}" # Test 2: Path-style (default behavior with s3_use_virtual_hosted_style=False) s3_logger_path = S3Logger( @@ -241,12 +232,8 @@ class TestS3V2UnitTests: call_args_path = s3_logger_path.async_httpx_client.put.call_args assert call_args_path is not None url_path = call_args_path[0][0] - expected_path_url = ( - "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" - ) - assert ( - url_path == expected_path_url - ), f"Expected path-style URL {expected_path_url}, got {url_path}" + expected_path_url = "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" + assert url_path == expected_path_url, f"Expected path-style URL {expected_path_url}, got {url_path}" # Test 3: Virtual-hosted-style with http protocol s3_logger_http = S3Logger( @@ -266,12 +253,10 @@ class TestS3V2UnitTests: call_args_http = s3_logger_http.async_httpx_client.put.call_args assert call_args_http is not None url_http = call_args_http[0][0] - expected_http_url = ( - "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" + expected_http_url = "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" + assert url_http == expected_http_url, ( + f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" ) - assert ( - url_http == expected_http_url - ), f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" # Test 4: Sync upload method with virtual-hosted-style s3_logger_sync_virtual = S3Logger( @@ -295,12 +280,10 @@ class TestS3V2UnitTests: call_args_sync = mock_sync_client.put.call_args assert call_args_sync is not None url_sync = call_args_sync[0][0] - expected_sync_url = ( - "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" + expected_sync_url = "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" + assert url_sync == expected_sync_url, ( + f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" ) - assert ( - url_sync == expected_sync_url - ), f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" # Test 5: Download method with virtual-hosted-style s3_logger_download_virtual = S3Logger( @@ -316,34 +299,27 @@ class TestS3V2UnitTests: mock_download_response.status_code = 200 mock_download_response.json = MagicMock(return_value={"downloaded": "data"}) s3_logger_download_virtual.async_httpx_client = AsyncMock() - s3_logger_download_virtual.async_httpx_client.get.return_value = ( - mock_download_response - ) + s3_logger_download_virtual.async_httpx_client.get.return_value = mock_download_response - result = asyncio.run( - s3_logger_download_virtual._download_object_from_s3( - "2025-09-14/download-test-key.json" - ) - ) + result = asyncio.run(s3_logger_download_virtual._download_object_from_s3("2025-09-14/download-test-key.json")) call_args_download = s3_logger_download_virtual.async_httpx_client.get.call_args assert call_args_download is not None url_download = call_args_download[0][0] expected_download_url = "https://download-bucket.download.endpoint.com/2025-09-14/download-test-key.json" - assert ( - url_download == expected_download_url - ), f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" + assert url_download == expected_download_url, ( + f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" + ) assert result == {"downloaded": "data"} @patch("asyncio.create_task") @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") - def test_s3_v2_put_url_encodes_spaces_in_object_key( - self, mock_periodic_flush, mock_create_task - ): - import requests + def test_s3_v2_put_url_encodes_spaces_in_object_key(self, mock_periodic_flush, mock_create_task): from unittest.mock import AsyncMock + import requests + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement mock_periodic_flush.return_value = None @@ -487,9 +463,7 @@ async def test_async_upload_exhausts_retries_on_persistent_503(): # All 3 attempts return 503 response_503 = MagicMock() response_503.status_code = 503 - response_503.raise_for_status = MagicMock( - side_effect=Exception("503 Service Unavailable") - ) + response_503.raise_for_status = MagicMock(side_effect=Exception("503 Service Unavailable")) logger.async_httpx_client = AsyncMock() logger.async_httpx_client.put = AsyncMock(return_value=response_503) @@ -528,12 +502,12 @@ async def test_async_upload_no_retry_on_4xx(): s3_object_download_filename="test-no-retry.json", ) - response_403 = MagicMock() - response_403.status_code = 403 - response_403.raise_for_status = MagicMock(side_effect=Exception("403 Forbidden")) + response_400 = MagicMock() + response_400.status_code = 400 + response_400.raise_for_status = MagicMock(side_effect=Exception("400 Bad Request")) logger.async_httpx_client = AsyncMock() - logger.async_httpx_client.put = AsyncMock(return_value=response_403) + logger.async_httpx_client.put = AsyncMock(return_value=response_400) with patch.object(logger, "handle_callback_failure") as mock_failure: await logger.async_upload_data_to_s3(test_element) @@ -543,6 +517,190 @@ async def test_async_upload_no_retry_on_4xx(): mock_failure.assert_called_once_with(callback_name="S3Logger") +_SIGV4_ACCESS_KEY = re.compile(r"Credential=(AKIA\d+)/") + + +@pytest.fixture +def rotating_profile(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str: + """ + A real botocore profile whose credential_process hands out a new key generation on every call and + expires inside the advisory refresh window, so RefreshableCredentials re-runs it on every property read. + """ + counter = tmp_path / "generation" + script = tmp_path / "rotate_credentials.py" + script.write_text( + textwrap.dedent( + f""" + import json, sys + from datetime import datetime, timedelta, timezone + from pathlib import Path + + counter = Path({str(counter)!r}) + generation = int(counter.read_text()) if counter.exists() else 0 + counter.write_text(str(generation + 1)) + expiry = (datetime.now(timezone.utc) + timedelta(minutes=12)).strftime("%Y-%m-%dT%H:%M:%SZ") + json.dump( + {{ + "Version": 1, + "AccessKeyId": f"AKIA{{generation}}", + "SecretAccessKey": f"secret-{{generation}}", + "SessionToken": f"token-{{generation}}", + "Expiration": expiry, + }}, + sys.stdout, + ) + """ + ) + ) + profile = f"rotating-{uuid.uuid4().hex}" + (tmp_path / "config").write_text(f"[profile {profile}]\ncredential_process = {sys.executable} {script}\n") + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + return profile + + +def _generation(request: httpx.Request) -> tuple[str, str]: + """(access key generation, session token generation) SigV4 baked into one request.""" + access_key = _SIGV4_ACCESS_KEY.search(request.headers["Authorization"]) + assert access_key is not None + return access_key.group(1).removeprefix("AKIA"), request.headers["X-Amz-Security-Token"].removeprefix("token-") + + +@asynccontextmanager +async def _s3_logger_on_production_handler(profile: str, statuses: list[int]): + """ + S3Logger wired to the real AsyncHTTPHandler over an httpx MockTransport that answers with the given + statuses in order, so the handler's own raise_for_status behaviour is exercised end to end. + """ + requests: list[httpx.Request] = [] + replies = iter(statuses) + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(next(replies), request=request, text="SignatureDoesNotMatch") + + handler = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_region_name="us-east-1", + s3_aws_profile_name=profile, + s3_flush_interval=3600, + ) + logger.async_httpx_client = handler + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + yield logger, requests, mock_sleep + await handler.client.aclose() + + +@pytest.mark.asyncio +async def test_async_upload_signs_with_one_frozen_credential_snapshot(rotating_profile: str, caplog): + """ + RefreshableCredentials refreshes on every property read once inside the advisory window, so signing + off the live object would mix the access key of one generation with the token of the next. + """ + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-frozen.json", + payload={"test": "frozen"}, + s3_object_download_filename="test-frozen.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [200]) as (logger, requests, _): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 1 + access_key_generation, token_generation = _generation(requests[0]) + assert access_key_generation == token_generation + assert "Error uploading to s3" not in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_retries_403_with_fresh_credentials_and_signature(rotating_profile: str, caplog): + """ + A 403 (SignatureDoesNotMatch after an IMDS rotation) must be retried, and the retry must fetch + credentials again and carry a signature computed from that newer generation. + """ + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-403.json", + payload={"test": "403"}, + s3_object_download_filename="test-403.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [403, 200]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 2 + first_key, first_token = _generation(requests[0]) + second_key, second_token = _generation(requests[1]) + assert first_key == first_token + assert second_key == second_token + assert int(second_key) > int(first_key) + assert requests[1].headers["Authorization"] != requests[0].headers["Authorization"] + mock_sleep.assert_awaited_once_with(1) + assert "Error uploading to s3" not in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_exhausts_403_retries_through_production_http_handler(rotating_profile: str, caplog): + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-403-exhausted.json", + payload={"test": "403-exhausted"}, + s3_object_download_filename="test-403-exhausted.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [403, 403, 403]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 3 + assert mock_sleep.await_args_list == [call(1), call(2)] + assert "Error uploading to s3" in caplog.text + + +@pytest.mark.asyncio +async def test_async_upload_does_not_retry_404_through_production_http_handler(rotating_profile: str, caplog): + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-404.json", + payload={"test": "404"}, + s3_object_download_filename="test-404.json", + ) + async with _s3_logger_on_production_handler(rotating_profile, [404]) as (logger, requests, mock_sleep): + await logger.async_upload_data_to_s3(test_element) + + assert len(requests) == 1 + mock_sleep.assert_not_awaited() + assert "Error uploading to s3" in caplog.text + + +def test_sync_upload_retries_403_with_fresh_signature(rotating_profile: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_PROFILE", rotating_profile) + logger = S3Logger(s3_bucket_name="test-bucket", s3_region_name="us-east-1", s3_flush_interval=3600) + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-sync-403.json", + payload={"test": "sync-403"}, + s3_object_download_filename="test-sync-403.json", + ) + requests: list[httpx.Request] = [] + replies = iter([403, 200]) + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(next(replies), request=request) + + handler = HTTPHandler() + handler.client = httpx.Client(transport=httpx.MockTransport(respond)) + with ( + patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it + "litellm.integrations.s3_v2._get_httpx_client", return_value=handler + ), + patch("time.sleep") as mock_sleep, + ): + logger.upload_data_to_s3(test_element) + + assert len(requests) == 2 + first_key, first_token = _generation(requests[0]) + second_key, second_token = _generation(requests[1]) + assert first_key == first_token + assert second_key == second_token + assert int(second_key) > int(first_key) + mock_sleep.assert_called_once_with(1) + + def test_sync_upload_retries_on_s3_503(): """ Test that the sync upload_data_to_s3 retries on transient S3 503. @@ -626,9 +784,7 @@ async def test_async_log_event_skips_when_standard_logging_object_missing(): # Nothing should have been queued (catches the case where code falls # through without returning and appends None to the queue) - assert ( - len(logger.log_queue) == 0 - ), "log_queue should be empty when standard_logging_object is missing" + assert len(logger.log_queue) == 0, "log_queue should be empty when standard_logging_object is missing" @pytest.mark.asyncio @@ -767,20 +923,18 @@ async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, # This should NOT be ignored - "s3_use_ssl": False, # This should also NOT be ignored - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, # This should NOT be ignored + "s3_use_ssl": False, # This should also NOT be ignored + }, ) with patch("asyncio.create_task"): - with patch( - "litellm.integrations.s3_v2.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client @@ -788,22 +942,16 @@ async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): logger = S3Logger() # Verify s3_verify is False, not None - assert ( - logger.s3_verify is False - ), f"Expected s3_verify=False, got {logger.s3_verify}" - assert ( - logger.s3_use_ssl is False - ), f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}" + assert logger.s3_verify is False, f"Expected s3_verify=False, got {logger.s3_verify}" + assert logger.s3_use_ssl is False, f"Expected s3_use_ssl=False, got {logger.s3_use_ssl}" # Verify that get_async_httpx_client was called with ssl_verify=False mock_get_client.assert_called_once() call_kwargs = mock_get_client.call_args.kwargs - assert ( - "params" in call_kwargs - ), "params should be passed to get_async_httpx_client" - assert call_kwargs["params"] == { - "ssl_verify": False - }, f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" + assert "params" in call_kwargs, "params should be passed to get_async_httpx_client" + assert call_kwargs["params"] == {"ssl_verify": False}, ( + f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" + ) @pytest.mark.asyncio @@ -820,17 +968,15 @@ async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_aws_access_key_id": "test-key", - "s3_aws_secret_access_key": "test-secret", - "s3_region_name": "us-east-1", - }, + "s3_bucket_name": "test-bucket", + "s3_aws_access_key_id": "test-key", + "s3_aws_secret_access_key": "test-secret", + "s3_region_name": "us-east-1", + }, ) with patch("asyncio.create_task"): - with patch( - "litellm.integrations.s3_v2.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.integrations.s3_v2.get_async_httpx_client") as mock_get_client: mock_client = AsyncMock() mock_get_client.return_value = mock_client @@ -838,9 +984,7 @@ async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): logger = S3Logger() # Verify s3_verify is None (default) - assert ( - logger.s3_verify is None - ), f"Expected s3_verify=None, got {logger.s3_verify}" + assert logger.s3_verify is None, f"Expected s3_verify=None, got {logger.s3_verify}" # Verify that get_async_httpx_client was called mock_get_client.assert_called_once() @@ -868,13 +1012,13 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatc litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, ) with patch("asyncio.create_task"): @@ -890,9 +1034,7 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatc httpx_client = logger.async_httpx_client.client # Check the _verify attribute (httpx internal) if hasattr(httpx_client, "_verify"): - assert ( - httpx_client._verify is False - ), f"Expected httpx client _verify=False, got {httpx_client._verify}" + assert httpx_client._verify is False, f"Expected httpx client _verify=False, got {httpx_client._verify}" @pytest.mark.asyncio @@ -910,13 +1052,13 @@ async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): litellm, "s3_callback_params", { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - }, + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, ) with patch("asyncio.create_task"): @@ -948,9 +1090,9 @@ async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): if hasattr(logger.async_httpx_client, "client"): httpx_client = logger.async_httpx_client.client if hasattr(httpx_client, "_verify"): - assert ( - httpx_client._verify is False - ), f"Expected async httpx client _verify=False, got {httpx_client._verify}" + assert httpx_client._verify is False, ( + f"Expected async httpx client _verify=False, got {httpx_client._verify}" + ) @pytest.mark.asyncio @@ -1017,9 +1159,7 @@ def patch_asyncio_create_task(): (True, True, None, None, ""), ], ) -def test_s3_object_key_prefix_combinations( - use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix -): +def test_s3_object_key_prefix_combinations(use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix): """ Validate correct S3 prefix composition for team alias + key alias combinations. """ @@ -1490,9 +1630,7 @@ def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): logger = S3Logger(s3_callback_params_override=override) assert logger.s3_bucket_name == "resolved-bucket" assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - assert ( - litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - ) + assert litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch): @@ -1520,9 +1658,7 @@ def _expected_content_md5(payload: dict) -> str: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps json_string = safe_dumps(payload) - return base64.b64encode( - hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest() - ).decode() + return base64.b64encode(hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()).decode() def _require_non_security_md5(monkeypatch): @@ -1658,9 +1794,9 @@ def test_s3_server_side_encryption_read_from_callback_params(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -1789,10 +1925,10 @@ def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") @@ -1863,10 +1999,10 @@ def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "AES256", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "AES256", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "AES256" @@ -1884,10 +2020,10 @@ def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": True, - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": True, + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -1902,10 +2038,10 @@ def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch): litellm, "s3_callback_params", { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": 12345, - }, + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": 12345, + }, ) logger = S3Logger() assert logger.s3_server_side_encryption == "aws:kms" @@ -2045,6 +2181,7 @@ async def test_download_signs_object_key_with_space_the_way_s3_does(): headers=call.kwargs["headers"], ) + _RESERVED_CHAR_KEYS = ( "2026-08-21/time-05-29-36_resp_bGl0ZWxsbTpjdXN0b20=.json", "session=logs/2026-08-21/time-05-29-36_abc.json", diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/test_litellm/litellm_core_utils/event_loop_lag.py new file mode 100644 index 00000000000..1cac0365547 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/event_loop_lag.py @@ -0,0 +1,40 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable +from typing import Final, TypeVar + +import litellm + +T = TypeVar("T") + + +def warm_tokenizer(model: str) -> None: + litellm.token_counter(model=model, text="load the tokenizer before anything is timed") + + +async def loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, float, tuple[float, ...]]: + finished: Final = asyncio.Event() + + async def timed() -> tuple[T, float]: + await asyncio.sleep(0) + started: Final = time.perf_counter() + try: + return await run(), time.perf_counter() - started + finally: + finished.set() + + (result, took), lags = await asyncio.gather(timed(), loop_wake_lags(finished)) + return result, took, lags + + +def assert_loop_stayed_free(took: float, lags: tuple[float, ...]) -> None: + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 65a6dd2a4ca..fbb9d178390 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1,9 +1,11 @@ import json +from datetime import datetime, timezone import pytest from fastapi.testclient import TestClient import litellm +from litellm._internal_context import pinned_billing_time from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) @@ -27,10 +29,10 @@ from litellm.types.utils import ( ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + BilledTokenRates, CostCalculatorUtils, PromptTokensDetailsResult, TokenRates, - TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, _is_off_peak, @@ -38,6 +40,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( apply_off_peak_pricing, calculate_cache_writing_cost, generic_cost_per_token, + get_billed_token_rates, get_token_type_cost_breakdown, ) from litellm.types.utils import CacheCreationTokenDetails, Usage @@ -3906,6 +3909,200 @@ def test_token_type_cost_breakdown_reconciles_with_generic_total(_local_model_co assert text_input_cost + breakdown.cache_read_cost == pytest.approx(prompt_cost) +def _custom_priced_usage() -> Usage: + return Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800, cache_creation_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ) + + +def test_token_type_cost_breakdown_prices_custom_pricing_from_its_flat_rates(): + """ + A custom-priced deployment, usually absent from the cost map, used to get zero cache and + reasoning lines while its total already billed cache tokens at the custom cache rates. + The lines must come from the same flat rates: a configured cache rate, else the input + rate for cache tokens and the output rate for reasoning tokens. + """ + from litellm.types.utils import CostPerToken + + breakdown = get_token_type_cost_breakdown( + model="openai/onprem-model", + custom_llm_provider="openai", + usage=_custom_priced_usage(), + custom_cost_per_token=CostPerToken( + input_cost_per_token=1e-6, output_cost_per_token=2e-6, cache_read_input_token_cost=1e-7 + ), + ) + + assert breakdown.cache_read_cost == pytest.approx(800 * 1e-7) + assert breakdown.cache_creation_cost == pytest.approx(100 * 1e-6) + assert breakdown.reasoning_cost == pytest.approx(200 * 2e-6) + + +def test_token_type_cost_breakdown_reconciles_with_custom_pricing_totals(): + from litellm.cost_calculator import cost_per_token + from litellm.types.utils import CostPerToken + + usage = _custom_priced_usage() + custom_cost_per_token = CostPerToken( + input_cost_per_token=1e-6, + output_cost_per_token=2e-6, + cache_read_input_token_cost=1e-7, + cache_creation_input_token_cost=1.25e-6, + ) + + prompt_cost, completion_cost = cost_per_token( + model="openai/onprem-model", + custom_llm_provider="openai", + prompt_tokens=1000, + completion_tokens=500, + usage_object=usage, + custom_cost_per_token=custom_cost_per_token, + ) + breakdown = get_token_type_cost_breakdown( + model="openai/onprem-model", + custom_llm_provider="openai", + usage=usage, + custom_cost_per_token=custom_cost_per_token, + ) + + assert 100 * 1e-6 + breakdown.cache_read_cost + breakdown.cache_creation_cost == pytest.approx(prompt_cost) + assert 300 * 2e-6 + breakdown.reasoning_cost == pytest.approx(completion_cost) + + +def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "tiered-cache-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + usage = Usage( + prompt_tokens=250_000, + completion_tokens=1_000, + total_tokens=251_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200_000, cache_creation_tokens=10_000), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ) + + rates = get_billed_token_rates(model="tiered-cache-model", custom_llm_provider="openai", usage=usage) + breakdown = get_token_type_cost_breakdown(model="tiered-cache-model", custom_llm_provider="openai", usage=usage) + + assert rates == BilledTokenRates( + input_cost_per_token=6e-6, + output_cost_per_token=3e-5, + cache_read_input_token_cost=6e-7, + cache_creation_input_token_cost=7.5e-6, + cache_creation_input_token_cost_above_1hr=0.0, + output_cost_per_reasoning_token=3e-5, + ) + assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost) + assert breakdown.cache_creation_cost == pytest.approx(10_000 * rates.cache_creation_input_token_cost) + assert breakdown.reasoning_cost == pytest.approx(200 * rates.output_cost_per_reasoning_token) + + +def test_a_pinned_billing_time_prices_the_totals_and_the_reported_rates_at_one_moment(monkeypatch): + """Totals and reported rates resolve off-peak pricing on separate paths that each read the + clock, so a window opening between the two reads used to leave them describing one request + at two different prices. Pinned, both must answer for the pinned moment.""" + monkeypatch.setitem( + litellm.model_cost, + "off-peak-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "off_peak_pricing": { + "hours_utc": "02:00-03:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + }, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + with pinned_billing_time(datetime(2026, 1, 1, 2, 30, tzinfo=timezone.utc)): + off_peak_prompt_cost, off_peak_completion_cost = generic_cost_per_token( + model="off-peak-model", usage=usage, custom_llm_provider="openai" + ) + off_peak_rates = get_billed_token_rates(model="off-peak-model", custom_llm_provider="openai", usage=usage) + with pinned_billing_time(datetime(2026, 1, 1, 12, 30, tzinfo=timezone.utc)): + peak_prompt_cost, peak_completion_cost = generic_cost_per_token( + model="off-peak-model", usage=usage, custom_llm_provider="openai" + ) + peak_rates = get_billed_token_rates(model="off-peak-model", custom_llm_provider="openai", usage=usage) + + assert off_peak_rates.input_cost_per_token == pytest.approx(1e-6) + assert peak_rates.input_cost_per_token == pytest.approx(3e-6) + assert off_peak_prompt_cost == pytest.approx(1000 * off_peak_rates.input_cost_per_token) + assert off_peak_completion_cost == pytest.approx(500 * off_peak_rates.output_cost_per_token) + assert peak_prompt_cost == pytest.approx(1000 * peak_rates.input_cost_per_token) + assert peak_completion_cost == pytest.approx(500 * peak_rates.output_cost_per_token) + + +def test_the_token_type_breakdown_carries_the_rates_it_billed_at(monkeypatch): + """Callers that report both the lines and the rates read the rates off the breakdown rather than + resolving them a second time, so the breakdown has to hand back exactly what it billed at.""" + monkeypatch.setitem( + litellm.model_cost, + "xai/tiered-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + usage = Usage( + prompt_tokens=200_000, + completion_tokens=1_000, + total_tokens=201_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100_000), + ) + + breakdown = get_token_type_cost_breakdown(model="xai/tiered-model", custom_llm_provider="xai", usage=usage) + + assert breakdown.rates == get_billed_token_rates( + model="xai/tiered-model", custom_llm_provider="xai", usage=usage + ) + assert breakdown.rates.cache_read_input_token_cost == pytest.approx(6e-7) + assert breakdown.cache_read_cost == pytest.approx(100_000 * breakdown.rates.cache_read_input_token_cost) + + +def test_the_token_type_breakdown_reports_no_rates_for_an_unpriced_model(): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + breakdown = get_token_type_cost_breakdown( + model="no-such-model-anywhere", custom_llm_provider="openai", usage=usage + ) + + assert breakdown.rates is None + + +def test_billed_token_rates_are_none_for_an_unpriced_model(): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + assert get_billed_token_rates(model="no-such-model-anywhere", custom_llm_provider="openai", usage=usage) is None + + def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost_map): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) @@ -3913,9 +4110,7 @@ def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost model="gpt-4o", custom_llm_provider="openai", usage=usage ) - assert breakdown == TokenTypeCostBreakdown( - reasoning_cost=0.0, cache_read_cost=0.0, cache_creation_cost=0.0 - ) + assert (breakdown.reasoning_cost, breakdown.cache_read_cost, breakdown.cache_creation_cost) == (0.0, 0.0, 0.0) @pytest.mark.parametrize( @@ -3987,9 +4182,7 @@ def test_token_type_cost_breakdown_handles_unknown_model_gracefully(): completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=5), ), ) - assert breakdown == TokenTypeCostBreakdown( - reasoning_cost=0.0, cache_read_cost=0.0, cache_creation_cost=0.0 - ) + assert (breakdown.reasoning_cost, breakdown.cache_read_cost, breakdown.cache_creation_cost) == (0.0, 0.0, 0.0) def test_token_type_cost_breakdown_applies_regional_uplift(_local_model_cost_map): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 6cc3dcceebc..bbb7b5f9c35 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,5 +1,6 @@ -import os +import json from collections.abc import Mapping, Sequence +from pathlib import Path import pytest @@ -11,8 +12,6 @@ from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebS from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams - - def test_web_search_cost_low(): web_search_options = WebSearchOptions(search_context_size="low") model_info = litellm.get_model_info("gpt-4o-search-preview") @@ -683,12 +682,13 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( def _openai_responses_with_web_search_calls(model, num_calls): - from litellm.types.llms.openai import ResponsesAPIResponse from openai.types.responses.response_function_web_search import ( ActionSearch, ResponseFunctionWebSearch, ) + from litellm.types.llms.openai import ResponsesAPIResponse + output = [ ResponseFunctionWebSearch( id=f"ws_{i}", @@ -859,11 +859,62 @@ def test_dated_search_preview_entries_carry_search_pricing(local_model_cost_map) custom_llm_provider="openai", standard_built_in_tools_params=None, ) - assert cost == pytest.approx(0.035), ( - f"dated search-preview id must bill the $0.035 search fee, got ${cost}" + assert cost == pytest.approx(0.025), ( + f"dated search-preview id must bill the $0.025 search fee, got ${cost}" ) +@pytest.mark.parametrize( + "web_search_options", + [ + None, + WebSearchOptions(search_context_size="low"), + WebSearchOptions(search_context_size="medium"), + WebSearchOptions(search_context_size="high"), + ], +) +def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias( + web_search_options: WebSearchOptions | None, local_model_cost_map: None +) -> None: + alias_info = litellm.get_model_info("gpt-4o-mini") + snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18") + + assert not snapshot_info["supports_web_search"] + assert not alias_info["supports_web_search"] + + snapshot_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( + web_search_options=web_search_options, model_info=snapshot_info + ) + alias_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( + web_search_options=web_search_options, model_info=alias_info + ) + + assert snapshot_cost == alias_cost == 0.025 + + +def test_gpt_4o_mini_web_search_price_matches_in_both_cost_maps(): + repo_root = Path(__file__).parents[4] + cost_maps = tuple( + json.loads((repo_root / path).read_text(encoding="utf-8")) + for path in ( + "model_prices_and_context_window.json", + "litellm/model_prices_and_context_window_backup.json", + ) + ) + canonical, backup = cost_maps + expected_search_price = { + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025, + "search_context_size_high": 0.025, + } + for model_name in ("gpt-4o-mini", "gpt-4o-mini-2024-07-18"): + canonical_entry = canonical[model_name] + backup_entry = backup[model_name] + assert canonical_entry["search_context_cost_per_query"] == expected_search_price + assert backup_entry["search_context_cost_per_query"] == expected_search_price + assert canonical_entry == backup_entry + + # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 304d732c518..8e46ae21de6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1,4 +1,6 @@ +from typing import Final +import pytest from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -99,3 +101,97 @@ def test_handle_invalid_parallel_tool_calls_skips_custom_tool_calls(): ) result = _handle_invalid_parallel_tool_calls([custom_tool_call, function_tool_call]) assert result == [custom_tool_call, function_tool_call] + + +def test_convert_empty_choices_response() -> None: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + ) + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + "vertex_ai_safety_results": ["blocked"], + } + result: Final = convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + assert result.choices == [] + assert getattr(result, "vertex_ai_safety_results") == ["blocked"] + + sync_stream: Final = list(convert_to_streaming_response(response_object=resp)) + assert len(sync_stream) == 1 + assert sync_stream[0].choices == [] + + +@pytest.mark.asyncio +async def test_convert_empty_choices_response_async() -> None: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response_async, + ) + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 0, "total_tokens": 10}, + } + async_chunks: Final = [chunk async for chunk in convert_to_streaming_response_async(response_object=resp)] + assert len(async_chunks) == 1 + assert async_chunks[0].choices == [] + + +def test_convert_missing_choices_raises_api_error() -> None: + from litellm.exceptions import APIError + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + } + with pytest.raises(APIError) as exc_info: + convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + assert "no 'choices'" in str(exc_info.value) + + +@pytest.mark.parametrize(("choices", "type_name"), [({}, "dict"), ("", "str"), (None, "NoneType"), (0, "int")]) +@pytest.mark.asyncio +async def test_convert_non_list_choices_raises_api_error(choices: object, type_name: str) -> None: + from litellm.exceptions import APIError + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, + ) + + resp: Final = { + "id": "x", + "created": 1, + "model": "gemini-3.5-flash", + "object": "chat.completion", + "choices": choices, + } + expected: Final = f"'choices' that is not a list \\({type_name}\\)" + with pytest.raises(APIError, match=expected): + convert_to_model_response_object( + response_object=resp, + model_response_object=ModelResponse(), + response_type="completion", + ) + with pytest.raises(APIError, match=expected): + list(convert_to_streaming_response(response_object=resp)) + with pytest.raises(APIError, match=expected): + async for _ in convert_to_streaming_response_async(response_object=resp): + pass diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index c037f928593..b5890d1a5b0 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1,23 +1,33 @@ +import copy import functools import json import os +import sys +from typing import Final from unittest.mock import MagicMock, patch import pytest - from litellm.litellm_core_utils.prompt_templates.common_utils import ( + ENCRYPTED_REASONING_SIGNATURE_PREFIX, TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, add_system_prompt_to_messages, + encrypted_content_from_signature, + encrypted_reasoning_signature, get_file_ids_from_messages, get_format_from_file_id, handle_any_messages_to_chat_completion_str_messages_conversion, hoist_images_from_tool_messages, + is_encrypted_reasoning_block, + responses_reasoning_items_from_thinking_blocks, split_concatenated_json_objects, + strip_encrypted_reasoning_from_messages, update_messages_with_model_file_ids, ) +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + def test_get_format_from_file_id(): unified_file_id = "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf" @@ -1435,7 +1445,7 @@ class TestFlattenTopLevelSchemaCombinators: assert schema == snapshot -class TestToolWithFlattenedParameters: +class TestToolWithSanitizedParameters: def _anyof_tool(self): return { "type": "function", @@ -1462,11 +1472,12 @@ class TestToolWithFlattenedParameters: def test_flattens_anyof_parameters_into_new_tool(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) tool = self._anyof_tool() - result = tool_with_flattened_parameters(tool) + result = tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) assert result is not tool parameters = result["function"]["parameters"] @@ -1479,7 +1490,8 @@ class TestToolWithFlattenedParameters: def test_clean_parameters_return_the_same_tool_object(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) tool = { @@ -1490,7 +1502,23 @@ class TestToolWithFlattenedParameters: }, } - assert tool_with_flattened_parameters(tool) is tool + assert tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) is tool + + def test_pattern_only_sanitizer_drops_the_regex_and_keeps_the_union(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + tool_with_sanitized_parameters, + ) + + tool = self._anyof_tool() + tool["function"]["parameters"]["properties"]["id"]["pattern"] = _ARTIFACT_FIELD_PATTERN + + result = tool_with_sanitized_parameters(tool, drop_non_python_regex_patterns) + + parameters = result["function"]["parameters"] + assert parameters["properties"]["id"] == {"type": "string"} + assert parameters["anyOf"] == self._anyof_tool()["function"]["parameters"]["anyOf"] + assert tool["function"]["parameters"]["properties"]["id"]["pattern"] == _ARTIFACT_FIELD_PATTERN @pytest.mark.parametrize( "tool", @@ -1503,10 +1531,127 @@ class TestToolWithFlattenedParameters: ) def test_non_dict_function_or_parameters_return_the_same_tool_object(self, tool): from litellm.litellm_core_utils.prompt_templates.common_utils import ( - tool_with_flattened_parameters, + flatten_combinators_and_drop_non_python_regex_patterns, + tool_with_sanitized_parameters, ) - assert tool_with_flattened_parameters(tool) is tool + assert tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) is tool + + +class TestDropNonPythonRegexPatterns: + """Claude Code's Artifact tool declares ECMA-262 ``\\p{..}`` escapes that OpenAI's + validator, which compiles ``pattern`` values and ``patternProperties`` keys with + Python ``re``, refuses as "not a 'regex'".""" + + def _schema(self, pattern): + return { + "type": "object", + "properties": { + "field": {"type": "string", "pattern": pattern}, + "writes": { + "type": "array", + "items": {"properties": {"doc_id": {"type": "string", "pattern": pattern}}}, + }, + "query": {"anyOf": [{"type": "string", "pattern": pattern}, {"type": "null"}]}, + "pair": {"type": "array", "prefixItems": [{"type": "string", "pattern": pattern}]}, + "extra": {"type": "object", "additionalProperties": {"type": "string", "pattern": pattern}}, + "tagged": { + "type": "object", + "patternProperties": {pattern: {"type": "string"}, "^x_": {"type": "integer"}}, + }, + }, + "$defs": {"segment": {"type": "string", "pattern": pattern}}, + "required": ["field"], + } + + def test_drops_every_regex_python_re_rejects_from_every_schema_position(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = self._schema(_ARTIFACT_FIELD_PATTERN) + + result = drop_non_python_regex_patterns(schema) + + assert '"pattern"' not in json.dumps(result) + properties = result["properties"] + assert properties["field"] == {"type": "string"} + assert properties["writes"]["items"]["properties"]["doc_id"] == {"type": "string"} + assert properties["query"]["anyOf"] == [{"type": "string"}, {"type": "null"}] + assert properties["pair"]["prefixItems"] == [{"type": "string"}] + assert properties["extra"]["additionalProperties"] == {"type": "string"} + assert properties["tagged"]["patternProperties"] == {"^x_": {"type": "integer"}} + assert result["$defs"]["segment"] == {"type": "string"} + assert result["required"] == ["field"] + assert schema == self._schema(_ARTIFACT_FIELD_PATTERN) + + def test_keeps_regexes_python_re_compiles_and_returns_the_same_object(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = self._schema(r'^(?!__.*__$)[^"\\./[\]]{1,200}$') + + assert drop_non_python_regex_patterns(schema) is schema + + def test_pattern_keys_inside_data_positions_are_not_regexes(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = { + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "template": {"type": "object", "default": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + "samples": {"type": "array", "examples": [{"pattern": _ARTIFACT_FIELD_PATTERN}]}, + "fixed": {"const": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + "vendor": {"type": "string", "x-litellm": {"pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + "required": ["pattern"], + } + + assert drop_non_python_regex_patterns(schema) is schema + + def test_regex_nested_past_what_python_re_can_parse_is_dropped_not_raised(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + schema = { + "type": "object", + "properties": {"deep": {"type": "string", "pattern": "(" * 2000 + "a" + ")" * 2000}}, + } + + assert drop_non_python_regex_patterns(schema)["properties"]["deep"] == {"type": "string"} + + def test_walks_schemas_deeper_than_the_interpreter_recursion_limit(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + depth = sys.getrecursionlimit() + leaf = {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN} + schema = functools.reduce( + lambda inner, _: {"type": "object", "properties": {"child": inner}}, range(depth), leaf + ) + + result = drop_non_python_regex_patterns(schema) + + assert functools.reduce(lambda node, _: node["properties"]["child"], range(depth), result) == {"type": "string"} + assert functools.reduce(lambda node, _: node["properties"]["child"], range(depth), schema) is leaf + + def test_leaves_levels_past_the_json_nesting_limit_alone(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_non_python_regex_patterns, + ) + + leaf = {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN} + schema = functools.reduce( + lambda inner, _: {"type": "object", "properties": {"child": inner}}, range(1100), leaf + ) + + assert drop_non_python_regex_patterns(schema) is schema class TestRequestContainsImageContent: @@ -1554,3 +1699,117 @@ class TestRequestContainsImageContent: for _ in range(50): nested = {"type": "tool_result", "content": [nested]} assert request_contains_image_content([{"role": "user", "content": [nested]}]) is False + + +class TestEncryptedReasoningReplay: + """Regression for https://github.com/BerriAI/litellm/issues/40288.""" + + def test_signature_round_trips_the_encrypted_content(self): + assert encrypted_content_from_signature(encrypted_reasoning_signature("gAAAA_bytes")) == "gAAAA_bytes" + + @pytest.mark.parametrize( + "signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7] + ) + def test_anything_else_is_not_encrypted_content(self, signature): + assert encrypted_content_from_signature(signature) is None + + def test_encrypted_thinking_block_replays_its_own_item(self): + items = responses_reasoning_items_from_thinking_blocks( + [{"type": "thinking", "thinking": "Plan.", "signature": encrypted_reasoning_signature("gAAAA_1")}] + ) + assert items == ( + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Plan."}], + "encrypted_content": "gAAAA_1", + }, + ) + + def test_encrypted_redacted_block_replays_with_an_empty_summary(self): + items = responses_reasoning_items_from_thinking_blocks( + [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_1")}] + ) + assert items == ({"type": "reasoning", "summary": [], "encrypted_content": "gAAAA_1"},) + + def test_plain_blocks_collapse_into_one_summary_item_around_encrypted_ones(self): + items = responses_reasoning_items_from_thinking_blocks( + [ + {"type": "thinking", "thinking": "A.", "signature": None}, + {"type": "thinking", "thinking": "B.", "signature": ""}, + {"type": "thinking", "thinking": "C.", "signature": encrypted_reasoning_signature("gAAAA_c")}, + {"type": "redacted_thinking", "data": "anthropic-minted-opaque-data"}, + {"type": "thinking", "thinking": "D."}, + ] + ) + assert items == ( + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}], + }, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "C."}], "encrypted_content": "gAAAA_c"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "D."}]}, + ) + assert all("id" not in item for item in items) + + def test_blocks_without_text_or_encrypted_content_produce_nothing(self): + assert responses_reasoning_items_from_thinking_blocks([{"type": "thinking", "thinking": ""}]) == () + assert responses_reasoning_items_from_thinking_blocks([]) == () + + @pytest.mark.parametrize( + ("block", "expected"), + [ + ({"type": "thinking", "thinking": "x", "signature": encrypted_reasoning_signature("g")}, True), + ({"type": "redacted_thinking", "data": encrypted_reasoning_signature("g")}, True), + ({"type": "thinking", "thinking": "x", "signature": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True), + ({"type": "redacted_thinking", "data": ENCRYPTED_REASONING_SIGNATURE_PREFIX}, True), + ({"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}, False), + ({"type": "redacted_thinking", "data": "EmwKAhgBEgy"}, False), + ({"type": "text", "text": encrypted_reasoning_signature("g")}, False), + ("not a block", False), + ], + ) + def test_is_encrypted_reasoning_block(self, block, expected): + assert is_encrypted_reasoning_block(block) is expected + + def test_strip_drops_every_bridge_tagged_block_and_leaves_no_unsigned_thinking_behind(self): + assistant_content = [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "thinking", "thinking": "packed by the bridge", "signature": encrypted_reasoning_signature("g1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("g2")}, + {"type": "thinking", "thinking": "", "signature": encrypted_reasoning_signature("g3")}, + {"type": "text", "text": "answer"}, + ] + messages = [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": assistant_content}, + {"role": "user", "content": [{"type": "text", "text": "follow-up"}]}, + ] + + strip_encrypted_reasoning_from_messages(messages) + + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"}, + {"type": "text", "text": "answer"}, + ] + assert all(block["signature"] for block in assistant_content if block["type"] == "thinking") + assert messages[0] == {"role": "user", "content": "question"} + assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]} + + @pytest.mark.parametrize( + "messages", + [ + "not a list", + None, + [{"role": "user", "content": None}], + [{"role": "user", "content": "plain string"}], + ["not a message"], + [{"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}]}], + ], + ) + def test_strip_leaves_history_without_bridge_reasoning_untouched(self, messages): + before = copy.deepcopy(messages) + + strip_encrypted_reasoning_from_messages(messages) + + assert messages == before diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index dd2d45f00c6..66d10fd1407 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -191,8 +191,16 @@ def test_bedrock_converse_assistant_with_empty_thinking_block_and_tool_calls(): {"type": "thinking", "thinking": "oss reasoning", "signature": None}, {"type": "thinking", "thinking": "oss reasoning", "signature": ""}, {"type": "thinking", "thinking": "oss reasoning"}, + {"type": "thinking", "thinking": "openai reasoning", "signature": "litellm_encrypted_reasoning:gAAAA"}, + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:gAAAA"}, + ], + ids=[ + "null_signature", + "empty_signature", + "missing_signature", + "encrypted_reasoning_signature", + "encrypted_reasoning_redacted_data", ], - ids=["null_signature", "empty_signature", "missing_signature"], ) def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block): """Open-source reasoning models (DeepSeek-R1, Qwen, etc.) emit thinking blocks @@ -219,7 +227,7 @@ def test_anthropic_messages_pt_drops_unsignable_thinking_block(thinking_block): assistant = next(m for m in result if m["role"] == "assistant") content = assistant["content"] assert all( - block.get("type") != "thinking" for block in content + block.get("type") not in ("thinking", "redacted_thinking") for block in content ), f"unsignable thinking block must be dropped, got {content!r}" assert any( block.get("type") == "text" and block.get("text") == "2+2 equals 4." diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 0495440c51c..c509c8399c9 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -6,6 +6,8 @@ count actual model entries, not reserved meta keys) and the extraction of the import json import os +import sys +import threading import pytest @@ -26,9 +28,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ( def _load_root_cost_map() -> dict: - path = os.path.join( - os.path.dirname(__file__), "../../../model_prices_and_context_window.json" - ) + path = os.path.join(os.path.dirname(__file__), "../../../model_prices_and_context_window.json") with open(path) as f: return json.load(f) @@ -44,9 +44,7 @@ def test_git_blob_id_is_what_git_hash_object_prints(): def _make_models(n: int) -> dict: - return { - f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n) - } + return {f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)} def test_count_model_entries_excludes_reserved_keys(): @@ -129,9 +127,7 @@ def test_finalize_pops_key_and_installs_rules(): def test_finalize_with_no_block_clears_rules(): previous = list(get_fallback_generalization_rules()) try: - set_fallback_generalizations( - [{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}] - ) + set_fallback_generalizations([{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}]) _finalize_model_cost_map(_make_models(2)) assert match_capability_generalizations("x-1") is None finally: @@ -317,9 +313,7 @@ def test_get_model_cost_map_stamps_loaded_at(): from litellm.litellm_core_utils import get_model_cost_map as module - client, _calls = _mock_client( - [httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client - ) + client, _calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) before = datetime.now(timezone.utc) module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client) @@ -328,6 +322,7 @@ def test_get_model_cost_map_stamps_loaded_at(): assert loaded_at is not None assert before <= loaded_at <= datetime.now(timezone.utc) + # --------------------------------------------------------------------------- # refetch_model_cost_map: retry/backoff behavior for runtime reloads # --------------------------------------------------------------------------- @@ -394,9 +389,7 @@ async def test_refetch_retries_429_honoring_retry_after(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert len(result.model_cost_map) > 100 assert calls["count"] == 3 @@ -408,9 +401,7 @@ async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff(): """All 429 without Retry-After: exponential backoff waits, then a failure value.""" client, calls = _mock_client([httpx.Response(429)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "429" in result.reason assert "after 3 attempts" in result.reason @@ -430,9 +421,7 @@ async def test_refetch_caps_retry_after_wait(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert sleeper.waits == [30.0] @@ -447,9 +436,7 @@ async def test_refetch_retries_transport_errors(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloaded) assert calls["count"] == 2 assert len(sleeper.waits) == 1 @@ -460,9 +447,7 @@ async def test_refetch_non_retryable_status_fails_immediately(): """A 404 is permanent: one attempt, no sleeps, failure value.""" client, calls = _mock_client([httpx.Response(404)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "404" in result.reason assert calls["count"] == 1 @@ -473,9 +458,7 @@ async def test_refetch_non_retryable_status_fails_immediately(): async def test_refetch_invalid_json_fails_immediately(): client, calls = _mock_client([httpx.Response(200, content=b"not json")]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map( - url=_URL, sleep=sleeper, rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "invalid JSON" in result.reason assert calls["count"] == 1 @@ -487,9 +470,7 @@ async def test_refetch_shrunk_map_fails_integrity_not_swapped_in(): """A drastically shrunk upstream file is rejected instead of being adopted.""" tiny = json.dumps(_make_models(60)).encode() client, _calls = _mock_client([httpx.Response(200, content=tiny)]) - result = await refetch_model_cost_map( - url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client - ) + result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) assert isinstance(result, ModelCostMapReloadUnavailable) assert "integrity validation" in result.reason @@ -586,69 +567,125 @@ from litellm.litellm_core_utils.get_model_cost_map import ( class _SyncSleepRecorder: """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" - def __init__(self): + def __init__(self, block=False): self.waits = [] + self.block = block + self.started = threading.Event() + self.release = threading.Event() def __call__(self, seconds: float) -> None: + if self.block: + self.started.set() + self.release.wait(timeout=10) self.waits.append(seconds) -def test_boot_load_retries_transient_failures_instead_of_falling_back(): - """A refused connection then a 503 at pod boot used to pin the process to the bundled - backup for its lifetime; both are transient and must be retried before giving up.""" +def _retry_threads(): + return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-retry"] + + +def test_boot_load_success_does_not_start_background_retry(): + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + ) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert _retry_threads() == [] + assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + assert get_model_cost_map_source_info()["source"] == "remote" + + +def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch): + import litellm + from litellm import utils as litellm_utils + from litellm.litellm_core_utils import get_model_cost_map as module + + original_model_cost = litellm.model_cost + monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost)) + for name, provider_models in tuple(vars(litellm).items()): + if name.endswith("_models") and isinstance(provider_models, set): + monkeypatch.setattr(litellm, name, set(provider_models)) + monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider)) + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + source_info = module._cost_map_source_info + for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): + monkeypatch.setattr(source_info, name, getattr(source_info, name)) + + remote_map = _load_root_cost_map() + remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} client, calls = _mock_client( [ httpx.ConnectError("connection refused"), - httpx.Response(503), - httpx.Response(200, content=_real_map_bytes()), + httpx.Response(200, content=json.dumps(remote_map).encode()), ], client_cls=httpx.Client, ) - sleeper = _SyncSleepRecorder() + sleeper = _SyncSleepRecorder(block=True) + litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}}) - cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - - assert calls["count"] == 3 - assert len(sleeper.waits) == 2 - assert 2.0 <= sleeper.waits[0] < 3.0 - assert 4.0 <= sleeper.waits[1] < 5.0 - source = get_model_cost_map_source_info() - assert source["source"] == "remote" - assert source["fallback_reason"] is None - assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} - - -def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts(): - """An outage longer than the retry budget still ends on the bundled backup, and the - recorded fallback reason says how many attempts were spent so operators can tell.""" - client, calls = _mock_client( - [httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client + cost_map = get_model_cost_map( + url=_URL, + max_attempts=3, + sleep=sleeper, + rng=random.Random(0), + client=client, ) - sleeper = _SyncSleepRecorder() - cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - - assert calls["count"] == 3 - assert sleeper.waits == [7.0, 7.0] - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert "after 3 attempts" in source["fallback_reason"] - assert len(cost_map) > 100 + assert calls["count"] == 1 + assert sleeper.waits == [] + assert sleeper.started.wait(timeout=10) + threads = _retry_threads() + try: + assert len(threads) == 1 + assert "claude-remote-only-test" not in cost_map + assert cost_map.keys() == _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys() + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed:") + sleeper.release.set() + for thread in threads: + thread.join(timeout=10) + assert all(not thread.is_alive() for thread in threads) + assert sleeper.waits and 2.0 <= sleeper.waits[0] < 3.0 + assert calls["count"] == 2 + assert "claude-remote-only-test" in litellm.model_cost + assert "claude-remote-only-test" in litellm.anthropic_models + assert "my-runtime-model" in litellm.model_cost + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + finally: + sleeper.release.set() + for thread in _retry_threads(): + thread.join(timeout=10) -def test_boot_load_does_not_retry_permanent_failures(): - """A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup.""" +def test_boot_load_does_not_retry_non_retryable_failure(): client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) sleeper = _SyncSleepRecorder() - get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + get_model_cost_map( + url=_URL, + sleep=sleeper, + rng=random.Random(0), + client=client, + ) assert calls["count"] == 1 assert sleeper.waits == [] - assert get_model_cost_map_source_info()["source"] == "local" - - get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0)) - assert sleeper.waits == [] - assert get_model_cost_map_source_info()["source"] == "local" + assert _retry_threads() == [] + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] is not None def test_boot_load_respects_local_env_override(monkeypatch): @@ -701,7 +738,9 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej ) get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=remote) shrunk_body = b'{"gpt-5.4-mini": {"mode": "chat", "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}}' - shrunk, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client) + shrunk, _ = _mock_client( + [httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client + ) get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=shrunk) @@ -711,3 +750,29 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej assert source["etag"] is None assert source["source_revision"] == _bundled_blob_id() assert source["source_revision"] != git_blob_id(shrunk_body) + + +@pytest.mark.parametrize( + ("argv0", "request_count"), + [ + ("/some/venv/bin/lite", 0), + ("/some/venv/bin/lite.exe", 0), + ("/some/venv/bin/python", 1), + ], +) +def test_boot_load_skips_remote_fetch_for_cli_processes( + monkeypatch: pytest.MonkeyPatch, argv0: str, request_count: int +) -> None: + monkeypatch.setattr(sys, "argv", [argv0, "--version"]) + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + + cost_map = get_model_cost_map(url=_URL, client=client) + + assert calls["count"] == request_count + assert cost_map + source = get_model_cost_map_source_info() + if request_count == 0: + assert source["source"] == "local" + else: + assert source["source"] == "remote" diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 0aa73833677..37e2031fdf4 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -6,7 +6,7 @@ import pytest import asyncio import traceback -from typing import Optional +from typing import Final, Optional import litellm from litellm import verbose_logger @@ -2633,6 +2633,48 @@ def test_dispatch_cached_response_extracts_delta( assert initialized_custom_stream_wrapper.response_id == "chatcmpl-cache-1" +def test_dispatch_cached_response_without_choices_is_an_empty_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """A cached completion with no choices replays as an empty, unfinished chunk + instead of raising IndexError on choices[0].""" + initialized_custom_stream_wrapper.custom_llm_provider = "cached_response" + chunk: Final = ModelResponseStream(id="chatcmpl-cache-empty", choices=[]) + + result, model_response, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, chunk + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] is None + assert initialized_custom_stream_wrapper.received_finish_reason is None + assert model_response.id == "chatcmpl-cache-empty" + + +@pytest.mark.asyncio +async def test_cached_response_without_choices_streams_a_single_stop_chunk( + logging_obj: Logging, +): + """A stream cache hit on a completion stored with choices == [] ends with one + finish_reason=stop chunk, the same shape the live empty stream produced.""" + + async def cached_chunks(): + yield ModelResponseStream(id="chatcmpl-cache-empty", choices=[]) + + wrapper: Final = CustomStreamWrapper( + completion_stream=cached_chunks(), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="cached_response", + ) + + chunks: Final = tuple([chunk async for chunk in wrapper]) + + assert len(chunks) == 1 + assert tuple(choice.finish_reason for chunk in chunks for choice in chunk.choices) == ("stop",) + assert all(choice.delta.content in (None, "") for chunk in chunks for choice in chunk.choices) + + def test_dispatch_vertex_ai_legacy_text_and_finish_reason( initialized_custom_stream_wrapper: CustomStreamWrapper, ): diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 4694fa8fbed..60f25c48443 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1,10 +1,16 @@ #### What this tests #### # This tests litellm.token_counter.token_counter() function +import asyncio +import base64 import importlib +import threading import time import traceback +from concurrent.futures import Future, wait +from typing import Final from unittest.mock import MagicMock +import anyio.to_thread import pytest import tiktoken @@ -14,9 +20,23 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import ( + _get_exact_count_function, + _get_extrapolating_count_function, + _get_tiktoken_count_function, + calculate_img_tokens, + high_detail_image_token_upper_bound, + offload_token_count, +) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text +from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, +) from tests.test_litellm.litellm_core_utils.messages_with_counts import ( MESSAGES_TEXT, MESSAGES_WITH_IMAGES, @@ -120,6 +140,135 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch): importlib.reload(litellm.constants) +async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): + warm_tokenizer("claude-fable-5") + + tokens, took, lags = await timed_with_loop_lags( + lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) + ) + + assert tokens > 0 + assert_loop_stayed_free(took, lags) + + +@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) +def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int): + count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) + front_heavy: Final = "a" * 1_000 + "b" * 4_000 + exact: Final = 1_000 + len(front_heavy) + + estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy) + + assert abs(estimate - exact) <= exact // 100 + assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars + + +def test_count_at_or_below_the_cap_is_exact(): + count_exactly: Final = MagicMock(side_effect=len) + + assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000 + assert count_exactly.call_args_list == [(("a" * 5_000,),)] + + +class _SlowEncoder: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.in_flight = 0 + self.peak_in_flight = 0 + + def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: + with self._lock: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + time.sleep(0.1) + with self._lock: + self.in_flight -= 1 + return [[0] * len(text) for text in texts] + + +@pytest.mark.asyncio +async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): + encoder: Final = _SlowEncoder() + count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) + shared_pool: Final = anyio.to_thread.current_default_thread_limiter() + burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]: + if counting.done(): + return () + await asyncio.sleep(0.01) + return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting)) + + counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst)))) + borrowed: Final = await shared_pool_borrowed_until_done(counting) + + assert await counting == [3] * burst + assert len(borrowed) > 1 and max(borrowed) == 0 + assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + +def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None: + def slow_count(counted: str) -> int: + time.sleep(0.1) + return len(counted) + + result.set_result(asyncio.run(offload_token_count(slow_count)(text))) + + +def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process(): + loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + results: Final = tuple(Future[int]() for _ in range(loops)) + threads: Final = tuple( + threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True) + for size, result in enumerate(results, start=1) + ) + for thread in threads: + thread.start() + + _, pending = wait(results, timeout=5) + + assert not pending + assert tuple(result.result() for result in results) == tuple(range(1, loops + 1)) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("8", 8), ("0", 4), ("not-an-int", 4)], +) +def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS") + importlib.reload(litellm.constants) + + +def test_token_counter_applies_the_default_cap(): + max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS + prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] + over_the_cap: Final = prose + "a" * 200_000 + exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap) + + estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap) + + assert estimate != exact + assert abs(estimate - exact) <= exact // 100 + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], +) +def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS") + importlib.reload(litellm.constants) + + def test_token_counter_with_prefix(): messages = [ {"role": "user", "content": "Who won the world cup in 2022?"}, @@ -1412,3 +1561,18 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): assert _count_user_content([prompt, named]) == _count_user_content( [prompt, {"type": "text", "text": "report.pdf"}] ) + + +def _png_data_url(width: int, height: int) -> str: + ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") + return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() + + +@pytest.mark.parametrize(("width", "height"), [(1, 1), (768, 768), (2000, 768), (768, 2000), (4096, 4096), (8000, 3072)]) +def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int, height: int) -> None: + assert calculate_img_tokens(_png_data_url(width, height), mode="high") <= high_detail_image_token_upper_bound() + + +def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None: + assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound() + assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index bf40f781fa3..9fe56f4dc65 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -315,6 +315,120 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert "event: message_start" in raw and "event: message_stop" in raw assert '"stop_reason": "end_turn"' in raw + @staticmethod + def _ended_tool_use_sse_chunks() -> list: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + return MaskArguments(guardrail_name="test") + + @staticmethod + def _partial_jsons(chunks: list) -> list: + return [ + json.loads(line[len("data:") :].strip())["delta"]["partial_json"] + for chunk in chunks + for line in chunk.decode().split("\n") + if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta" + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_input_back_into_sse_chunks(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', "", ""] + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert "persim" not in raw + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_use_name_back_into_sse_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call.function.name = "lookup_fruit_reviewed" + return inputs + + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + raw = b"".join(chunks).decode() + assert '"name": "lookup_fruit_reviewed"' in raw and '"id": "toolu_1"' in raw + assert '"name": "lookup_fruit"' not in raw + assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"} + + @pytest.mark.asyncio + async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_tool_use_sse_chunks() + original = [bytes(chunk) for chunk in chunks] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert chunks == original + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + server_tool_use = [ + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ] + tool_use = self._ended_tool_use_sse_chunks() + chunks = ( + tool_use[:1] + + [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in server_tool_use] + + [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]] + ) + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() @@ -2156,3 +2270,29 @@ class TestAnthropicMessagesHandlerStreamingScanKey: assert open_key == StreamingScanKey(texts=("hi",)) assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] assert ended_key != open_key + + +class TestAnthropicMessagesHandlerPostCallHookResponse: + def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self): + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + assembled = ModelResponse( + id="msg_1", + model="claude", + choices=[Choices(message=Message(role="assistant", content="hello world"), finish_reason="stop")], + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + ) + + hook_response = AnthropicMessagesHandler().post_call_hook_response(assembled) + + assert hook_response["type"] == "message" + assert hook_response["role"] == "assistant" + assert hook_response["content"] == [{"type": "text", "text": "hello world"}] + assert hook_response["stop_reason"] == "end_turn" + assert hook_response["usage"]["input_tokens"] == 1 + assert hook_response["usage"]["output_tokens"] == 2 + + def test_anything_else_reaches_the_hook_untouched(self): + native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]} + + assert AnthropicMessagesHandler().post_call_hook_response(native) is native diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index c59ec70b015..f6ee1cd71c0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -10,6 +10,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.litellm_core_utils.prompt_templates.factory import ( THOUGHT_SIGNATURE_SEPARATOR, @@ -41,6 +42,21 @@ from litellm.types.utils import ( ) +def test_translate_openai_response_to_anthropic_empty_choices() -> None: + response: Final = ModelResponse( + id="chatcmpl-empty", + model="gemini-3.5-flash", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10), + ) + + result: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "end_turn" + assert result["usage"]["input_tokens"] == 10 + + def test_translate_chat_refusal_to_anthropic_response(): response = ModelResponse( id="chatcmpl-refusal", @@ -408,6 +424,43 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks(): assert result[1]["tool_calls"][0]["id"] == "toolu_01234" +def test_translate_anthropic_messages_to_openai_drops_bridge_encrypted_reasoning_blocks(): + """A session that moves from an OpenAI reasoning model to a chat provider replays reasoning only OpenAI can read. + + Gemini rejects the whole request when such a block reaches it as a thought_signature, so the + adapter drops those blocks and keeps the provider-signed ones. + """ + + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "text", "text": "Who drinks water?"}], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The Norwegian."}, + ], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_signed"}, + {"type": "text", "text": "Still the Norwegian."}, + ], + ), + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages) + + assert [m["role"] for m in result] == ["user", "assistant", "assistant"] + assert not result[1].get("thinking_blocks") + assert result[1]["content"] == "The Norwegian." + assert [b["signature"] for b in result[2]["thinking_blocks"]] == ["EqQBCkYIAxgCIkA_signed"] + + def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py new file mode 100644 index 00000000000..c64e9d392e5 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py @@ -0,0 +1,50 @@ +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) + + +def _transform(messages): + return AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-5", + messages=messages, + anthropic_messages_optional_request_params={"max_tokens": 1024}, + litellm_params={}, + headers={}, + ) + + +def test_reasoning_replayed_from_the_responses_bridge_never_reaches_anthropic(): + """Claude Code resumed on a Claude model echoes the thinking blocks a gpt turn produced.""" + messages = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The answer."}, + ], + }, + {"role": "user", "content": "And the next one?"}, + ] + request = _transform(messages) + assert request["messages"][1]["content"] == [{"type": "text", "text": "The answer."}] + assert len(messages[1]["content"]) == 3 + + +def test_anthropic_signed_thinking_blocks_are_forwarded_untouched(): + messages = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."}, + ], + }, + ] + assert _transform(messages)["messages"] == messages diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py index 16e8cf0e90e..b66075f691b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py @@ -67,6 +67,39 @@ def test_build_responses_kwargs_prefers_explicit_prompt_cache_key_over_derived() assert responses_kwargs["prompt_cache_key"] == "explicit-key" +def test_build_responses_kwargs_asks_openai_for_encrypted_reasoning_without_thinking(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + extra_kwargs={"custom_llm_provider": "openai"}, + ) + assert responses_kwargs["include"] == ["reasoning.encrypted_content"] + assert "reasoning" not in responses_kwargs + + +def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_rejects_it(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="perplexity/sonar", + thinking={"type": "enabled", "budget_tokens": 4096}, + extra_kwargs={"custom_llm_provider": "perplexity"}, + ) + assert "include" not in responses_kwargs + assert "reasoning" in responses_kwargs + + +def test_build_responses_kwargs_keeps_the_deployment_include_next_to_encrypted_reasoning(): + responses_kwargs = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model="openai/gpt-5.6-luna", + extra_kwargs={"custom_llm_provider": "openai", "include": ["file_search_call.results"]}, + ) + assert responses_kwargs["include"] == ["reasoning.encrypted_content", "file_search_call.results"] + + def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key(): responses_kwargs = _build_responses_kwargs( max_tokens=1024, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index d9df5df426a..bfe2d6b7cea 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -10,6 +10,9 @@ from types import SimpleNamespace sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, +) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, ) @@ -114,7 +117,7 @@ class TestReasoningItemWithoutSummaryText: """ @staticmethod - def _gpt_turn(reasoning_summary_deltas: list) -> list: + def _gpt_turn(reasoning_summary_deltas: list, encrypted_content: str | None = None) -> list: return [ {"type": "response.created"}, {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, @@ -122,7 +125,10 @@ class TestReasoningItemWithoutSummaryText: {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": delta} for delta in reasoning_summary_deltas ), - {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + { + "type": "response.output_item.done", + "item": {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted_content}, + }, {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, @@ -171,6 +177,70 @@ class TestReasoningItemWithoutSummaryText: assert not [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] +_ENCRYPTED_REASONING = "gAAAAABp_encrypted_reasoning_bytes_only_openai_can_read" + + +class TestEncryptedReasoningIsStreamedForReplay: + """Regression for https://github.com/BerriAI/litellm/issues/40288. + + The client echoes a thinking block's signature (or a redacted block's data) back on the + next turn, so the item's ``encrypted_content`` has to reach it through one of those. + """ + + def test_encrypted_content_is_streamed_as_the_signature_before_the_block_closes(self): + chunks = _drain_async( + TestReasoningItemWithoutSummaryText._gpt_turn( + reasoning_summary_deltas=["Weighing options"], encrypted_content=_ENCRYPTED_REASONING + ) + ) + + assert [(c["type"], c.get("index"), c.get("delta", {}).get("type")) for c in chunks[1:5]] == [ + ("content_block_start", 0, None), + ("content_block_delta", 0, "thinking_delta"), + ("content_block_delta", 0, "signature_delta"), + ("content_block_stop", 0, None), + ] + assert chunks[3]["delta"]["signature"] == encrypted_reasoning_signature(_ENCRYPTED_REASONING) + + def test_reasoning_without_summary_streams_a_redacted_thinking_block(self): + chunks = _drain_async( + TestReasoningItemWithoutSummaryText._gpt_turn( + reasoning_summary_deltas=[], encrypted_content=_ENCRYPTED_REASONING + ) + ) + + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_stop", 0), + ("content_block_start", 1), + ("content_block_delta", 1), + ("content_block_stop", 1), + ] + assert chunks[1]["content_block"] == { + "type": "redacted_thinking", + "data": encrypted_reasoning_signature(_ENCRYPTED_REASONING), + } + + def test_summary_parts_are_separated_inside_the_one_thinking_block(self): + """Two summary parts read as two paragraphs, not as one run-on sentence.""" + events = [ + {"type": "response.created"}, + {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, + {"type": "response.reasoning_summary_part.added", "item_id": "rs_1", "summary_index": 0}, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "First."}, + {"type": "response.reasoning_summary_part.added", "item_id": "rs_1", "summary_index": 1}, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "Second."}, + {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + ] + chunks = _process_all(events) + + thinking = "".join( + c["delta"]["thinking"] for c in chunks if c.get("delta", {}).get("type") == "thinking_delta" + ) + assert thinking == "First.\n\nSecond." + assert [c["type"] for c in chunks].count("content_block_start") == 1 + + class TestToolUseBlockClosedExactlyOnce: """Regression for https://github.com/BerriAI/litellm/issues/37273. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index edcc7adddb7..4ad559aa547 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -19,6 +19,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, @@ -566,6 +567,66 @@ class TestTranslateMessagesToResponsesInput: result = _translate_messages(messages) assert "id" not in result[0] + def test_thinking_block_with_encrypted_signature_replays_the_encrypted_content(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288 (inbound fault site).""" + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "Private reasoning.", + "signature": encrypted_reasoning_signature("gAAAA_turn_one"), + } + ], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Private reasoning."}], + "encrypted_content": "gAAAA_turn_one", + } + ] + + def test_redacted_thinking_with_encrypted_data_replays_the_encrypted_content(self): + messages = [ + { + "role": "assistant", + "content": [{"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_turn_one")}], + } + ] + result = _translate_messages(messages) + assert result == [{"type": "reasoning", "summary": [], "encrypted_content": "gAAAA_turn_one"}] + + def test_each_encrypted_thinking_block_stays_its_own_reasoning_item(self): + """Two upstream items must not be merged into one, or the encrypted content of one is lost.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "First.", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "thinking", "thinking": "Second.", "signature": encrypted_reasoning_signature("gAAAA_2")}, + ], + } + ] + result = _translate_messages(messages) + assert [item["encrypted_content"] for item in result] == ["gAAAA_1", "gAAAA_2"] + + def test_anthropic_signed_thinking_block_replays_as_a_summary_only_item(self): + """A real Anthropic signature is opaque here, so it never masquerades as encrypted content.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "Private reasoning.", "signature": "ErcBCkgIValid"}], + } + ] + result = _translate_messages(messages) + assert result == [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "Private reasoning."}]} + ] + def test_consecutive_thinking_blocks_become_one_reasoning_item(self): """Summary parts of one upstream reasoning item are regrouped into that item.""" messages = [ @@ -1102,6 +1163,23 @@ class TestTranslateRequestBroaderCoverage: kwargs = _ADAPTER.translate_request(req) assert "reasoning" not in kwargs + def test_thinking_asks_for_the_encrypted_reasoning(self): + """The documented way to get reasoning that survives store=false is to ask for it.""" + req = _make_request(thinking={"type": "enabled", "budget_tokens": 12000}) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["include"] == ["reasoning.encrypted_content"] + + def test_encrypted_reasoning_is_asked_for_without_a_thinking_block(self): + """A reasoning model reasons whether or not the client sent `thinking`, so the replay needs it either way.""" + kwargs = _ADAPTER.translate_request(_make_request()) + assert kwargs["include"] == ["reasoning.encrypted_content"] + + def test_encrypted_reasoning_is_not_asked_for_when_the_provider_rejects_include(self): + req = _make_request(thinking={"type": "enabled", "budget_tokens": 12000}) + kwargs = _ADAPTER.translate_request(req, include_encrypted_reasoning=False) + assert kwargs["reasoning"] == {"effort": "high"} + assert "include" not in kwargs + def test_metadata_user_id_mapped_to_user(self): req = _make_request(metadata={"user_id": "user-42"}) kwargs = _ADAPTER.translate_request(req) @@ -1246,7 +1324,9 @@ def _make_function_call_item(call_id: str, name: str, arguments: str) -> MagicMo return item -def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> MagicMock: +def _make_reasoning_item( + summaries: List[str], item_id: str = "rs_test_1", encrypted_content: str | None = None +) -> MagicMock: """Build a mock ResponseReasoningItem.""" from openai.types.responses import ResponseReasoningItem # type: ignore[import] @@ -1259,9 +1339,13 @@ def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> Ma item = MagicMock(spec=ResponseReasoningItem) item.id = item_id item.summary = summary_mocks + item.encrypted_content = encrypted_content return item +_ENCRYPTED_REASONING = "gAAAAABp_encrypted_reasoning_bytes_only_openai_can_read" + + class TestTranslateResponse: """Responses API -> AnthropicMessagesResponse conversion.""" @@ -1386,7 +1470,81 @@ class TestTranslateResponse: reasoning = _make_reasoning_item(["Part one.", "Part two."], item_id="rs_abc123") response = _make_mock_response(output=[reasoning]) result: Any = _ADAPTER.translate_response(response) - assert [block["signature"] for block in result["content"]] == [None, None] + assert [block["signature"] for block in result["content"]] == [None] + assert "rs_abc123" not in json.dumps(result["content"]) + + def test_summary_parts_join_into_one_thinking_block(self): + """One reasoning item is one block, so its signature is echoed back exactly once.""" + reasoning = _make_reasoning_item(["Part one.", "Part two."]) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert [block["thinking"] for block in result["content"]] == ["Part one.\n\nPart two."] + + def test_encrypted_content_rides_the_thinking_signature(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288 (outbound fault site).""" + reasoning = _make_reasoning_item(["Part one."], item_id="rs_abc123", encrypted_content=_ENCRYPTED_REASONING) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + { + "type": "thinking", + "thinking": "Part one.", + "signature": encrypted_reasoning_signature(_ENCRYPTED_REASONING), + } + ] + + def test_reasoning_without_summary_becomes_redacted_thinking(self): + """With summaries off the encrypted reasoning still has to reach the client to be replayed.""" + reasoning = _make_reasoning_item([], encrypted_content=_ENCRYPTED_REASONING) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature(_ENCRYPTED_REASONING)} + ] + + def test_dict_reasoning_item_carries_its_encrypted_content(self): + response = _make_mock_response( + output=[ + { + "type": "reasoning", + "id": "rs_dict_1", + "encrypted_content": _ENCRYPTED_REASONING, + "summary": [{"type": "summary_text", "text": "Weighing the options."}], + } + ] + ) + result: Any = _ADAPTER.translate_response(response) + assert result["content"][0]["signature"] == encrypted_reasoning_signature(_ENCRYPTED_REASONING) + + def test_reasoning_item_round_trip_is_byte_stable(self): + """Regression for https://github.com/BerriAI/litellm/issues/40288. + + The reasoning item the next turn replays must be the one OpenAI produced, with its + encrypted reasoning intact, and identical on every later turn so the prompt cache + prefix keeps matching. + """ + reasoning = _make_reasoning_item(["Part one.", "Part two."], encrypted_content=_ENCRYPTED_REASONING) + turn: Any = _ADAPTER.translate_response(_make_mock_response(output=[reasoning])) + history = [{"role": "assistant", "content": turn["content"]}] + + replayed_items = [_translate_messages(history) for _ in range(2)] + + assert replayed_items[0] == replayed_items[1] + assert replayed_items[0] == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Part one.\n\nPart two."}], + "encrypted_content": _ENCRYPTED_REASONING, + } + ] + + def test_redacted_reasoning_round_trip_replays_the_encrypted_content(self): + reasoning = _make_reasoning_item([], encrypted_content=_ENCRYPTED_REASONING) + turn: Any = _ADAPTER.translate_response(_make_mock_response(output=[reasoning])) + + replayed = _translate_messages([{"role": "assistant", "content": turn["content"]}]) + + assert replayed == [{"type": "reasoning", "summary": [], "encrypted_content": _ENCRYPTED_REASONING}] def test_dict_reasoning_item_becomes_thinking_block(self): """A reasoning item arriving as a plain dict is kept, not dropped.""" @@ -1402,14 +1560,26 @@ class TestTranslateResponse: result: Any = _ADAPTER.translate_response(response) assert result["content"] == [{"type": "thinking", "thinking": "Weighing the options.", "signature": None}] - def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self): + @pytest.mark.parametrize( + ("summaries", "encrypted_content"), + [ + (["Part one."], None), + (["Part one."], _ENCRYPTED_REASONING), + ([], _ENCRYPTED_REASONING), + ], + ids=["unsigned_thinking", "encrypted_thinking", "encrypted_redacted_thinking"], + ) + def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self, summaries, encrypted_content): """Replaying this turn to an Anthropic model must not send a signature it cannot verify.""" from litellm.litellm_core_utils.prompt_templates.factory import ( _drop_unsignable_thinking_blocks, ) - response = _make_mock_response(output=[_make_reasoning_item(["Part one."], item_id="rs_abc123")]) + response = _make_mock_response( + output=[_make_reasoning_item(summaries, item_id="rs_abc123", encrypted_content=encrypted_content)] + ) result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 assert _drop_unsignable_thinking_blocks(result["content"]) == [] def test_usage_mapped_correctly(self): 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 ae620fdd6dc..e1b39c4ba13 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -42,12 +42,15 @@ FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789" def test_is_claude_code_one_shot_subagent_request(messages, system, expected): from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request - assert is_claude_code_one_shot_subagent_request( - messages=messages, - system=system, - tools=None, - user_agent="claude-cli/2.1.263 (external, cli)", - ) is expected + assert ( + is_claude_code_one_shot_subagent_request( + messages=messages, + system=system, + tools=None, + user_agent="claude-cli/2.1.263 (external, cli)", + ) + is expected + ) class TestOptionallyHandleAnthropicOAuth: @@ -1541,6 +1544,71 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_content_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["thinking"] + def test_strip_keeps_encrypted_reasoning_blocks_for_the_responses_bridge(self): + """The /v1/messages handler runs this before dispatch, so the bridge must still see the replay.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, + ) + from litellm.llms.anthropic.common_utils import ( + strip_empty_content_blocks_from_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + {"type": "text", "text": "The answer."}, + ], + } + ] + assert strip_empty_content_blocks_from_anthropic_messages(msgs) == msgs + + def test_strip_encrypted_reasoning_drops_only_the_bridge_tagged_blocks(self): + """A session resumed on an Anthropic model replays reasoning only OpenAI can verify.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + encrypted_reasoning_signature, + ) + from litellm.llms.anthropic.common_utils import ( + strip_encrypted_reasoning_blocks_from_anthropic_messages, + ) + + msgs = [ + {"role": "user", "content": "Solve it."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_1")}, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_2")}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan", "signature": encrypted_reasoning_signature("gAAAA_3")}, + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."}, + ], + }, + ] + out = strip_encrypted_reasoning_blocks_from_anthropic_messages(msgs) + assert [m["role"] for m in out] == ["user", "assistant"] + assert [b["type"] for b in out[1]["content"]] == ["thinking", "redacted_thinking", "text"] + assert out[1]["content"][0]["signature"] == "EqQBCkYIAxgCIkA_anthropic_signed" + assert len(msgs[1]["content"]) == 2 + assert len(msgs[2]["content"]) == 4 + + def test_strip_encrypted_reasoning_leaves_malformed_messages_for_the_provider_to_reject(self): + """A bare string in messages must reach Anthropic as a 400, not die in the stripper as a 500.""" + from litellm.llms.anthropic.common_utils import ( + strip_encrypted_reasoning_blocks_from_anthropic_messages, + ) + + msgs = ["hi", {"role": "user", "content": "hello"}] + assert strip_encrypted_reasoning_blocks_from_anthropic_messages(msgs) == msgs + def test_strip_empty_text_blocks_treats_null_text_as_empty(self): from litellm.llms.anthropic.common_utils import ( strip_empty_content_blocks_from_anthropic_messages, @@ -2199,3 +2267,25 @@ def test_create_anthropic_model_list_response_empty(): assert response["has_more"] is False assert response["first_id"] is None assert response["last_id"] is None + + +def test_create_anthropic_model_list_response_lists_ids_as_told(): + """listed_ids renames an entry for the caller while display_name and every other field stay keyed to the served + id, and the envelope's first/last ids follow the renamed entries.""" + from litellm.llms.anthropic.common_utils import ( + create_anthropic_model_list_response, + ) + + response = create_anthropic_model_list_response( + [ + {"id": "gpt-4o", "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": 1000000}, + {"id": "claude-haiku-4-5", "object": "model", "created": 0, "owned_by": "openai"}, + ], + display_names={"gpt-4o": "GPT 4o"}, + listed_ids={"gpt-4o": "claude-router-gpt-4o[1m]"}, + ) + + gpt, haiku = response["data"] + assert (gpt["id"], gpt["display_name"], gpt["max_input_tokens"]) == ("claude-router-gpt-4o[1m]", "GPT 4o", 1000000) + assert (haiku["id"], haiku["display_name"]) == ("claude-haiku-4-5", "claude-haiku-4-5") + assert (response["first_id"], response["last_id"]) == ("claude-router-gpt-4o[1m]", "claude-haiku-4-5") diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index 4e6b9ed0188..bc6cb0c0fed 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -198,6 +198,9 @@ def test_azure_gpt_5_takes_the_reasoning_path() -> None: assert "reasoning_effort" in supported +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestAzureToolSchemaCombinatorFlattening: """ Regression tests for LIT-6510: Azure's chat completions validator rejects @@ -259,6 +262,26 @@ class TestAzureToolSchemaCombinatorFlattening: self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) assert tool == self._anyof_tool() + def test_transform_request_drops_non_python_regex_pattern(self): + tool = { + "type": "function", + "function": { + "name": "Artifact", + "parameters": { + "type": "object", + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + }, + } + + request = self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) + + assert request["tools"][0]["function"]["parameters"] == { + "type": "object", + "properties": {"field": {"type": "string"}}, + } + assert tool["function"]["parameters"]["properties"]["field"]["pattern"] == _ARTIFACT_FIELD_PATTERN + def test_clean_object_schema_passes_through_as_same_object(self): tool = { "type": "function", diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index 29b74c2ee4a..c7e86616ee2 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -1,11 +1,20 @@ import json +from datetime import datetime from unittest.mock import MagicMock import httpx +import pytest - -from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig -from litellm.types.utils import ModelResponse +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound +from litellm.llms.azure.passthrough.transformation import ( + AzurePassthroughConfig, + azure_router_model_in_endpoint, + foreign_azure_deployment, +) +from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.utils import EmbeddingResponse, ModelResponse def _azure_chat_completion_body(): @@ -73,22 +82,408 @@ def test_azure_passthrough_logging_non_streaming_response_chat_completions(): assert result.usage.total_tokens == 18 -def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): - """ - Endpoints other than chat/completions (responses, messages, images) fall - through to None — matches base-class behavior and Bedrock's "unknown - endpoint" handling. Not a regression; just scoping. - """ - config = AzurePassthroughConfig() - logging_obj = MagicMock() - - result = config.logging_non_streaming_response( - model="gpt-4.1-mini", +def _relay_logging_obj(model: str) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": "https://my-resource.openai.azure.com", "custom_llm_provider": "azure"}, + optional_params={}, custom_llm_provider="azure", - httpx_response=_make_httpx_response(_azure_chat_completion_body()), + ) + return logging_obj + + +def _relay_logging_result(model: str, endpoint: str, body, status_code: int = 200): + logging_obj = _relay_logging_obj(model) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request( + "POST", f"https://my-resource.openai.azure.com/{endpoint}?api-version=2025-04-01-preview" + ), + ) + result = AzurePassthroughConfig().logging_non_streaming_response( + model=model, + custom_llm_provider="azure", + httpx_response=response, request_data={}, logging_obj=logging_obj, - endpoint="openai/responses", + endpoint=endpoint, + ) + return result, logging_obj + + +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1000, "total_tokens": 1000}, +} + +RESPONSES_BODY = { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4.1-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, +} + + +def test_azure_passthrough_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", "openai/deployments/text-embedding-3-small/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure/text-embedding-3-small")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1000 * per_token) + + +def test_azure_passthrough_responses_relay_is_costed_per_token(): + result, logging_obj = _relay_logging_result("gpt-4.1-mini", "openai/responses", RESPONSES_BODY) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(result, ResponsesAPIResponse) + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_failed_embeddings_relay_is_not_costed(): + result, logging_obj = _relay_logging_result( + "text-embedding-3-small", + "openai/deployments/text-embedding-3-small/embeddings", + {"error": {"code": "429", "message": "rate limited"}}, + status_code=429, ) assert result is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_azure_passthrough_logging_non_streaming_response_unknown_endpoint_returns_none(): + result, logging_obj = _relay_logging_result( + "gpt-4o-mini-tts", "openai/deployments/gpt-4o-mini-tts/audio/speech", {"audio": "..."} + ) + + assert result is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def _sse_line(payload: dict) -> str: + return "data: " + json.dumps(payload) + + +def _azure_chat_completion_chunks() -> list[str]: + head = {"id": "chatcmpl-abc123", "object": "chat.completion.chunk", "created": 1700000000, "model": "gpt-4.1-mini"} + return [ + _sse_line( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}], + } + ), + _sse_line( + {**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]} + ), + _sse_line({**head, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), + _sse_line({**head, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}}), + "data: [DONE]", + ] + + +def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens == 10 + assert response.usage.completion_tokens == 8 + + +def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request(): + messages = [{"role": "user", "content": "Say hi in three words"}] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter(model="gpt-4.1-mini", messages=messages) + assert response.usage.completion_tokens > 0 + + +def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_without_fetching_the_image(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png", "detail": "high"}}, + ], + } + ] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + text_only_messages = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] + assert isinstance(response, ModelResponse) + assert response.usage.prompt_tokens == ( + litellm.token_counter(model="gpt-4.1-mini", messages=text_only_messages) + high_detail_image_token_upper_bound() + ) + + +def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_chat_completion_chunks(), + litellm_logging_obj=MagicMock(), + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/embeddings", + ) + + assert response is None + + +def _azure_responses_stream_chunks(terminal_event: str | None = "response.completed") -> list[str]: + in_progress = {**RESPONSES_BODY, "status": "in_progress", "output": [], "usage": None} + events = [ + ("response.created", {"type": "response.created", "sequence_number": 0, "response": in_progress}), + ( + "response.output_text.delta", + {"type": "response.output_text.delta", "sequence_number": 1, "item_id": "msg_1", "delta": "hi"}, + ), + ] + ( + [(terminal_event, {"type": terminal_event, "sequence_number": 2, "response": RESPONSES_BODY})] + if terminal_event + else [] + ) + return [line for name, payload in events for line in (f"event: {name}", _sse_line(payload))] + + +def test_azure_passthrough_streaming_responses_chunks_are_costed_per_token(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + info = litellm.get_model_info("azure/gpt-4.1-mini") + + assert isinstance(response, ResponseCompletedEvent) + assert response.response.usage.input_tokens == 1000 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +def test_azure_passthrough_streaming_responses_without_a_terminal_event_are_not_costed(): + logging_obj = _relay_logging_obj("gpt-4.1-mini") + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=_azure_responses_stream_chunks(terminal_event=None), + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/responses", + ) + + assert response is None + assert logging_obj.call_type == "allm_passthrough_route" + + +def _complete_url(request_query_params: dict, litellm_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params=request_query_params, + litellm_params=litellm_params, + ) + return url + + +def test_azure_passthrough_url_forwards_the_callers_api_version(): + url = _complete_url(request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={}) + + assert url.path == "/openai/deployments/gpt-4.1-mini/chat/completions" + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_prefers_the_callers_api_version_over_the_deployments(): + url = _complete_url( + request_query_params={"api-version": "2025-04-01-preview"}, litellm_params={"api_version": "2024-10-21"} + ) + + assert url.params["api-version"] == "2025-04-01-preview" + + +def test_azure_passthrough_url_fills_in_the_deployments_api_version_when_the_caller_sends_none(): + url = _complete_url(request_query_params={}, litellm_params={"api_version": "2024-10-21"}) + + assert url.params["api-version"] == "2024-10-21" + + +FULL_URL_API_BASE = ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" +) + + +def _full_url_complete_url(request_query_params: dict) -> httpx.URL: + url, _ = AzurePassthroughConfig().get_complete_url( + api_base=FULL_URL_API_BASE, + api_key="key", + model="gpt-4.1-mini", + endpoint="chat/completions", + request_query_params=request_query_params, + litellm_params={}, + ) + return url + + +def test_azure_passthrough_url_prefers_the_callers_api_version_over_a_full_url_api_bases(): + url = _full_url_complete_url(request_query_params={"api-version": "2025-04-01-preview"}) + + assert str(url) == ( + "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions" + "?api-version=2025-04-01-preview" + ) + + +def test_azure_passthrough_url_keeps_a_full_url_api_bases_api_version_when_the_caller_sends_none(): + url = _full_url_complete_url(request_query_params={}) + + assert url.params["api-version"] == "2024-10-21" + + +def test_azure_passthrough_url_strips_the_leading_router_model_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt-4.1-mini/openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) + + +def test_azure_passthrough_url_rewrites_the_model_group_only_as_a_whole_segment(): + url, _ = AzurePassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com", + api_key="key", + model="gpt-4.1-mini", + endpoint="gpt/openai/deployments/gpt-4.1-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "gpt"}}, + ) + + assert ( + str(url) + == "https://my-resource.openai.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-10-21" + ) + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], +) +def test_azure_passthrough_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert ( + AzurePassthroughConfig().is_streaming_request( + endpoint="openai/deployments/x/chat/completions", request_data=request_data + ) + is expected + ) + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("gpt/openai/deployments/gpt/chat/completions", None), + ("openai/deployments/gpt/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4-mini/chat/completions", None), + ("gpt/openai/deployments/GPT-5.4-MINI/chat/completions", None), + ("gpt/openai/deployments/Gpt/chat/completions", "Gpt"), + ("gpt/models/chat/completions", None), + ("gpt/openai/deployments/gpt-5.4/chat/completions", "gpt-5.4"), + ("gpt/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/victim/gpt/chat/completions", "victim"), + ("gpt/openai/deployments/GPT-5.4/chat/completions", "GPT-5.4"), + ], +) +def test_foreign_azure_deployment_names_a_segment_outside_the_group(endpoint, expected): + assert foreign_azure_deployment(endpoint, "gpt", lambda: frozenset({"gpt-5.4-mini"})) == expected + + +def test_foreign_azure_deployment_skips_the_router_when_the_segment_is_the_group_itself(): + def served_models(): + raise AssertionError("the router must not be consulted for the group's own name") + + assert foreign_azure_deployment("gpt/openai/deployments/gpt/chat/completions", "gpt", served_models) is None + + +@pytest.mark.parametrize( + "endpoint, expected", + [ + ("other-group/openai/deployments/other-group/chat/completions", "other-group"), + ("openai/deployments/gpt/chat/completions", "gpt"), + ("openai/deployments/my-azure-deployment/chat/completions", None), + ("gpt", None), + ], +) +def test_azure_router_model_in_endpoint_picks_the_first_router_model_segment(endpoint, expected): + assert azure_router_model_in_endpoint(endpoint, frozenset({"gpt", "other-group"})) == expected diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index 0cac2705ab0..726c9f65681 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -1,4 +1,5 @@ from copy import deepcopy +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -243,6 +244,9 @@ def test_provider_config_manager_o_series_selection(): assert not isinstance(default_config, AzureOpenAIOSeriesResponsesAPIConfig) +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestAzureResponsesAPIConfig: def setup_method(self): self.config = AzureOpenAIResponsesAPIConfig() @@ -599,6 +603,31 @@ class TestAzureResponsesAPIConfig: assert result["tools"][0] is tool assert "anyOf" in result["tools"][0]["parameters"] + def test_azure_drops_non_python_regex_pattern_while_keeping_gpt5_combinators(self): + tool = { + "type": "function", + "name": "Artifact", + "parameters": { + "type": "object", + "anyOf": [{"properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}}], + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + }, + } + + result = self.config.transform_responses_api_request( + model="my-eastus-deployment", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(model_info={"base_model": "azure/gpt-5.4-mini"}), + headers={}, + ) + + assert result["tools"][0]["parameters"] == { + "type": "object", + "anyOf": [{"properties": {"field": {"type": "string"}}}], + "properties": {"field": {"type": "string"}}, + } + def test_azure_keeps_combinators_for_unrecognized_deployment_without_base_model(self): tool = self._anyof_tool() diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index f000abb4c9a..c959c201ccb 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -385,6 +385,58 @@ def test_select_azure_base_url_called(setup_mocks): setup_mocks["select_url"].assert_called_once() +def test_initialize_defaults_max_retries_to_litellm_default(setup_mocks): + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={}, + api_key="test-api-key", + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version="2023-06-01", + is_async=False, + ) + + assert result["max_retries"] == litellm.constants.DEFAULT_MAX_RETRIES + + +@pytest.mark.parametrize( + "configured, expected", + [(0, 0), (5, 5), (None, litellm.constants.DEFAULT_MAX_RETRIES)], +) +def test_initialize_honors_explicit_max_retries(setup_mocks, configured, expected): + result = BaseAzureLLM().initialize_azure_sdk_client( + litellm_params={"max_retries": configured}, + api_key="test-api-key", + api_base="https://test.openai.azure.com", + model_name="gpt-4", + api_version="2023-06-01", + is_async=False, + ) + + assert result["max_retries"] == expected + + +def test_default_max_retries_env_var_reaches_azure_sdk_client(): + import subprocess + import sys + + code = ( + "from litellm.llms.azure.common_utils import BaseAzureLLM\n" + "client = BaseAzureLLM().get_azure_openai_client(" + "api_key='test-api-key', api_base='https://test.openai.azure.com', api_version='2024-02-01'," + " client=None, _is_async=True, litellm_params={}, model='gpt-4')\n" + "print(client.max_retries)" + ) + completed = subprocess.run( + [sys.executable, "-c", code], + env={**os.environ, "DEFAULT_MAX_RETRIES": "0"}, + capture_output=True, + text=True, + check=True, + ) + + assert completed.stdout.strip() == "0" + + @pytest.mark.parametrize( "call_type", [ diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index 284a912d9a4..75e046825a3 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -6,6 +6,7 @@ import pytest import litellm +from litellm.images.utils import ImageEditRequestUtils from litellm.llms.azure_ai.image_edit import ( AzureFoundryMAIImageEditConfig, get_azure_ai_image_edit_config, @@ -70,44 +71,48 @@ class TestAzureMAIImageEdit: assert "/mai/v1/images/edits" in url assert "api-version=preview" in url - def test_map_openai_params_keeps_size(self): - config = AzureFoundryMAIImageEditConfig() - optional_params = config.map_openai_params( - image_edit_optional_params={"size": "1792x1024", "n": 1}, + def test_get_optional_params_image_edit_size_raises_400(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError, match="size") as exc_info: + ImageEditRequestUtils.get_optional_params_image_edit( + model="MAI-Image-2.5", + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={"size": "1024x1024", "n": 1}, + ) + assert exc_info.value.status_code == 400 + + def test_get_optional_params_image_edit_size_dropped_with_drop_params(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + optional_params = ImageEditRequestUtils.get_optional_params_image_edit( model="MAI-Image-2.5", + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={"size": "1024x1024", "n": 1}, drop_params=True, ) - assert optional_params["size"] == "1792x1024" + assert "size" not in optional_params assert optional_params["n"] == 1 - assert "width" not in optional_params - assert "height" not in optional_params - def test_map_openai_params_defaults_size(self): - config = AzureFoundryMAIImageEditConfig() - optional_params = config.map_openai_params( - image_edit_optional_params={}, + def test_get_optional_params_image_edit_without_size_forwards_nothing_extra(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + optional_params = ImageEditRequestUtils.get_optional_params_image_edit( model="MAI-Image-2.5", - drop_params=True, + image_edit_provider_config=AzureFoundryMAIImageEditConfig(), + image_edit_optional_params={}, ) - assert optional_params["size"] == "1024x1024" + assert optional_params == {} - def test_map_openai_params_unsupported_size_raises(self): - config = AzureFoundryMAIImageEditConfig() - with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): - config.map_openai_params( - image_edit_optional_params={"size": "auto"}, - model="MAI-Image-2.5", - drop_params=True, - ) - - def test_map_openai_params_invalid_size_format_raises(self): - config = AzureFoundryMAIImageEditConfig() - with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): - config.map_openai_params( - image_edit_optional_params={"size": "1024xabc"}, - model="MAI-Image-2.5", - drop_params=True, + def test_image_edit_size_surfaces_as_400(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_edit( + model="azure_ai/MAI-Image-2.5", + image=io.BytesIO(b"fake-image-bytes"), + prompt="Turn this into a studio product shot", + size="1024x1024", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", ) + assert exc_info.value.status_code == 400 def test_transform_image_edit_request_uses_image_field(self): config = AzureFoundryMAIImageEditConfig() @@ -117,14 +122,14 @@ class TestAzureMAIImageEdit: model="MAI-Image-2.5", prompt="Turn this into a studio product shot", image=image_bytes, - image_edit_optional_request_params={"size": "1024x1024", "n": 1}, + image_edit_optional_request_params={"n": 1}, litellm_params={}, headers={}, ) assert data["model"] == "MAI-Image-2.5" assert data["prompt"] == "Turn this into a studio product shot" - assert data["size"] == "1024x1024" + assert "size" not in data assert data["n"] == 1 assert len(files) == 1 assert files[0][0] == "image" diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 9bdc79919d2..55656b97c57 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -3,8 +3,8 @@ from unittest.mock import MagicMock import httpx import pytest - import litellm +from litellm.exceptions import UnsupportedParamsError from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.image_generation import get_azure_image_generation_config from litellm.llms.azure.image_generation.http_utils import ( @@ -29,9 +29,7 @@ from litellm.utils import get_optional_params_image_gen class TestAzureMAIImageGeneration: def test_is_mai_model(self): assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5") - assert AzureFoundryMAIImageGenerationConfig.is_mai_model( - "azure_ai/MAI-Image-2.5" - ) + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("azure_ai/MAI-Image-2.5") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") @@ -42,16 +40,10 @@ class TestAzureMAIImageGeneration: api_base="https://my-resource.services.ai.azure.com", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_mai_image_generation_url_preserves_full_path(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( api_base=api, api_version="preview", @@ -63,10 +55,7 @@ class TestAzureMAIImageGeneration: api_base="https://my-resource.services.ai.azure.com/mai/v1", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_azure_ai_image_generation_config_returns_mai(self): config = get_azure_ai_image_generation_config("MAI-Image-2.5") @@ -104,13 +93,13 @@ class TestAzureMAIImageGeneration: config = AzureFoundryMAIImageGenerationConfig() optional_params = get_optional_params_image_gen( model="MAI-Image-2.5", - size="1792x1024", + size="1024x1024", n=1, custom_llm_provider="azure_ai", provider_config=config, drop_params=True, ) - assert optional_params["width"] == 1792 + assert optional_params["width"] == 1024 assert optional_params["height"] == 1024 assert "size" not in optional_params @@ -127,10 +116,7 @@ class TestAzureMAIImageGeneration: assert "api-version=preview" in url def test_mai_json_body_keeps_model(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" data = { "model": "MAI-Image-2.5", "prompt": "A photograph of a red fox", @@ -176,7 +162,7 @@ class TestAzureMAIImageGeneration: def test_map_openai_params_unsupported_size_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): + with pytest.raises(UnsupportedParamsError, match="Unsupported size value: 'auto'"): config.map_openai_params( non_default_params={"size": "auto"}, optional_params={}, @@ -186,7 +172,7 @@ class TestAzureMAIImageGeneration: def test_map_openai_params_invalid_custom_size_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): + with pytest.raises(UnsupportedParamsError, match="Invalid size format: '1024xabc'"): config.map_openai_params( non_default_params={"size": "1024xabc"}, optional_params={}, @@ -194,9 +180,138 @@ class TestAzureMAIImageGeneration: drop_params=True, ) + @pytest.mark.parametrize("size", ["512x512", "256x256", "700x1400"]) + def test_map_openai_params_size_below_minimum_dimension_raises(self, size): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="at least 768 pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + @pytest.mark.parametrize("size", ["1792x1024", "1024x1792"]) + def test_map_openai_params_size_over_total_pixel_budget_raises(self, size): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + @pytest.mark.parametrize("size", ["1032x1024", "1376x768"]) + def test_map_openai_params_size_at_live_pixel_cap_passes_through(self, size): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert optional_params["width"] * optional_params["height"] == 1_056_768 + + def test_map_openai_params_size_one_pixel_over_live_cap_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"): + config.map_openai_params( + non_default_params={"size": "1033x1024"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_map_openai_params_explicit_width_height_not_range_checked(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"width": 1792, "height": 1024}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == 1024 + + @pytest.mark.parametrize("n", [2, 4, "2", 0, -1]) + def test_map_openai_params_n_other_than_one_raises(self, n): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"): + config.map_openai_params( + non_default_params={"n": n}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_map_openai_params_non_numeric_n_raises_400(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(UnsupportedParamsError, match="not a whole number of images") as exc_info: + config.map_openai_params( + non_default_params={"n": "abc"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert exc_info.value.status_code == 400 + + def test_get_optional_params_image_gen_global_drop_params_drops_multi_image_n(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", True) + optional_params = get_optional_params_image_gen( + model="MAI-Image-2.5", + n=4, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryMAIImageGenerationConfig(), + ) + assert "n" not in optional_params + assert optional_params["width"] == 1024 + + def test_get_optional_params_image_gen_without_any_drop_params_still_raises(self, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(UnsupportedParamsError, match="returns exactly 1 image per request"): + get_optional_params_image_gen( + model="MAI-Image-2.5", + n=4, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryMAIImageGenerationConfig(), + ) + + def test_map_openai_params_multi_image_n_dropped_with_drop_params(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 4}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert "n" not in optional_params + + def test_map_openai_params_single_image_n_still_passes_through(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 1}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert optional_params["n"] == 1 + + @pytest.mark.parametrize("params", [{"n": 2}, {"n": "abc"}, {"size": "512x512"}, {"size": "1792x1024"}]) + def test_image_generation_rejected_params_surface_as_400(self, params): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_generation( + model="azure_ai/MAI-Image-2.5", + prompt="A photograph of a red fox", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", + **params, + ) + assert exc_info.value.status_code == 400 + def test_map_openai_params_unsupported_param_raises(self): config = AzureFoundryMAIImageGenerationConfig() - with pytest.raises(ValueError, match="Parameter quality is not supported"): + with pytest.raises(UnsupportedParamsError, match="Parameter quality is not supported"): config.map_openai_params( non_default_params={"quality": "hd"}, optional_params={}, @@ -343,16 +458,12 @@ class TestAzureMAIImageGeneration: litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") - image_response = ImageResponse( - data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] - ) + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) cost = azure_ai_image_cost_calculator( model=model, image_response=image_response, ) - assert ( - cost == len(image_response.data or []) * model_info["output_cost_per_image"] - ) + assert cost == len(image_response.data or []) * model_info["output_cost_per_image"] assert cost > 0 diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py new file mode 100644 index 00000000000..c8007acf70f --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -0,0 +1,617 @@ +import json +from datetime import datetime +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.azure_ai.passthrough.transformation import AzureAIPassthroughConfig +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.types.rerank import RerankResponse +from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, ModelResponse +from litellm.utils import ProviderConfigManager + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" +RESPONSES_COMPLETED_EVENT = { + "type": "response.completed", + "sequence_number": 2, + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100}, + }, +} + + +class _SpendProbe(CustomLogger): + logged_call_type: str | None = None + logged_cost: float | None = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_call_type = kwargs["call_type"] + self.logged_cost = kwargs["response_cost"] + + +@pytest.fixture(autouse=True) +def clear_azure_ai_env(monkeypatch): + for env_var in ("AZURE_AI_API_BASE", "AZURE_AI_API_KEY", "AZURE_AD_TOKEN", "AZURE_API_KEY"): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "api_key", None) + + +def test_provider_config_manager_resolves_azure_ai_passthrough_config(): + config = ProviderConfigManager.get_provider_passthrough_config( + model="Cohere-parse-v5", provider=LlmProviders.AZURE_AI + ) + + assert isinstance(config, AzureAIPassthroughConfig) + + +def test_router_model_prefix_is_stripped_and_native_path_kept_verbatim(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert base == FOUNDRY_BASE + + +def test_model_group_prefix_is_stripped_when_router_metadata_names_it(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key=None, + model="Cohere-parse-v5", + endpoint="/parse-alias/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "parse-alias"}}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_model_inside_the_path_stays_and_query_params_are_forwarded(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/", + api_key=None, + model="gpt-5.4-mini", + endpoint="openai/deployments/gpt-5.4-mini/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2024-10-21" + + +def test_api_base_that_already_ends_in_models_is_cut_back_to_the_foundry_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert base == FOUNDRY_BASE + + +def test_full_url_api_base_that_already_ends_with_the_native_path_is_not_doubled(): + model_router_url = ( + "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router/chat/completions" + ) + + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{model_router_url}?api-version=2025-01-01-preview", + api_key="key", + model="model_router/model-router", + endpoint="model-router/chat/completions", + request_query_params=None, + litellm_params={"litellm_metadata": {"model_group": "model-router"}}, + ) + + assert str(url) == f"{model_router_url}?api-version=2025-01-01-preview" + assert base == "https://my-resource.cognitiveservices.azure.com/openai/deployments/model-router" + + +@pytest.mark.parametrize("relayed_deployment", ["gpt-4o", "GPT-4o"]) +def test_deployment_root_api_base_is_not_repeated_when_the_relay_carries_the_deployment_path(relayed_deployment): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/gpt-4o", + api_key="key", + model="gpt-4o", + endpoint=f"aoai-gpt-4o/openai/deployments/{relayed_deployment}/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-gpt-4o"}}, + ) + + assert str(url) == ( + f"https://my-resource.openai.azure.com/openai/deployments/{relayed_deployment}/chat/completions" + "?api-version=2024-10-21" + ) + assert base == "https://my-resource.openai.azure.com" + + +def test_deployment_named_like_the_first_native_segment_keeps_its_deployment_root(): + url, base = AzureAIPassthroughConfig().get_complete_url( + api_base="https://my-resource.openai.azure.com/openai/deployments/chat", + api_key="key", + model="chat", + endpoint="aoai-chat/chat/completions", + request_query_params={"api-version": "2024-10-21"}, + litellm_params={"litellm_metadata": {"model_group": "aoai-chat"}}, + ) + + assert str(url) == "https://my-resource.openai.azure.com/openai/deployments/chat/chat/completions?api-version=2024-10-21" + assert base == "https://my-resource.openai.azure.com/openai/deployments/chat" + + +def test_parse_relay_under_a_models_api_base_targets_the_foundry_root(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models", + api_key="key", + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + + +def test_deployment_api_version_fills_in_when_the_caller_sends_none(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + +def test_callers_api_version_beats_the_deployments(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=FOUNDRY_BASE, + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params={"api-version": "2025-04-01-preview"}, + litellm_params={"api_version": "2024-05-01-preview"}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2025-04-01-preview" + + +def test_api_version_on_the_configured_api_base_is_the_last_fallback(): + url, _ = AzureAIPassthroughConfig().get_complete_url( + api_base=f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview", + api_key="key", + model="gpt-5.4-mini", + endpoint="gpt-5.4-mini/models/chat/completions", + request_query_params=None, + litellm_params={}, + ) + + assert str(url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + + +def test_missing_api_base_raises_instead_of_building_a_relative_url(): + with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): + AzureAIPassthroughConfig().get_complete_url( + api_base=None, + api_key=None, + model="Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + request_query_params=None, + litellm_params={}, + ) + + +def _auth_headers(api_key: str | None, api_base: str, litellm_params: dict | None = None) -> dict: + return AzureAIPassthroughConfig().validate_environment( + headers={"content-type": "application/json"}, + model="Cohere-parse-v5", + messages=[], + optional_params={}, + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + ) + + +def test_foundry_host_gets_the_api_key_header(): + headers = _auth_headers(api_key="deployment-key", api_base=FOUNDRY_BASE) + + assert headers == {"content-type": "application/json", "api-key": "deployment-key"} + + +def test_serverless_host_gets_a_bearer_token(): + headers = _auth_headers(api_key="deployment-key", api_base="https://cohere-parse.eastus.models.ai.azure.com") + + assert headers["Authorization"] == "Bearer deployment-key" + assert "api-key" not in headers + + +def test_entra_token_is_used_when_the_deployment_has_no_api_key(): + headers = _auth_headers(api_key=None, api_base=FOUNDRY_BASE, litellm_params={"azure_ad_token": "entra-token"}) + + assert headers["Authorization"] == "Bearer entra-token" + + +def test_no_credentials_at_all_raises(): + with pytest.raises(ValueError, match="Missing Azure AI credentials"): + _auth_headers(api_key=None, api_base=FOUNDRY_BASE) + + +@pytest.mark.parametrize( + "request_data, expected", + [({"stream": True}, True), ({"stream": 1}, True), ({"stream": False}, False), ({}, False)], +) +def test_is_streaming_request_reads_the_stream_flag(request_data, expected): + assert ( + AzureAIPassthroughConfig().is_streaming_request(endpoint="models/chat/completions", request_data=request_data) + is expected + ) + + +def _chat_completion_response() -> httpx.Response: + body = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + +def test_chat_completions_relay_yields_a_model_response_for_cost_tracking(): + result = AzureAIPassthroughConfig().logging_non_streaming_response( + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + httpx_response=_chat_completion_response(), + request_data={"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hi"}]}, + logging_obj=MagicMock(), + endpoint="models/chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "hi" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 8 + + +def _non_chat_logging_result(content: bytes, content_type: str): + parse_response = httpx.Response( + status_code=200, + headers={"content-type": content_type}, + content=content, + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + return AzureAIPassthroughConfig().logging_non_streaming_response( + model="Cohere-parse-v5", + custom_llm_provider="azure_ai", + httpx_response=parse_response, + request_data={"model": "Cohere-parse-v5"}, + logging_obj=MagicMock(), + endpoint="providers/cohere/v2/parse", + ) + + +def test_non_chat_relay_with_a_non_json_body_logs_the_raw_text(): + assert _non_chat_logging_result(b"page one", "text/plain") == {"response": "page one"} + + +def _relay_logging_obj( + model: str, + api_base: str, + stream: bool = False, + callbacks: list[CustomLogger] | None = None, + endpoint: str = "", +) -> Logging: + logging_obj = Logging( + model=model, + messages=[], + stream=stream, + call_type="allm_passthrough_route", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + dynamic_async_success_callbacks=callbacks, + ) + logging_obj.update_environment_variables( + model=model, + litellm_params={"api_base": api_base, "custom_llm_provider": "azure_ai"}, + optional_params={}, + custom_llm_provider="azure_ai", + endpoint=endpoint, + ) + return logging_obj + + +def _relay_logging_result( + config: AzureAIPassthroughConfig, + model: str, + native_path: str, + body, + api_base: str = FOUNDRY_BASE, + status_code: int = 200, +): + relayed_url = f"{FOUNDRY_BASE}/{native_path}?api-version=2024-05-01-preview" + logging_obj = _relay_logging_obj(model, api_base) + response = httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", relayed_url), + ) + result = config.logging_non_streaming_response( + model=model, + custom_llm_provider="azure_ai", + httpx_response=response, + request_data={"model": model}, + logging_obj=logging_obj, + endpoint=f"{model}/{native_path}", + ) + return result, logging_obj + + +MISTRAL_OCR_BODY = { + "pages": [{"index": 0, "markdown": "page one"}, {"index": 1, "markdown": "page two"}], + "model": "mistral-document-ai-2512", + "usage_info": {"pages_processed": 2, "doc_size_bytes": 4321}, +} + + +def test_mistral_document_ai_relay_is_costed_per_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + per_page = litellm.get_model_info("azure_ai/mistral-document-ai-2512")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 2 + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_page) + + +def test_ocr_route_under_a_models_api_base_is_still_recognised(): + result, _ = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + MISTRAL_OCR_BODY, + api_base=f"{FOUNDRY_BASE}/models", + ) + + assert isinstance(result, OCRResponse) + + +def test_relay_to_a_non_ocr_route_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "mistral-document-ai-2512", "models/info", {"name": "mistral-document-ai-2512"} + ) + + assert result == {"response": {"name": "mistral-document-ai-2512"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +COHERE_PARSE_BODY = {"id": "parse-1", "pages": [], "meta": {"billed_units": {"pages": 3}}} + + +def test_cohere_parse_relay_is_costed_per_billed_page(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "Cohere-parse-v5", "providers/cohere/v2/parse", COHERE_PARSE_BODY + ) + per_page = litellm.get_model_info("azure_ai/Cohere-parse-v5")["ocr_cost_per_page"] + + assert isinstance(result, OCRResponse) + assert result.usage_info.pages_processed == 3 + assert logging_obj.call_type == "aocr" + assert per_page > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(3 * per_page) + + +def test_deployment_without_an_ocr_config_is_never_costed_as_ocr(): + config = AzureAIPassthroughConfig(ocr_config_for=lambda model: None) + result, logging_obj = _relay_logging_result( + config, "mistral-document-ai-2512", "providers/mistral/azure/ocr", MISTRAL_OCR_BODY + ) + + assert result == {"response": MISTRAL_OCR_BODY} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_accepted_ocr_job_without_a_result_body_is_not_costed(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + {"status": "running"}, + status_code=202, + ) + + assert result == {"response": {"status": "running"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_unparseable_ocr_body_falls_back_to_the_passthrough_object(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "mistral-document-ai-2512", + "providers/mistral/azure/ocr", + ["not", "an", "ocr", "body"], + ) + + assert result == {"response": '["not", "an", "ocr", "body"]'} + assert logging_obj.call_type == "allm_passthrough_route" + + +EMBEDDINGS_BODY = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "embed-v-4-0", + "usage": {"prompt_tokens": 1200, "total_tokens": 1200}, +} + +RERANK_BODY = { + "id": "rerank-1", + "results": [{"index": 1, "relevance_score": 0.9}, {"index": 0, "relevance_score": 0.2}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 2}}, +} + +IMAGE_BODY = {"created": 1, "data": [{"b64_json": "AAAA"}]} + + +def test_foundry_embeddings_relay_is_costed_per_input_token(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "embed-v-4-0", "models/embeddings", EMBEDDINGS_BODY + ) + per_token = litellm.get_model_info("azure_ai/embed-v-4-0")["input_cost_per_token"] + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(1200 * per_token) + + +def test_cohere_rerank_relay_is_costed_per_search_unit(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "cohere-rerank-v4.0-fast", "providers/cohere/v2/rerank", RERANK_BODY + ) + per_query = litellm.get_model_info("azure_ai/cohere-rerank-v4.0-fast")["input_cost_per_query"] + + assert isinstance(result, RerankResponse) + assert logging_obj.call_type == "arerank" + assert per_query > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(2 * per_query) + + +def test_image_generation_relay_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "openai/deployments/FLUX.2-pro/images/generations", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert per_image > 0 + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + +def test_flux_2_relay_through_the_provider_route_is_costed_per_image(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), "FLUX.2-pro", "providers/blackforestlabs/v1/flux-2-pro", IMAGE_BODY + ) + per_image = litellm.get_model_info("azure_ai/FLUX.2-pro")["output_cost_per_image"] + + assert isinstance(result, ImageResponse) + assert logging_obj.call_type == "aimage_generation" + assert logging_obj._response_cost_calculator(result=result) == pytest.approx(per_image) + + +def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): + result, logging_obj = _relay_logging_result( + AzureAIPassthroughConfig(), + "cohere-rerank-v4.0-fast", + "providers/cohere/v2/rerank", + {"message": "invalid request"}, + status_code=400, + ) + + assert result == {"response": {"message": "invalid request"}} + assert logging_obj.call_type == "allm_passthrough_route" + + +def test_streaming_chat_completion_chunks_are_costed_like_azure(): + head = {"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "gpt-5.4-mini"} + chunks = [ + "data: " + + json.dumps( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + } + ), + "data: " + + json.dumps({**head, "choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}}), + "data: [DONE]", + ] + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=MagicMock(), + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "hi" + assert response.usage.total_tokens == 4 + + +def test_streaming_responses_chunks_through_a_router_relay_are_costed_like_azure(): + logging_obj = _relay_logging_obj("gpt-5.4-mini", FOUNDRY_BASE) + + response = AzureAIPassthroughConfig().handle_logging_collected_chunks( + all_chunks=["event: response.completed", "data: " + json.dumps(RESPONSES_COMPLETED_EVENT)], + litellm_logging_obj=logging_obj, + model="gpt-5.4-mini", + custom_llm_provider="azure_ai", + endpoint="gpt/openai/responses", + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert response is not None + assert response.response.usage.output_tokens == 100 + assert logging_obj.call_type == "aresponses" + assert logging_obj._response_cost_calculator(result=response.response) == pytest.approx( + 1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"] + ) + + +async def test_streaming_responses_relay_flush_reaches_the_success_callbacks_with_a_price(): + probe = _SpendProbe() + logging_obj = _relay_logging_obj( + "gpt-5.4-mini", FOUNDRY_BASE, stream=True, callbacks=[probe], endpoint="gpt/openai/responses" + ) + stream = "event: response.completed\ndata: " + json.dumps(RESPONSES_COMPLETED_EVENT) + "\n\n" + + await logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=[stream.encode()], provider_config=AzureAIPassthroughConfig() + ) + info = litellm.get_model_info("azure_ai/gpt-5.4-mini") + + assert probe.logged_call_type == "allm_passthrough_route" + assert probe.logged_cost == pytest.approx(1000 * info["input_cost_per_token"] + 100 * info["output_cost_per_token"]) diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index c2c448cd7e2..96a2fa6ec67 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -1,5 +1,7 @@ import json +from unittest.mock import MagicMock +import httpx import pytest @@ -190,3 +192,45 @@ def test_get_error_class_preserves_provider_headers(): assert isinstance(error, BedrockError) assert error.headers == {"x-amzn-RequestId": "req-invoke-500"} assert error.response.headers["x-amzn-requestid"] == "req-invoke-500" + + +def test_transform_response_hands_json_mode_to_nova(): + """The invoke dispatcher forwards its json_mode argument to Nova instead of dropping it.""" + from litellm.types.utils import ModelResponse + + response_json = { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_nova_json", + "name": "json_tool_call", + "input": {"city": "Paris", "temperature": 21}, + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9}, + } + raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock")) + + result = AmazonInvokeConfig().transform_response( + model="invoke/amazon.nova-lite-v1:0", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "weather"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=True, + ) + + assert result.choices[0].message.tool_calls is None + assert json.loads(result.choices[0].message.content) == {"city": "Paris", "temperature": 21} diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index f34b8eb1fb9..8831fd9061c 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -6,6 +6,7 @@ extension, and AWS credential resolution is stubbed so nothing reaches STS. from __future__ import annotations +import asyncio from unittest.mock import MagicMock, patch import httpx @@ -16,6 +17,7 @@ from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe RUST_RESPONSE = { "created": 1_700_000_000, @@ -308,7 +310,9 @@ CONVERSE_RESPONSE = { } -async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): +async def _drive_async_completion( + *, skip_pre_call_logging: bool, logging_obj, credentials: Credentials = RESOLVED_CREDENTIALS +): """Run the real `async_completion` with a stubbed transport.""" import httpx as _httpx @@ -335,7 +339,7 @@ async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): stream=None, optional_params={"maxTokens": 16}, litellm_params={"aws_region_name": "us-west-2"}, - credentials=RESOLVED_CREDENTIALS, + credentials=credentials, headers={}, client=client, skip_pre_call_logging=skip_pre_call_logging, @@ -357,6 +361,23 @@ async def test_async_completion_logs_pre_call_by_default(): assert logging_obj.pre_call.call_count == 1 +@pytest.mark.asyncio +async def test_async_completion_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: botocore refreshes expiring credentials inside SigV4 signing with a + blocking HTTP call, so `async_completion` must sign on a worker thread to keep the loop serving.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + response = await _drive_async_completion( + skip_pre_call_logging=False, logging_obj=MagicMock(), credentials=probe.credentials() + ) + await release + + assert response.choices[0].message.content == "hi" + assert probe.served_during_refresh is True + + def _sync_client_returning_converse_response(): client = MagicMock() client.post.side_effect = lambda **_kwargs: httpx.Response( diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index f0e361ceb88..2e9ea90f3b8 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -382,6 +382,8 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): "us.openai.gpt-5.6-sol", "global.openai.gpt-5.6-terra", "bedrock/converse/us.openai.gpt-5.6-luna", + "us.openai.gpt-6-astra", + "bedrock/converse/global.openai.gpt-6-astra", ], ) def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(model, local_model_cost_map): @@ -412,6 +414,7 @@ def test_reasoning_effort_maps_to_reasoning_effort_for_openai_gpt5_converse(mode [ "us.openai.gpt-5.6-sol", "bedrock/converse/global.openai.gpt-5.6-luna", + "us.openai.gpt-6-astra", ], ) def test_openai_gpt5_converse_never_forwards_thinking(model, local_model_cost_map): @@ -863,6 +866,191 @@ def test_get_supported_openai_params(): assert "reasoning_effort" in supported_params +@pytest.mark.parametrize( + "model", + [ + "bedrock/us.deepseek.r1-v1:0", + "bedrock/converse/us.deepseek.r1-v1:0", + "bedrock/deepseek.v3-v1:0", + "bedrock/deepseek.v3.2", + ], +) +def test_bedrock_deepseek_does_not_advertise_thinking(model): + """DeepSeek reasons natively on Bedrock and does not take the Anthropic-shaped `thinking` + field (R1 400s on it, V3 ignores it), so it must not be advertised as supported.""" + config = AmazonConverseConfig() + supported_params = config.get_supported_openai_params(model=model) + assert "thinking" not in supported_params + assert "output_config" not in supported_params + + +@pytest.mark.parametrize("model", ["bedrock/us.deepseek.r1-v1:0", "bedrock/converse/us.deepseek.r1-v1:0"]) +def test_bedrock_deepseek_r1_does_not_advertise_reasoning_effort(model): + """DeepSeek R1 always reasons and returns a 400 for any reasoning_effort shape.""" + config = AmazonConverseConfig() + assert "reasoning_effort" not in config.get_supported_openai_params(model=model) + + +@pytest.mark.parametrize("model", ["bedrock/deepseek.v3-v1:0", "bedrock/deepseek.v3.2", "bedrock/us.deepseek.v3.2"]) +def test_bedrock_deepseek_v3_advertises_reasoning_effort(model): + """DeepSeek V3 on Bedrock accepts a raw reasoning_effort in additionalModelRequestFields.""" + config = AmazonConverseConfig() + assert "reasoning_effort" in config.get_supported_openai_params(model=model) + + +@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"]) +def test_bedrock_deepseek_thinking_raises_without_drop_params(model): + """Passing `thinking` to Bedrock DeepSeek must fail client-side with a clear + UnsupportedParamsError instead of leaking through to Bedrock.""" + with pytest.raises(litellm.UnsupportedParamsError): + litellm.utils.get_optional_params( + model=model, + custom_llm_provider="bedrock", + thinking={"type": "enabled", "budget_tokens": 1024}, + ) + + +def test_bedrock_deepseek_r1_reasoning_effort_raises_without_drop_params(): + with pytest.raises(litellm.UnsupportedParamsError): + litellm.utils.get_optional_params( + model="us.deepseek.r1-v1:0", + custom_llm_provider="bedrock", + reasoning_effort="high", + ) + + +@pytest.mark.parametrize("model", ["us.deepseek.r1-v1:0", "deepseek.v3.2"]) +def test_bedrock_deepseek_thinking_dropped_does_not_leak_into_request(model): + """With drop_params, `thinking` is dropped rather than forwarded into + additionalModelRequestFields for Bedrock DeepSeek.""" + optional_params = litellm.utils.get_optional_params( + model=model, + custom_llm_provider="bedrock", + thinking={"type": "enabled", "budget_tokens": 1024}, + drop_params=True, + ) + assert "thinking" not in optional_params + + config = AmazonConverseConfig() + request = config._transform_request( + model=f"bedrock/converse/{model}", + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert "thinking" not in (request.get("additionalModelRequestFields") or {}) + + +@pytest.mark.parametrize("param", ["thinking", "reasoning_effort"]) +def test_bedrock_deepseek_r1_reasoning_params_not_forwarded_by_map(param): + """Even when map_openai_params is called directly (bypassing the supported-params + gate), DeepSeek R1 must not forward thinking/reasoning_effort into + additionalModelRequestFields, since Bedrock rejects both with a 400.""" + config = AmazonConverseConfig() + model = "bedrock/converse/us.deepseek.r1-v1:0" + value = {"type": "enabled", "budget_tokens": 1024} if param == "thinking" else "high" + + optional_params = config.map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert "thinking" not in optional_params + assert "reasoning_effort" not in optional_params + + request = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert request.get("additionalModelRequestFields") is None + + +def test_bedrock_deepseek_v3_reasoning_effort_forwarded_raw(): + """DeepSeek V3 takes reasoning_effort verbatim in additionalModelRequestFields, never + converted into the Anthropic `thinking` block that Claude models get.""" + config = AmazonConverseConfig() + model = "bedrock/deepseek.v3.2" + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high", "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert "thinking" not in optional_params + + request = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Say hi in one word."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + assert request["additionalModelRequestFields"] == {"reasoning_effort": "high"} + + +def test_bedrock_deepseek_v3_thinking_dropped_by_map(): + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}, "max_tokens": 100}, + optional_params={}, + model="bedrock/deepseek.v3.2", + drop_params=False, + ) + assert "thinking" not in optional_params + assert "reasoning_effort" not in optional_params + + +@pytest.mark.parametrize( + "model, param, value, kept_key", + [ + ( + "bedrock/us.anthropic.claude-opus-4-20250514-v1:0", + "thinking", + {"type": "enabled", "budget_tokens": 1024}, + "thinking", + ), + ( + "bedrock/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + "thinking", + {"type": "enabled", "budget_tokens": 1024}, + "thinking", + ), + ( + "bedrock/openai.gpt-oss-safeguard-20b-1:0", + "reasoning_effort", + "high", + "reasoning_effort", + ), + ( + "bedrock/us.amazon.nova-2-lite-v1:0", + "reasoning_effort", + "high", + "reasoningConfig", + ), + ], +) +def test_bedrock_non_deepseek_reasoning_params_preserved(model, param, value, kept_key): + """The DeepSeek leak fix must only drop reasoning request params for DeepSeek. + + Claude behind an application-inference-profile ARN, gpt-oss-safeguard (absent from the + cost map so `supports_reasoning` is False), and Nova 2 all reason via a request param and + must keep it. Regression guard against gating the drop on a positive allowlist, which + silently degraded reasoning for anything the allowlist/ARN introspection missed.""" + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={param: value, "max_tokens": 100}, + optional_params={}, + model=model, + drop_params=False, + ) + assert kept_key in optional_params + + def test_get_supported_openai_params_bedrock_converse(): """ Test that all documented bedrock converse models have the same set of supported openai params when using @@ -6727,3 +6915,41 @@ def test_forced_tool_choice_forwarded_on_converse_models_that_support_it( ) assert result == {"any": {}} + + +def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it(): + response_json = { + "metrics": {"latencyMs": 900}, + "output": { + "message": { + "content": [ + { + "toolUse": { + "input": {"city": "Paris", "population": 2100000}, + "name": "json_tool_call", + "toolUseId": "tooluse_invoke_nova_json", + } + } + ], + "role": "assistant", + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 40, "outputTokens": 20, "totalTokens": 60}, + } + raw_response = httpx.Response(200, json=response_json, request=httpx.Request("POST", "https://bedrock.test")) + logging_obj = MagicMock() + result = AmazonConverseConfig().transform_response( + model="bedrock/invoke/us.amazon.nova-micro-v1:0", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[], + optional_params={"tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": {}}}]}, + litellm_params={}, + encoding=None, + json_mode=True, + ) + assert result.choices[0].message.tool_calls is None + assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py new file mode 100644 index 00000000000..3622ce7f212 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -0,0 +1,54 @@ +import asyncio +from unittest.mock import AsyncMock + +import httpx +import pytest +from botocore.credentials import RefreshableCredentials + +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe + + +class _ProbedCountTokensHandler(BedrockCountTokensHandler): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_handle_count_tokens_request_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the count_tokens handler signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = AsyncMock(spec=AsyncHTTPHandler) + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json={"inputTokens": 7}, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + result = await _ProbedCountTokensHandler(probe).handle_count_tokens_request( + request_data={ + "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": "hi"}], + }, + litellm_params={"aws_region_name": "us-west-2"}, + resolved_model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + client=client, + ) + await release + + assert result == {"input_tokens": 7} + assert client.post.call_args.kwargs["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index ddbd3a2e9ba..18f4b0f6ced 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -1,11 +1,15 @@ import json +import asyncio from unittest.mock import Mock, patch +import httpx import pytest +import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.base import HiddenParams +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock async invoke responses async_invoke_response = { @@ -422,3 +426,34 @@ class TestBedrockAsyncInvokeEmbedding: async_endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/async-invoke" ) + + +@pytest.mark.asyncio +async def test_async_invoke_status_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the GetAsyncInvoke poll is a signed GET, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + embedder = BedrockEmbedding() + probe = EventLoopProbe() + + with ( + patch.object(embedder, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + respx.mock, + ): + route = respx.get(url__regex=r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/async-invoke/.*").mock( + return_value=httpx.Response(200, json=async_invoke_status_response) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + status = await embedder._get_async_invoke_status( + invocation_arn=async_invoke_status_response["invocationArn"], aws_region_name="us-east-1" + ) + await release + + assert status["status"] == "InProgress" + assert "Authorization" in route.calls.last.request.headers + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index bcd1a29d0e8..4aad2846537 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1,12 +1,17 @@ import json +import asyncio import os from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, MagicMock import pytest +import httpx import litellm from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.bedrock.embed.embedding import BedrockEmbedding +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock responses for different embedding models titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} @@ -1062,6 +1067,41 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" +@pytest.mark.asyncio +async def test_async_single_func_embeddings_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: Titan, Nova, and TwelveLabs embeddings sign one SigV4 request per + input, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call, + so each signing must run on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = MagicMock() + client.__class__ = AsyncHTTPHandler + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json=titan_embedding_response, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + ) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await BedrockEmbedding()._async_single_func_embeddings( + client=client, + timeout=None, + batch_data=[{"inputText": test_input}], + credentials=probe.credentials(), + extra_headers=None, + endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-text-v1/invoke", + aws_region_name="us-west-2", + model="amazon.titan-embed-text-v1", + logging_obj=MagicMock(), + provider="amazon", + ) + await release + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + assert "Authorization" in client.post.call_args.kwargs["headers"] + assert probe.served_during_refresh is True marengo_3_embedding_response = {"data": [{"embedding": [0.01 * i for i in range(512)]}]} MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw==" diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/test_litellm/llms/bedrock/event_loop_probe.py new file mode 100644 index 00000000000..c347247ec32 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/event_loop_probe.py @@ -0,0 +1,57 @@ +"""Refreshable credentials whose refresh only completes while the event loop keeps serving.""" + +from __future__ import annotations + +import asyncio +import threading +import time +from datetime import datetime, timedelta, timezone +from typing import Final + +from botocore.credentials import RefreshableCredentials + +REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0 +REFRESH_START_TIMEOUT_SECONDS: Final = 10.0 + + +class EventLoopProbe: + """Blocks inside botocore's credential refresh until a coroutine on the loop releases it. + + Signing on the event loop thread can never be released, so `served_during_refresh` reads False there + and True only when the refresh ran on another thread while the loop stayed responsive. + """ + + def __init__(self) -> None: + self.refresh_started: Final = threading.Event() + self.loop_served: Final = threading.Event() + self.served_during_refresh: bool | None = None + + def refresh(self) -> dict[str, str | None]: + self.refresh_started.set() + served: Final = self.loop_served.wait(timeout=REFRESH_RELEASE_TIMEOUT_SECONDS) + if self.served_during_refresh is None: + self.served_during_refresh = served + return { + "access_key": "AKIAREFRESHED", + "secret_key": "refreshed-secret", + "token": None, + "expiry_time": (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat(), + } + + def credentials(self) -> RefreshableCredentials: + return RefreshableCredentials( + access_key="AKIASTALE", + secret_key="stale-secret", + token=None, + expiry_time=datetime.now(timezone.utc) + timedelta(seconds=60), + refresh_using=self.refresh, + method="event-loop-probe", + ) + + async def release_refresh_from_the_loop(self) -> None: + deadline: Final = time.monotonic() + REFRESH_START_TIMEOUT_SECONDS + while not self.refresh_started.is_set(): + if time.monotonic() > deadline: + raise TimeoutError("signing finished without ever starting a credential refresh") + await asyncio.sleep(0.005) + self.loop_served.set() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index f854d806bdc..2c7e476e9a8 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1,4 +1,6 @@ +import asyncio import json +from concurrent.futures import ThreadPoolExecutor import os import threading import time @@ -22,7 +24,10 @@ from litellm.llms.bedrock.base_aws_llm import ( AwsAuthError, BaseAWSLLM, Boto3CredentialsInfo, + run_aws_signing, + sign_request_off_loop_if_aws, ) +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Global variable for the base_aws_llm.py file path @@ -38,6 +43,14 @@ def flush_shared_bedrock_iam_cache(): yield +@pytest.fixture(autouse=True) +def _clean_ssl_env(monkeypatch): + """get_ssl_verify reads these, so the sts client's verify= would otherwise depend on + the ambient environment. The published images set SSL_CERT_FILE.""" + for env_var in ("SSL_CERT_FILE", "SSL_VERIFY"): + monkeypatch.delenv(env_var, raising=False) + + def test_base_aws_llm_instances_share_process_wide_iam_cache(): """Regression LIT-2662: new instances must reuse iam_cache (Bedrock passthrough is per-request).""" first = BaseAWSLLM() @@ -3215,3 +3228,53 @@ class TestGetRequestHeadersResign: extra_headers={"Authorization": "Bearer foo"}, ) assert prepped.headers["Authorization"] == "Bearer foo" + + +@pytest.mark.asyncio +async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credentials_refresh(): + """Regression for issue #40165: an AWS provider's signing (and the botocore credential refresh + inside it) must run off the event loop, so other requests keep being served meanwhile.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-runtime.us-west-2.amazonaws.com/", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-west-2").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws(BaseAWSLLM(), sign, headers={"Content-Type": "application/json"}) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True + + +def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): + """A signing parked on botocore's refresh lock must not hold a default-executor thread, since every + other provider's async entry point hops through that same executor. The scenario runs on its own loop + so the one-thread default executor it pins never leaks into the session loop.""" + + async def scenario() -> tuple[str, str]: + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + signing_parked = asyncio.Event() + refresh_done = threading.Event() + + def sign() -> str: + loop.call_soon_threadsafe(signing_parked.set) + refresh_done.wait() + return threading.current_thread().name + + signing = asyncio.create_task(run_aws_signing(sign)) + try: + await asyncio.wait_for(signing_parked.wait(), timeout=5) + other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) + finally: + refresh_done.set() + return other_provider.name, await signing + + other_provider, signing_thread = asyncio.run(scenario()) + assert other_provider != signing_thread + assert signing_thread.startswith("aws-signing") diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 1be94d4daa2..e83a844c87e 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -6,15 +6,20 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.ht """ import json +import asyncio from unittest.mock import patch import httpx import pytest +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest import litellm from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig +from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.types.utils import LlmProviders +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.fixture @@ -710,3 +715,26 @@ def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id): resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name) assert provider == "bedrock_mantle" assert resolved_model == model_id + + +@pytest.mark.asyncio +async def test_mantle_signing_runs_off_the_event_loop(): + """Regression for issue #40165: Mantle signs with SigV4 through a composed BaseAWSLLM, so the + off-loop gate must recognise it too, or its credential refresh blocks the loop like Bedrock's did.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-mantle.us-east-1.api.aws/v1/responses", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-east-1").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws( + BedrockMantleChatConfig(), sign, headers={"Content-Type": "application/json"} + ) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 8e0415d50de..a7520bd5955 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -10,18 +10,23 @@ from unittest.mock import MagicMock, patch import httpx import pytest - +import litellm +from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError +from litellm.main import responses_api_bridge_check from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager -from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig class TestChatGPTResponsesAPITransformation: @pytest.mark.parametrize( "model_name", [ + "chatgpt/gpt-5.5", + "chatgpt/gpt-5.6-luna", + "chatgpt/gpt-5.6-sol", + "chatgpt/gpt-5.6-terra", "chatgpt/gpt-5.4", "chatgpt/gpt-5.4-pro", "chatgpt/gpt-5.3-chat-latest", @@ -40,6 +45,52 @@ class TestChatGPTResponsesAPITransformation: assert isinstance(config, ChatGPTResponsesAPIConfig) assert config.custom_llm_provider == LlmProviders.CHATGPT + @pytest.mark.parametrize( + "model_name", + [ + "chatgpt/gpt-5.5", + "chatgpt/gpt-5.6-luna", + "chatgpt/gpt-5.6-sol", + "chatgpt/gpt-5.6-terra", + ], + ) + def test_chatgpt_responses_model_metadata(self, model_name: str, local_model_cost_map: None) -> None: + model_info = litellm.get_model_info(model_name) + + assert model_info["litellm_provider"] == "chatgpt" + assert model_info["mode"] == "responses" + assert model_info["supported_endpoints"] == [ + "/v1/chat/completions", + "/v1/responses", + ] + assert model_info["max_input_tokens"] == 1050000 + assert model_info["max_output_tokens"] == 128000 + + @pytest.mark.parametrize( + "model_name", + [ + "gpt-5.5", + "gpt-5.6-luna", + "gpt-5.6-sol", + "gpt-5.6-terra", + ], + ) + def test_chatgpt_models_bridge_chat_completions_to_responses( + self, model_name: str, local_model_cost_map: None + ) -> None: + """A chat completions request for these models must take the Responses bridge. + + `gpt-5.6-*` also exists as an openai chat model, so an unregistered + chatgpt model resolves to mode "chat" here and never reaches the bridge. + """ + model_info, resolved_model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="chatgpt", + ) + + assert model_info["mode"] == "responses" + assert resolved_model == model_name + @patch("litellm.llms.chatgpt.responses.transformation.Authenticator") def test_chatgpt_responses_endpoint_url(self, mock_authenticator_class): mock_auth_instance = MagicMock() diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 98a5f4b2db5..cea2d439198 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock, patch import httpx import pytest +from botocore.credentials import RefreshableCredentials import litellm from litellm._logging import verbose_logger @@ -19,6 +20,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( @@ -29,11 +31,15 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, +) from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -813,6 +819,65 @@ async def test_anthropic_messages_streaming_response_aclose_closes_agentic_upstr assert tracker.closed is True +class _ProbedBedrockMessagesConfig(AmazonAnthropicClaudeMessagesConfig): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_signs_bedrock_off_the_event_loop(monkeypatch): + """Regression for issue #40165: /v1/messages on Bedrock signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + handler = BaseLLMHTTPHandler() + upstream_response = httpx.Response( + 200, + json={ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": "claude-haiku-4-5-20251001", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=upstream_response) + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.dynamic_success_callbacks = None + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + await handler.async_anthropic_messages_handler( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=_ProbedBedrockMessagesConfig(probe), + anthropic_messages_optional_request_params={"max_tokens": 16}, + custom_llm_provider="bedrock", + litellm_params=GenericLiteLLMParams(aws_region_name="us-west-2"), + logging_obj=mock_logging_obj, + client=mock_client, + stream=False, + kwargs={}, + ) + await release + + sent_headers = mock_client.post.call_args.kwargs["headers"] + assert sent_headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_passes_litellm_metadata(): """Ensure litellm_metadata from kwargs is forwarded via update_from_kwargs. @@ -3367,6 +3432,26 @@ async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_tran assert captured["body"] == {"transformed_by": "async"} assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + assert not any(thread.name.startswith("aws-signing") for thread in config.sign_threads + pre_call_threads) + + +class _AWSTransformRecordingConfig(SignsRequestsWithAWS, _TransformRecordingConfig): + pass + + +async def test_completion_signs_aws_configs_on_the_aws_signing_pool_after_the_async_transform(): + config = _AWSTransformRecordingConfig(transform_async=True) + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread.name.startswith("aws-signing") for thread in config.sign_threads) + assert pre_call_threads and all(thread.name.startswith("aws-signing") for thread in pre_call_threads) async def test_completion_keeps_sync_transform_request_before_returning_by_default(): diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 02655fb7f77..caf7bed7385 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -255,6 +255,19 @@ def test_transform_messages_sanitizes_empty_content(): assert result[1]["content"] == "Hi" +def test_transform_request_preserves_unity_model_service_name(): + config = DatabricksConfig() + result = config.transform_request( + model="system.ai.kimi-k3", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["model"] == "system.ai.kimi-k3" + + def test_transform_request_strips_thinking_blocks_and_reasoning_content(): """Regression for LIT-6762: replaying an assistant turn that litellm decorated with `thinking_blocks` / `reasoning_content` made Databricks 400 with @@ -590,3 +603,87 @@ def test_chunk_parser_without_usage_still_parses_content(): assert result.id == "chatcmpl-test" assert result.model == "databricks-claude-sonnet-5" assert result.choices[0]["delta"]["content"] == "hi" + + +@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"]) +def test_transform_choices_surfaces_top_level_reasoning_content(reasoning_key: str) -> None: + config = DatabricksConfig() + databricks_choices = [ + { + "message": { + "role": "assistant", + "content": "391", + reasoning_key: "We need answer just number. 17*23=391.", + }, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.content == "391" + assert choices[0].message.reasoning_content == "We need answer just number. 17*23=391." + assert getattr(choices[0].message, "thinking_blocks", None) is None + + +def test_transform_choices_parses_think_tags_in_string_content(): + config = DatabricksConfig() + databricks_choices = [ + { + "message": {"role": "assistant", "content": "17 times 23391"}, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.content == "391" + assert choices[0].message.reasoning_content == "17 times 23" + + +def test_transform_choices_prefers_reasoning_blocks_over_top_level_field(): + config = DatabricksConfig() + databricks_choices = [ + { + "message": { + "role": "assistant", + "content": [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "from block"}]}, + {"type": "text", "text": "391"}, + ], + "reasoning_content": "from field", + }, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.reasoning_content == "from block" + assert choices[0].message.content == "391" + + +@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"]) +def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) -> None: + iterator = DatabricksChatResponseIterator(None, sync_stream=True) + chunk = { + "id": "1", + "object": "chat.completion.chunk", + "created": 0, + "model": "lit-qa-deepseek-v4-flash", + "choices": [ + { + "delta": {"role": "assistant", "content": None, reasoning_key: "We need answer"}, + "index": 0, + "finish_reason": None, + } + ], + } + + parsed = iterator.chunk_parser(chunk) + + assert parsed.choices[0].delta.reasoning_content == "We need answer" + assert parsed.choices[0].delta.content is None diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py index 39198bb20f3..c6ad78366f7 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -657,6 +657,77 @@ class TestEndpointURLConstruction: assert api_base.endswith("/chat/completions") + def test_chat_gateway_endpoint_for_unity_model_on_legacy_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/serving-endpoints", + api_key="test-key", + model="system.ai.kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + + def test_chat_gateway_endpoint_preserves_explicit_gateway_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/ai-gateway/mlflow/v1/", + api_key="test-key", + model="system.ai.kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + + def test_chat_gateway_preserves_unity_model_service_name_with_explicit_base(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + config = DatabricksConfig() + request = config.transform_request( + model="catalog.schema.kimi-k3", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert config.get_complete_url( + api_base="https://test.net/ai-gateway/mlflow/v1", + api_key="test-key", + model="catalog.schema.kimi-k3", + optional_params={}, + litellm_params={}, + ) == "https://test.net/ai-gateway/mlflow/v1/chat/completions" + assert request["model"] == "catalog.schema.kimi-k3" + + def test_chat_legacy_endpoint_remains_default(self, monkeypatch): + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + url = DatabricksConfig().get_complete_url( + api_base="https://test.net/serving-endpoints", + api_key="test-key", + model="databricks-kimi-k3", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://test.net/serving-endpoints/chat/completions" + def test_embeddings_endpoint(self, monkeypatch): """Embeddings endpoint is correctly appended.""" monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py new file mode 100644 index 00000000000..bc0ea23e249 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py @@ -0,0 +1,155 @@ +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config +from litellm.llms.hosted_vllm.image_edit.transformation import HostedVLLMImageEditConfig +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +PNG_BYTES = b"\x89PNG\r\n\x1a\nfakepng" +MODEL = "Qwen/Qwen-Image-Edit-2511" + + +@pytest.fixture(autouse=True) +def _clear_hosted_vllm_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False) + monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False) + + +def test_provider_config_registration(): + config = ProviderConfigManager.get_provider_image_edit_config( + model=f"hosted_vllm/{MODEL}", + provider=LlmProviders.HOSTED_VLLM, + ) + + assert isinstance(config, HostedVLLMImageEditConfig) + assert isinstance(get_hosted_vllm_image_edit_config(MODEL), HostedVLLMImageEditConfig) + + +@pytest.mark.parametrize( + "api_base", + ["http://localhost:8091", "http://localhost:8091/", "http://localhost:8091/v1", "http://localhost:8091/v1/"], +) +def test_get_complete_url_appends_images_edits(api_base: str): + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=api_base, litellm_params={}) + == "http://localhost:8091/v1/images/edits" + ) + + +def test_get_complete_url_falls_back_to_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://vllm-omni:8000/v1") + config = HostedVLLMImageEditConfig() + + assert ( + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + == "http://vllm-omni:8000/v1/images/edits" + ) + + +def test_get_complete_url_requires_api_base(): + config = HostedVLLMImageEditConfig() + + with pytest.raises(ValueError, match="api_base not set"): + config.get_complete_url(model=MODEL, api_base=None, litellm_params={}) + + +def test_validate_environment_defaults_to_fake_api_key(): + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers == {"Authorization": "Bearer fake-api-key"} + + +def test_validate_environment_uses_provided_api_key_and_keeps_headers(): + headers = HostedVLLMImageEditConfig().validate_environment( + headers={"X-Test": "1"}, + model=MODEL, + api_key="my-custom-key", + ) + + assert headers == {"X-Test": "1", "Authorization": "Bearer my-custom-key"} + + +def test_validate_environment_falls_back_to_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOSTED_VLLM_API_KEY", "env-key") + + headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL) + + assert headers["Authorization"] == "Bearer env-key" + + +def test_image_edit_posts_multipart_to_vllm_omni(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + response = litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + api_key="test-key", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + seed=42, + ) + + assert response.data + assert len(captured) == 1 + request = captured[0] + assert str(request.url) == "http://localhost:8091/v1/images/edits" + assert request.headers["authorization"] == "Bearer test-key" + assert request.headers["content-type"].startswith("multipart/form-data") + assert b'name="image[]"' in request.content + assert PNG_BYTES in request.content + assert f'name="model"\r\n\r\n{MODEL}'.encode() in request.content + assert b'name="prompt"\r\n\r\nadd a hat' in request.content + assert b'name="seed"\r\n\r\n42' in request.content + + +@pytest.mark.parametrize("param", ["mask", "quality", "input_fidelity"]) +def test_params_vllm_omni_ignores_are_not_advertised(param: str): + supported = HostedVLLMImageEditConfig().get_supported_openai_params(MODEL) + + assert param not in supported + assert {"image", "prompt", "n", "size", "response_format", "background", "user"} <= set(supported) + + +def test_image_edit_rejects_quality_unless_dropped(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + with pytest.raises(litellm.UnsupportedParamsError, match="quality"): + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + ) + assert captured == [] + + litellm.image_edit( + model=f"hosted_vllm/{MODEL}", + image=PNG_BYTES, + prompt="add a hat", + api_base="http://localhost:8091", + client=client, + quality="low", + drop_params=True, + ) + + assert len(captured) == 1 + assert b'name="quality"' not in captured[0].content + assert b'name="prompt"\r\n\r\nadd a hat' in captured[0].content diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index aff0530ee8b..5a29a96829f 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1113,6 +1113,102 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert chunks[1].choices[0].delta.content in (None, "") assert chunks[1].choices[0].finish_reason == "stop" + @staticmethod + def _ended_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None): + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(fragment("", name="lookup_fruit", call_id="call_1")), + chunk(fragment('{"fruit":')), + chunk(fragment(' "persimmon"}')), + chunk(None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_arguments_back_into_chunks(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""] + assert fragments[0][0].function.name == "lookup_fruit" + assert fragments[0][0].id == "call_1" + assert chunks[3].choices[0].delta.tool_calls is None + assert chunks[3].choices[0].finish_reason == "tool_calls" + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_tool_call_name_back_into_chunks(self): + class RenameTool(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for tool_call in inputs.get("tool_calls", []): + tool_call["function"]["name"] = "lookup_fruit_reviewed" + return inputs + + handler = OpenAIChatCompletionsHandler() + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=RenameTool(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]] + assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None] + assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"} + assert fragments[0].id == "call_1" + + @pytest.mark.asyncio + async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_tool_call_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] + assert [fragment[0].function.arguments for fragment in fragments] == ["", '{"fruit":', ' "persimmon"}'] + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): handler = OpenAIChatCompletionsHandler() @@ -1179,6 +1275,62 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + @staticmethod + def _two_choice_tool_call_stream_chunks() -> list: + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk( + choice_index: int, tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None + ) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=choice_index, + delta=Delta(tool_calls=[tool_call] if tool_call else None), + finish_reason=finish_reason, + ) + ], + ) + + def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None): + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + return [ + chunk(0, fragment("", name="lookup_fruit", call_id="call_1")), + chunk(1, fragment("", name="lookup_fruit", call_id="call_2")), + chunk(0, fragment('{"fruit": "persimmon"}')), + chunk(1, fragment('{"fruit": "durian"}')), + chunk(0, None, finish_reason="tool_calls"), + chunk(1, None, finish_reason="tool_calls"), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 9737d63cc26..b110586ae5b 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -4,6 +4,7 @@ Tests for OpenAI GPT transformation (litellm/llms/openai/chat/gpt_transformation import pytest +from typing import Final import litellm @@ -1168,6 +1169,9 @@ class TestOpenAIPromptCacheBreakpointChatPath: assert "prompt_cache_options" not in request +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + + class TestToolSchemaCombinatorFlatteningForOpenAI: """ Regression tests for LIT-6488: OpenAI's chat completions validator rejects @@ -1281,3 +1285,50 @@ class TestToolSchemaCombinatorFlatteningForOpenAI: parameters = request["tools"][0]["function"]["parameters"] assert "anyOf" not in parameters assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + + @staticmethod + def _artifact_tool(): + return { + "type": "function", + "function": { + "name": "Artifact", + "parameters": { + "type": "object", + "properties": {"field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}}, + "required": ["field"], + }, + }, + } + + def test_drops_non_python_regex_pattern_for_hosted_openai(self): + tool = self._artifact_tool() + + request = self._transform(self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [tool]) + + assert request["tools"][0]["function"]["parameters"] == { + "type": "object", + "properties": {"field": {"type": "string"}}, + "required": ["field"], + } + assert tool == self._artifact_tool() + + def test_custom_api_base_drops_non_python_regex_pattern_but_keeps_union(self): + tool = self._anyof_tool() + tool["function"]["parameters"]["properties"]["id"]["pattern"] = _ARTIFACT_FIELD_PATTERN + + request = self._transform( + self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": "http://localhost:8000/v1"}, [tool] + ) + + parameters = request["tools"][0]["function"]["parameters"] + assert parameters["properties"]["id"] == {"type": "string"} + assert parameters["anyOf"] == self._anyof_tool()["function"]["parameters"]["anyOf"] + + def test_non_openai_provider_keeps_non_python_regex_pattern(self): + tool = self._artifact_tool() + + request = self._transform( + self.config, "some-oss-model", {"custom_llm_provider": "groq", "api_base": None}, [tool] + ) + + assert request["tools"][0] is tool diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 33c8c97fea7..a4f0a77a9b6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -10,23 +10,33 @@ from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock +import logging + import pytest from fastapi import HTTPException -from openai.types.responses import ResponseFunctionToolCall +from pydantic import BaseModel +from openai.types.responses import ( + ResponseCustomToolCall, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, + ResponseFunctionToolCall, +) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.responses.main import GenericResponseOutputItem, OutputText +from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -56,6 +66,60 @@ class MockGuardrail(CustomGuardrail): return inputs +class PersimmonMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + tool_calls = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": tool_call["function"]["arguments"].replace("persimmon", "[MASKED]"), + }, + } + for tool_call in inputs.get("tool_calls", []) + ] + return {**inputs, "tool_calls": tool_calls} + + +class FlatShapeGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + flat_tool_calls = [{"name": "exec", "input": "rm -rf /"} for _ in inputs.get("tool_calls", [])] + return {**inputs, "tool_calls": flat_tool_calls} + + +class DroppingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "tool_calls": []} + + +CUSTOM_TOOL_CALL_ITEM = { + "type": "custom_tool_call", + "id": "ctc_1", + "call_id": "call_exec_1", + "name": "exec", + "input": "echo persimmon", + "status": "completed", +} + + class TestOpenAIResponsesHandlerDiscovery: """Test that the handler is properly discovered by the guardrail system""" @@ -556,7 +620,7 @@ class TestOpenAIResponsesHandlerToolCallExtraction: texts_to_check: List[str] = [] images_to_check: List[str] = [] - tool_calls_to_check: List[Any] = [] + tool_calls_to_check: List[ChatCompletionToolCallChunk] = [] task_mappings: List[Tuple[int, int]] = [] # Extract tool calls @@ -627,6 +691,123 @@ class TestOpenAIResponsesHandlerToolCallExtraction: == '{"location":"Boston, MA","unit":"celsius"}' ) + @pytest.mark.parametrize( + "output_item", + [ + dict(CUSTOM_TOOL_CALL_ITEM), + CustomToolCallOutputItem(**CUSTOM_TOOL_CALL_ITEM), + ResponseCustomToolCall(**{key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "status"}), + ], + ids=["dict", "litellm_typed", "openai_typed"], + ) + def test_extract_custom_tool_call_input_as_arguments(self, output_item): + handler = OpenAIResponsesHandler() + texts_to_check: List[str] = [] + tool_calls_to_check: List[Any] = [] + + handler._extract_output_text_and_images( + output_item=output_item, + output_idx=2, + texts_to_check=texts_to_check, + images_to_check=[], + task_mappings=[], + tool_calls_to_check=tool_calls_to_check, + ) + + assert texts_to_check == [] + assert tool_calls_to_check == [ + { + "id": "call_exec_1", + "type": "function", + "function": {"name": "exec", "arguments": "echo persimmon"}, + "index": 2, + } + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("typed", [False, True], ids=["dict", "typed"]) + async def test_process_output_response_writes_tool_call_rewrites_back(self, typed): + handler = OpenAIResponsesHandler() + function_call = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_fn_1", + "name": "lookup_fruit", + "arguments": '{"fruit": "persimmon"}', + "status": "completed", + } + message = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "running persimmon", "annotations": []}], + } + payload = { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [message, function_call, dict(CUSTOM_TOOL_CALL_ITEM)], + } + response = ResponsesAPIResponse.model_validate(payload) if typed else payload + + result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask")) + + output = result.output if typed else result["output"] + function_item, custom_item = output[1], output[2] + assert (function_item.arguments if typed else function_item["arguments"]) == '{"fruit": "[MASKED]"}' + assert (custom_item.input if typed else custom_item["input"]) == "echo [MASKED]" + assert (custom_item.name if typed else custom_item["name"]) == "exec" + assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon" + + @staticmethod + def _custom_tool_call_response(item: dict) -> dict: + return { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [item], + } + + @pytest.mark.asyncio + async def test_process_output_response_ignores_tool_call_rewrites_in_another_shape(self): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + result = await handler.process_output_response(response, FlatShapeGuardrail(guardrail_name="flat")) + + assert result["output"][0]["input"] == "echo persimmon" + assert result["output"][0]["name"] == "exec" + + @pytest.mark.asyncio + async def test_process_output_response_warns_when_guardrail_drops_tool_calls(self, caplog): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await handler.process_output_response(response, DroppingGuardrail(guardrail_name="dropper")) + + assert result["output"][0]["input"] == "echo persimmon" + assert any( + "dropper" in record.getMessage() and "0 tool calls for the 1 scanned" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_process_output_response_keeps_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + nameless_item = {key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "name"} + response = self._custom_tool_call_response(nameless_item) + + result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask")) + + assert result["output"][0]["input"] == "echo [MASKED]" + assert "name" not in result["output"][0] + @pytest.mark.asyncio async def test_process_output_response_with_tool_calls(self): """Test processing output response containing function tool calls""" @@ -1195,6 +1376,352 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @staticmethod + def _ended_function_call_stream_events() -> List[dict]: + def item(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_123", + "call_id": "call_123", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_123", + "output_index": 0, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "created_at": 1, + "model": "gpt-4o", + "output": [item('{"fruit": "persimmon"}', "completed")], + "status": "completed", + }, + }, + ] + + @staticmethod + def _argument_masking_guardrail() -> CustomGuardrail: + class MaskArguments(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + tool_calls = [ + {**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}} + for tool_call in inputs.get("tool_calls", []) + ] + return {**inputs, "tool_calls": tool_calls} + + return MaskArguments(guardrail_name="test-mask-arguments") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_function_call_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["item"]["arguments"] == "" + assert events[1]["delta"] == '{"fruit": "[MASKED]"}' + assert events[2]["delta"] == "" + assert events[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["response"]["output"][0]["name"] == "lookup_fruit" + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_typed_function_call_events(self): + from litellm.types.llms.openai import ( + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ) + + handler = OpenAIResponsesHandler() + typed_events: List[Any] = [ + model.model_validate(event) + for model, event in zip( + ( + OutputItemAddedEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ), + self._ended_function_call_stream_events(), + ) + ] + completed_event = typed_events[5] + assert isinstance(completed_event, ResponseCompletedEvent) + assert isinstance(completed_event.response, ResponsesAPIResponse) + assert isinstance(completed_event.response.output[0], ResponseFunctionToolCall) + + await handler.process_output_streaming_response( + responses_so_far=typed_events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert typed_events[1].delta == '{"fruit": "[MASKED]"}' + assert typed_events[2].delta == "" + assert typed_events[3].arguments == '{"fruit": "[MASKED]"}' + assert typed_events[4].item.arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}' + assert completed_event.response.output[0].name == "lookup_fruit" + + @staticmethod + def _ended_custom_tool_call_stream_events() -> List[dict]: + def item(input_text: str, status: str) -> dict: + return {**CUSTOM_TOOL_CALL_ITEM, "input": input_text, "status": status} + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "echo "}, + {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "persimmon"}, + {"type": "response.custom_tool_call_input.done", "item_id": "ctc_1", "output_index": 0, "input": "echo persimmon"}, + {"type": "response.output_item.done", "output_index": 0, "item": item("echo persimmon", "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "created_at": 1, + "model": "gpt-5.6", + "output": [item("echo persimmon", "completed")], + "status": "completed", + }, + }, + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_custom_tool_call_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["item"]["input"] == "" + assert events[1]["delta"] == "echo [MASKED]" + assert events[2]["delta"] == "" + assert events[3]["input"] == "echo [MASKED]" + assert events[4]["item"]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["name"] == "exec" + assert "arguments" not in events[5]["response"]["output"][0] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keep_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + items = [events[0]["item"], events[4]["item"], events[5]["response"]["output"][0]] + for item in items: + del item["name"] + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[3]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]" + assert all("name" not in item for item in items) + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self): + from litellm.types.llms.openai import ( + OutputItemAddedEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ) + + handler = OpenAIResponsesHandler() + typed_events: List[BaseModel] = [ + model.model_validate({**event, "sequence_number": sequence_number}) + for sequence_number, (model, event) in enumerate( + zip( + ( + OutputItemAddedEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, + OutputItemDoneEvent, + ResponseCompletedEvent, + ), + self._ended_custom_tool_call_stream_events(), + ) + ) + ] + completed_event = typed_events[5] + assert isinstance(completed_event.response.output[0], CustomToolCallOutputItem) + + await handler.process_output_streaming_response( + responses_so_far=typed_events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert typed_events[1].delta == "echo [MASKED]" + assert typed_events[2].delta == "" + assert typed_events[3].input == "echo [MASKED]" + assert typed_events[4].item.input == "echo [MASKED]" + assert completed_event.response.output[0].input == "echo [MASKED]" + assert completed_event.response.output[0].name == "exec" + + @pytest.mark.asyncio + async def test_deliver_ended_stream_custom_tool_call_rewrite_without_matching_events_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + events[5]["response"]["output"] = [{**events[5]["response"]["output"][0], "call_id": "call_999"}] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @staticmethod + def _bridged_function_call_stream_events() -> List[dict]: + reasoning = {"type": "reasoning", "id": "rs_1", "summary": []} + text = {"type": "output_text", "text": "Looking that up", "annotations": []} + message = {"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [text]} + + def function_call(arguments: str, status: str) -> dict: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning)}, + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "status": "in_progress", "content": []}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": "Looking that up"}, + {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": [dict(text)]}}, + {"type": "response.output_item.added", "output_index": 1, "item": function_call("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.done", + "item_id": "fc_1", + "output_index": 1, + "arguments": '{"fruit": "persimmon"}', + }, + {"type": "response.output_item.done", "output_index": 1, "item": function_call('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "model": "claude-haiku-4-5", + "output": [ + dict(reasoning), + {**message, "content": [dict(text)]}, + function_call('{"fruit": "persimmon"}', "completed"), + ], + }, + }, + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keys_bridged_function_call_events_by_call_id(self): + handler = OpenAIResponsesHandler() + events = self._bridged_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[6]["delta"] == '{"fruit": "[MASKED]"}' + assert events[7]["delta"] == "" + assert events[8]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[5]["item"]["name"] == "lookup_fruit" + assert events[9]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[10]["response"]["output"][2]["arguments"] == '{"fruit": "[MASKED]"}' + assert events[3]["delta"] == "Looking that up" + assert events[4]["item"]["content"][0]["text"] == "Looking that up" + assert events[10]["response"]["output"][1]["content"][0]["text"] == "Looking that up" + assert events[1]["item"] == {"type": "reasoning", "id": "rs_1", "summary": []} + + @pytest.mark.asyncio + @pytest.mark.parametrize("mismatch", ["orphan_call_id", "duplicate_call_id"]) + async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(self, mismatch): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + envelope_item = events[5]["response"]["output"][0] + if mismatch == "orphan_call_id": + events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}] + else: + events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self): + handler = OpenAIResponsesHandler() + events = self._ended_function_call_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._argument_masking_guardrail(), + litellm_logging_obj=None, + ) + + assert events[1]["delta"] == '{"fruit":' + assert events[3]["arguments"] == '{"fruit": "persimmon"}' + assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "persimmon"}' + @pytest.mark.asyncio @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type): @@ -2522,8 +3049,21 @@ class TestOpenAIResponsesHandlerStreamingScanKey: assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] assert ended_key != open_key + def test_completed_event_with_a_custom_tool_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + ended_key = handler.get_streaming_scan_key( + [self._delta(0, "hi"), self._completed(1, [message, dict(CUSTOM_TOOL_CALL_ITEM)])] + ) + rewritten_key = handler.get_streaming_scan_key( + [self._delta(0, "hi"), self._completed(1, [message, {**CUSTOM_TOOL_CALL_ITEM, "input": "echo kumquat"}])] + ) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "echo persimmon" in ended_key.tool_calls[0] + assert rewritten_key != ended_key + def test_completed_event_reads_every_output_text_part(self): - from litellm.types.responses.main import GenericResponseOutputItem, OutputText + from litellm.types.responses.main import CustomToolCallOutputItem, GenericResponseOutputItem, OutputText item = GenericResponseOutputItem( type="message", diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index c5902b32a06..4cf8767764b 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -1,5 +1,6 @@ import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -20,6 +21,8 @@ from litellm.types.llms.openai import ( ) from litellm.types.router import GenericLiteLLMParams +_ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' + class TestOpenAIResponsesAPIConfig: def setup_method(self): @@ -2022,6 +2025,86 @@ class TestFlattenToolSchemaCombinatorsWiring: assert "anyOf" not in result["tools"][1]["parameters"] +class TestToolSchemaRegexPatternWiring: + """Claude Code's Artifact tool reaches /v1/responses (the /v1/messages bridge) with an + ECMA-262 ``pattern``; OpenAI compiles patterns with Python ``re`` and 400s + "'...' is not a 'regex'" for every model family, so the keyword is dropped. + """ + + def _artifact_tool(self): + return { + "type": "function", + "name": "Artifact", + "parameters": { + "type": "object", + "properties": { + "field": {"type": "string", "pattern": _ARTIFACT_FIELD_PATTERN}, + "doc_id": {"type": "string", "pattern": r"^(?!\.\.?(?:/|$))[A-Za-z0-9_\-.~:@+]{1,200}$"}, + }, + "required": ["field"], + }, + } + + @pytest.mark.parametrize("model", ["gpt-5.6", "gpt-4o", "o3"]) + def test_openai_drops_only_the_pattern_python_re_rejects_for_every_family(self, model): + tool = self._artifact_tool() + + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model=model, + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + properties = result["tools"][0]["parameters"]["properties"] + assert properties["field"] == {"type": "string"} + assert properties["doc_id"] == tool["parameters"]["properties"]["doc_id"] + assert result["tools"][0]["parameters"]["required"] == ["field"] + assert tool["parameters"]["properties"]["field"]["pattern"] == _ARTIFACT_FIELD_PATTERN + assert json.loads(json.dumps(result["tools"])) == result["tools"] + + def test_openai_drops_patterns_inside_codex_namespace_tools(self): + namespace = {"type": "namespace", "name": "mcp__claude", "tools": [self._artifact_tool()]} + + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-5.6", + input="hi", + response_api_optional_request_params={"tools": [namespace]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0]["tools"][0]["parameters"]["properties"]["field"] == {"type": "string"} + + def test_openai_compact_request_drops_patterns(self): + _, data = OpenAIResponsesAPIConfig().transform_compact_response_api_request( + model="gpt-5.6", + input="hi", + response_api_optional_request_params={"tools": [self._artifact_tool()]}, + api_base="https://api.openai.com/v1/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["tools"][0]["parameters"]["properties"]["field"] == {"type": "string"} + + def test_non_openai_subclass_keeps_patterns(self): + from litellm.llms.hosted_vllm.responses.transformation import HostedVLLMResponsesAPIConfig + + tool = self._artifact_tool() + + result = HostedVLLMResponsesAPIConfig().transform_responses_api_request( + model="hosted_vllm/qwen", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0] is tool + + class TestReasoningFollowsModelSupport: """Responses API clients like Codex send `reasoning` on every request, and OpenAI 400s it on non-reasoning models like gpt-4o. drop_params must strip it there, the same way the diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py index eadd87d9c92..08e46b1ffac 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -322,8 +322,8 @@ class TestModelCostEntry: entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-preview"] assert entry["mode"] == "audio_transcription" assert entry["litellm_provider"] == "vertex_ai" - assert entry["input_cost_per_audio_token"] == pytest.approx(2.5e-06) - assert entry["input_cost_per_token"] == pytest.approx(2.5e-06) + assert entry["input_cost_per_audio_token"] == pytest.approx(2e-06) + assert entry["input_cost_per_token"] == pytest.approx(2e-06) assert entry["output_cost_per_token"] == pytest.approx(1.2e-05) assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 1950c37a12e..546cff18b5d 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -873,3 +873,118 @@ def test_llm_passthrough_route_propagates_allm_passthrough_route_to_logging_obj( assert captured_litellm_params.get("allm_passthrough_route") is True assert LitellmLogging._is_sync_litellm_request(captured_litellm_params) is False + + +FOUNDRY_BASE = "https://my-resource.services.ai.azure.com" + + +def _foundry_parse_response() -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id":"parse-1","pages":[]}', + request=httpx.Request("POST", f"{FOUNDRY_BASE}/providers/cohere/v2/parse"), + ) + + +def test_azure_ai_relay_reaches_the_deployment_with_its_own_credential(): + """ + Regression for LIT-7022: azure_ai had no passthrough config, so every + /azure_ai// relay raised "Provider azure_ai not found" + before a request was built. + """ + client = HTTPHandler() + + with patch.object(client.client, "send", return_value=_foundry_parse_response()) as mock_send: + response = llm_passthrough_route( + model="azure_ai/Cohere-parse-v5", + endpoint="Cohere-parse-v5/providers/cohere/v2/parse", + method="POST", + custom_llm_provider="azure_ai", + api_base=FOUNDRY_BASE, + api_key="deployment-key", + json={"model": "Cohere-parse-v5", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=client, + litellm_logging_obj=MagicMock(), + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_azure_ai_model_through_the_deployment_api_base(): + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-parse", + "litellm_params": { + "model": "azure_ai/Cohere-parse-v5", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + + with patch.object(async_client.client, "send", AsyncMock(return_value=_foundry_parse_response())) as mock_send: + response = await router.allm_passthrough_route( + model="foundry-parse", + method="POST", + endpoint="foundry-parse/providers/cohere/v2/parse", + json={"model": "foundry-parse", "document": {"type": "image_url", "image_url": "https://x/y.png"}}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/providers/cohere/v2/parse" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "Cohere-parse-v5" + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_router_relays_an_openai_model_on_a_foundry_base_as_azure_ai(monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + router = litellm.Router( + model_list=[ + { + "model_name": "foundry-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": FOUNDRY_BASE, + "api_key": "deployment-key", + }, + } + ] + ) + async_client = AsyncHTTPHandler() + upstream = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=( + b'{"id":"chatcmpl-1","object":"chat.completion","model":"gpt-5.4-mini",' + b'"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"hi"}}],' + b'"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' + ), + request=httpx.Request("POST", f"{FOUNDRY_BASE}/models/chat/completions"), + ) + + with patch.object(async_client.client, "send", AsyncMock(return_value=upstream)) as mock_send: + await router.allm_passthrough_route( + model="foundry-gpt", + method="POST", + endpoint="foundry-gpt/models/chat/completions", + request_query_params={"api-version": "2024-05-01-preview"}, + json={"model": "foundry-gpt", "messages": [{"role": "user", "content": "hi"}]}, + client=async_client, + ) + + sent = mock_send.call_args.kwargs["request"] + assert str(sent.url) == f"{FOUNDRY_BASE}/models/chat/completions?api-version=2024-05-01-preview" + assert sent.headers["api-key"] == "deployment-key" + assert json.loads(sent.content)["model"] == "gpt-5.4-mini" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 20f4719e4bf..e6c8d4ee039 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7193,14 +7193,13 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}" - async def test_per_server_challenge_for_gateway_managed_oauth2(self): - """Anonymous request to a per-server path whose single target is a gateway-managed - oauth2 server: 401 plus the RFC 9728 challenge advertising the PER-SERVER - protected-resource metadata in the same URL spelling the request used, so a keyless - DCR client configured with either per-server spelling discovers the gateway as the - authorization server (LIT-4864). Covers interactive and M2M, which the gateway can - both serve end to end.""" - from litellm.types.mcp import MCPAuth + @pytest.mark.parametrize( + "auth_type", + (None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token", "oauth2"), + ) + @pytest.mark.parametrize("bearer_presented", (False, True)) + async def test_per_server_challenge_for_gateway_owned_auth(self, auth_type, bearer_presented): + """Gateway admission challenges are independent of upstream authentication.""" from litellm.types.mcp_server.mcp_server_manager import MCPServer server = MCPServer( @@ -7209,7 +7208,7 @@ class TestAggregateGatewayDcrChallenge: server_name="github", url="https://upstream.example/mcp", transport="http", - auth_type=MCPAuth.oauth2, + auth_type=auth_type, ) for path, expected_metadata_path in ( ("/mcp/github", "/.well-known/oauth-protected-resource/mcp/github"), @@ -7223,10 +7222,16 @@ class TestAggregateGatewayDcrChallenge: ): mock_mgr.get_mcp_server_by_name.return_value = server with pytest.raises(HTTPException) as exc_info: - await MCPRequestHandler.process_mcp_request(self._scope(path=path)) + await MCPRequestHandler.process_mcp_request( + self._scope( + path=path, + extra_headers=((b"authorization", b"Bearer invalid-key"),) if bearer_presented else (), + ) + ) assert exc_info.value.status_code == 401 www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] - assert www_authenticate == f'Bearer resource_metadata="http://testserver{expected_metadata_path}"' + error = 'error="invalid_token", ' if bearer_presented else "" + assert www_authenticate == f'Bearer {error}resource_metadata="http://testserver{expected_metadata_path}"' async def test_per_server_challenge_keeps_spelling_under_server_root_path(self): """On a sub-path deployment the challenge must still advertise the spelling the client @@ -7303,10 +7308,7 @@ class TestAggregateGatewayDcrChallenge: ) def test_challenge_target_excludes_every_non_gateway_managed_mode(self): - """Unit pin of the challenge-target owner: only a resolved gateway-managed oauth2 - target (interactive or M2M) yields a per-server challenge; delegate-auth oauth2 - (whose keyless flow is upstream PKCE via the relay), every client-forwarded auth - type, OBO, api_key, unknown names, and CSV paths yield None (LIT-4864).""" + """Gateway challenges exclude unresolved, delegated, and client-forwarded targets.""" from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( _gateway_dcr_challenge_target, ) @@ -7333,7 +7335,11 @@ class TestAggregateGatewayDcrChallenge: (_server(MCPAuth.true_passthrough), None), (_server(MCPAuth.oauth_delegate), None), (_server(MCPAuth.oauth_delegate, dcr_bridge=True), None), - (_server(MCPAuth.api_key), None), + (_server(MCPAuth.api_key), "srv"), + (_server(MCPAuth.none, extra_headers=["Authorization"]), None), + (_server(None, extra_headers=["X-API-Key"]), None), + (_server(MCPAuth.none, extra_headers=["Authorization"], oauth_passthrough=True), None), + (_server(MCPAuth.oauth2_id_jag), None), (None, None), ] for resolved, expected in cases: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 6b3098d9e60..5fab4ceec72 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -19,6 +19,7 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials import ( ApiKeyConfig, + AuthConfig, AuthorizationCodeConfig, AwsSigV4Config, Byok, @@ -1203,3 +1204,71 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot(): assert isinstance(result, Ok) headers, _ = await _emitted_async(result.ok) assert headers["Authorization"] == "caller-token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("config", "subject", "expected_source", "expected_header"), + [ + (NoneConfig(), _SUBJECT, "no-auth", None), + (PassthroughConfig(), _SUBJECT, "no-auth", None), + (PassthroughConfig(), _with_inbound("Bearer caller-token"), "oauth2-passthrough", "Bearer caller-token"), + (ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), _SUBJECT, "static-token", "Bearer static-key"), + (AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"), + ], +) +async def test_resolved_source_matches_the_credential_sent_upstream( + config: AuthConfig, subject: Subject, expected_source: str, expected_header: str | None +) -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="stored-alice")}) + provider = UpstreamCredentialProvider(oauth_token_store=store) + result = await resolve_credentials_with_source(provider, subject, _spec(config)) + assert isinstance(result, Ok) + assert result.ok.source.value == expected_source + assert _emitted(result.ok.auth).get("Authorization") == expected_header + assert "stored-alice" not in repr(result.ok) + assert "static-key" not in repr(result.ok) + + +@pytest.mark.asyncio +async def test_resolved_source_preserves_missing_user_token_error() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + +@pytest.mark.asyncio +async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) + m2m = await resolve_credentials_with_source(UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M)) + assert isinstance(m2m, Ok) + headers, _ = await _emitted_async(m2m.ok.auth) + assert headers["Authorization"] == "Bearer m2m-at" + assert m2m.ok.source.value == "m2m-client-credentials" + assert source.gets == ["s"] + + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) + exchanged = await resolve_credentials_with_source(UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO)) + assert isinstance(exchanged, Ok) + assert _emitted(exchanged.ok.auth)["Authorization"] == "Bearer exchanged-at" + assert exchanged.ok.source.value == "token-exchange" + assert len(exchanger.calls) == 1 + + +@pytest.mark.asyncio +async def test_id_jag_source_describes_final_token_after_both_exchanges() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source + + endpoint = _FakeTokenEndpoint(_two_leg_ok("resource-token")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await resolve_credentials_with_source(provider, _with_inbound("identity-token"), _spec(_id_jag_config())) + assert isinstance(result, Ok) + assert result.ok.source.value == "id-jag" + assert _emitted(result.ok.auth)["Authorization"] == "Bearer resource-token" + assert len(endpoint.calls) == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 763200c3709..a7a65bfe466 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7651,12 +7651,6 @@ async def test_token_endpoint_client_secret_basic_without_secret_returns_400(): assert exc_info.value.status_code == 400 -# ------------------------------------------------------------------- -# Non-oauth2 (auth_type=none, access-group gated) servers must not be -# driven through the gateway OAuth authorize/token/register/discovery -# flow, and must not be advertised as OAuth-protected in discovery docs. -# ------------------------------------------------------------------- - def _access_group_none_server(server_name="access_group_server"): """A non-oauth2, access-group gated MCP server: no client_id, no OAuth.""" @@ -7794,35 +7788,38 @@ async def test_register_client_rejects_non_oauth2_server(): @pytest.mark.asyncio -async def test_oauth_protected_resource_404_for_non_oauth2_server(): - """Discovery must not advertise a none-auth server as an OAuth-protected resource.""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_protected_resource_response, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - except ImportError: - pytest.skip("MCP discoverable endpoints not available") +@pytest.mark.parametrize( + "auth_type", (None, "none", "api_key", "bearer_token", "basic", "aws_sigv4", "authorization", "token") +) +@pytest.mark.parametrize("use_standard_pattern", (False, True)) +async def test_oauth_protected_resource_for_gateway_owned_auth(auth_type, use_standard_pattern): + from starlette.requests import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_protected_resource_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _access_group_none_server().model_copy(update={"auth_type": auth_type}) + request = Request( + {"type": "http", "scheme": "https", "path": "/", "headers": [(b"host", b"litellm.example.com")]} + ) global_mcp_server_manager.registry.clear() - server = _access_group_none_server() global_mcp_server_manager.registry[server.server_id] = server - - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - try: - with pytest.raises(HTTPException) as exc_info: - await _build_oauth_protected_resource_response( - request=mock_request, - mcp_server_name="access_group_server", - use_standard_pattern=False, - ) - assert exc_info.value.status_code == 404 - assert "not an OAuth-protected resource" in str(exc_info.value.detail) + response = await _build_oauth_protected_resource_response( + request=request, + mcp_server_name="access_group_server", + use_standard_pattern=use_standard_pattern, + ) + resource_path = "/mcp/access_group_server" if use_standard_pattern else "/access_group_server/mcp" + assert response == { + "resource": f"https://litellm.example.com{resource_path}", + "authorization_servers": ["https://litellm.example.com/mcp"], + "scopes_supported": [], + } finally: global_mcp_server_manager.registry.clear() @@ -7914,9 +7911,7 @@ async def test_oauth_protected_resource_passthrough_none_auth_not_404(): @pytest.mark.asyncio async def test_oauth_protected_resource_404_for_unknown_server_name(): - """A discovery request for an unknown server name returns the same 404 as a non-oauth2 - server (not a 200 metadata doc with broken URLs), so the well-known paths cannot be used - to enumerate non-OAuth server names.""" + """Unknown server names must not produce metadata advertising nonexistent resources.""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _build_oauth_protected_resource_response, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 73a52a8d2e8..8aa4cfc5619 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -833,7 +833,7 @@ async def test_manual_delivery_page_renders_the_url_as_data_never_as_a_shell_com assert 'value="' in body -def _scoped_mcp_server(name="github", **kw): +def _scoped_mcp_server(name="github", auth_type="oauth2", **kw): from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -844,7 +844,7 @@ def _scoped_mcp_server(name="github", **kw): alias=name, url="https://upstream.example/mcp", transport="http", - auth_type=MCPAuth.oauth2, + auth_type=MCPAuth(auth_type) if auth_type is not None else None, **kw, ) @@ -2044,3 +2044,45 @@ async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage(): status, body = await _introspect(minted.token.get_secret_value(), master_key=None) assert (status, body["error"]) == (500, "server_error") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_type", [None, "none", "api_key", "bearer_token", "basic", "authorization", "token", "aws_sigv4"] +) +@pytest.mark.parametrize("resource", ["https://llm.example.com/mcp/github", "https://llm.example.com/github/mcp"]) +async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(auth_type, resource): + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server(auth_type=auth_type) + vendor = _VendorCredential("absent") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, resource) + described = await _describe_page(response, scoped_server=server, vendor=vendor) + assert json.loads(described.body) == { + "state": "m2m", + "client_origin": "https://claude.ai", + "server_id": "github-id", + "server_name": "github", + "connected": True, + } + unreachable = await _complete_page(response, scoped_server=server, reachable=_ServerReachability(False)) + assert unreachable.status_code == 400 + cache = DualCache() + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + assert vendor.calls == [] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + redeemed = await _redeem(code, client_id, cache=cache, resource=resource) + assert redeemed.status_code == 200 + payload = json.loads(redeemed.body) + assert _opened_principal(payload).resource_server_id == "github-id" + renewed = await _redeem( + None, client_id, cache=cache, grant_type="refresh_token", refresh_token=payload["refresh_token"] + ) + assert renewed.status_code == 200 + assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index d299239f68e..7d46bb0237a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -3,11 +3,17 @@ Tests for MCPDebug — MCP OAuth2 debug response headers. """ import asyncio -from unittest.mock import MagicMock +from typing import Final + +import pytest +from starlette.types import Message + +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution from litellm.proxy._experimental.mcp_server.mcp_debug import ( MCP_DEBUG_REQUEST_HEADER, MCPDebug, + MCPAuthDiagnostics, ) @@ -166,77 +172,6 @@ class TestBuildDebugHeaders: assert set(headers.keys()) == expected_keys -class TestResolveAuthResolution: - def _make_server(self, **kwargs): - server = MagicMock() - server.alias = kwargs.get("alias", "test") - server.server_name = kwargs.get("server_name", "test") - server.has_client_credentials = kwargs.get("has_client_credentials", False) - server.authentication_token = kwargs.get("authentication_token", None) - server.auth_type = kwargs.get("auth_type", None) - return server - - def test_per_request_header(self): - server = self._make_server() - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header="Bearer xxx", - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "per-request-header" - - def test_server_specific_header(self): - server = self._make_server(alias="atlas") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers={"atlas": {"Authorization": "Bearer xxx"}}, - oauth2_headers=None, - ) - assert result == "per-request-header" - - def test_m2m(self): - server = self._make_server(has_client_credentials=True) - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "m2m-client-credentials" - - def test_static_token(self): - server = self._make_server(authentication_token="static-tok") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "static-token" - - def test_oauth2_passthrough(self): - server = self._make_server(auth_type="oauth2") - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers={"Authorization": "Bearer eyJ..."}, - ) - assert result == "oauth2-passthrough" - - def test_no_auth(self): - server = self._make_server() - result = MCPDebug.resolve_auth_resolution( - server, - mcp_auth_header=None, - mcp_server_auth_headers=None, - oauth2_headers=None, - ) - assert result == "no-auth" - - class TestWrapSendWithDebugHeaders: def test_injects_headers(self): captured = [] @@ -269,3 +204,90 @@ class TestWrapSendWithDebugHeaders: asyncio.run(wrapped(body_msg)) assert captured[0] == body_msg + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", tuple(AuthResolution)) +@pytest.mark.parametrize("method", ("GET", "DELETE", "POST")) +async def test_debug_defers_resolution_until_first_frame_only_for_post(source: AuthResolution, method: str) -> None: + captured: Final[list[Message]] = [] + diagnostics: Final = MCPAuthDiagnostics() + + async def send(message: Message) -> None: + captured.append(message) + + wrapped: Final = MCPDebug.wrap_send_with_debug_headers( + send, diagnostics.headers(), diagnostics.headers, request_method=method + ) + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + assert len(captured) == (0 if method == "POST" else 1) + diagnostics.record("s1", source) + body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True} + await wrapped(body) + assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == ( + source.value.encode() if method == "POST" else b"unresolved" + ) + assert captured[1] == body + + +@pytest.mark.asyncio +async def test_early_stream_frame_reports_unresolved_without_waiting() -> None: + captured: Final[list[Message]] = [] + diagnostics: Final = MCPAuthDiagnostics() + + async def send(message: Message) -> None: + captured.append(message) + + wrapped: Final = MCPDebug.wrap_send_with_debug_headers(send, {}, diagnostics.headers, request_method="POST") + await wrapped({"type": "http.response.start", "status": 200, "headers": []}) + await wrapped({"type": "http.response.body", "body": b": ping\n\n", "more_body": True}) + diagnostics.record("s1", AuthResolution.stored_user_token) + await wrapped({"type": "http.response.body", "body": b"data: pong\n\n", "more_body": False}) + assert len(captured) == 3 + assert dict(captured[0]["headers"])[b"x-mcp-debug-auth-resolution"] == b"unresolved" + + +def test_diagnostics_keep_requests_separate_and_do_not_collapse_multiple_servers() -> None: + alice: Final = MCPAuthDiagnostics() + bob: Final = MCPAuthDiagnostics() + alice.record("s1", AuthResolution.stored_user_token) + assert bob.resolution() == "unresolved" + alice.record("s1", AuthResolution.token_exchange) + assert alice.resolution() == "token-exchange" + alice.record("s2", AuthResolution.static_token) + assert alice.resolution() == "multiple" + assert alice.headers()["x-mcp-debug-auth-resolutions"] == '{"s1":"token-exchange","s2":"static-token"}' + + +@pytest.mark.asyncio +async def test_concurrent_mcp_messages_record_on_their_own_http_scope() -> None: + from unittest.mock import MagicMock + + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_debug import ( + MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, + record_auth_resolution, + ) + + session: Final = MagicMock() + first: Final = MCPAuthDiagnostics() + second: Final = MCPAuthDiagnostics() + + async def record(diagnostics: MCPAuthDiagnostics, source: AuthResolution) -> None: + context: Final = RequestContext( + request_id=1, meta=None, session=session, lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + ) + token: Final = request_ctx.set(context) + try: + await asyncio.sleep(0) + record_auth_resolution("same-server", source) + finally: + request_ctx.reset(token) + + await asyncio.gather(record(first, AuthResolution.stored_user_token), record(second, AuthResolution.per_request_header)) + assert first.resolution() == "stored-user-token" + assert second.resolution() == "per-request-header" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 9b6eaba3177..12805e355f9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -3,6 +3,7 @@ import contextvars import os from datetime import datetime, timedelta from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1866,96 +1867,85 @@ async def test_streamable_http_session_manager_is_stateless(): @pytest.mark.asyncio -async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(): - """ - Test that routing correctly sends: - - initialize (no mcp-session-id) → stateful manager (so client gets mcp-session-id) - - tools/list (no mcp-session-id) → stateless manager (curl, Inspector) - """ - try: - from litellm.proxy._experimental.mcp_server.server import ( - handle_streamable_http_mcp, - session_manager_stateful, - session_manager_stateless, +@pytest.mark.parametrize("debug", (False, True)) +@pytest.mark.parametrize( + ("method", "request_body", "stateful"), + ( + ("POST", b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', True), + ("POST", b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', False), + ("GET", b"", False), + ("DELETE", b"", False), + ), +) +async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( + debug: bool, method: str, request_body: bytes, stateful: bool +) -> None: + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + from starlette.types import Message, Receive, Scope, Send + + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.mcp_debug import record_auth_resolution + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution + + scope: Final[Scope] = {"type": "http", "method": method, "path": "/mcp", "headers": []} + receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False}) + send: Final = AsyncMock() + observe_start: Final = AsyncMock() + body: Final[Message] = {"type": "http.response.body", "body": b"data: pong\n\n", "more_body": True} + + async def handle_request(request_scope: Scope, receive: Receive, outgoing: Send) -> None: + await outgoing({"type": "http.response.start", "status": 200, "headers": []}) + await observe_start(send.await_count) + context: Final = RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, request=Request(request_scope) ) - except ImportError: - pytest.skip("MCP server not available") + token: Final = request_ctx.set(context) + try: + record_auth_resolution("s1", AuthResolution.stored_user_token) + finally: + request_ctx.reset(token) + await outgoing(body) - async def make_request(method_body: bytes, path: str = "/mcp/progress_test"): - scope = { - "type": "http", - "method": "POST", - "path": path, - "headers": [ - (b"content-type", b"application/json"), - (b"authorization", b"Bearer test-key"), - ], - } - receive = AsyncMock( - return_value={ - "type": "http.request", - "body": method_body, - "more_body": False, - } - ) - send = AsyncMock() - - stateless_called = [] - stateful_called = [] - - async def stateless_handle(s, r, se): - stateless_called.append(1) - - async def stateful_handle(s, r, se): - stateful_called.append(1) - - with ( - patch( - "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", - new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + stateless_handle: Final = AsyncMock(side_effect=handle_request) + stateful_handle: Final = AsyncMock(side_effect=handle_request) + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + UserAPIKeyAuth(user_id="debug-user"), + None, + None, + None, + None, + {"x-litellm-mcp-debug": "true"} if debug else {}, ), - patch( - "litellm.proxy._experimental.mcp_server.server.set_auth_context", - ), - patch( - "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", - True, - ), - patch.object( - session_manager_stateless, - "handle_request", - side_effect=stateless_handle, - ), - patch.object( - session_manager_stateful, - "handle_request", - side_effect=stateful_handle, - ), - patch.object( - session_manager_stateless, - "_server_instances", - {}, - ), - patch.object( - session_manager_stateful, - "_server_instances", - {}, - ), - ): - await handle_streamable_http_mcp(scope, receive, send) + ), + patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) - return bool(stateless_called), bool(stateful_called) - - # initialize → stateful - init_body = b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}' - stateless_called, stateful_called = await make_request(init_body) - assert stateful_called and not stateless_called, "initialize (no session) should route to stateful, not stateless" - - # tools/list → stateless - tools_body = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' - stateless_called, stateful_called = await make_request(tools_body) - assert stateless_called and not stateful_called, "tools/list (no session) should route to stateless, not stateful" + assert stateful_handle.await_count == (1 if stateful else 0) + assert stateless_handle.await_count == (0 if stateful else 1) + observe_start.assert_awaited_once_with(0 if debug and method == "POST" else 1) + assert send.await_count == 2 + assert send.call_args_list[0].args[0]["status"] == 200 + assert send.call_args_list[1].args[0] == body + headers: Final = dict(send.call_args_list[0].args[0]["headers"]) + if debug: + assert headers[b"x-mcp-debug-auth-resolution"] == (b"stored-user-token" if method == "POST" else b"unresolved") + else: + assert not any(name.startswith(b"x-mcp-debug") for name in headers) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 46fef83092d..adf985e9a21 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,7 +5,7 @@ import logging import os import sys from datetime import datetime -from typing import Any, Dict, Final, Optional +from typing import Any, Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12567,3 +12567,112 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, with pytest.raises(HTTPException) as exc_info: await call assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("config", "extra_headers", "expected_source", "expected_authorization"), + [ + ("stored", None, "stored-user-token", "Bearer stored-token"), + ("stored", {"aUtHoRiZaTiOn": "Bearer injected"}, "stored-user-token", "Bearer stored-token"), + ("static", None, "static-token", "Bearer static-token"), + ("static", {"authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), + ("none", None, "no-auth", None), + ("none", {"Authorization": "Bearer injected"}, "extra-headers", "Bearer injected"), + ], +) +async def test_debug_resolution_matches_final_header_conflict_winner( + config: Literal["stored", "static", "none"], + extra_headers: dict[str, str] | None, + expected_source: str, + expected_authorization: str | None, +) -> None: + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + from pydantic import SecretStr + + from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import MCPAuthenticatedUser + from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics + from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + ApiKeyConfig, AuthorizationCodeConfig, NoneConfig, ServerSpec, SharedKey, UpstreamCredentialProvider, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + class Store: + def __init__(self) -> None: + self.calls = 0 + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + self.calls += 1 + return OAuthToken(access_token="stored-token") if user_id == "alice" else None + + store = Store() + context = MCPAuthenticatedUser(UserAPIKeyAuth(user_id="alice")) + diagnostics = MCPAuthDiagnostics() + token = request_ctx.set(RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + )) + selected = { + "stored": AuthorizationCodeConfig(), + "static": ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-token"))), + "none": NoneConfig(), + }[config] + try: + auth, remaining = await MCPServerManager()._resolve_v2_auth( + server=MCPServer( + server_id="s", name="s", transport="http", url="https://up.example/mcp", + static_headers={"Authorization": "Bearer configured"}, + ), + spec=ServerSpec(server_id="s", resource="https://up.example/mcp", config=selected), + provider=UpstreamCredentialProvider(oauth_token_store=store), + subject_token=None, + user_api_key_auth=context.user_api_key_auth, + extra_headers=extra_headers, + ) + request = httpx.Request("GET", "https://up.example/mcp", headers=remaining) + if auth is not None: + next(auth.auth_flow(request)) + assert diagnostics.resolution() == expected_source + assert request.headers.get("Authorization") == expected_authorization + assert store.calls == (1 if config == "stored" else 0) + finally: + request_ctx.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "stdio"]) +async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Literal["http", "stdio"]) -> None: + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + diagnostics = MCPAuthDiagnostics() + token = request_ctx.set(RequestContext( + request_id=1, meta=None, session=MagicMock(), lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + )) + try: + server = MCPServer( + server_id="signed", name="signed", transport=transport, + url="https://up.example/mcp", auth_type="aws_sigv4", + aws_access_key_id="AKIDEXAMPLE", aws_secret_access_key="test-signing-secret", + aws_region_name="us-east-1", aws_service_name="execute-api", + command="python", args=["-c", "pass"], + ) + client = await MCPServerManager()._create_mcp_client(server) + if transport == "stdio": + assert diagnostics.resolution() == "not-applicable" + else: + assert diagnostics.resolution() == "aws-sigv4" + request = httpx.Request("POST", "https://up.example/mcp", content=b"{}") + next(client._aws_auth.auth_flow(request)) + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 ") + assert "Credential=AKIDEXAMPLE/" in request.headers["Authorization"] + finally: + request_ctx.reset(token) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index bd692776c82..16c6aa128d0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -3,7 +3,7 @@ import inspect import json import sys from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 @@ -87,6 +87,125 @@ def _route_has_dependency(route, dependency) -> bool: class TestExecuteWithMcpClient: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("auth_type", "auth_value", "expected_auth"), + ( + (MCPAuth.none, None, {}), + (MCPAuth.basic, "preview:correct", {"Authorization": "Basic cHJldmlldzpjb3JyZWN0"}), + (MCPAuth.basic, None, {"Authorization": "Basic cHJldmlldzpzdG9yZWQ="}), + (MCPAuth.bearer_token, "edited", {"Authorization": "Bearer edited"}), + (MCPAuth.api_key, "edited", {"X-API-Key": "edited"}), + (MCPAuth.token, "edited", {"Authorization": "token edited"}), + (MCPAuth.authorization, "Custom edited", {"Authorization": "Custom edited"}), + ), + ) + async def test_static_preview_uses_edited_connection_instead_of_registered_server( + self, + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuth, + auth_value: str | None, + expected_auth: dict[str, str], + ) -> None: + from starlette.datastructures import Headers + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview-server", + name="saved", + url="https://stored.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token="preview:stored", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, + server_name="edited", + url="https://stored.example/corrected-mcp", + transport=MCPTransport.sse, + auth_type=auth_type, + credentials={"auth_value": auth_value} if auth_value is not None else None, + static_headers={"X-Preview": "edited"}, + ) + staged: Final = rest_endpoints._stage_server_test(payload, Headers()) + + async def inspect_connection(client: MCPClient) -> dict[str, object]: + return {"url": client.server_url, "transport": client.transport_type, "headers": client._get_auth_headers()} + + result: Final = await rest_endpoints._execute_with_mcp_client( + staged.request, + inspect_connection, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, + ) + assert result == { + "url": "https://stored.example/corrected-mcp", + "transport": MCPTransport.sse, + "headers": {"X-Preview": "edited", **expected_auth}, + } + assert manager.get_mcp_server_by_id(saved.server_id) is saved + assert saved.url == "https://stored.example/mcp" + + @pytest.mark.parametrize( + ("saved_url", "url", "same_origin"), + ( + ("https://stored.example/mcp", "https://other.example/mcp", False), + ("https://stored.example/mcp", "http://stored.example/mcp", False), + ("https://stored.example/mcp", "https://stored.example:8443/mcp", False), + ("https://stored.example/mcp", "https://stored.example:443/mcp", True), + ("http://stored.example/mcp", "http://stored.example:80/edited", True), + ("https://stored.example/mcp", "HTTPS://STORED.EXAMPLE/edited", True), + ("https://[::1]/mcp", "https://[::1]/edited", True), + ("https://[::1]/mcp", "https://[::1]:443/edited", True), + ("https://[::1]/mcp", "https://[::2]/edited", False), + ("https://stored.example/mcp", "https://stored.example:invalid/mcp", False), + ), + ) + @pytest.mark.parametrize("explicit_credential", (None, "preview:explicit")) + def test_static_preview_respects_origin_when_inheriting_credentials( + self, + monkeypatch: pytest.MonkeyPatch, + saved_url: str, + url: str, + same_origin: bool, + explicit_credential: str | None, + ) -> None: + from starlette.datastructures import Headers + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview-server", + name="saved", + url=saved_url, + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + authentication_token="preview:stored", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, + url=url, + transport=MCPTransport.http, + auth_type=MCPAuth.basic, + credentials={"auth_value": explicit_credential} if explicit_credential else None, + ) + staged: Final = rest_endpoints._stage_server_test(payload, Headers()) + expected: Final = explicit_credential or ("preview:stored" if same_origin else None) + assert staged.mcp_auth_header == expected + assert staged.request.credentials == ({"auth_value": expected} if expected else None) + @pytest.mark.asyncio async def test_redacts_stack_trace(self, monkeypatch): async def fake_create_client(*args, **kwargs): @@ -113,7 +232,7 @@ class TestExecuteWithMcpClient: assert "stack_trace" not in result @pytest.mark.asyncio - async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def test_timeout_caps_hanging_operation_and_names_origin(self, monkeypatch): async def fake_create_client(*args, **kwargs): return object() @@ -138,7 +257,7 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] @pytest.mark.asyncio async def test_timeout_covers_client_creation(self, monkeypatch): @@ -166,15 +285,15 @@ class TestExecuteWithMcpClient: ) assert result["error"] is True - assert "https://mcp.example.com/mcp/" in result["message"] + assert "https://mcp.example.com" in result["message"] def test_timeout_defaults_to_tool_listing_timeout(self): default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default assert default == MCP_TOOL_LISTING_TIMEOUT - def test_connection_error_message_timeout_names_url_and_budget(self): + def test_connection_error_message_timeout_names_origin_and_budget(self): message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) - assert "https://api.example.com/mcp/" in message + assert "https://api.example.com" in message assert "30s" in message def test_connection_error_message_hides_arbitrary_http_exception_detail(self): @@ -592,7 +711,7 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert result["error"] is True - assert "Failed to connect to MCP server" in result["message"] + assert "reference" in result["message"] # Error message must not leak raw exception details assert "cancel scope" not in result["message"] @@ -3427,10 +3546,191 @@ class TestConnectionErrorMessage: message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message + @pytest.mark.parametrize( + "error_type", [httpx.ReadError, httpx.WriteError, httpx.RemoteProtocolError, ConnectionResetError] + ) + def test_interrupted_connection_message_is_safe(self, error_type: type[Exception]) -> None: + message: Final = rest_endpoints._connection_error_message( + error_type("secret-transport-detail"), "https://example.com/?token=secret-query", 30 + ) + assert "connection was interrupted" in message + assert "secret" not in message + + def test_closed_connection_explains_incomplete_request(self) -> None: + from mcp import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=-32000, message="Connection closed", data="secret-data")), None, 30 + ) + assert "connection was closed before the request completed" in message + assert "secret" not in message + + def test_timeout_does_not_claim_the_server_sent_nothing(self) -> None: + message: Final = rest_endpoints._connection_error_message(TimeoutError(), None, 30) + assert "no valid MCP response received" in message + + @pytest.mark.asyncio + @pytest.mark.parametrize("sdk_timeout", [True, False]) + @pytest.mark.parametrize("read_timeout", [0, 1]) + async def test_timeout_message_uses_the_deadline_that_expired(self, sdk_timeout: bool, read_timeout: int) -> None: + from mcp import McpError + from mcp.types import ErrorData + + async def operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + try: + raise TimeoutError("secret-timeout") + except TimeoutError as elapsed: + if not sdk_timeout: + raise + try: + raise McpError(ErrorData(code=408, message="secret-sdk-timeout")) from elapsed + except McpError as sdk_error: + raise TimeoutError() from sdk_error + + payload: Final = NewMCPServerRequest( + server_name="timeout", url="https://example.com", auth_type=MCPAuth.none, timeout=read_timeout + ) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, operation, timeout_seconds=30) + assert (f"within {read_timeout}s" if sdk_timeout else "within 30s") in result["message"] + assert "secret" not in result["message"] + def test_unknown_error_falls_back_to_generic(self): message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message - assert "proxy logs" in message.lower() + assert "reference" in message.lower() + + def test_sdk_session_terminated_explains_endpoint_and_retry(self) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=32600, message="Session terminated")), "https://example.com/mcp", 30.0 + ) + + assert "session was terminated" in message + assert "MCP endpoint" in message + assert "transport" in message + assert "retry" in message + assert "404" not in message + + @pytest.mark.parametrize("code", [-32700, -32601, -32602, -32603, -32000, 32600, 408]) + def test_rpc_errors_include_code_without_echoing_upstream_data(self, code: int) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + message: Final = rest_endpoints._connection_error_message( + McpError(ErrorData(code=code, message="secret-message", data={"token": "secret-data"})), + "https://example.com/secret-path?token=secret-query", + 30.0, + ) + + assert f"JSON-RPC code {code}" in message + assert "secret" not in message + assert "timed out" not in message + assert "session was terminated" not in message + + @pytest.mark.parametrize("status_code", [401, 403, 404, 405, 429, 503]) + def test_wrapped_http_failures_preserve_status(self, status_code: int) -> None: + response: Final = httpx.Response(status_code, text="secret-body") + upstream: Final = httpx.HTTPStatusError( + "secret-exception", + request=httpx.Request("POST", "https://example.com/?token=secret-query"), + response=response, + ) + wrapped: Final = BaseExceptionGroup( + "secret-group", [asyncio.CancelledError(), BaseExceptionGroup("nested", [upstream])] + ) + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert f"HTTP {status_code}" in message + assert "secret" not in message + + def test_explicit_cause_is_classified_before_incidental_context(self) -> None: + wrapped: Final = RuntimeError("secret-wrapper") + wrapped.__cause__ = httpx.ConnectError("secret-cause") + wrapped.__context__ = TimeoutError("secret-context") + + message: Final = rest_endpoints._connection_error_message(wrapped, "https://example.com", 30.0) + + assert "unreachable" in message + assert "secret" not in message + + def test_timeout_url_redacts_credentials_path_query_and_fragment(self) -> None: + message: Final = rest_endpoints._connection_error_message( + TimeoutError("secret-error"), + "https://secret-user:secret-pass@example.com:8443/secret-path?token=secret-query#secret-fragment", + 30.0, + ) + + assert "https://example.com:8443" in message + assert "30s" in message + assert "secret" not in message + + def test_unknown_failure_reference_matches_safe_diagnostics(self, caplog: pytest.LogCaptureFixture) -> None: + import re + + try: + raise RuntimeError("secret-exception-body") + except RuntimeError as exc: + message: Final = rest_endpoints._connection_error_message( + exc, "https://secret-user:secret-password@example.com/secret-path?token=secret-query", 30.0 + ) + + reference: Final = re.search(r"reference ([a-f0-9]{32})", message) + assert reference is not None + diagnostics: Final = tuple( + record for record in caplog.records if "MCP connection test failed" in record.message + ) + assert len(diagnostics) == 1 + assert reference.group(1) in diagnostics[0].message + assert "RuntimeError" in diagnostics[0].message + assert "test_unknown_failure_reference_matches_safe_diagnostics" in diagnostics[0].message + assert diagnostics[0].exc_info is None + assert "secret" not in message + diagnostics[0].message + + @pytest.mark.parametrize("exc", [ValueError("secret-config"), HTTPException(500, "secret-detail")]) + def test_unrelated_errors_are_not_misreported_as_invalid_mcp(self, exc: Exception) -> None: + message: Final = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) + + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + def test_configuration_validation_error_uses_unknown_fallback(self) -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError) as caught: + NewMCPServerRequest.model_validate({"server_name": "example", "transport": "secret-invalid-transport"}) + + message: Final = rest_endpoints._connection_error_message(caught.value, "https://example.com", 30.0) + assert "reference" in message + assert "invalid MCP response" not in message + assert "secret" not in message + + @pytest.mark.asyncio + async def test_connection_test_preserves_cancellation(self) -> None: + async def cancelled_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise asyncio.CancelledError + + payload: Final = NewMCPServerRequest(server_name="cancelled", url="https://example.com", auth_type=MCPAuth.none) + with pytest.raises(asyncio.CancelledError): + await rest_endpoints._execute_with_mcp_client(payload, cancelled_operation) + + @pytest.mark.asyncio + async def test_unknown_failure_preserves_response_contract(self) -> None: + async def failing_operation(client: rest_endpoints.MCPClient) -> dict[str, object]: + raise RuntimeError("secret-operation") + + payload: Final = NewMCPServerRequest(server_name="unknown", url="https://example.com", auth_type=MCPAuth.none) + result: Final = await rest_endpoints._execute_with_mcp_client(payload, failing_operation) + + assert result["error"] is True + assert result["status"] == "error" + assert "reference" in result["message"] + assert "secret" not in result["message"] + assert "stack_trace" not in result class TestGetServerAuthHeaderGroupDefault: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 2284a05b2e9..8777e24e209 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json from types import SimpleNamespace -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: @@ -37,6 +37,7 @@ from litellm.proxy.auth.auth_checks import ( _can_object_call_vector_stores, _check_end_user_budget, _check_team_member_budget, + _fetch_key_object_from_db_with_reconnect, _get_fuzzy_user_object, _get_team_db_check, _log_budget_lookup_failure, @@ -55,6 +56,7 @@ from litellm.caching.redis_cache import RedisCache from litellm.constants import ( DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + PROXY_DB_LOOKUP_MAX_CONCURRENCY, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) @@ -564,6 +566,43 @@ async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_e assert mock_prisma_client.get_data.await_count == 1 +class _InFlightCountingPrisma: + def __init__(self) -> None: + self.in_flight = 0 + self.max_in_flight = 0 + + async def get_data( + self, token: str, table_name: str, parent_otel_span: None, proxy_logging_obj: None + ) -> UserAPIKeyAuth: + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + await asyncio.sleep(0.001) + self.in_flight -= 1 + return UserAPIKeyAuth(token=token) + + +@pytest.mark.asyncio +async def test_fetch_key_object_from_db_bounds_in_flight_prisma_requests(): + prisma: Final = _InFlightCountingPrisma() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 5 + + results: Final = await asyncio.gather( + *( + _fetch_key_object_from_db_with_reconnect( + hashed_token=f"hashed-token-{i}", + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + parent_otel_span=None, + proxy_logging_obj=None, + ) + for i in range(burst) + ) + ) + + assert len(results) == burst + assert {r.token for r in results if r is not None} == {f"hashed-token-{i}" for i in range(burst)} + assert prisma.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY + + def _fake_redis_cache(): fake_redis = MagicMock() fake_redis.async_get_cache = AsyncMock(return_value=None) @@ -2374,6 +2413,44 @@ def _mock_prisma_for_team_lookup(find_unique): return mock_prisma_client +_TEAM_ALIAS_TABLE_ROW = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"} + + +def _prisma_team_row(include): + """Mimics Prisma: the `litellm_model_table` relation rides on the row only when the query `include`s it.""" + columns = {"team_id": "team-aliases", "team_alias": "aliases", "models": ["gpt-4o"]} + row = ( + {**columns, "litellm_model_table": _TEAM_ALIAS_TABLE_ROW} + if (include or {}).get("litellm_model_table") + else columns + ) + return SimpleNamespace(dict=lambda: row, model_dump=lambda: row) + + +@pytest.mark.asyncio +async def test_get_team_object_loads_model_aliases_relation(): + """LIT-5858: the auth path read teams without `include`ing `litellm_model_table`, so every JWT + team came back with `model_aliases=None` and alias requests 403'd.""" + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.auth.team_grants import team_model_aliases + + async def find_unique(where, include=None): + return _prisma_team_row(include) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + team = await get_team_object( + team_id="team-aliases", + prisma_client=_mock_prisma_for_team_lookup(AsyncMock(side_effect=find_unique)), + user_api_key_cache=mock_cache, + check_db_only=True, + ) + + assert team_model_aliases(team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio async def test_get_team_object_distinguishes_absent_team_from_unreadable_row(): """A deleted team and a database that would not answer both surface as a 404, @@ -6195,6 +6272,32 @@ async def test_get_team_object_by_alias_db_fetch_returns_cached_obj(): assert result.models == ["gpt-4"] +@pytest.mark.asyncio +async def test_get_team_object_by_alias_loads_model_aliases_relation(): + """LIT-5858: same regression as `test_get_team_object_loads_model_aliases_relation`, for the + `team_alias_jwt_field` lookup.""" + from litellm.proxy.auth.auth_checks import get_team_object_by_alias + from litellm.proxy.auth.team_grants import team_model_aliases + + async def find_many(where, include=None): + return [_prisma_team_row(include)] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + team = await get_team_object_by_alias( + team_alias="aliases", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + assert team_model_aliases(team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio async def test_get_org_object_by_alias_db_fetch_returns_validated_org(): from litellm.proxy._types import LiteLLM_OrganizationTable diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index aaf630ad29b..3c725148d3f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -565,6 +565,38 @@ def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): ) +def _azure_relay_router(): + from litellm.router import Router + + return Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://a.services.ai.azure.com", "api_key": "k"}, + }, + { + "model_name": "other-group", + "litellm_params": {"model": "azure/gpt-5.4", "api_base": "https://b.openai.azure.com", "api_key": "k"}, + }, + ] + ) + + +@pytest.mark.parametrize( + "route, request_data, expected", + [ + ("/azure_ai/other-group/openai/deployments/other-group/chat/completions", {"model": "gpt"}, "other-group"), + ("/azure_ai/other-group/models/chat/completions", {}, "other-group"), + ("/azure/openai/deployments/gpt/chat/completions", {"model": "other-group"}, "gpt"), + ("/azure/openai/deployments/gpt/chat/completions", {}, "gpt"), + ("/azure/openai/deployments/my-azure-deployment/chat/completions", {"model": "gpt"}, "gpt"), + ("/azure_ai/gpt", {"model": "other-group"}, "other-group"), + ], +) +def test_get_model_from_request_azure_relay_routes_use_the_model_group_in_the_path(route, request_data, expected): + assert get_model_from_request(request_data=request_data, route=route, llm_router=_azure_relay_router()) == expected + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 99a0a4c0a8b..94226b5404d 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -13,6 +13,7 @@ from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, JWTLiteLLMRoleMap, LiteLLM_JWTAuth, + LiteLLM_ModelTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -1255,6 +1256,57 @@ async def test_find_team_with_model_access_model_group(monkeypatch): assert team_obj.team_id == "team-1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_aliases", + ['{"fast": "gpt-4o"}', {"fast": "gpt-4o"}], + ids=["json-string", "dict"], +) +async def test_find_team_with_model_access_resolves_team_model_alias(monkeypatch, model_aliases): + """LIT-5858: a JWT team that grants `gpt-4o` under the alias `fast` must resolve a request + for `fast`. The JWT path used to pass `team_model_aliases=None`, so every alias request 403'd.""" + import sys + import types + + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + router = Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}]) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + team = LiteLLM_TeamTable( + team_id="team-aliases", + models=["gpt-4o"], + litellm_model_table=LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin"), + ) + + async def mock_get_team_object(*args, **kwargs): + return team + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + user_api_key_cache = DualCache() + + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids={"team-aliases"}, + requested_model="fast", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), + ) + + assert team_id == "team-aliases" + assert team_obj is team + + @pytest.mark.asyncio async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatch): """Regression for #31189: a single-team JWT that grants the requested model diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 8d93d801bfd..e209a491b0a 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.login_utils import ( LoginResult, authenticate_user, get_ui_credentials, + is_env_credential_login_enabled, ) @@ -185,6 +186,7 @@ async def test_authenticate_user_invalid_credentials(): assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "401" assert "Invalid credentials" in exc_info.value.message + assert "UI_USERNAME" in exc_info.value.message @pytest.mark.asyncio @@ -799,3 +801,158 @@ class TestDisablePasswordLoginWhenSSOEnabled: assert isinstance(result, LoginResult) assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestDisableEnvCredentialLogin: + """`disable_env_credential_login` must reject a login with the env + credentials (UI_USERNAME/UI_PASSWORD, or the master-key fallback when + UI_PASSWORD is unset) while leaving database-user password logins + untouched, so admins with real accounts keep a way in.""" + + @pytest.mark.asyncio + async def test_rejects_correct_env_credentials_when_disabled(self): + master_key = "sk-1234" + ui_username = "admin" + ui_password = "env-only-password" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "401" + assert "UI_USERNAME" not in exc_info.value.message + assert "UI_PASSWORD" not in exc_info.value.message + + @pytest.mark.asyncio + async def test_rejects_master_key_fallback_when_disabled(self): + """With UI_PASSWORD unset, the master key IS the env password, so the + setting must reject it too or it protects nothing by default.""" + master_key = "sk-1234" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": "admin"}, clear=True): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username="admin", + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert exc_info.value.code == "401" + + @pytest.mark.asyncio + async def test_db_user_login_still_works_when_disabled(self): + master_key = "sk-1234" + user_email = "admin@example.com" + password = "Str0ng!Passw0rd" + + mock_user = LiteLLM_UserTable( + user_id="db-admin-1", + user_email=user_email, + password=hash_token(token=password), + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user) + + with patch.dict( + os.environ, + { + "UI_USERNAME": "admin", + "UI_PASSWORD": "env-password", + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + stack.enter_context( + patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "db-user-token"}, + ) + ) + result = await authenticate_user( + username=user_email, + password=password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_env_credential_login": True}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == "db-admin-1" + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + @pytest.mark.asyncio + async def test_env_login_still_works_when_setting_absent(self): + """Env-credential login is the bootstrap path on a fresh install and + must stay on by default.""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": master_key, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + _patch_successful_admin_login_deps(stack) + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + +class TestIsEnvCredentialLoginEnabled: + """Drives the Admin UI warning banner: it must be True exactly when a + login with the env credentials could actually succeed.""" + + def test_enabled_by_default(self): + assert is_env_credential_login_enabled({}) is True + + def test_disabled_by_dedicated_setting(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": True}) is False + + def test_explicit_false_keeps_it_enabled(self): + assert is_env_credential_login_enabled({"disable_env_credential_login": False}) is True + + def test_disabled_when_sso_gate_blocks_all_password_logins(self): + """`disable_password_login_when_sso_enabled` with SSO configured + rejects every username/password login before the env comparison runs, + so the banner must not nag about an already-unreachable path.""" + with ExitStack() as stack: + _patch_sso_configured(stack, configured=True) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is False + + def test_enabled_when_sso_gate_is_set_but_sso_not_configured(self): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=False) + assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is True diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py new file mode 100644 index 00000000000..447fc1c93a1 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -0,0 +1,129 @@ +import pytest + +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_VerificationTokenView, + Member, + UserAPIKeyAuth, +) +from litellm.models.team import LiteLLM_ModelTable +from litellm.proxy.auth.team_grants import team_grants, team_model_aliases + +TEAM_ID = "team-grants" +USER_ID = "user-in-team" +ALIASES = {"fast": "gpt-4o-mini", "smart": "gpt-4o"} + + +def _alias_table(model_aliases) -> LiteLLM_ModelTable: + return LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin") + + +def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=TEAM_ID, + team_alias="grants-team", + tpm_limit=1000, + rpm_limit=10, + max_budget=50.0, + soft_budget=25.0, + spend=12.5, + models=["gpt-4o", "gpt-4o-mini"], + blocked=True, + metadata={"tier": "gold"}, + litellm_model_table=_alias_table(model_aliases), + object_permission_id="op-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", mcp_servers=["mcp-a"]), + members_with_roles=[ + Member(user_id="someone-else", role="user"), + Member(user_id=USER_ID, role="admin"), + ], + ) + + +def _membership() -> LiteLLM_TeamMembership: + return LiteLLM_TeamMembership( + user_id=USER_ID, + team_id=TEAM_ID, + spend=3.25, + litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5), + ) + + +def test_team_grants_cover_every_team_field_the_key_path_gets(): + """Class guard for LIT-5858 and its siblings: every ``team_*`` column the combined-view SQL hands the + virtual-key path must come out of the projection too, with the team's actual value, so adding a column + to ``LiteLLM_VerificationTokenView`` without teaching ``team_grants`` fails here instead of in prod.""" + team = _full_team() + grants = team_grants(team_object=team, team_membership=_membership(), user_id=USER_ID) + token = UserAPIKeyAuth(team_id=TEAM_ID, **grants) + + view_team_fields = {name for name in LiteLLM_VerificationTokenView.model_fields if name.startswith("team_")} + assert view_team_fields - {"team_id"} <= set(grants) + assert all(grants[name] is not None for name in view_team_fields - {"team_id"}) + + assert token.team_alias == "grants-team" + assert token.team_tpm_limit == 1000 + assert token.team_rpm_limit == 10 + assert token.team_max_budget == 50.0 + assert token.team_soft_budget == 25.0 + assert token.team_spend == 12.5 + assert token.team_models == ["gpt-4o", "gpt-4o-mini"] + assert token.team_blocked is True + assert token.team_metadata == {"tier": "gold"} + assert token.team_model_aliases == ALIASES + assert token.team_object_permission_id == "op-1" + assert token.team_object_permission is not None + assert token.team_object_permission.mcp_servers == ["mcp-a"] + assert token.team_member == Member(user_id=USER_ID, role="admin") + assert token.team_member_spend == 3.25 + assert token.team_member_tpm_limit == 500 + assert token.team_member_rpm_limit == 5 + + +def test_team_grants_without_team_leave_token_defaults(): + token = UserAPIKeyAuth(**team_grants(team_object=None, team_membership=None, user_id=USER_ID)) + assert token == UserAPIKeyAuth() + + +@pytest.mark.parametrize( + "stored_aliases", + [ALIASES, '{"fast": "gpt-4o-mini", "smart": "gpt-4o"}'], + ids=["json-object", "json-string-as-written-by-team-new"], +) +def test_team_model_aliases_decode_both_storage_shapes(stored_aliases): + team = _full_team(model_aliases=stored_aliases) + assert team_model_aliases(team) == ALIASES + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] == ALIASES + + +@pytest.mark.parametrize("stored_aliases", [None, "not json", '["a", "b"]', {"fast": 3}], ids=str) +def test_team_model_aliases_treat_unusable_column_as_no_aliases(stored_aliases): + team = _full_team(model_aliases=stored_aliases) + assert team_model_aliases(team) is None + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] is None + + +def test_team_model_aliases_none_without_relation_loaded(): + team = _full_team() + team.litellm_model_table = None + assert team_model_aliases(team) is None + assert team_model_aliases(None) is None + + +def test_team_member_is_the_callers_row_only(): + team = _full_team() + assert team_grants(team_object=team, team_membership=None, user_id="someone-else")["team_member"] == Member( + user_id="someone-else", role="user" + ) + assert team_grants(team_object=team, team_membership=None, user_id="stranger")["team_member"] is None + assert team_grants(team_object=team, team_membership=None, user_id=None)["team_member"] is None + + +def test_membership_limits_absent_without_membership_row(): + grants = team_grants(team_object=_full_team(), team_membership=None, user_id=USER_ID) + assert grants["team_member_spend"] is None + assert grants["team_member_tpm_limit"] is None + assert grants["team_member_rpm_limit"] is None diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index fd289e33ea6..6cce6d0316b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -7133,3 +7133,162 @@ def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(t assert report["outcomes"] == ["accepted", "rejected"] auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard-return", "proxy-admin-return"]) +async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_admin): + """LIT-5858: the team-based JWT path hand-built ``UserAPIKeyAuth`` from a short list of team fields, so the + team's model aliases (and on the admin return, its object permission) never reached the token and alias + requests 403'd. Both returns now go through ``team_grants``; pin the fields that used to be dropped.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.models.team import LiteLLM_ModelTable + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + Member, + ) + + class _AcceptEveryJwt(JWTHandler): + def is_jwt(self, token: str) -> bool: + return True + + jwt_handler = _AcceptEveryJwt() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + team = LiteLLM_TeamTable( + team_id="team-jwt-aliases", + team_alias="jwt-aliases", + models=["gpt-4o"], + max_budget=40.0, + spend=4.0, + blocked=False, + metadata={"tier": "gold"}, + litellm_model_table=LiteLLM_ModelTable( + model_aliases='{"fast": "gpt-4o"}', created_by="admin", updated_by="admin" + ), + object_permission_id="op-jwt", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-jwt", mcp_servers=["mcp-a"]), + members_with_roles=[Member(user_id="jwt-user", role="admin")], + ) + membership = LiteLLM_TeamMembership(user_id="jwt-user", team_id="team-jwt-aliases", spend=1.5) + builder_result = { + "is_proxy_admin": is_proxy_admin, + "team_object": team, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": "jwt", + "team_id": "team-jwt-aliases", + "user_id": "jwt-user", + "user_email": "jwt-user@example.com", + "end_user_id": None, + "org_id": None, + "team_membership": membership, + "jwt_claims": {"sub": "jwt-user"}, + } + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": DualCache(), + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {"enable_jwt_auth": True}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": jwt_handler, + "premium_user": True, + "litellm_proxy_admin_name": "admin", + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http", "headers": [], "method": "POST"}) + request._url = URL(url="/chat/completions") + with patch( # test-quality-ok: auth_builder is the claim-resolution seam; the regression is how its result is projected onto the token + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=builder_result, + ): + token = await _user_api_key_auth_builder( + request=request, + api_key="Bearer header.payload.signature", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert token.team_id == "team-jwt-aliases" + assert token.user_role == (LitellmUserRoles.PROXY_ADMIN if is_proxy_admin else LitellmUserRoles.INTERNAL_USER) + assert token.team_model_aliases == {"fast": "gpt-4o"} + assert token.team_object_permission is not None + assert token.team_object_permission.mcp_servers == ["mcp-a"] + assert token.team_object_permission_id == "op-jwt" + assert token.team_alias == "jwt-aliases" + assert token.team_models == ["gpt-4o"] + assert token.team_max_budget == 40.0 + assert token.team_spend == 4.0 + assert token.team_metadata == {"tier": "gold"} + assert token.team_member == Member(user_id="jwt-user", role="admin") + assert token.team_member_spend == 1.5 + assert token.jwt_claims == {"sub": "jwt-user"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"]) +async def test_claude_view_normalizes_before_model_access(monkeypatch, route): + from starlette.requests import Request + from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access + + source = "foo[1m]" + encoded = "claude-router-" + source.encode().hex() + "[1m]" + router = litellm.Router(model_list=[{"model_name": source, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}]) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + data = {"model": encoded, "messages": [{"role": "user", "content": "hi"}]} + request = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + token = UserAPIKeyAuth(models=[source]) + await _enforce_key_and_fallback_model_access(valid_token=token, request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + assert data["model"] == source + assert (await request.json())["model"] == source + assert json.loads(await request.body())["model"] == source + assert request.scope["parsed_body"][1]["model"] == source + with pytest.raises(ProxyException): + await _enforce_key_and_fallback_model_access(valid_token=UserAPIKeyAuth(models=["other"]), request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "hierarchical", "unclaimed"]) +async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer): + from starlette.requests import Request + from litellm.proxy.auth.user_api_key_auth import _normalize_claude_model + + encoded = "claude-router-666f6f" + names = ("foo", "other", encoded) if layer == "literal" else ("foo", "other") + alias = {encoded: "other"} + router = litellm.Router(model_list=[{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names], model_group_alias=alias if layer == "router" else None) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) + token = UserAPIKeyAuth(aliases=alias if layer == "key" else {}, router_settings={"model_group_alias": alias} if layer == "hierarchical" else None) + data = {"model": encoded} + request = Request({"type": "http", "method": "POST", "path": "/v1/messages", "headers": [], "query_string": b""}) + await _normalize_claude_model(data, token, request, "/v1/messages") + assert data["model"] == ("foo" if layer == "unclaimed" else encoded) + await _normalize_claude_model(data, token, request, "/v1/messages") + assert data["model"] == ("foo" if layer == "unclaimed" else encoded) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py index 4a3b3ef22c4..77742ea9f9f 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -159,6 +159,10 @@ class TestUpCommand: assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" assert "apiKeyHelper" not in captured["settings"] + # The ephemeral proxy serves only the autorouter, so a starting model left by + # `lite configure claude --model` or a user pin would 400 on the first message. + assert captured["settings"]["model"] == "autorouter" + assert captured["settings"]["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" assert captured["settings_mode"] == 0o600 assert terminate_calls == [99999] diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py deleted file mode 100644 index 87a33c79a79..00000000000 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py +++ /dev/null @@ -1,63 +0,0 @@ -from litellm.proxy.client.cli.commands.autoroute.settings import ( - ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, - merge_claude_settings_static_token, -) - - -def test_preserves_unrelated_top_level_keys(): - merged = merge_claude_settings_static_token({"theme": "dark"}, "http://127.0.0.1:4000", "token-abc") - assert merged["theme"] == "dark" - - -def test_preserves_unrelated_env_keys(): - settings = {"env": {"SOME_OTHER_VAR": "value"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["SOME_OTHER_VAR"] == "value" - - -def test_sets_base_url_and_auth_token(): - merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") - assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" - assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" - - -def test_preserves_existing_tool_search(): - settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" - - -def test_drops_stray_api_key(): - settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert "ANTHROPIC_API_KEY" not in merged["env"] - - -def test_removes_existing_api_key_helper(): - settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token"} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert "apiKeyHelper" not in merged - - -def test_does_not_mutate_input(): - settings = {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} - merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert settings == {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} - - -def test_forces_all_claude_code_default_model_tiers_to_the_autorouter(): - # A bare "*" model_name deployment looks like the obvious way to catch every request - # regardless of which model Claude Code thinks it's using, but Router's auto-router - # registry is keyed by the literal requested model string with no wildcard resolution - # (litellm/router.py:10711-10717) -- so the only reliable way to make every one of Claude - # Code's own tiers hit the auto-router is to override the env vars it reads per tier. - merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000", "token-abc") - for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: - assert merged["env"][key] == "autorouter" - - -def test_overrides_a_preexisting_default_model_env_var(): - settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} - merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" diff --git a/tests/test_litellm/proxy/client/cli/conftest.py b/tests/test_litellm/proxy/client/cli/conftest.py new file mode 100644 index 00000000000..50c76d3f125 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/conftest.py @@ -0,0 +1,32 @@ +import os +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest + +REAL_CLAUDE_SETTINGS: Final = Path(os.path.expanduser("~")) / ".claude" / "settings.json" + + +def _current_bytes() -> bytes | None: + return REAL_CLAUDE_SETTINGS.read_bytes() if REAL_CLAUDE_SETTINGS.exists() else None + + +@pytest.fixture(autouse=True) +def isolated_claude_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Path]: + before: Final = _current_bytes() + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(tmp_path / ".claude")) + yield tmp_path + after: Final = _current_bytes() + if after == before: + return + if before is None: + REAL_CLAUDE_SETTINGS.unlink() + else: + REAL_CLAUDE_SETTINGS.write_bytes(before) + pytest.fail( + f"this test wrote the developer's real {REAL_CLAUDE_SETTINGS}; the original bytes were restored. " + "Resolve the Claude settings path at call time (never Path.home() at import) and point the test at tmp_path" + ) diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index a8a6659fe9a..bebea285edc 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -28,6 +28,7 @@ from litellm.proxy.client.cli.commands.agents import ( ) AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" +CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" def _agent_command(name): @@ -163,6 +164,27 @@ class TestBuildAgentEnv: assert env["PATH"] == "/usr/bin" assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} + def test_anthropic_profile_leaves_the_bearer_to_the_api_key_helper(self): + env = build_agent_env( + {"ANTHROPIC_AUTH_TOKEN": "stale-token", "ANTHROPIC_API_KEY": "real-key"}, + "http://localhost:4000/", + "sk-key", + frozenset({"anthropic"}), + export_anthropic_token=False, + ) + assert "ANTHROPIC_AUTH_TOKEN" not in env + assert "ANTHROPIC_API_KEY" not in env + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["ENABLE_TOOL_SEARCH"] == "true" + assert env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + + def test_helper_mode_still_exports_the_openai_key(self): + env = build_agent_env( + {}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}), export_anthropic_token=False + ) + assert "ANTHROPIC_AUTH_TOKEN" not in env + assert env["OPENAI_API_KEY"] == "sk-key" + class TestAgentLaunchArgs: def test_claude_and_opencode_get_no_extra_args(self): @@ -509,6 +531,25 @@ class TestRunAgent: assert "ANTHROPIC_API_KEY" not in env assert "OPENAI_BASE_URL" not in env + def test_helper_supplied_token_never_reaches_the_launch_env(self): + calls = {} + verified = [] + + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={"PATH": "/usr/bin", "ANTHROPIC_AUTH_TOKEN": "stale-token"}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda base_url, api_key: verified.append(api_key), + launcher=lambda p, a, e: calls.update(env=dict(e)), + export_anthropic_token=False, + ) + + assert verified == ["sk-key"] + assert "ANTHROPIC_AUTH_TOKEN" not in calls["env"] + assert calls["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + def test_codex_gets_openai_env(self): calls = {} run_agent( @@ -1052,6 +1093,101 @@ class TestAgentCommands: in result.output ) + def _invoke_claude_with_settings(self, tmp_path, settings, obj, *, default_settings=None): + config_dir = tmp_path / "claude-config" + config_dir.mkdir() + if settings is not None: + (config_dir / "settings.json").write_text(json.dumps(settings)) + default_path = tmp_path / "home-claude" / "settings.json" + default_path.parent.mkdir() + if default_settings is not None: + default_path.write_text(json.dumps(default_settings)) + captured = {} + with ( + patch(f"{CLAUDE_SETTINGS_MODULE}.CLAUDE_SETTINGS_PATH", default_path), + patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"), + patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), + ): + result = self.runner.invoke( + _agent_command("claude"), [], obj=obj, env={"CLAUDE_CONFIG_DIR": str(config_dir)} + ) + assert result.exit_code == 0, result.output + return captured, result.output + + def test_helper_is_read_from_the_config_dir_claude_code_uses(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is False + assert str(tmp_path / "claude-config" / "settings.json") in output + + def test_helper_only_in_the_default_file_keeps_the_env_token_when_config_dir_points_elsewhere(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + None, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + default_settings={"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + ) + + assert captured["export_anthropic_token"] is True + assert "apiKeyHelper" not in output + + def test_stored_login_with_a_matching_helper_leaves_the_token_to_the_helper(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is False + assert "reads its key from the apiKeyHelper" in output + + def test_explicit_key_is_exported_even_when_a_helper_matches(self, tmp_path): + captured, output = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": False}, + ) + + assert captured["export_anthropic_token"] is True + assert "apiKeyHelper" not in output + + def test_helper_for_another_proxy_keeps_the_env_token(self, tmp_path): + captured, _ = self._invoke_claude_with_settings( + tmp_path, + {"apiKeyHelper": "/usr/local/bin/lite --base-url https://other.example.com auth print-token"}, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is True + + def test_no_claude_settings_keeps_the_env_token(self, tmp_path): + captured, _ = self._invoke_claude_with_settings( + tmp_path, + None, + {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert captured["export_anthropic_token"] is True + + def test_codex_never_consults_claude_settings(self): + captured = {} + with ( + patch(f"{AGENTS_MODULE}.lite_api_key_helper_configured", side_effect=AssertionError("consulted")), + patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), + ): + result = self.runner.invoke( + _agent_command("codex"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, + ) + + assert result.exit_code == 0, result.output + assert captured["export_anthropic_token"] is True + def test_codex_shows_friendly_name(self): captured = {} with patch( diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 821323e722c..1f314e0c9d8 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -6,7 +6,6 @@ from pathlib import Path from unittest.mock import Mock, patch - import pytest from click.testing import CliRunner @@ -27,6 +26,7 @@ from litellm.proxy.client.cli.commands.auth import ( print_token, whoami, ) +from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner @@ -84,7 +84,7 @@ class TestPollingErrorSurfacing: } with patch("requests.get", return_value=mock_response) as mock_get, patch("time.sleep"): - with pytest.raises(ValueError, match='Your litellm CLI is out of date and uses a login flow') as exc_info: + with pytest.raises(ValueError, match="Your litellm CLI is out of date and uses a login flow") as exc_info: _poll_for_ready_data("http://test/sso/cli/poll/sk-legacy") assert mock_get.call_count == 1 @@ -151,7 +151,7 @@ class TestStartCliSsoFlowErrors: mock_response.status_code = 404 with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError, match='Either --base-url is wrong, or the proxy is older than') as exc_info: + with pytest.raises(ValueError, match="Either --base-url is wrong, or the proxy is older than") as exc_info: _start_cli_sso_flow("https://old-proxy.example.com") message = str(exc_info.value) @@ -167,7 +167,7 @@ class TestStartCliSsoFlowErrors: mock_response.json.return_value = {"detail": "Too many CLI login attempts. Try again later."} with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError, match='Too many CLI login attempts\\. Try again later\\.') as exc_info: + with pytest.raises(ValueError, match="Too many CLI login attempts\\. Try again later\\.") as exc_info: _start_cli_sso_flow("https://test.example.com") assert "HTTP 429" in str(exc_info.value) @@ -183,7 +183,7 @@ class TestStartCliSsoFlowErrors: mock_response.text = "Sign in to corporate VPN" with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError, match='A proxy, load balancer, or auth gateway in front of') as exc_info: + with pytest.raises(ValueError, match="A proxy, load balancer, or auth gateway in front of") as exc_info: _start_cli_sso_flow("https://test.example.com") message = str(exc_info.value) @@ -197,7 +197,7 @@ class TestStartCliSsoFlowErrors: from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow with patch("requests.post", side_effect=requests.ConnectionError("Connection refused")): - with pytest.raises(ValueError, match='Connection refused\\. Check that the proxy is running') as exc_info: + with pytest.raises(ValueError, match="Connection refused\\. Check that the proxy is running") as exc_info: _start_cli_sso_flow("https://unreachable.example.com") message = str(exc_info.value) @@ -584,13 +584,9 @@ class TestLogoutCommand: assert "could not be checked" in result.output assert DISABLE_KEYRING_ENV_VAR in result.output - def test_logout_warns_when_the_keychain_refuses_to_release_the_entry( - self, isolated_home, secret_vault_factory - ): + def test_logout_warns_when_the_keychain_refuses_to_release_the_entry(self, isolated_home, secret_vault_factory): """A locked keychain leaves a live credential behind that the user believes is gone.""" - vault = secret_vault_factory( - blob=_secret_blob("https://test.example.com", "sk-stored"), erasable=False - ) + vault = secret_vault_factory(blob=_secret_blob("https://test.example.com", "sk-stored"), erasable=False) _write_token_file(isolated_home, key=None) result = self.runner.invoke(logout, obj={"secret_vault": vault}) @@ -1210,9 +1206,7 @@ class TestKeychainBackedCommands: assert str(token_file) in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" - def test_login_points_a_user_missing_the_keyring_package_at_the_install( - self, isolated_home, secret_vault_factory - ): + def test_login_points_a_user_missing_the_keyring_package_at_the_install(self, isolated_home, secret_vault_factory): """`lite` ships with every install, the keyring package only with the cli extra. Telling that user their machine has no keychain sends them looking for a problem they do not have.""" result = self._login(secret_vault_factory(available=False, failure=KeyringNotInstalled())) @@ -1223,9 +1217,7 @@ class TestKeychainBackedCommands: assert "No OS keychain available" not in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" - def test_login_keeps_the_credential_when_the_backend_keeps_nothing( - self, isolated_home, secret_vault_factory - ): + def test_login_keeps_the_credential_when_the_backend_keeps_nothing(self, isolated_home, secret_vault_factory): """A backend that accepts writes and stores nothing must not be reported as keychain storage, because the file is then told to drop the only remaining copy.""" result = self._login(secret_vault_factory(discards=True)) @@ -1236,9 +1228,7 @@ class TestKeychainBackedCommands: assert "keyring --enable" in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" - def test_login_names_the_kill_switch_instead_of_blaming_the_machine( - self, isolated_home, secret_vault_factory - ): + def test_login_names_the_kill_switch_instead_of_blaming_the_machine(self, isolated_home, secret_vault_factory): result = self._login(secret_vault_factory(available=False, failure=KeyringDisabled())) assert result.exit_code == 0 @@ -1297,9 +1287,7 @@ class TestKeychainBackedCommands: assert "could not be read" in result.output assert "lite login" in result.output - def test_whoami_does_not_call_a_credential_it_cannot_read_authenticated( - self, isolated_home, secret_vault_factory - ): + def test_whoami_does_not_call_a_credential_it_cannot_read_authenticated(self, isolated_home, secret_vault_factory): """A login whose secret is stuck in an unreachable keychain authenticates nothing. Leading with "Authenticated" and a token age reads as a working session, and sends the user looking for the problem somewhere other than the keychain the notice underneath names.""" @@ -1316,9 +1304,7 @@ class TestKeychainBackedCommands: assert "the credential cannot be read" in result.output assert "could not be read" in result.output - def test_whoami_names_the_kill_switch_rather_than_a_missing_package( - self, isolated_home, secret_vault_factory - ): + def test_whoami_names_the_kill_switch_rather_than_a_missing_package(self, isolated_home, secret_vault_factory): """Every unreachable keychain used to be described as a locked one needing the keyring package installed. Someone who set the kill switch has the package and an unlocked keychain, so that advice sends them to fix two things that were never wrong.""" @@ -1330,9 +1316,7 @@ class TestKeychainBackedCommands: assert DISABLE_KEYRING_ENV_VAR in result.output assert "pip install" not in result.output - def test_print_token_points_an_install_without_keyring_at_the_package( - self, isolated_home, secret_vault_factory - ): + def test_print_token_points_an_install_without_keyring_at_the_package(self, isolated_home, secret_vault_factory): _write_token_file(isolated_home, key=None) vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) obj = {"base_url": "https://test.example.com", "secret_vault": vault} @@ -1398,9 +1382,22 @@ class TestLoginConfigClaude: def setup_method(self): self.runner = CliRunner() - def _run_login(self, tmp_path, args, base_url="https://test.example.com"): - settings_path = tmp_path / "claude" / "settings.json" + def _isolate_default_settings(self, tmp_path, monkeypatch): + """The default file, its `lite up` backup and its configure receipt all live under tmp_path.""" backup_path = tmp_path / "claude_settings_backup.json" + monkeypatch.setattr( + claude_settings_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup_path, "lite up", "lite down"),) + ) + monkeypatch.setattr( + claude_settings_module, "CLAUDE_SETTINGS_PATH", tmp_path / "default-home" / ".claude" / "settings.json" + ) + monkeypatch.setattr(claude_settings_module, "CONFIGURE_STATE_PATH", tmp_path / "claude_configure_state.json") + return backup_path + + def _run_login(self, tmp_path, monkeypatch, args, base_url="https://test.example.com", *, config_dir_env=None): + settings_path = tmp_path / "claude" / "settings.json" + backup_path = self._isolate_default_settings(tmp_path, monkeypatch) + env = {"CLAUDE_CONFIG_DIR": str(settings_path.parent)} if config_dir_env is None else config_dir_env poll_response = Mock() poll_response.status_code = 200 poll_response.json.return_value = { @@ -1416,55 +1413,102 @@ class TestLoginConfigClaude: patch("requests.get", return_value=poll_response), patch("litellm.proxy.client.cli.commands.auth.save_cli_token"), patch("litellm.proxy.client.cli.interface.show_commands"), - patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path), - patch( - "litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS", - (SettingsFileOwner(backup_path, "lite up", "lite down"),), - ), patch( "litellm.proxy.client.cli.commands.claude_settings.shutil.which", return_value="/usr/local/bin/lite", ), ): - result = self.runner.invoke(login, args, obj={"base_url": base_url}) + result = self.runner.invoke(login, args, obj={"base_url": base_url}, env=env) return result, settings_path, backup_path - def test_default_login_does_not_touch_claude_settings(self, tmp_path): - result, settings_path, _backup_path = self._run_login(tmp_path, []) + def test_default_login_does_not_touch_claude_settings(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, []) assert result.exit_code == 0 assert "Login successful!" in result.output assert not settings_path.exists() assert "Configured Claude Code" not in result.output - def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path): - result, settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code == 0 written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" - assert "Configured Claude Code" in result.output + assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output + assert "pins a proxy model for every tier" not in result.output + assert "the model Claude Code starts on" in result.output - def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path): + def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) settings_path.write_text(json.dumps({"theme": "dark", "env": {"KEEP": "me"}})) - result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code == 0 written = json.loads(settings_path.read_text()) assert written["theme"] == "dark" assert written["env"]["KEEP"] == "me" - def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path): + def _run_login_refused_before_the_sso_flow(self, tmp_path, monkeypatch, config_dir): + self._isolate_default_settings(tmp_path, monkeypatch).write_text("{}") + with patch("requests.post") as post, patch("webbrowser.open") as browser: + result = self.runner.invoke( + login, + ["--config-claude"], + obj={"base_url": "https://test.example.com"}, + env={"CLAUDE_CONFIG_DIR": config_dir}, + ) + assert result.exit_code != 0 + assert "not logging in" in result.output and "lite down" in result.output + assert "`lite up` is currently managing" in result.output + assert "Login successful!" not in result.output + post.assert_not_called() + browser.assert_not_called() + assert not (tmp_path / "default-home" / ".claude" / "settings.json").exists() + + def test_refuses_before_logging_in_while_lite_up_holds_the_default_settings_file(self, tmp_path, monkeypatch): + self._run_login_refused_before_the_sso_flow(tmp_path, monkeypatch, config_dir="") + + def test_refuses_before_logging_in_while_lite_up_holds_the_default_file_reached_through_a_symlink( + self, tmp_path, monkeypatch + ): + default_config_dir = tmp_path / "default-home" / ".claude" + default_config_dir.mkdir(parents=True) + alias = tmp_path / "claude-alias" + alias.symlink_to(default_config_dir, target_is_directory=True) + + self._run_login_refused_before_the_sso_flow(tmp_path, monkeypatch, config_dir=str(alias)) + + def test_flag_writes_an_alternate_config_dir_even_while_lite_up_holds_the_default_file(self, tmp_path, monkeypatch): + (tmp_path / "claude_settings_backup.json").write_text("{}") + + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" + assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output + + def test_flag_keeps_a_config_dir_receipt_apart_from_the_default_file_receipt(self, tmp_path, monkeypatch): + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + + assert result.exit_code == 0, result.output + default_receipt = tmp_path / "claude_configure_state.json" + assert not default_receipt.exists() + receipts = list((tmp_path / "claude_configure_state").glob("*.json")) + assert len(receipts) == 1 + assert json.loads(receipts[0].read_text())["file_existed"] is False + + def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) settings_path.write_text("not json at all {{{") - result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"]) + result, _settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"]) assert result.exit_code != 0 assert "Login successful!" in result.output @@ -1853,7 +1897,10 @@ class TestPkcePrintToken: assert result.stdout == "" assert sum(len(session.posts) for session in _FakeSession.instances) == 1 assert result.output.count("Could not renew the key") == 1 - assert "Could not renew the key: token request failed with 400: the refresh token was already used" in result.output + assert ( + "Could not renew the key: token request failed with 400: the refresh token was already used" + in result.output + ) assert "Key expired. Run 'lite login --pkce' again." in result.output save.assert_not_called() diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index e5f2a9d95bd..fc2d98f2264 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -1,7 +1,9 @@ import json +import os import shlex import stat import time +from pathlib import Path from unittest.mock import patch import pytest @@ -9,14 +11,30 @@ from click.testing import CliRunner from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord from litellm.proxy.client.cli import cli +from litellm.litellm_core_utils.private_json import commit_staged_json from litellm.proxy.client.cli.commands.claude_settings import ( + ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, AUTOROUTE_BACKUP_PATH, BACKUP_PATH, + CLAUDE_SETTINGS_PATH, + CONFIGURE_STATE_PATH, + OWNED_ENV_KEYS, + OWNED_TOP_LEVEL_KEYS, SETTINGS_FILE_OWNERS, + ApiKeyHelper, ClaudeSettingsError, + KeepModel, SettingsFileOwner, + StartOn, + StaticToken, + UnpinModel, + claude_settings_path, + configure_claude_settings, + configure_state_path, + lite_api_key_helper_configured, + merge_claude_settings, resolve_api_key_helper, - write_claude_settings, + unconfigure_claude_settings, ) @@ -24,6 +42,7 @@ def _owners(*backup_paths): """Stand-in owners for the real `lite up` / `lite autoroute up` registry.""" return tuple(SettingsFileOwner(path, "lite up", "lite down") for path in backup_paths) + CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" @@ -97,17 +116,28 @@ def lite_on_path(): yield -class TestWriteClaudeSettings: +def _helper_configure(base_url, settings_path, owners, state_path=None): + """`lite login --config-claude`'s shape: the login credential behind apiKeyHelper, no pinned model.""" + state = state_path if state_path is not None else settings_path.parent.parent / "state.json" + root = base_url.rstrip("/") + configure_claude_settings( + root, ApiKeyHelper(resolve_api_key_helper(root)), KeepModel(), settings_path, state, owners + ) + + +class TestConfigureWithTheLoginHelper: def test_creates_the_file_and_its_parent_when_missing(self, paths, lite_on_path): settings_path, backup_path = paths assert not settings_path.parent.exists() - write_claude_settings("https://proxy.example.com/", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com/", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" + assert "model" not in written def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path): settings_path, backup_path = paths @@ -123,7 +153,7 @@ class TestWriteClaudeSettings: ) ) - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["theme"] == "dark" @@ -135,26 +165,31 @@ class TestWriteClaudeSettings: def test_rerunning_against_a_new_proxy_refreshes_both_base_url_and_helper(self, paths, lite_on_path): settings_path, backup_path = paths - write_claude_settings("https://first.example.com", settings_path, _owners(backup_path)) - write_claude_settings("https://second.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://first.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://second.example.com", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://second.example.com" assert "second.example.com" in written["apiKeyHelper"] assert "first.example.com" not in written["apiKeyHelper"] - def test_drops_a_stray_static_api_key_so_the_helper_token_wins(self, paths, lite_on_path): + def test_drops_stray_static_credentials_so_the_helper_token_wins(self, paths, lite_on_path): + # Claude Code prefers ANTHROPIC_AUTH_TOKEN over apiKeyHelper, so a virtual key left behind + # by an earlier `lite configure claude --api-key` would silently keep winning. settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) - settings_path.write_text(json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked"}})) + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked", "ANTHROPIC_AUTH_TOKEN": "sk-old"}}) + ) - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - assert "ANTHROPIC_API_KEY" not in json.loads(settings_path.read_text())["env"] + env = json.loads(settings_path.read_text())["env"] + assert "ANTHROPIC_API_KEY" not in env and "ANTHROPIC_AUTH_TOKEN" not in env def test_written_file_is_owner_only(self, paths, lite_on_path): settings_path, backup_path = paths - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert stat.S_IMODE(settings_path.stat().st_mode) == 0o600 def test_refuses_while_lite_up_holds_a_backup(self, paths, lite_on_path): @@ -162,7 +197,7 @@ class TestWriteClaudeSettings: backup_path.write_text("{}") with pytest.raises(ClaudeSettingsError, match="lite down"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert not settings_path.exists() @@ -172,7 +207,7 @@ class TestWriteClaudeSettings: settings_path.write_text("not json at all {{{") with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert settings_path.read_text() == "not json at all {{{" @@ -180,7 +215,7 @@ class TestWriteClaudeSettings: settings_path, backup_path = paths with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert not settings_path.exists() @@ -196,7 +231,7 @@ class TestWriteClaudeSettings: settings_path.write_bytes(b'{"theme": "\xff\xfe"}') with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths, lite_on_path): """An unreadable settings file must not surface as "Authentication failed". @@ -210,16 +245,18 @@ class TestWriteClaudeSettings: settings_path.mkdir() with pytest.raises(ClaudeSettingsError, match="Could not read"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths, lite_on_path): settings_path, backup_path = paths - with patch( - f"{CLAUDE_SETTINGS_MODULE}.write_private_json", - side_effect=OSError("Read-only file system"), - ): - with pytest.raises(ClaudeSettingsError, match="Read-only file system"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + settings_path.parent.mkdir(parents=True) + settings_path.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not write"): + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + finally: + settings_path.parent.chmod(0o700) + assert not settings_path.exists() class TestApiKeyHelperIsActuallyInvocable: @@ -303,7 +340,7 @@ class TestConflictingOwnersOfTheSettingsFile: backup.write_text("{}") stand_in = SettingsFileOwner(backup, owner.start_command, owner.stop_command) with pytest.raises(ClaudeSettingsError, match="currently managing"): - write_claude_settings("https://proxy.example.com", settings_path, (stand_in,)) + _helper_configure("https://proxy.example.com", settings_path, (stand_in,)) backup.unlink() assert not settings_path.exists() @@ -314,9 +351,9 @@ class TestConflictingOwnersOfTheSettingsFile: autoroute = SettingsFileOwner(backup, "lite autoroute up", "lite autoroute down") with pytest.raises(ClaudeSettingsError, match="`lite autoroute up` is currently managing"): - write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + _helper_configure("https://proxy.example.com", settings_path, (autoroute,)) with pytest.raises(ClaudeSettingsError, match="Run `lite autoroute down` first"): - write_claude_settings("https://proxy.example.com", settings_path, (autoroute,)) + _helper_configure("https://proxy.example.com", settings_path, (autoroute,)) def test_the_registry_matches_the_paths_the_commands_actually_use(self): """A second definition of the autoroute dir must not drift from this one.""" @@ -341,7 +378,7 @@ class TestDoesNotDestroyUserOwnedStructure: link.parent.mkdir() link.symlink_to(real) - write_claude_settings("https://proxy.example.com", link, ()) + _helper_configure("https://proxy.example.com", link, ()) assert link.is_symlink() assert json.loads(real.read_text())["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" @@ -354,6 +391,567 @@ class TestDoesNotDestroyUserOwnedStructure: settings_path.write_text(json.dumps({"theme": "dark", "env": "not-an-object"})) with pytest.raises(ClaudeSettingsError, match="non-object"): - write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path)) + _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert json.loads(settings_path.read_text())["env"] == "not-an-object" + + +class TestClaudeSettingsPath: + def test_defaults_to_the_home_settings_file(self): + assert claude_settings_path({}) == CLAUDE_SETTINGS_PATH + assert claude_settings_path({"CLAUDE_CONFIG_DIR": ""}) == CLAUDE_SETTINGS_PATH + + def test_follows_claude_config_dir_like_claude_code_does(self, tmp_path): + assert claude_settings_path({"CLAUDE_CONFIG_DIR": str(tmp_path)}) == tmp_path / "settings.json" + + def test_expands_a_tilde_in_claude_config_dir(self): + assert claude_settings_path({"CLAUDE_CONFIG_DIR": "~/.claude-work"}) == ( + Path.home() / ".claude-work" / "settings.json" + ) + + +class TestConfigureStatePath: + """Each settings file gets its own undo receipt: the default file keeps the long-standing path, and + a CLAUDE_CONFIG_DIR file gets one keyed by its resolved location, so `lite unconfigure claude` + under one config dir never restores the other file's history.""" + + @pytest.fixture + def default_paths(self, tmp_path): + default_settings = tmp_path / "home" / ".claude" / "settings.json" + default_state = tmp_path / "home" / ".litellm" / "claude_configure_state.json" + with ( + patch(f"{CLAUDE_SETTINGS_MODULE}.CLAUDE_SETTINGS_PATH", default_settings), + patch(f"{CLAUDE_SETTINGS_MODULE}.CONFIGURE_STATE_PATH", default_state), + ): + yield default_settings, default_state + + def test_the_default_file_keeps_the_default_receipt(self, default_paths): + default_settings, default_state = default_paths + assert configure_state_path(default_settings) == default_state + + def test_a_symlink_alias_of_the_default_file_shares_its_receipt(self, default_paths): + default_settings, default_state = default_paths + default_settings.parent.mkdir(parents=True) + alias = default_settings.parent.parent / "claude-alias" + alias.symlink_to(default_settings.parent, target_is_directory=True) + assert configure_state_path(alias / "settings.json") == default_state + + def test_another_settings_file_gets_a_receipt_of_its_own_beside_the_default_one(self, default_paths, tmp_path): + _default_settings, default_state = default_paths + work_state = configure_state_path(tmp_path / "work" / "settings.json") + play_state = configure_state_path(tmp_path / "play" / "settings.json") + assert work_state != default_state and play_state != default_state + assert work_state != play_state + assert work_state.parent == play_state.parent == default_state.parent / "claude_configure_state" + assert work_state == configure_state_path(tmp_path / "work" / "settings.json") + + def test_configure_and_unconfigure_under_a_config_dir_leave_the_default_receipt_alone( + self, default_paths, tmp_path, lite_on_path + ): + _default_settings, default_state = default_paths + work_settings = tmp_path / "work" / "settings.json" + work_state = configure_state_path(work_settings) + configure_claude_settings( + "https://proxy.example.com", + ApiKeyHelper(resolve_api_key_helper("https://proxy.example.com")), + KeepModel(), + work_settings, + work_state, + (), + ) + assert work_state.exists() and not default_state.exists() + outcome = unconfigure_claude_settings(work_settings, work_state, ()) + assert outcome.file_removed and not work_settings.exists() + assert not work_state.exists() + + +class TestLiteApiKeyHelperConfigured: + def _settings(self, tmp_path, payload): + settings_path = tmp_path / "settings.json" + settings_path.write_text(payload) + return settings_path + + def test_recognises_the_helper_lite_login_wrote_for_this_proxy(self, tmp_path, lite_on_path): + settings_path = tmp_path / "settings.json" + _helper_configure("https://proxy.example.com/", settings_path, (), tmp_path / "state.json") + + assert lite_api_key_helper_configured("https://proxy.example.com/", settings_path) is True + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is True + + def test_a_helper_for_another_proxy_does_not_count(self, tmp_path, lite_on_path): + settings_path = tmp_path / "settings.json" + _helper_configure("https://other.example.com", settings_path, (), tmp_path / "state.json") + + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_a_hand_written_helper_does_not_count(self, tmp_path, lite_on_path): + settings_path = self._settings(tmp_path, json.dumps({"apiKeyHelper": "cat ~/.my-proxy-key"})) + + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_missing_or_helperless_settings_do_not_count(self, tmp_path, lite_on_path): + assert lite_api_key_helper_configured("https://proxy.example.com", tmp_path / "absent.json") is False + helperless = json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://proxy.example.com"}}) + settings_path = self._settings(tmp_path, helperless) + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_unreadable_settings_fall_back_to_false(self, tmp_path, lite_on_path): + settings_path = self._settings(tmp_path, "{not json") + + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + def test_lite_missing_from_path_falls_back_to_false(self, tmp_path): + helper = "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" + settings_path = self._settings(tmp_path, json.dumps({"apiKeyHelper": helper})) + with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): + assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False + + +class TestMergeClaudeSettings: + """One merge for every way Claude Code gets wired: `lite up`, `lite login --config-claude`, + `lite configure claude` and `lite autoroute up`.""" + + def test_a_static_token_lands_in_env_and_the_helper_slot_is_cleared(self): + settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "leaked"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000/", StaticToken("token-abc")) + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert "ANTHROPIC_API_KEY" not in merged["env"] + assert "apiKeyHelper" not in merged + assert "model" not in merged + assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) + + def test_a_helper_lands_top_level_and_the_static_slots_are_cleared(self): + settings = {"env": {"ANTHROPIC_AUTH_TOKEN": "sk-old", "ANTHROPIC_API_KEY": "leaked"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", ApiKeyHelper("lite auth print-token")) + assert merged["apiKeyHelper"] == "lite auth print-token" + assert "ANTHROPIC_AUTH_TOKEN" not in merged["env"] and "ANTHROPIC_API_KEY" not in merged["env"] + + def test_keeps_existing_switch_values_and_unrelated_keys_without_mutating_the_input(self): + settings = {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", StaticToken("token-abc")) + assert merged["theme"] == "dark" + assert merged["env"]["SOME_OTHER_VAR"] == "value" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" + assert settings == {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} + + def test_a_default_model_sets_only_the_row_claude_code_starts_on(self): + merged = merge_claude_settings( + {}, "http://127.0.0.1:4000", StaticToken("token-abc"), default_model="claude-auto" + ) + assert merged["model"] == "claude-auto" + assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) + + def test_a_tier_model_forces_every_claude_code_tier_as_autoroute_needs(self): + # Router's auto-router registry is keyed by the literal requested model string with no + # wildcard resolution, so `lite autoroute up` overrides the env var each tier reads. + settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} + merged = merge_claude_settings( + settings, "http://127.0.0.1:4000", StaticToken("token-abc"), tier_model="autorouter" + ) + assert {merged["env"][key] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS} == {"autorouter"} + assert "model" not in merged + + def test_touches_exactly_the_declared_owned_keys(self): + # The receipt and unconfigure restore exactly OWNED_*_KEYS, so a key the merge writes outside + # that table would be written by configure and never undone. + settings = { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "old", "ENABLE_TOOL_SEARCH": "false"}, + "apiKeyHelper": "old-helper", + "model": "old-model", + } + for credential in (StaticToken("token-abc"), ApiKeyHelper("helper")): + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", credential, default_model="claude-auto") + changed_top_level = {key for key in set(settings) | set(merged) if settings.get(key) != merged.get(key)} + assert changed_top_level - {"env"} <= set(OWNED_TOP_LEVEL_KEYS) + changed_env = { + key + for key in set(settings["env"]) | set(merged["env"]) + if settings["env"].get(key) != merged["env"].get(key) + } + assert changed_env <= set(OWNED_ENV_KEYS) + assert merged["permissions"] == {"allow": ["Bash"]} + assert merged["env"]["KEEP_ME"] == "1" + + +PROXY = "http://127.0.0.1:4000" +ANTHROPIC = "https://api.anthropic.com" +HELPER = ApiKeyHelper("lite auth print-token") +ORIGINAL = { + "theme": "dark", + "permissions": {"allow": ["Bash"]}, + "env": {"KEEP_ME": "1", "ANTHROPIC_API_KEY": "sk-ant-mine", "ANTHROPIC_BASE_URL": ANTHROPIC}, + "apiKeyHelper": "/usr/local/bin/lite auth print-token", + "model": "claude-opus-5", +} + + +def _set(path, value): + """A user edit: set (or with `_ABSENT`, remove) the key at a dotted path in the settings file.""" + + def edit(settings): + section, _, key = path.rpartition(".") + container = settings.setdefault(section, {}) if section else settings + if value is _ABSENT: + container.pop(key, None) + else: + container[key] = value + return settings + + return edit + + +_ABSENT = object() + + +class _Rig: + """One settings file plus receipt under tmp_path, driven through the public functions only.""" + + def __init__(self, tmp_path, initial): + self.settings = tmp_path / "claude" / "settings.json" + self.state = tmp_path / "state" / "claude_configure_state.json" + if initial is not None: + self.settings.parent.mkdir(parents=True) + self.settings.write_text(json.dumps(initial)) + + def read(self): + return json.loads(self.settings.read_text()) if self.settings.exists() else None + + def configure(self, credential=StaticToken("sk-virtual-key"), model=StartOn("claude-auto"), **kwargs): + configure_claude_settings(PROXY, credential, model, self.settings, self.state, (), **kwargs) + + def edit(self, *edits): + settings = self.read() + for apply in edits: + settings = apply(settings) + self.settings.write_text(json.dumps(settings)) + + def unconfigure(self): + return unconfigure_claude_settings(self.settings, self.state, ()) + + +# Each row: initial file, steps (configure kwargs dicts or edit callables) between the first configure +# and unconfigure, the expected file afterwards, and the expected outcome fields. Sequences that used +# to be one test each; the receipt's rules are what make them all come out right. +UNDO_SCENARIOS = { + "plain round trip": (ORIGINAL, [], ORIGINAL, {"kept": ()}), + "no file before": (None, [], None, {"file_removed": True}), + "no env before": ({"theme": "dark"}, [], {"theme": "dark"}, {}), + "null env before": ({"theme": "dark", "env": None}, [], {"theme": "dark", "env": None}, {}), + "empty env before": ({"theme": "dark", "env": {}}, [], {"theme": "dark", "env": {}}, {}), + "user edits stay and are named": ( + ORIGINAL, + [_set("env.ENABLE_TOOL_SEARCH", "false"), _set("model", "claude-sonnet-4-6")], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "claude-sonnet-4-6"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}, "withheld": ()}, + ), + "user filled an env configure created": (None, [_set("env.MY_VAR", "mine")], {"env": {"MY_VAR": "mine"}}, {}), + "user deleted the file": (None, [lambda s: None], None, {"file_removed": True, "restored": (), "kept": ()}), + "user removed our key: neither restored nor kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_AUTH_TOKEN", _ABSENT)], + ORIGINAL, + {"not_restored": {"env.ANTHROPIC_AUTH_TOKEN"}, "kept": ()}, + ), + "restored names only what changed": ( + {"model": "claude-opus-5"}, + [], + {"model": "claude-opus-5"}, + { + "restored": { + "env.ANTHROPIC_BASE_URL", + "env.ENABLE_TOOL_SEARCH", + "env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", + "apiKeyHelper", + }, + "kept": (), + }, + {"credential": HELPER, "model": KeepModel()}, + ), + "repeat across credential kinds keeps the first snapshot": ( + ORIGINAL, + [ + {"credential": HELPER, "model": UnpinModel()}, + {"credential": StaticToken("sk-rotated"), "model": StartOn("claude-sonnet-4-6")}, + ], + ORIGINAL, + {}, + ), + "repeat without a model lets go of our pin, user had none": ({}, [{"model": UnpinModel()}], {}, {}), + "repeat without a model lets go of our pin, user had one": ( + {"model": "claude-opus-5"}, + [{"model": UnpinModel()}], + {"model": "claude-opus-5"}, + {}, + ), + "re-login keeps our pin": (None, [{"credential": HELPER, "model": KeepModel()}], None, {"file_removed": True}), + "edit between configures survives an unpin repeat": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": HELPER, "model": UnpinModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures survives a re-login": ( + ORIGINAL, + [ + _set("model", "my-favourite"), + _set("env.ENABLE_TOOL_SEARCH", "false"), + {"credential": HELPER, "model": KeepModel()}, + ], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, + ), + "edit between configures: a same-model repeat displaces it, so it is what comes back": ( + ORIGINAL, + [_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": HELPER}], + {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, + {"kept": {"env.ENABLE_TOOL_SEARCH"}, "restored_includes": {"model"}}, + ), + "base URL changed since: credentials withheld, receipt kept": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {**ORIGINAL, "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, "apiKeyHelper": _ABSENT}, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC), ("apiKeyHelper", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL"}, + "receipt_kept": True, + }, + ), + "base URL changed and back: judged against the URL the restored file holds": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", ANTHROPIC)], + ORIGINAL, + {"withheld": ()}, + ), + "credential captured beside no URL goes back only beside no URL": ( + {"env": {"ANTHROPIC_API_KEY": "sk-default-endpoint"}}, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")], + {"env": {"ANTHROPIC_BASE_URL": "http://other-proxy:4000"}}, + { + "withheld": {("env.ANTHROPIC_API_KEY", "no ANTHROPIC_BASE_URL (Anthropic's default endpoint)")}, + "receipt_kept": True, + }, + ), + "restored document empty while a credential is withheld: file goes, receipt stays": ( + None, + [ + _set("env.ANTHROPIC_API_KEY", "sk-user"), + {"credential": HELPER, "model": KeepModel()}, + _set("env.ANTHROPIC_BASE_URL", _ABSENT), + ], + None, + {"withheld": {("env.ANTHROPIC_API_KEY", PROXY)}, "file_removed": True, "receipt_kept": True}, + {"credential": HELPER, "model": KeepModel()}, + ), + "a credential the user changed is kept, never also withheld": ( + ORIGINAL, + [_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000"), _set("apiKeyHelper", "/opt/mine/helper")], + { + **ORIGINAL, + "env": {"KEEP_ME": "1", "ANTHROPIC_BASE_URL": "http://other-proxy:4000"}, + "apiKeyHelper": "/opt/mine/helper", + }, + { + "withheld": {("env.ANTHROPIC_API_KEY", ANTHROPIC)}, + "kept": {"env.ANTHROPIC_BASE_URL", "apiKeyHelper"}, + "receipt_kept": True, + }, + ), +} + + +def _expected_file(expected): + if expected is None: + return None + return {k: v for k, v in expected.items() if v is not _ABSENT} + + +class TestConfigureAndUnconfigure: + """`configure_claude_settings` records how to undo itself; `unconfigure_claude_settings` undoes only that.""" + + @pytest.mark.parametrize("scenario", UNDO_SCENARIOS.values(), ids=UNDO_SCENARIOS.keys()) + def test_undo_matrix(self, tmp_path, scenario): + initial, steps, expected, outcome_expectations, *first = scenario + rig = _Rig(tmp_path, initial) + rig.configure(**(first[0] if first else {})) + for step in steps: + if isinstance(step, dict): + rig.configure(**step) + elif rig.settings.exists() and step(json.loads(rig.settings.read_text())) is None: + rig.settings.unlink() + else: + rig.edit(step) + + outcome = rig.unconfigure() + + assert rig.read() == _expected_file(expected) + assert rig.state.exists() == outcome_expectations.get("receipt_kept", False) + for field, want in outcome_expectations.items(): + if field == "withheld": + assert {(item.key, item.endpoint) for item in outcome.withheld} == set(want) + elif field == "not_restored": + assert not set(want) & set(outcome.restored) and not set(want) & set(outcome.kept) + elif field == "restored_includes": + assert set(want) <= set(outcome.restored) + elif field in ("restored", "kept"): + assert set(getattr(outcome, field)) == set(want) + elif field != "receipt_kept": + assert getattr(outcome, field) == want + assert not {item.key for item in outcome.withheld} & set(outcome.kept) + + def test_configure_writes_owner_only_and_the_receipt_never_holds_the_key(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure(credential=StaticToken("sk-virtual-key-never-on-disk-twice")) + configured = rig.read() + assert configured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key-never-on-disk-twice" + assert configured["env"]["ANTHROPIC_BASE_URL"] == PROXY and configured["model"] == "claude-auto" + assert "ANTHROPIC_API_KEY" not in configured["env"] and "apiKeyHelper" not in configured + assert stat.S_IMODE(rig.settings.stat().st_mode) == 0o600 == stat.S_IMODE(rig.state.stat().st_mode) + assert "sk-virtual-key-never-on-disk-twice" not in rig.state.read_text() + + def test_withheld_credentials_come_back_once_the_url_points_at_their_server_again(self, tmp_path): + # The kept receipt owns only the withheld slots: the second unconfigure restores exactly those. + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + rig.edit(_set("env.ANTHROPIC_BASE_URL", "http://other-proxy:4000")) + rig.unconfigure() + rig.edit(_set("env.ANTHROPIC_BASE_URL", ANTHROPIC), _set("theme", "light")) + outcome = rig.unconfigure() + assert rig.read() == {**ORIGINAL, "theme": "light"} + assert set(outcome.restored) == {"env.ANTHROPIC_API_KEY", "apiKeyHelper"} + assert outcome.kept == () and outcome.withheld == () and not rig.state.exists() + + @pytest.mark.parametrize( + ("path", "value", "repeat_credential"), + [ + ("env.ANTHROPIC_API_KEY", "sk-user-added-later", HELPER), + ("env.ANTHROPIC_AUTH_TOKEN", "sk-users-own-token", HELPER), + ("apiKeyHelper", "/opt/mine/helper", StaticToken("sk-rotated")), + ], + ids=["user-adds-api-key", "user-replaces-our-token", "user-sets-own-helper"], + ) + def test_a_credential_the_user_set_between_two_configures_is_what_comes_back( + self, tmp_path, path, value, repeat_credential + ): + # The repeat's merge clears the slot, so the displaced value is snapshotted and is what returns; + # it was set while the file pointed at the proxy, so it returns once the file points there again. + rig = _Rig(tmp_path, {"theme": "dark"}) + rig.configure(credential=HELPER, model=KeepModel()) + rig.edit(_set(path, value)) + rig.configure(credential=repeat_credential, model=KeepModel()) + assert not _lookup(rig.read(), path) + + outcome = rig.unconfigure() + assert [(item.key, item.endpoint) for item in outcome.withheld] == [(path, PROXY)] + assert rig.read() == {"theme": "dark"} and rig.state.exists() + rig.settings.write_text(json.dumps({"theme": "dark", "env": {"ANTHROPIC_BASE_URL": PROXY}})) + outcome = rig.unconfigure() + assert _lookup(rig.read(), path) == value + assert outcome.restored == (path,) and outcome.withheld == () and not rig.state.exists() + + def test_a_receipt_commit_that_fails_leaves_no_staged_token_behind(self, tmp_path): + rig = _Rig(tmp_path, {}) + + def commit_receipt_fails(staged, path): + if path == str(rig.state): + os.unlink(staged) + raise OSError("receipt rename failed") + commit_staged_json(staged, path) + + with pytest.raises(ClaudeSettingsError, match=r"Could not write .*receipt rename failed"): + rig.configure(credential=StaticToken("sk-never-left-in-a-temp-file"), commit=commit_receipt_fails) + assert not list(rig.settings.parent.glob(".tmp-*")) and not list(rig.state.parent.glob(".tmp-*")) + assert rig.read() == {} and not rig.state.exists() + + @pytest.mark.parametrize("configured_before", [False, True], ids=["first-configure", "repeat-configure"]) + def test_a_settings_commit_that_fails_after_the_receipt_landed_puts_the_receipt_back( + self, tmp_path, configured_before + ): + # The two renames are not atomic: a settings rename that fails after the receipt landed must + # not leave a receipt describing settings that were never written. + rig = _Rig(tmp_path, ORIGINAL) + if configured_before: + rig.configure() + receipt_before = rig.state.read_text() if configured_before else None + settings_before = rig.settings.read_text() + + def commit_settings_fails(staged, path): + if path == str(rig.settings): + os.unlink(staged) + raise OSError("rename failed") + commit_staged_json(staged, path) + + with pytest.raises(ClaudeSettingsError, match="rename failed"): + rig.configure(credential=StaticToken("sk-rotated"), commit=commit_settings_fails) + assert rig.settings.read_text() == settings_before + assert (rig.state.read_text() if rig.state.exists() else None) == receipt_before + if configured_before: + rig.unconfigure() + assert rig.read() == ORIGINAL + + def test_a_failed_repeat_configure_leaves_the_earlier_undo_intact(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + receipt_before = rig.state.read_text() + rig.settings.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not write"): + rig.configure(credential=StaticToken("sk-rotated")) + finally: + rig.settings.parent.chmod(0o700) + assert rig.state.read_text() == receipt_before and not list(rig.state.parent.glob(".tmp-*")) + assert rig.read()["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" + rig.unconfigure() + assert rig.read() == ORIGINAL + + def test_unconfigure_reports_a_receipt_it_cannot_remove_as_a_settings_error(self, tmp_path): + rig = _Rig(tmp_path, ORIGINAL) + rig.configure() + rig.state.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not remove"): + rig.unconfigure() + finally: + rig.state.parent.chmod(0o700) + + def test_configure_writes_through_a_symlinked_settings_file(self, tmp_path): + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text(json.dumps({"theme": "dark"})) + link = tmp_path / "settings.json" + link.symlink_to(target) + configure_claude_settings(PROXY, StaticToken("sk-virtual-key"), UnpinModel(), link, tmp_path / "state.json", ()) + assert link.is_symlink() + assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" + + @pytest.mark.parametrize("operation", ["configure", "unconfigure"]) + def test_refuses_while_a_temporary_owner_holds_a_backup(self, paths, tmp_path, operation): + settings_path, backup_path = paths + backup_path.write_text("{}") + owners = _owners(backup_path) + state = tmp_path / "state.json" + attempt = ( + (lambda: configure_claude_settings(PROXY, StaticToken("k"), UnpinModel(), settings_path, state, owners)) + if operation == "configure" + else (lambda: unconfigure_claude_settings(settings_path, state, owners)) + ) + with pytest.raises(ClaudeSettingsError, match="lite down"): + attempt() + assert not settings_path.exists() + + def test_unconfigure_without_a_receipt_is_an_error_not_a_silent_no_op(self, tmp_path): + with pytest.raises(ClaudeSettingsError, match="nothing to undo"): + _Rig(tmp_path, None).unconfigure() + + +def _lookup(settings, path): + section, _, key = path.rpartition(".") + return (settings.get(section) or {}).get(key) if section else settings.get(key) diff --git a/tests/test_litellm/proxy/client/cli/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py new file mode 100644 index 00000000000..ac7408339fe --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -0,0 +1,431 @@ +import json +import os +import stat + +import click +import pytest +import requests +import responses +from click.testing import CliRunner + +from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module +from litellm.proxy.client.cli.commands import configure as configure_module +from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner +from litellm.proxy.client.cli.commands.configure import configure_claude, configure_group, interactive_configure + +PROXY = "http://proxy.test:4000" +VALID_KEY = "sk-virtual-key" +LISTED_MODELS = ("claude-auto", "gpt-5.6-luna") + + +def _mock_models(): + responses.get( + f"{PROXY}/v1/models", + json={"data": [{"id": model, "object": "model"} for model in LISTED_MODELS]}, + match=[responses.matchers.header_matcher({"Authorization": f"Bearer {VALID_KEY}"})], + ) + responses.get(f"{PROXY}/v1/models", status=401) + + +@pytest.fixture +def paths(monkeypatch, tmp_path): + """The default settings file, reached the way Claude Code reaches it: CLAUDE_CONFIG_DIR names its directory.""" + settings_path = tmp_path / "claude" / "settings.json" + state_path = tmp_path / "litellm" / "claude_configure_state.json" + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(settings_path.parent)) + monkeypatch.setattr(claude_settings_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude_settings_module, "CONFIGURE_STATE_PATH", state_path) + return settings_path, state_path + + +@pytest.fixture +def lite_on_path(monkeypatch, tmp_path): + """A real `lite` executable on PATH, so the apiKeyHelper command resolves without patching.""" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + lite = bin_dir / "lite" + lite.write_text("#!/bin/sh\nexit 0\n") + lite.chmod(lite.stat().st_mode | stat.S_IXUSR) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ.get('PATH', '')}") + return str(lite) + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def lite_up_backup(monkeypatch, tmp_path): + """A `lite up` session holding its backup, the local precondition every settings write refuses on.""" + backup = tmp_path / "claude_settings_backup.json" + backup.write_text("{}") + monkeypatch.setattr( + claude_settings_module, "SETTINGS_FILE_OWNERS", (SettingsFileOwner(backup, "lite up", "lite down"),) + ) + return backup + + +def _configure(runner, *args): + return runner.invoke(cli, ["--base-url", PROXY, "configure", "claude", *args]) + + +class TestConfigureClaudeWithAVirtualKey: + @responses.activate + def test_writes_settings_and_reports_without_echoing_the_key(self, runner, paths): + _mock_models() + settings_path, state_path = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto") + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_BASE_URL"] == PROXY + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" + assert written["model"] == "claude-auto" + assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in written["env"] + assert state_path.exists() + assert VALID_KEY not in result.output + assert "Starting model: claude-auto" in result.output + assert "1 of the proxy's 2 models" in result.output + assert "lite unconfigure claude" in result.output + assert [call.request.headers.get("x-gateway-client") for call in responses.calls] == ["claude-code"] + + @responses.activate + def test_takes_the_key_from_the_global_option_and_keeps_claude_codes_default(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = runner.invoke(cli, ["--base-url", PROXY, "--api-key", VALID_KEY, "configure", "claude"]) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert "model" not in written + assert "Starting model: not pinned" in result.output + + @responses.activate + def test_refuses_a_model_the_proxy_does_not_list(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-nope") + assert result.exit_code != 0 + assert "'claude-nope' is not served" in result.output + assert "claude-auto, gpt-5.6-luna" in result.output + assert not settings_path.exists() + + @responses.activate + def test_refuses_a_key_the_proxy_rejects(self, runner, paths): + _mock_models() + settings_path, _ = paths + result = _configure(runner, "--api-key", "sk-wrong") + assert result.exit_code != 0 + assert "rejected your key (HTTP 401)" in result.output + assert not settings_path.exists() + + @responses.activate + @pytest.mark.parametrize( + ("mock", "expected", "unexpected"), + [ + ( + lambda: responses.get(f"{PROXY}/v1/models", body=requests.ConnectionError("refused")), + "Is the proxy at", + "answered", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", status=500), + "The proxy at http://proxy.test:4000 answered", + "Is the proxy at", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", body="not json"), + "answered, so check that it is a LiteLLM proxy", + "Is the proxy at", + ), + ( + lambda: responses.get(f"{PROXY}/v1/models", json={"data": []}), + "Claude Code would have nothing to run", + "Is the proxy at", + ), + ], + ids=["unreachable", "http-500", "non-json-body", "empty-list"], + ) + def test_the_listing_hint_matches_how_the_listing_failed(self, runner, paths, mock, expected, unexpected): + # Only a proxy that never answered gets the "is it running" question; a 500, a non-JSON body or an + # empty list prove it is up, and the hint says so instead. + mock() + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code != 0 + assert expected in result.output and unexpected not in result.output + assert not settings_path.exists() + + @responses.activate + @pytest.mark.parametrize("entry", ["virtual-key", "login", "interactive"]) + def test_refuses_while_lite_up_holds_a_backup_before_any_login_or_request( + self, runner, paths, monkeypatch, lite_up_backup, entry + ): + _mock_models() + + def login_must_not_run(ctx): + raise AssertionError("the local precondition must be checked before a login is attempted") + + monkeypatch.setattr(configure_module, "ensure_fresh_login", login_must_not_run) + if entry == "interactive": + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": None}) + with pytest.raises(click.ClickException, match="lite down"): + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=lambda listed: None) + else: + args = ["--api-key", VALID_KEY] if entry == "virtual-key" else [] + result = runner.invoke(configure_claude, args, obj={"base_url": PROXY, "api_key": None}) + assert result.exit_code != 0 and "lite down" in result.output + assert len(responses.calls) == 0 + assert not paths[0].exists() + + @responses.activate + def test_says_so_when_the_key_is_written_through_a_symlink(self, runner, paths, tmp_path): + _mock_models() + settings_path, _ = paths + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text("{}") + settings_path.parent.mkdir(parents=True) + settings_path.symlink_to(target) + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code == 0, result.output + assert "keep it out of version control" in result.output + assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + + +class TestConfigureClaudeWithTheLogin: + def _stored_login(self, monkeypatch): + monkeypatch.setattr(configure_module, "ensure_fresh_login", lambda ctx: None) + monkeypatch.setattr(configure_module, "get_stored_api_key", lambda expected_base_url, vault: VALID_KEY) + + @responses.activate + def test_uses_the_login_through_the_helper_and_writes_no_secret(self, runner, paths, monkeypatch, lite_on_path): + _mock_models() + self._stored_login(monkeypatch) + settings_path, _ = paths + result = runner.invoke( + configure_claude, + ["--model", "claude-auto"], + obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": True}, + ) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["apiKeyHelper"] == f"{lite_on_path} --base-url {PROXY} auth print-token" + assert "ANTHROPIC_AUTH_TOKEN" not in written["env"] + assert written["model"] == "claude-auto" + assert VALID_KEY not in settings_path.read_text() + assert "read through apiKeyHelper" in result.output + + @responses.activate + def test_an_explicit_key_still_wins_over_a_stored_login(self, runner, paths, monkeypatch, lite_on_path): + _mock_models() + self._stored_login(monkeypatch) + settings_path, _ = paths + result = runner.invoke( + configure_claude, + ["--api-key", VALID_KEY], + obj={"base_url": PROXY, "api_key": "sk-login-jwt", "api_key_from_token_file": True}, + ) + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY and "apiKeyHelper" not in written + + +class TestInteractiveConfigure: + @responses.activate + def test_asks_for_targets_and_a_starting_model_then_configures(self, paths): + _mock_models() + settings_path, _ = paths + asked = {} + + def pick_model(listed): + asked["listed"] = tuple(listed) + return "claude-auto" + + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + assert asked["listed"] == LISTED_MODELS + assert json.loads(settings_path.read_text())["model"] == "claude-auto" + + def test_does_nothing_when_claude_code_is_not_picked(self, paths): + settings_path, _ = paths + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + interactive_configure(ctx, pick_targets=lambda: (), pick_model=lambda listed: None) + assert not settings_path.exists() + + def test_bare_configure_without_a_terminal_names_the_non_interactive_command(self, runner, paths): + result = runner.invoke(cli, ["--base-url", PROXY, "configure"]) + assert result.exit_code != 0 + assert "lite configure claude --api-key" in result.output + + +class TestUnconfigureClaude: + @responses.activate + def test_restores_the_original_file_and_removes_the_receipt(self, runner, paths): + _mock_models() + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + original = {"theme": "dark", "model": "claude-opus-5"} + settings_path.write_text(json.dumps(original)) + assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text()) == original + assert not state_path.exists() + assert "Restored in" in result.output and "model" in result.output + assert "ANTHROPIC_API_KEY" not in result.output, "a key that never existed was not restored" + + @responses.activate + def test_a_file_only_configure_created_is_reported_removed_not_restored(self, runner, paths): + _mock_models() + settings_path, _ = paths + assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert not settings_path.exists() + assert "No settings file remains" in result.output and "Restored" not in result.output + + @responses.activate + def test_says_when_nothing_was_still_ours_and_names_what_it_kept(self, runner, paths): + _mock_models() + settings_path, _ = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"theme": "dark"})) + assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + edited = json.loads(settings_path.read_text()) + edited["env"] = {key: f"{value}-edited" for key, value in edited["env"].items()} + edited["model"] = "mine" + settings_path.write_text(json.dumps(edited)) + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert "Nothing in" in result.output and "was still ours to restore" in result.output + assert "Left as you changed them since:" in result.output and "model" in result.output + + @responses.activate + def test_names_the_server_a_withheld_credential_was_captured_with_and_keeps_the_receipt(self, runner, paths): + _mock_models() + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://api.anthropic.com", "ANTHROPIC_API_KEY": "sk-ant"}}) + ) + assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + edited = json.loads(settings_path.read_text()) + edited["env"]["ANTHROPIC_BASE_URL"] = "http://other-proxy:4000" + settings_path.write_text(json.dumps(edited)) + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert "env.ANTHROPIC_API_KEY (captured with https://api.anthropic.com)" in result.output + assert str(state_path) in result.output and state_path.exists() + assert "sk-ant" not in result.output + + def test_refuses_while_lite_up_holds_a_backup(self, runner, paths, lite_up_backup): + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code != 0 and "lite down" in result.output + + @responses.activate + def test_a_config_dir_is_configured_and_undone_apart_from_the_default_file( + self, runner, paths, monkeypatch, tmp_path, lite_up_backup + ): + _mock_models() + default_settings, default_state = paths + work_dir = tmp_path / "claude-work" + work_dir.mkdir() + original = {"theme": "dark"} + (work_dir / "settings.json").write_text(json.dumps(original)) + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(work_dir)) + + configured = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto") + assert configured.exit_code == 0, configured.output + assert f"Configured Claude Code: {work_dir / 'settings.json'}" in configured.output + assert json.loads((work_dir / "settings.json").read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + assert not default_settings.exists() and not default_state.exists() + + undone = runner.invoke(cli, ["unconfigure", "claude"]) + assert undone.exit_code == 0, undone.output + assert json.loads((work_dir / "settings.json").read_text()) == original + assert not default_settings.exists() and not default_state.exists() + assert runner.invoke(cli, ["unconfigure", "claude"]).exit_code != 0, "the receipt is gone with the undo" + + def test_without_a_receipt_it_fails_loudly(self, runner, paths): + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code != 0 + assert "nothing to undo" in result.output + + +class TestClaudeCodeView: + VIEW = {"anthropic-version": "2023-06-01", "x-gateway-client": "claude-code"} + + def _mock(self, rows): + responses.get( + f"{PROXY}/v1/models", + json={"data": rows}, + match=[responses.matchers.header_matcher({"Authorization": f"Bearer {VALID_KEY}", **self.VIEW})], + ) + + @responses.activate + @pytest.mark.parametrize( + "model, pinned", + [ + ("literal-claude-router-source", "emitted-literal"), + ("marked-sibling", "emitted-marked[1m]"), + ("emitted-collision", "emitted-source-priority"), + ("emitted-only", "emitted-only"), + ], + ) + def test_pins_source_identity_before_emitted_id(self, runner, paths, model, pinned): + self._mock( + [ + {"id": "emitted-collision", "source_model": "other-source"}, + {"id": "emitted-source-priority", "source_model": "emitted-collision"}, + {"id": "emitted-marked[1m]", "source_model": "marked-sibling"}, + {"id": "emitted-literal", "source_model": "literal-claude-router-source"}, + {"id": "emitted-only"}, + ] + ) + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", model) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text())["model"] == pinned + assert f"Starting model: {pinned}" in result.output + assert len(responses.calls) == 1 + + @responses.activate + def test_refuses_unknown_short_suffix(self, runner, paths): + self._mock([{"id": "emitted-router-source", "source_model": "literal-router-source"}]) + settings_path, _ = paths + result = _configure(runner, "--api-key", VALID_KEY, "--model", "source") + assert result.exit_code != 0 + assert "'source' is not served" in result.output + assert not settings_path.exists() + + @responses.activate + def test_interactive_picker_uses_source_names(self, paths): + self._mock([{"id": "emitted", "source_model": "source"}]) + settings_path, _ = paths + asked = {} + ctx = click.Context( + configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} + ) + + def pick_model(listed): + asked["listed"] = tuple(listed) + return "source" + + interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + assert asked["listed"] == ("source",) + assert json.loads(settings_path.read_text())["model"] == "emitted" + + @responses.activate + def test_counts_what_an_older_proxy_lets_the_picker_show(self, runner, paths): + _mock_models() + result = _configure(runner, "--api-key", VALID_KEY) + assert result.exit_code == 0, result.output + assert "/model will list 1 of the proxy's 2 models: Claude Code shows only ids containing" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 0dd388919a5..b73d1acc6e3 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -7,8 +7,6 @@ from unittest.mock import Mock, patch import pytest from click.testing import CliRunner - - import litellm.proxy.client.cli from litellm._version import version as litellm_version from litellm.proxy.client.cli import cli diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 68c0ac70064..03e5d7dd197 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -4,13 +4,16 @@ import stat from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import pytest import requests from litellm.proxy.client.cli.commands.pi import ( + ListingFailure, ModelLimits, PiSyncError, fetch_model_ids, fetch_model_limits, + fetch_model_listing, models_json_path, provider_block, sync_models_json, @@ -28,6 +31,10 @@ class _FakeResponse: return self._payload +def _refused(*args, **kwargs): + raise requests.ConnectionError("refused") + + class TestFetchModelIds: def test_returns_ids_in_proxy_order_deduped(self): captured = {} @@ -44,6 +51,43 @@ class TestFetchModelIds: assert captured["url"] == "http://localhost:4000/v1/models" assert captured["headers"] == {"Authorization": "Bearer sk-key"} + def test_returns_rows_with_optional_source_model_and_dedups_identical_rows(self): + result = fetch_model_listing( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse( + 200, + {"data": [{"id": "emitted", "source_model": "source"}, {"id": "emitted", "source_model": "source"}]}, + ), + ) + assert not isinstance(result, PiSyncError) + assert tuple((model.id, model.source_model) for model in result) == (("emitted", "source"),) + + @pytest.mark.parametrize( + "entry", + [ + {"id": ""}, + {"id": "emitted", "source_model": ""}, + {"id": "emitted", "source_model": 1}, + ], + ) + def test_rejects_invalid_model_identity(self, entry): + result = fetch_model_listing( + "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(200, {"data": [entry]}) + ) + assert isinstance(result, PiSyncError) and result.kind is ListingFailure.BAD_BODY + + def test_rejects_conflicting_emitted_id_mappings(self): + result = fetch_model_listing( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse( + 200, + {"data": [{"id": "emitted", "source_model": "one"}, {"id": "emitted", "source_model": "two"}]}, + ), + ) + assert isinstance(result, PiSyncError) and result.kind is ListingFailure.BAD_BODY + def test_network_error_is_a_value(self): def boom(*a, **k): raise requests.ConnectionError("refused") @@ -53,9 +97,7 @@ class TestFetchModelIds: assert "Could not list models" in result.message def test_non_200_is_a_value(self): - result = fetch_model_ids( - "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500) - ) + result = fetch_model_ids("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)) assert isinstance(result, PiSyncError) assert "HTTP 500" in result.message @@ -75,6 +117,22 @@ class TestFetchModelIds: ) assert isinstance(result, PiSyncError) assert "no models" in result.message + assert result.kind is ListingFailure.EMPTY + + @pytest.mark.parametrize( + ("get", "kind"), + [ + (_refused, ListingFailure.UNREACHABLE), + (lambda *a, **k: _FakeResponse(401), ListingFailure.REJECTED), + (lambda *a, **k: _FakeResponse(403), ListingFailure.REJECTED), + (lambda *a, **k: _FakeResponse(500), ListingFailure.OTHER), + (lambda *a, **k: _FakeResponse(200), ListingFailure.BAD_BODY), + ], + ids=["unreachable", "401", "403", "500", "bad-body"], + ) + def test_the_failure_kind_is_decided_where_the_response_is_classified(self, get, kind): + result = fetch_model_ids("http://localhost:4000", "sk-key", get=get) + assert isinstance(result, PiSyncError) and result.kind is kind class TestFetchModelLimits: diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index aead1764b0e..111d2f3d682 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -11,11 +11,11 @@ from click.testing import CliRunner from litellm.proxy.client.cli.commands import up as up_module from litellm.proxy.client.cli.commands.agents import AgentRunError -from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError +from litellm.proxy.client.cli.commands.claude_settings import ApiKeyHelper, ClaudeSettingsError from litellm.proxy.client.cli.commands.up import ( BackupRecord, UpError, - _ensure_fresh_login, + ensure_fresh_login, down, load_json_or_empty, merge_claude_settings, @@ -40,12 +40,12 @@ def _patch_paths(monkeypatch, tmp_path): class TestMergeClaudeSettings: def test_preserves_unrelated_top_level_keys(self): - merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", "helper") + merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["theme"] == "dark" def test_preserves_unrelated_env_keys(self): settings = {"env": {"SOME_OTHER_VAR": "value"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"]["SOME_OTHER_VAR"] == "value" def test_overrides_base_url_and_helper(self): @@ -53,7 +53,7 @@ class TestMergeClaudeSettings: "env": {"ANTHROPIC_BASE_URL": "https://old.example.com"}, "apiKeyHelper": "old-helper", } - merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") + merged = merge_claude_settings(settings, "http://localhost:4000/", ApiKeyHelper("new-helper")) assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" @@ -61,21 +61,21 @@ class TestMergeClaudeSettings: def test_preserves_existing_gateway_model_discovery(self): settings = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "0"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "0" def test_preserves_existing_tool_search(self): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} - merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert "ANTHROPIC_API_KEY" not in merged["env"] def test_works_from_empty_settings(self): - merged = merge_claude_settings({}, "http://localhost:4000", "helper") + merged = merge_claude_settings({}, "http://localhost:4000", ApiKeyHelper("helper")) assert merged["env"] == { "ANTHROPIC_BASE_URL": "http://localhost:4000", "ENABLE_TOOL_SEARCH": "true", @@ -85,7 +85,7 @@ class TestMergeClaudeSettings: def test_does_not_mutate_input(self): settings = {"env": {"FOO": "bar"}} - merge_claude_settings(settings, "http://localhost:4000", "helper") + merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) assert settings == {"env": {"FOO": "bar"}} @@ -327,7 +327,7 @@ class TestEnsureFreshLogin: monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [] @@ -339,7 +339,7 @@ class TestEnsureFreshLogin: monkeypatch, on_login=lambda: store.log_in({"key": "sk-b", "base_url": "http://proxy-b:4000"}, "sk-b") ) - _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) assert login_calls == [("http://proxy-b:4000", False)] assert store.key_requests == ["http://proxy-b:4000", "http://proxy-b:4000"] @@ -353,7 +353,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in({"key": "sk-a", "base_url": "http://proxy-a:4000"}, "sk-a"), ) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [("http://proxy-a:4000", False)] @@ -363,7 +363,7 @@ class TestEnsureFreshLogin: monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) with pytest.raises(UpError, match="Run `lite login` first"): - _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) def test_trusts_a_pkce_credential_that_was_renewed_on_the_way_in(self, monkeypatch): """A --pkce key inside its freshness buffer is renewed by `get_stored_api_key`, so `lite up` @@ -377,7 +377,7 @@ class TestEnsureFreshLogin: ) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [] assert store.key_requests == ["http://proxy-a:4000"] @@ -390,7 +390,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"), ) - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) assert login_calls == [("http://proxy-a:4000", True)] @@ -399,7 +399,7 @@ class TestEnsureFreshLogin: _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=-10), {}) with pytest.raises(UpError, match="Run `lite login --pkce` first"): - _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + ensure_fresh_login(_make_ctx("http://proxy-a:4000")) def test_trusts_the_key_the_cli_group_already_resolved_instead_of_reading_the_token_file_again( self, monkeypatch @@ -409,7 +409,7 @@ class TestEnsureFreshLogin: store = _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=86_400), {}) login_calls = _capture_login(monkeypatch) - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group")) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key="sk-pkce-renewed-by-the-group")) assert login_calls == [] assert store.key_requests == [] @@ -423,7 +423,7 @@ class TestEnsureFreshLogin: ) with pytest.raises(UpError, match="Run `lite login --pkce` first"): - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) assert store.key_requests == [] @@ -435,7 +435,7 @@ class TestEnsureFreshLogin: on_login=lambda: store.log_in(_pkce_record("http://proxy-a:4000", seconds_left=86_400), "sk-pkce-fresh"), ) - _ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) + ensure_fresh_login(_make_group_ctx("http://proxy-a:4000", api_key=None)) assert login_calls == [("http://proxy-a:4000", True)] assert store.key_requests == ["http://proxy-a:4000"] @@ -620,10 +620,7 @@ class TestUpCanInvokeTheRealLoginCommand: ctx.obj = {"base_url": "http://127.0.0.1:9"} ctx.invoke(real_login, pkce=False) - with ( - patch(f"{AUTH_MODULE}.CLAUDE_SETTINGS_PATH", settings_path), - patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")), - ): - CliRunner().invoke(driver, [], standalone_mode=False) + with patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")): + CliRunner().invoke(driver, [], standalone_mode=False, env={"CLAUDE_CONFIG_DIR": str(tmp_path)}) assert not settings_path.exists() diff --git a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py new file mode 100644 index 00000000000..7ef03140093 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py @@ -0,0 +1,109 @@ +"""Model identity survives Claude Code presentation, filtering and configured alias precedence.""" + +from itertools import combinations + +import pytest + +from litellm import Router +from litellm.proxy.common_utils.model_listing_utils import ( + ClaudeCodeRoutingNames, + claude_code_group_name, + claude_code_model_id, + claude_code_requested_group, + claude_code_view_ids, +) + + +def _encoded(name): + return "claude-router-" + name.encode().hex() + + +def _marked(name): + return f"{_encoded(name)}[1m]" + + +def _row(name, limit=1000000): + return {"id": name, "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": limit} + + +def _router(*names, aliases=None): + return Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} + for name in names + ], + model_group_alias=aliases, + ) + + +@pytest.mark.parametrize("limit", [None, 999999, 1000000]) +@pytest.mark.parametrize("name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"]) +def test_listing_round_trips_entire_source_name(name, limit): + names = frozenset({name}) + view = claude_code_model_id(name, limit, names) + assert (claude_code_group_name(view, names) or view) == name + if "claude" not in name: + assert view.startswith(_encoded(name)) + assert ("[1m]" in view.lower()) == (limit == 1000000 or "[1m]" in name.lower() and "claude" in name) + + +def test_collision_matrix_round_trips_without_duplicate_ids(): + universe = ("foo", "foo[1m]", "claude-router-foo", _encoded("foo"), _encoded("foo") + "[1m]", "claude-opus-5", "claude-opus-5[1m]") + for pair in combinations(universe, 2): + for visible in (pair, pair[:1], pair[1:]): + names = frozenset(pair) + view = claude_code_view_ids(tuple(_row(n) for n in visible), {"user-agent": "claude-code/2.1.267"}, names) + assert len(set(view.values())) == len(visible) + assert all((claude_code_group_name(shown, names) or shown) == source for source, shown in view.items()) + + +@pytest.mark.parametrize("spelling", ["claude-router-foo", "claude-router-ff", "claude-router-66 6f6f", "claude-router-666F6F", "claude-router-", _encoded("missing")]) +def test_unknown_or_noncanonical_ids_are_never_guessed(spelling): + assert claude_code_group_name(spelling, frozenset({"foo"})) is None + + +@pytest.mark.parametrize("headers,enabled", [ + ({"user-agent": "claude-code/2.1.267"}, True), + ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), + ({"x-gateway-client": "Claude-Code"}, True), + ({"user-agent": "anthropic-sdk-python/0.40"}, False), + ({}, False), +]) +def test_only_claude_code_gets_the_view(headers, enabled): + rows = (_row("foo"), _row("claude-opus-5")) + view = claude_code_view_ids(rows, headers, frozenset(row["id"] for row in rows)) + assert dict(view) == ({"foo": _encoded("foo") + "[1m]", "claude-opus-5": "claude-opus-5[1m]"} if enabled else {}) + + +@pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "team", "wildcard"]) +def test_configured_names_outrank_generated_ids_even_when_hidden_from_listing(monkeypatch, layer): + import litellm + + encoded = _encoded("foo") + alias = {encoded: "other"} + monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) + router = _router("foo", "other", *( (encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), aliases=alias if layer == "router" else None) + maps = (alias,) if layer in ("key", "team") else () + names = ClaudeCodeRoutingNames(router, None, maps) + assert claude_code_requested_group(encoded, router, None, maps) is None + assert claude_code_view_ids((_row("foo", None),), {"x-gateway-client": "claude-code"}, names)["foo"] == "foo" + + +@pytest.mark.parametrize("source", ["foo", "foo[1m]", "世界"]) +def test_mutation_breaking_the_hex_name_cannot_route_to_the_source(source): + router = _router(source) + encoded = _encoded(source) + malformed = encoded[:-1] + ("0" if encoded[-1] != "0" else "1") + assert claude_code_requested_group(malformed, router, None) is None + assert claude_code_requested_group(_marked(source), router, None) == source + + +def test_team_public_name_uses_the_same_scope_at_list_and_request(): + router = Router(model_list=[{ + "model_name": "model_name_team-a_id", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, + }]) + shown = claude_code_view_ids((_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a"))["shared"] + assert claude_code_requested_group(shown, router, "team-a") == "shared" + assert claude_code_requested_group(shown, router, "team-b") is None diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index ba342342366..875dca4bee3 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -11,15 +11,30 @@ clobber a pre-existing ``DATABASE_URL_READ_REPLICA``. A pre-existing ``DATABASE_URL`` (password auth) is likewise left untouched. """ +import datetime +import hashlib import os +import socket +import ssl +import tempfile +import threading import urllib.parse +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final from unittest.mock import patch import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec from pydantic import ValidationError from litellm.proxy.db.db_url_settings import ( + PG_SSL_REQUEST, DatabaseURLSettings, + translate_libpq_ssl_params, unsupported_db_scheme, unsupported_db_scheme_message, ) @@ -381,9 +396,7 @@ def test_writer_password_is_percent_encoded(monkeypatch): def test_writer_url_not_clobbered_when_already_set(monkeypatch): """An operator-pinned DATABASE_URL (e.g. helm's $(VAR) assembly) always wins over the discrete fields.""" - monkeypatch.setenv( - "DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db" - ) + monkeypatch.setenv("DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db") monkeypatch.setenv("DATABASE_HOST", "writer.example.com") monkeypatch.setenv("DATABASE_USER", "litellm") monkeypatch.setenv("DATABASE_NAME", "litellm_db") @@ -515,9 +528,7 @@ def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch): def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") - monkeypatch.setenv( - "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" - ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db") with pytest.raises(RuntimeError, match=r"DATABASE_URL_READ_REPLICA.*mysql"): _apply() @@ -542,15 +553,11 @@ def test_reader_inherits_writer_connection_params(monkeypatch): "DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&pool_timeout=20&pgbouncer=true", ) - monkeypatch.setenv( - "DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db" - ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db") _apply() - query = urllib.parse.parse_qs( - urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query - ) + query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query) assert query["connection_limit"] == ["3"] assert query["pool_timeout"] == ["20"] assert query["pgbouncer"] == ["true"] @@ -568,9 +575,7 @@ def test_reader_keeps_its_own_pinned_connection_params(monkeypatch): _apply() - query = urllib.parse.parse_qs( - urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query - ) + query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query) assert query["connection_limit"] == ["50"] assert query["pool_timeout"] == ["20"] @@ -776,6 +781,170 @@ def test_libpq_verify_full_and_sslrootcert_become_prisma_strict_sslcert(monkeypa } +def _issue_cert( + subject: str, issuer: x509.Certificate | None, issuer_key: ec.EllipticCurvePrivateKey | None, ca: bool +) -> tuple[x509.Certificate, ec.EllipticCurvePrivateKey]: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name((x509.NameAttribute(x509.NameOID.COMMON_NAME, subject),)) + now: Final = datetime.datetime.now(datetime.timezone.utc) + builder: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(issuer.subject if issuer else name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=5)) + .not_valid_after(now + datetime.timedelta(days=1)) + .add_extension(x509.BasicConstraints(ca=ca, path_length=None), critical=True) + .add_extension(x509.SubjectAlternativeName((x509.DNSName("localhost"),)), critical=False) + ) + return builder.sign(issuer_key or key, hashes.SHA256()), key + + +def _pem(cert: x509.Certificate) -> bytes: + return cert.public_bytes(serialization.Encoding.PEM) + + +class _TlsPostgresStub: + """Answers one libpq ``SSLRequest`` with ``S`` and serves ``leaf + intermediate``.""" + + def __init__(self, chain_pem: Path, key_pem: Path) -> None: + self.context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + self.context.load_cert_chain(str(chain_pem), str(key_pem)) + self.listener: Final = socket.create_server(("127.0.0.1", 0)) + self.port: Final[int] = self.listener.getsockname()[1] + self.thread: Final = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + + def _serve(self) -> None: + with self.listener: + while True: + try: + conn: socket.socket = self.listener.accept()[0] + except OSError: + return + with conn: + try: + if conn.recv(8) == PG_SSL_REQUEST: + conn.sendall(b"S") + with self.context.wrap_socket(conn, server_side=True) as tls: + tls.recv(1) + except OSError: + continue + + +@dataclass(frozen=True, slots=True) +class _RdsLikePki: + bundle: Path + wrong_bundle: Path + root: Path + port: int + + +@pytest.fixture +def rds_like_pki(tmp_path: Path) -> Iterator[_RdsLikePki]: + """An RDS-shaped trust setup: the server sends leaf + intermediate, the + bundle holds only self-signed roots, and the right root is not first.""" + root, root_key = _issue_cert("Real Root CA", None, None, ca=True) + decoys: Final = tuple(_issue_cert(f"Decoy Root CA {i}", None, None, ca=True)[0] for i in range(3)) + intermediate, intermediate_key = _issue_cert("Intermediate CA", root, root_key, ca=True) + leaf, leaf_key = _issue_cert("localhost", intermediate, intermediate_key, ca=False) + chain_pem: Final = tmp_path / "server-chain.pem" + chain_pem.write_bytes(_pem(leaf) + _pem(intermediate)) + key_pem: Final = tmp_path / "server.key" + key_pem.write_bytes( + leaf_key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ) + ) + bundle: Final = tmp_path / "global-bundle.pem" + bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys) + _pem(root)) + wrong_bundle: Final = tmp_path / "wrong-bundle.pem" + wrong_bundle.write_bytes(b"".join(_pem(decoy) for decoy in decoys)) + root_pem: Final = tmp_path / "root.pem" + root_pem.write_bytes(_pem(root)) + stub: Final = _TlsPostgresStub(chain_pem, key_pem) + yield _RdsLikePki(bundle=bundle, wrong_bundle=wrong_bundle, root=root_pem, port=stub.port) + stub.listener.close() + + +def _params(url: str) -> tuple[tuple[str, str], ...]: + return tuple(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True)) + + +def test_multi_root_bundle_is_pinned_to_the_root_the_server_chains_to( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki +): + """Prisma's ``sslcert`` loads only the first certificate of the file, so + handing it the whole RDS bundle trusts one region's root and fails with + "unable to get local issuer certificate" everywhere else. The URL Prisma + receives must point at a single-certificate file holding the server's root.""" + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}", + ) + + _apply() + + (sslmode, sslcert, sslaccept, _) = _params(os.environ["DATABASE_URL"]) + assert (sslmode, sslaccept) == (("sslmode", "require"), ("sslaccept", "strict")) + assert sslcert[0] == "sslcert" and sslcert[1] != str(rds_like_pki.bundle) + assert Path(sslcert[1]).read_bytes() == rds_like_pki.root.read_bytes() + + +def test_pinned_root_replaces_a_planted_symlink_instead_of_writing_through_it( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki, tmp_path: Path +): + """The pinned file has a predictable name in a shared temp dir, so a symlink + planted there must not redirect the write onto its target.""" + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + root_der: Final = x509.load_pem_x509_certificate(rds_like_pki.root.read_bytes()).public_bytes( + serialization.Encoding.DER + ) + pinned: Final = tmp_path / f"litellm-sslcert-{hashlib.sha256(root_der).hexdigest()[:16]}.pem" + victim: Final = tmp_path / "victim.txt" + victim.write_text("untouched") + pinned.symlink_to(victim) + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db?sslmode=verify-full&sslrootcert={rds_like_pki.bundle}", + ) + + _apply() + + assert ("sslcert", str(pinned)) in _params(os.environ["DATABASE_URL"]) + assert victim.read_text() == "untouched" + assert not pinned.is_symlink() and pinned.read_bytes() == rds_like_pki.root.read_bytes() + + +def test_bundle_without_the_servers_root_is_passed_through_unchanged( + monkeypatch: pytest.MonkeyPatch, rds_like_pki: _RdsLikePki +): + """Nothing in the bundle verifies the server, so no root is pinned and + Prisma keeps rejecting the connection instead of trusting a root the + operator never shipped.""" + monkeypatch.setenv( + "DATABASE_URL", + f"postgresql://u:p@localhost:{rds_like_pki.port}/litellm_db" + f"?sslmode=verify-full&sslrootcert={rds_like_pki.wrong_bundle}", + ) + + _apply() + + assert ("sslcert", str(rds_like_pki.wrong_bundle)) in _params(os.environ["DATABASE_URL"]) + + +def test_root_cert_resolver_receives_the_urls_host_and_default_port(): + def resolver(cert_path: str, host: str, port: int) -> str: + return f"/pinned/{host}/{port}{cert_path}" + + url: Final = translate_libpq_ssl_params( + "postgresql://u:p@db.example.com/litellm_db?sslmode=verify-full&sslrootcert=/certs/bundle.pem", resolver + ) + + assert ("sslcert", "/pinned/db.example.com/5432/certs/bundle.pem") in _params(url) + + def test_libpq_verify_ca_becomes_prisma_strict(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=verify-ca") diff --git a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py index 93a11a914cb..045261e2d53 100644 --- a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py +++ b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py @@ -8,8 +8,11 @@ from datetime import datetime, timezone import pytest +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY from litellm.proxy.db.gateway_request_tracking import ( + GATEWAY_REQUESTS_JOB_NAME, GatewayRequestAccumulator, + GatewayRequestRedisBuffer, commit_gateway_requests_to_db, flush_gateway_requests, ) @@ -83,40 +86,47 @@ def test_drain_snapshot_is_not_mutated_by_later_records(): # ── commit ──────────────────────────────────────────────────────────────────── -class FakeTable: - def __init__(self) -> None: - self.upserts: list[dict] = [] - - def upsert(self, *, where: dict, data: dict) -> None: - self.upserts.append({"where": where, "data": data}) - - -class FakeBatcher: - def __init__(self, table: FakeTable) -> None: - self.litellm_dailygatewayrequests = table - - async def __aenter__(self) -> "FakeBatcher": - return self - - async def __aexit__(self, *args: object) -> bool: - return False - - class FakeDB: - def __init__(self, table: FakeTable) -> None: - self._table = table + def __init__(self) -> None: + self.statements: list[tuple[str, tuple[object, ...]]] = [] - def batch_(self) -> FakeBatcher: - return FakeBatcher(self._table) + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + return len(args) // 5 class FakePrismaClient: def __init__(self) -> None: - self.table = FakeTable() - self.db = FakeDB(self.table) + self.db = FakeDB() -def test_commit_upserts_one_incrementing_row_per_key(): +def _rows_written(client: FakePrismaClient) -> list[tuple[object, ...]]: + """Every (date, category, route, successful, failed) tuple the database received, in statement order.""" + return [params[i : i + 5] for _, params in client.db.statements for i in range(0, len(params), 5)] + + +def test_commit_increments_with_a_single_statement_for_the_whole_snapshot(): + """One statement per flush is the whole point: the previous per-key upsert cost + the primary (workers x routes) statements per interval.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route=route): ( + GatewayRequestCounts(successful_requests=7, failed_requests=2) + ) + for route in ("/chat/completions", "/embeddings", "/responses", "/v1/messages", "/mcp") + } + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + assert len(client.db.statements) == 1 + sql, params = client.db.statements[0] + assert sql.count("ON CONFLICT") == 1 + assert sql.count("(NOW() AT TIME ZONE 'UTC'))") == 5 + assert len(params) == 25 + + +def test_commit_sql_adds_to_the_existing_row_instead_of_replacing_it(): + """A worker only knows its own share; the SQL must add EXCLUDED onto the stored count.""" client = FakePrismaClient() snapshot = { GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( @@ -126,20 +136,36 @@ def test_commit_upserts_one_incrementing_row_per_key(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - assert len(client.table.upserts) == 1 - written = client.table.upserts[0] - assert written["where"] == { - "date_category_route": { - "date": "2026-08-01", - "category": "llm", - "route": "/chat/completions", - } + sql, params = client.db.statements[0] + assert 'INSERT INTO "LiteLLM_DailyGatewayRequests"' in sql + assert 'ON CONFLICT ("date", "category", "route") DO UPDATE SET' in sql + assert ( + '"successful_requests" = "LiteLLM_DailyGatewayRequests"."successful_requests" + EXCLUDED."successful_requests"' + in sql + ) + assert '"failed_requests" = "LiteLLM_DailyGatewayRequests"."failed_requests" + EXCLUDED."failed_requests"' in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 7, 2) + + +def test_commit_placeholders_line_up_with_params(): + """$n positions are generated per row; a drift here silently swaps a route for a count.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=0) + ), + GatewayRequestKey(date="2026-08-01", category="mcp", route="/mcp"): ( + GatewayRequestCounts(successful_requests=0, failed_requests=3) + ), } - assert written["data"]["update"] == { - "successful_requests": {"increment": 7}, - "failed_requests": {"increment": 2}, - } - assert written["data"]["create"]["successful_requests"] == 7 + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + sql, params = client.db.statements[0] + assert "($1::text, $2::text, $3::text, $4::bigint, $5::bigint," in sql + assert "($6::text, $7::text, $8::text, $9::bigint, $10::bigint," in sql + assert "$11" not in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 1, 0, "2026-08-01", "mcp", "/mcp", 0, 3) def test_commit_is_deterministically_ordered(): @@ -154,17 +180,14 @@ def test_commit_is_deterministically_ordered(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - written_order = [ - (row["where"]["date_category_route"]["date"], row["where"]["date_category_route"]["category"]) - for row in client.table.upserts - ] + written_order = [(row[0], row[1]) for row in _rows_written(client)] assert written_order == [("2026-08-01", "llm"), ("2026-08-01", "mcp"), ("2026-08-02", "llm")] def test_commit_skips_the_database_entirely_when_nothing_accumulated(): client = FakePrismaClient() asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot={})) - assert client.table.upserts == [] + assert client.db.statements == [] # ── flush ───────────────────────────────────────────────────────────────────── @@ -177,12 +200,12 @@ def test_flush_drains_and_commits(): asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 + assert len(client.db.statements) == 1 assert acc.drain() == {} class ExplodingDB: - def batch_(self): + async def execute_raw(self, query: str, *args: object) -> int: raise RuntimeError("db gone") @@ -208,10 +231,7 @@ def test_failed_flush_keeps_counts_for_the_next_attempt(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert client.table.upserts[0]["data"]["update"] == { - "successful_requests": {"increment": 1}, - "failed_requests": {"increment": 1}, - } + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] def test_restored_counts_merge_with_requests_recorded_meanwhile(): @@ -223,5 +243,272 @@ def test_restored_counts_merge_with_requests_recorded_meanwhile(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 - assert client.table.upserts[0]["data"]["update"]["successful_requests"] == {"increment": 2} + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingDBWithInFlightRequest: + """Fails the write after a request has been recorded while it was in flight.""" + + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.accumulator = accumulator + + async def execute_raw(self, query: str, *args: object) -> int: + _record(self.accumulator, 500) + raise RuntimeError("db gone") + + +class ExplodingClientWithInFlightRequest: + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.db = ExplodingDBWithInFlightRequest(accumulator) + + +def test_restore_keeps_requests_recorded_while_the_failed_write_was_in_flight(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(ExplodingClientWithInFlightRequest(acc), acc)) + + client = FakePrismaClient() + asyncio.run(flush_gateway_requests(client, acc)) + + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] + + +# ── redis buffer ────────────────────────────────────────────────────────────── + + +class FakeRedis: + def __init__(self) -> None: + self.lists: dict[str, list[str]] = {} + + async def async_rpush(self, key: str, values: list[str]) -> int: + self.lists.setdefault(key, []).extend(values) + return len(self.lists[key]) + + async def async_lpop(self, key: str, count: int) -> list[str] | None: + queue = self.lists.get(key, []) + if not queue: + return None + popped, self.lists[key] = queue[:count], queue[count:] + return popped + + +class FakePodLock: + def __init__(self, *, leader: bool) -> None: + self.leader = leader + self.held: list[str] = [] + self.released: list[str] = [] + + async def acquire_lock(self, cronjob_id: str) -> bool: + self.held.append(cronjob_id) + return self.leader + + async def release_lock(self, cronjob_id: str) -> None: + self.released.append(cronjob_id) + + +class FakeLease: + """Redis-side view of the job lock: SET NX by pod id, re-entrant for the holder, freed only by release or TTL.""" + + def __init__(self) -> None: + self.holder: str | None = None + + +class FakeLeasePodLock: + def __init__(self, lease: FakeLease, pod_id: str) -> None: + self.lease = lease + self.pod_id = pod_id + + async def acquire_lock(self, cronjob_id: str) -> bool: + if self.lease.holder is None: + self.lease.holder = self.pod_id + return self.lease.holder == self.pod_id + + async def release_lock(self, cronjob_id: str) -> None: + if self.lease.holder == self.pod_id: + self.lease.holder = None + + +def _buffer(redis: FakeRedis, *, leader: bool) -> tuple[GatewayRequestRedisBuffer, FakePodLock]: + lock = FakePodLock(leader=leader) + return GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=lock), lock # pyright: ignore[reportArgumentType] # duck-typed fakes + + +def test_non_leader_workers_push_to_redis_and_never_touch_the_database(): + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(3): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + assert client.db.statements == [] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_folds_every_workers_snapshot_into_one_statement(): + """Fifty workers each flushing the same routes must cost the primary one statement, not fifty.""" + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(50): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500, route="/responses") + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader_acc = GatewayRequestAccumulator() + _record(leader_acc, 200) + leader, lock = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [ + (_today(), "llm", "/chat/completions", 51, 0), + (_today(), "llm", "/responses", 0, 50), + ] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + assert lock.held == [GATEWAY_REQUESTS_JOB_NAME] + assert lock.released == [] + + +def test_leader_keeps_the_lease_so_staggered_pods_cost_one_statement_per_interval(): + """Pods flush on their own clocks; without the lease each one would win the lock in turn and commit alone.""" + redis = FakeRedis() + client = FakePrismaClient() + lease = FakeLease() + pods = tuple( + GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=FakeLeasePodLock(lease, f"pod-{i}")) # pyright: ignore[reportArgumentType] # duck-typed fakes + for i in range(4) + ) + + for _interval in range(3): + for pod in pods: + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(client, acc, pod)) + + assert lease.holder == "pod-0" + assert len(client.db.statements) == 3 + assert [row[3] for row in _rows_written(client)] == [1, 4, 4] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_drains_a_backlog_deeper_than_one_capped_pop(): + """More workers than MAX_REDIS_BUFFER_DEQUEUE_COUNT must not leave a growing tail queued behind the cap.""" + redis = FakeRedis() + client = FakePrismaClient() + workers = MAX_REDIS_BUFFER_DEQUEUE_COUNT * 2 + 1 + for _ in range(workers): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", workers, 0)] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + +def test_leader_with_nothing_buffered_writes_nothing(): + redis = FakeRedis() + client = FakePrismaClient() + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert client.db.statements == [] + assert lock.released == [] + + +def test_leader_requeues_to_redis_when_the_database_commit_fails(): + """Counts popped from Redis are gone from every worker; a failed commit must put them back.""" + redis = FakeRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 200) + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(ExplodingClient(), acc, leader)) + + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + assert lock.released == [] + assert acc.drain() == {} + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingRedis(FakeRedis): + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis gone") + + +class UnreadableRedis(FakeRedis): + async def async_lpop(self, key: str, count: int) -> list[str] | None: + raise RuntimeError("redis gone mid-flush") + + +class UnwritableRedis(FakeRedis): + """Pops succeed, pushes fail: a Redis that went read-only between the leader's pop and its re-queue.""" + + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis read-only") + + +def test_leader_keeps_popped_counts_in_memory_when_both_the_database_and_the_requeue_fail(): + """The pop removed the only copy; if Redis will not take it back the leader itself must carry it.""" + redis = FakeRedis() + worker_acc = GatewayRequestAccumulator() + _record(worker_acc, 200) + _record(worker_acc, 200) + worker, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(FakePrismaClient(), worker_acc, worker)) + + degraded = UnwritableRedis() + degraded.lists = redis.lists + leader_acc = GatewayRequestAccumulator() + leader, _ = _buffer(degraded, leader=True) + asyncio.run(flush_gateway_requests(ExplodingClient(), leader_acc, leader)) + assert degraded.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +def test_leader_whose_redis_read_fails_leaves_the_pushed_rows_for_the_next_flush(): + """The scheduler job must not raise, and nothing is popped so nothing needs restoring anywhere.""" + redis = UnreadableRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + client = FakePrismaClient() + leader, _ = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, acc, leader)) + + assert client.db.statements == [] + assert acc.drain() == {} + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + + +def test_failed_redis_push_keeps_counts_locally_for_the_next_flush(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500) + buffer, lock = _buffer(ExplodingRedis(), leader=True) + + asyncio.run(flush_gateway_requests(FakePrismaClient(), acc, buffer)) + + assert lock.held == [] + assert acc.drain() == { + GatewayRequestKey(date=_today(), category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=1) + ) + } diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 8cb3fc665eb..0361d5cfe8f 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -7,6 +7,7 @@ allowed to run: only when the row is missing or belongs to an older window. from __future__ import annotations +import asyncio from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -14,6 +15,7 @@ from typing import Final import pytest from litellm.caching.dual_cache import DualCache +from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) @@ -42,6 +44,19 @@ class _FakeSpendLogsTable: return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}] +class _InFlightCountingTable: + def __init__(self) -> None: + self.in_flight = 0 + self.max_in_flight = 0 + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace: + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + await asyncio.sleep(0.001) + self.in_flight -= 1 + return SimpleNamespace(token=where["token"], spend=1.0) + + class _FakePrismaClient: def __init__( self, @@ -55,6 +70,7 @@ class _FakePrismaClient: litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error), litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error), + litellm_verificationtoken=_InFlightCountingTable(), ) @@ -306,6 +322,21 @@ async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_err ) +@pytest.mark.asyncio +async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): + """Per-counter singleflight only collapses duplicates of one key. A cold-cache burst + over many distinct keys must still not flood the prisma engine HTTP pool (LIT-6435).""" + prisma: Final = _FakePrismaClient() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 5 + + results: Final = await asyncio.gather( + *(SpendCounterReseed.from_db(prisma_client=prisma, counter_key=f"spend:key:hashed-{i}") for i in range(burst)) + ) + + assert results == [1.0] * burst + assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth 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 7da55f22bda..c0762edec92 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 @@ -3,6 +3,8 @@ Unit tests for Bedrock Guardrails """ import json +import asyncio +from datetime import datetime, timezone import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -28,6 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockTextContent, ) from litellm.types.utils import CallTypes, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.mark.asyncio @@ -5842,3 +5845,36 @@ async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): assert response["action"] == "NONE" assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the ApplyGuardrail request is signed with SigV4, and botocore + refreshes expiring credentials inside that signing with a blocking HTTP call, so it must run + on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"action": "NONE", "outputs": [], "assessments": []}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=allowed)): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await guardrail._post_apply_guardrail_content( + content=[{"text": {"text": "hello"}}], + base_request_data={"source": "INPUT"}, + credentials=probe.credentials(), + aws_region_name="us-east-1", + api_key=None, + request_data={}, + event_type=GuardrailEventHooks.pre_call, + start_time=datetime.now(timezone.utc), + completed_chunk_usages=[], + ) + await release + + assert response["action"] == "NONE" + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index f4af77d2e40..bbc8fd539a3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -5,10 +5,12 @@ All Bedrock HTTP calls are mocked; no real AWS calls are made. """ import json +import asyncio import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest +import httpx from fastapi import HTTPException @@ -21,6 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailResponse, ) from litellm.types.utils import Choices, Message, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe CONTENT_FILTER_CHECKS = {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}} @@ -861,3 +864,33 @@ async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeyp {"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8} ] assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_invoke_guardrail_checks_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the checks request is signed with SigV4, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + g = BedrockGuardrail(checks=CONTENT_FILTER_CHECKS, content_filter_threshold=0.5) + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.1}]}}}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with ( + patch.object(g, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + patch.object(g.async_handler, "post", new=AsyncMock(return_value=allowed)), + ): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await g.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"messages": []}, + ) + await release + + assert response == BedrockGuardrailResponse() + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index db2f94306fb..0d723d6671f 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -56,13 +56,13 @@ class TestPolicyFromLitellmParams: class _FakeRouter: - """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + """Minimal stand-in for litellm.Router.deployments_for_request, for policy_for_model.""" def __init__(self, deployments: list[dict[str, Any]]): self._deployments = deployments - def get_model_list(self, model_name, team_id=None): - return [d for d in self._deployments if d.get("model_name") == model_name] + def deployments_for_request(self, model, request_kwargs): + return [d for d in self._deployments if d.get("model_name") == model] def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: @@ -78,23 +78,23 @@ def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[ class TestPolicyForModel: def test_no_router_returns_none(self): - assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=None, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_no_marker_deployment_returns_none(self): router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_without_policy_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_with_policy_is_found(self): router = _FakeRouter( [_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) def test_picks_the_marker_whose_tags_the_request_carries(self): @@ -107,8 +107,8 @@ class TestPolicyForModel: ] ) - eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) - us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + eu = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) @@ -116,7 +116,7 @@ class TestPolicyForModel: def test_untagged_marker_matches_any_request(self): router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) policy = policy_for_model( - llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("anything",) ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) @@ -128,14 +128,14 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-default"}), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) def test_no_untagged_fallback_means_no_policy(self): """No matching marker means no policy, not an unrelated slice's compression.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) assert ( - policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None + policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) is None ) def test_tag_scoped_marker_takes_precedence_over_untagged(self): @@ -147,7 +147,7 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e9e58347337..b06d87ac67e 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1301,6 +1301,33 @@ def test_health_readiness_details_returns_diagnostic_fields(monkeypatch): assert "cache" in response_data +@pytest.mark.parametrize( + "general_settings, expected_warning", + [ + ({}, True), + ({"disable_env_credential_login": True}, False), + ], +) +def test_health_readiness_details_reports_env_credential_login_warning(monkeypatch, general_settings, expected_warning): + """ + The Admin UI banner is driven by this flag: it must be True while + env-credential login is possible and False once + `disable_env_credential_login` turns that login path off. + """ + app = FastAPI() + app.include_router(_health_endpoints_module.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + client = TestClient(app) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + response = client.get("/health/readiness/details") + + assert response.status_code == 200, response.text + assert response.json()["show_env_credential_login_warning"] is expected_warning + + def test_health_readiness_allows_explicit_legacy_public_details(monkeypatch): """ Operators can explicitly preserve the legacy public readiness payload. diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 0ff8b67b1a7..0cd6b4ede9c 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1861,3 +1861,60 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch): ) assert capacity_blocked.value.status_code == 429 assert "Model capacity reached" in capacity_blocked.value.detail["error"] + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_priority_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + RateLimitResponse, + RateLimitStatus, + get_or_create_request_stash, + ) + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=75, + limit_remaining=74, + rate_limit_type="requests", + descriptor_key="priority_model", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75 + assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74 + assert additional_headers["x-litellm-priority"] == "premium" + assert additional_headers["x-litellm-rate-limiter-version"] == "v3" + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 4003286d887..6d382370f5f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6171,3 +6171,68 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses(): data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5)) ) assert get_request_stash().batch_enqueued_reservation == reservation + + +@pytest.mark.asyncio +async def test_post_call_success_hook_attaches_ratelimit_headers_to_dict_response(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}}, + } + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-dict-response")), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + assert additional_headers["x-litellm-attempted-retries"] == 0 + assert additional_headers["x-ratelimit-model_saturation_check-limit-requests"] == 100 + assert additional_headers["x-ratelimit-model_saturation_check-remaining-requests"] == 99 + + +@pytest.mark.asyncio +async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + get_or_create_request_stash().rate_limit_response = RateLimitResponse( + overall_code="OK", + statuses=[ + RateLimitStatus( + code="OK", + current_limit=100, + limit_remaining=99, + rate_limit_type="requests", + descriptor_key="model_saturation_check", + ) + ], + ) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-raw-dict")), + response=response, + ) + + assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 2839acab6b0..bdaca9ffc2d 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3670,5 +3670,42 @@ async def test_post_call_success_hook_contains_header_merge_failures( ) +@pytest.mark.asyncio +async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loop(rate_limiter): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + handler, _cache = rate_limiter + stash = get_or_create_request_stash() + warm_tokenizer("claude-fable-5") + data: dict[str, object] = { + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + itpm_descriptor = { + "key": PROJECT_ITPM_DESCRIPTOR_KEY, + "value": "proj-loop:claude-fable-5", + "rate_limit": {"tokens_per_unit": 10_000_000, "window_size": 60}, + } + + _, took, lags = await timed_with_loop_lags( + lambda: handler._reserve_project_io_tokens_or_raise( + descriptors=[itpm_descriptor], + data=data, + requested_model="claude-fable-5", + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-itpm-loop"), project_id="proj-loop"), + tpm_reservation_scopes=[], + tpm_reservation_amount=0, + ) + ) + + assert stash.rate_limit_response is not None + assert_loop_stayed_free(took, lags) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index 93dc429168f..bb71d67f24e 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -265,7 +265,7 @@ class TestSuggesterRejectsModelsWithoutToolCalling: def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map): supported_params = litellm.get_supported_openai_params( - model="amazon.nova-pro-v1:0", + model="meta.llama4-scout-17b-instruct-v1:0", custom_llm_provider="bedrock", ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index ec62cc47018..7ece35ceedf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -4,13 +4,17 @@ Tests for cost tracking settings management endpoints. Tests the GET and PATCH endpoints for managing cost discount configuration. """ +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from pydantic import ValidationError import litellm +from litellm._internal_context import pinned_billing_time +from litellm.proxy._types import CostEstimateRequest from litellm.proxy.management_endpoints.cost_tracking_settings import router from litellm.proxy.proxy_server import app @@ -789,13 +793,13 @@ INPUT_TOKENS = 1000 OUTPUT_TOKENS = 500 -def _router_pricing(**pricing: float) -> MagicMock: +def _router_pricing(model: str = AN_UNDERLYING_MODEL, **pricing: float) -> MagicMock: mock_router = MagicMock() mock_router.get_model_list.return_value = [ { "model_name": AN_ALIAS, "litellm_params": { - "model": AN_UNDERLYING_MODEL, + "model": model, "custom_llm_provider": "openai", **pricing, }, @@ -811,9 +815,7 @@ async def _estimate(mock_router: MagicMock | None, model: str = AN_ALIAS, **over request = CostEstimateRequest( model=model, - input_tokens=INPUT_TOKENS, - output_tokens=OUTPUT_TOKENS, - **overrides, + **{"input_tokens": INPUT_TOKENS, "output_tokens": OUTPUT_TOKENS, **overrides}, ) with patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point "litellm.proxy.proxy_server.llm_router", mock_router @@ -909,3 +911,299 @@ class TestEstimateCostPeriodTotals: assert response.cost_per_request == pytest.approx(0.0022) assert response.daily_margin_cost == pytest.approx(0.02) assert response.daily_cost == pytest.approx(0.22) + + +CACHE_READ_TOKENS = 800 +CACHE_CREATION_TOKENS = 100 +REASONING_TOKENS = 200 +TEXT_INPUT_TOKENS = INPUT_TOKENS - CACHE_READ_TOKENS - CACHE_CREATION_TOKENS +TEXT_OUTPUT_TOKENS = OUTPUT_TOKENS - REASONING_TOKENS + + +async def _estimate_with_cache_and_reasoning(mock_router: MagicMock | None, model: str = AN_ALIAS, **overrides: int): + return await _estimate( + mock_router, + model=model, + cache_read_input_tokens=CACHE_READ_TOKENS, + cache_creation_input_tokens=CACHE_CREATION_TOKENS, + reasoning_tokens=REASONING_TOKENS, + **overrides, + ) + + +class TestEstimateCostCacheAndReasoningTokens: + @pytest.mark.asyncio + async def test_a_mapped_model_bills_cache_and_reasoning_tokens_at_their_own_rates(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "output_cost_per_reasoning_token": 1e-5, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate_with_cache_and_reasoning(None, model=A_MAPPED_MODEL, num_requests_per_day=10) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 3e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 3.75e-6) + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 1e-5) + assert response.input_cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 3e-6 + CACHE_READ_TOKENS * 3e-7 + CACHE_CREATION_TOKENS * 3.75e-6 + ) + assert response.output_cost_per_request == pytest.approx(TEXT_OUTPUT_TOKENS * 15e-6 + REASONING_TOKENS * 1e-5) + assert response.cost_per_request == pytest.approx( + response.input_cost_per_request + response.output_cost_per_request + ) + assert response.daily_cache_read_cost == pytest.approx(10 * CACHE_READ_TOKENS * 3e-7) + assert response.daily_cache_creation_cost == pytest.approx(10 * CACHE_CREATION_TOKENS * 3.75e-6) + assert response.daily_reasoning_cost == pytest.approx(10 * REASONING_TOKENS * 1e-5) + assert response.monthly_cache_read_cost is None + assert response.cache_read_input_token_cost == pytest.approx(3e-7) + assert response.cache_creation_input_token_cost == pytest.approx(3.75e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(1e-5) + assert ( + response.cache_read_input_tokens, + response.cache_creation_input_tokens, + response.reasoning_tokens, + ) == (CACHE_READ_TOKENS, CACHE_CREATION_TOKENS, REASONING_TOKENS) + + @pytest.mark.asyncio + async def test_a_model_without_cache_or_reasoning_prices_estimates_what_the_proxy_bills(self, monkeypatch): + """The cost calculator bills cache tokens of a cost-map model without cache prices at zero + and its reasoning tokens at the output rate. The estimate reports those effective rates.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + {"input_cost_per_token": 5e-6, "output_cost_per_token": 6e-6, "litellm_provider": "openai", "mode": "chat"}, + ) + + response = await _estimate_with_cache_and_reasoning(None, model=A_MAPPED_MODEL) + + assert response.cache_read_cost_per_request == 0.0 + assert response.cache_creation_cost_per_request == 0.0 + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 6e-6) + assert response.input_cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6) + assert response.cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6 + OUTPUT_TOKENS * 6e-6) + assert response.cache_read_input_token_cost == 0.0 + assert response.cache_creation_input_token_cost == 0.0 + assert response.output_cost_per_reasoning_token == pytest.approx(6e-6) + + @pytest.mark.asyncio + async def test_a_request_without_cache_or_reasoning_tokens_estimates_as_before(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "output_cost_per_reasoning_token": 1e-5, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate(None, model=A_MAPPED_MODEL, num_requests_per_day=10) + + assert response.cost_per_request == pytest.approx(INPUT_TOKENS * 3e-6 + OUTPUT_TOKENS * 15e-6) + assert response.cache_read_cost_per_request == 0.0 + assert response.cache_creation_cost_per_request == 0.0 + assert response.reasoning_cost_per_request == 0.0 + assert response.daily_cache_read_cost == 0.0 + assert response.daily_reasoning_cost == 0.0 + + @pytest.mark.asyncio + async def test_a_custom_priced_deployment_bills_cache_and_reasoning_tokens_from_its_flat_rates(self): + response = await _estimate_with_cache_and_reasoning( + _router_pricing(input_cost_per_token=1e-6, output_cost_per_token=2e-6, cache_read_input_token_cost=1e-7) + ) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 1e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 1e-6) + assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 2e-6) + assert response.cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 1e-6 + CACHE_READ_TOKENS * 1e-7 + CACHE_CREATION_TOKENS * 1e-6 + OUTPUT_TOKENS * 2e-6 + ) + assert response.cache_read_input_token_cost == pytest.approx(1e-7) + assert response.cache_creation_input_token_cost == pytest.approx(1e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(2e-6) + + @pytest.mark.asyncio + async def test_a_custom_priced_deployment_of_a_mapped_model_inherits_its_built_in_cache_rates(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 5e-6, + "output_cost_per_token": 6e-6, + "cache_read_input_token_cost": 5e-7, + "cache_creation_input_token_cost": 6.25e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate_with_cache_and_reasoning( + _router_pricing(model=A_MAPPED_MODEL, input_cost_per_token=1e-6, output_cost_per_token=2e-6) + ) + + assert response.cache_read_cost_per_request == pytest.approx(CACHE_READ_TOKENS * 5e-7) + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 6.25e-6) + assert response.input_cost_per_request == pytest.approx( + TEXT_INPUT_TOKENS * 1e-6 + CACHE_READ_TOKENS * 5e-7 + CACHE_CREATION_TOKENS * 6.25e-6 + ) + assert response.cache_read_input_token_cost == pytest.approx(5e-7) + assert response.cache_creation_input_token_cost == pytest.approx(6.25e-6) + + @pytest.mark.asyncio + async def test_a_tiered_model_reports_the_rates_its_lines_were_billed_at(self, monkeypatch): + """Above a token tier the calculator bills every line at the tier's rate, so the reported + rates must be the tier's too: each line equals its token count times the rate next to it.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "cache_creation_input_token_cost": 3.75e-6, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-6, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + response = await _estimate( + None, + model=A_MAPPED_MODEL, + input_tokens=250_000, + cache_read_input_tokens=200_000, + cache_creation_input_tokens=10_000, + output_tokens=1_000, + reasoning_tokens=200, + ) + + assert response.input_cost_per_token == pytest.approx(6e-6) + assert response.output_cost_per_token == pytest.approx(3e-5) + assert response.cache_read_input_token_cost == pytest.approx(6e-7) + assert response.cache_creation_input_token_cost == pytest.approx(7.5e-6) + assert response.output_cost_per_reasoning_token == pytest.approx(3e-5) + assert response.cache_read_cost_per_request == pytest.approx(200_000 * response.cache_read_input_token_cost) + assert response.cache_creation_cost_per_request == pytest.approx( + 10_000 * response.cache_creation_input_token_cost + ) + assert response.reasoning_cost_per_request == pytest.approx(200 * response.output_cost_per_reasoning_token) + assert response.input_cost_per_request == pytest.approx( + 40_000 * response.input_cost_per_token + + response.cache_read_cost_per_request + + response.cache_creation_cost_per_request + ) + assert response.output_cost_per_request == pytest.approx(1_000 * response.output_cost_per_token) + + @pytest.mark.asyncio + async def test_a_quote_prices_its_totals_and_its_rates_at_the_same_moment(self, monkeypatch): + """The totals and the reported rates resolve off-peak pricing on separate paths. A quote + taken as a window opens must not bill on one side of it and report rates from the other.""" + monkeypatch.setitem( + litellm.model_cost, + A_MAPPED_MODEL, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "off_peak_pricing": { + "hours_utc": "02:00-03:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + }, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + + with pinned_billing_time(datetime(2026, 1, 1, 2, 30, tzinfo=timezone.utc)): + response = await _estimate(None, model=A_MAPPED_MODEL) + + assert response.input_cost_per_token == pytest.approx(1e-6) + assert response.output_cost_per_token == pytest.approx(5e-6) + assert response.input_cost_per_request == pytest.approx(INPUT_TOKENS * response.input_cost_per_token) + assert response.output_cost_per_request == pytest.approx(OUTPUT_TOKENS * response.output_cost_per_token) + + + @pytest.mark.asyncio + async def test_an_unrouted_model_reports_the_rates_of_the_provider_the_calculator_inferred(self, monkeypatch): + """The cost calculator infers a provider this endpoint never resolved, and the provider decides + whether a tier threshold is inclusive. xai bills a request sitting exactly on the 200k threshold + at the tier rate, so the reported rates have to be the tier's rather than the sub-tier base.""" + an_xai_model = "xai/tiered-model" + monkeypatch.setitem( + litellm.model_cost, + an_xai_model, + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + + response = await _estimate( + None, + model=an_xai_model, + input_tokens=200_000, + cache_read_input_tokens=100_000, + output_tokens=1_000, + ) + + assert response.input_cost_per_token == pytest.approx(6e-6) + assert response.output_cost_per_token == pytest.approx(3e-5) + assert response.cache_read_input_token_cost == pytest.approx(6e-7) + assert response.cache_read_cost_per_request == pytest.approx(100_000 * response.cache_read_input_token_cost) + assert response.input_cost_per_request == pytest.approx( + 100_000 * response.input_cost_per_token + response.cache_read_cost_per_request + ) + assert response.output_cost_per_request == pytest.approx(1_000 * response.output_cost_per_token) + + +class TestCostEstimateRequestTokenSubsets: + def test_cache_tokens_beyond_the_input_tokens_are_rejected(self): + with pytest.raises(ValidationError, match="cannot exceed input_tokens"): + CostEstimateRequest( + model=AN_ALIAS, + input_tokens=INPUT_TOKENS, + output_tokens=OUTPUT_TOKENS, + cache_read_input_tokens=INPUT_TOKENS, + cache_creation_input_tokens=1, + ) + + def test_reasoning_tokens_beyond_the_output_tokens_are_rejected(self): + with pytest.raises(ValidationError, match="cannot exceed output_tokens"): + CostEstimateRequest( + model=AN_ALIAS, + input_tokens=INPUT_TOKENS, + output_tokens=OUTPUT_TOKENS, + reasoning_tokens=OUTPUT_TOKENS + 1, + ) + + def test_the_endpoint_answers_422_when_cache_tokens_exceed_input_tokens(self): + response = client.post( + "/cost/estimate", + headers={"Authorization": "Bearer sk-1234"}, + json={"model": AN_ALIAS, "input_tokens": 1000, "output_tokens": 100, "cache_read_input_tokens": 8000}, + ) + + assert response.status_code == 422 + assert "cannot exceed input_tokens" in response.text diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a873a367eab..65cc23ea67f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5085,6 +5085,104 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch): assert len(deleted_keys) == 2 +class _JWTMappingRow: + def __init__(self, token, jwt_claim_name, jwt_claim_value): + self.token = token + self.jwt_claim_name = jwt_claim_name + self.jwt_claim_value = jwt_claim_value + + +class _CascadingJWTMappingTable: + """Mapping rows that LiteLLM_JWTKeyMapping_token_fkey drops when their key is deleted.""" + + def __init__(self, rows): + self.rows = rows + + async def find_many(self, where, **kwargs): + return [row for row in self.rows if row.token == where["token"]] + + def cascade(self, deleted_tokens): + self.rows = [row for row in self.rows if row.token not in deleted_tokens] + + +class _RecordingEvict: + def __init__(self): + self.cache_keys = () + + async def __call__(self, cache_keys, user_api_key_cache): + self.cache_keys = tuple(cache_keys) + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypatch): + """Deleting a key must evict its jwt_key_mapping cache entries (LIT-5380). + + The FK cascade removes the mapping rows, so a surviving cache entry would keep + resolving the deleted token hash and 401 every JWT call from that identity until + virtual_key_mapping_cache_ttl expires, instead of auto-registering again. + """ + jwt_table = _CascadingJWTMappingTable( + [_JWTMappingRow("hashed-token-1", "email", "user@example.com")] + ) + + key1 = LiteLLM_VerificationToken( + token="hashed-token-1", + user_id="user-123", + team_id=None, + key_alias="jwt-mapped-key", + spend=0.0, + max_budget=None, + models=[], + aliases={}, + config={}, + permissions={}, + metadata={}, + model_max_budget={}, + model_spend={}, + soft_budget_cooldown=False, + allowed_routes=[], + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[key1] + ) + mock_prisma_client.db.litellm_jwtkeymapping = jwt_table + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + + async def cascading_delete_data(tokens): + jwt_table.cascade(tokens) + return list(tokens) + + mock_prisma_client.delete_data = AsyncMock(side_effect=cascading_delete_data) + + recording_evict = _RecordingEvict() + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.evict_and_broadcast", + recording_evict, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", + lambda token: token, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ) + + await delete_verification_tokens( + tokens=["hashed-token-1"], + user_api_key_cache=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + ) + + assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",) + + @pytest.mark.asyncio async def test_delete_key_fn_persists_deleted_keys(monkeypatch): from litellm.proxy._types import KeyRequest @@ -17975,6 +18073,32 @@ def test_key_request_blank_organization_id_is_unset(): assert UpdateKeyRequest(key="sk-1", organization_id="org-1").organization_id == "org-1" +def test_update_key_request_blank_team_id_is_not_a_team_change(): + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + is_different_team, + ) + + blank = UpdateKeyRequest(key="sk-1", team_id="", key_alias="renamed") + assert blank.team_id is None + assert "team_id" not in blank.model_dump(exclude_unset=True) + assert blank.model_dump(exclude_unset=True) == {"key": "sk-1", "key_alias": "renamed"} + assert is_different_team(data=blank, existing_key_row=LiteLLM_VerificationToken(token="hashed")) is False + assert ( + is_different_team(data=blank, existing_key_row=LiteLLM_VerificationToken(token="hashed", team_id="team-1")) + is False + ) + assert "team_id" in UpdateKeyRequest(key="sk-1", team_id=None).model_dump(exclude_unset=True) + assert UpdateKeyRequest(key="sk-1", team_id="team-1").team_id == "team-1" + assert ( + is_different_team( + data=UpdateKeyRequest(key="sk-1", team_id="team-1"), + existing_key_row=LiteLLM_VerificationToken(token="hashed"), + ) + is True + ) + + def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch): """key_generation_check with team_id="" must take the personal-key path instead of failing the team lookup with "Unable to find team object" (LIT-3925).""" 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 051e6bed4fd..2f6561046b1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2205,6 +2205,50 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): assert update_call_kwargs.get("include", {}).get("object_permission") is True +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint_name", ["team_model_add", "team_model_delete"]) +async def test_team_model_add_delete_keep_model_aliases_in_team_cache(endpoint_name, monkeypatch): + """LIT-5858: Prisma only returns `litellm_model_table` when the `update` asks for it, so the refreshed + cache entry lost the team's model aliases and JWT alias requests 403'd until the next DB read.""" + from litellm.proxy._types import TeamModelAddRequest, TeamModelDeleteRequest + from litellm.proxy.auth.team_grants import team_model_aliases + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_model_add, team_model_delete + + columns = {"team_id": "team-1234", "models": ["gpt-4o", "openai/*"]} + alias_table = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"} + + async def update(where, data, include=None): + row = {**columns, "litellm_model_table": alias_table} if (include or {}).get("litellm_model_table") else columns + return SimpleNamespace(team_id="team-1234", model_dump=lambda: row) + + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=SimpleNamespace(model_dump=lambda: columns)) + prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=update) + prisma_client.db.execute_raw = AsyncMock(return_value=None) + cache = UserApiKeyCache() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + if endpoint_name == "team_model_add": + await team_model_add( + data=TeamModelAddRequest(team_id="team-1234", models=["team-byok-1"]), + http_request=MagicMock(), + user_api_key_dict=admin, + ) + else: + await team_model_delete( + data=TeamModelDeleteRequest(team_id="team-1234", models=["openai/*"]), + http_request=MagicMock(), + user_api_key_dict=admin, + ) + + cached_team = await cache.async_get_cache(key="team_id:team-1234", model_type=LiteLLM_TeamTableCachedObj) + assert team_model_aliases(cached_team) == {"fast": "gpt-4o"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "endpoint_name", diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 6f9142c85df..a3d3ae32169 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -9,8 +9,10 @@ import pytest import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, + count_relayed_prompt_tokens, ) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, @@ -2037,3 +2039,56 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: if __name__ == "__main__": pytest.main([__file__]) + + +ONE_PIXEL_PNG_DATA_URL = ( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" +) +UNREACHABLE_IMAGE_URL = "http://127.0.0.1:9/doc.png" +TEXT_ONLY_MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "Describe this"}]}] + + +def _image_messages(url: str, detail: str) -> list[dict]: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": url, "detail": detail}}, + ], + } + ] + + +def test_count_relayed_prompt_tokens_counts_a_data_url_image_exactly(): + messages = _image_messages(ONE_PIXEL_PNG_DATA_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + + +def test_count_relayed_prompt_tokens_keeps_a_low_detail_remote_image_at_the_base_count(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "low") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == litellm.token_counter( + model="gpt-4.1-mini", messages=messages + ) + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) < high_detail_image_token_upper_bound() + + +def test_count_relayed_prompt_tokens_charges_only_the_remote_high_detail_image_at_the_upper_bound(): + messages = _image_messages(UNREACHABLE_IMAGE_URL, "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() + ) + + +@pytest.mark.parametrize("scheme", ["HTTPS://", "Http://"]) +def test_count_relayed_prompt_tokens_charges_an_uppercase_scheme_remote_high_detail_image_at_the_upper_bound(scheme): + messages = _image_messages(scheme + UNREACHABLE_IMAGE_URL.split("://", 1)[1], "high") + + assert count_relayed_prompt_tokens("gpt-4.1-mini", messages) == ( + litellm.token_counter(model="gpt-4.1-mini", messages=TEXT_ONLY_MESSAGES) + high_detail_image_token_upper_bound() + ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index d9969dd1dc9..addc952af14 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1,3 +1,4 @@ +import asyncio import base64 import contextlib import json @@ -19,6 +20,7 @@ from starlette.datastructures import FormData import litellm from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, @@ -1852,11 +1854,11 @@ class TestBedrockAgentRuntimePassthroughToggle: return request @contextlib.contextmanager - def _patched_dispatch(self, general_settings: Mapping[str, object]): + def _patched_dispatch(self, general_settings: Mapping[str, object], credentials: object | None = None): from botocore.credentials import Credentials bedrock_llm: Final = Mock() - bedrock_llm.get_credentials = Mock(return_value=Credentials("ak", "sk")) + bedrock_llm.get_credentials = Mock(return_value=credentials or Credentials("ak", "sk")) forwarder: Final = AsyncMock(return_value="forwarded") with ( @@ -1891,6 +1893,27 @@ class TestBedrockAgentRuntimePassthroughToggle: forwarder.assert_awaited_once() assert "bedrock-agent-runtime.us-east-1.amazonaws.com" in create_route.call_args.kwargs["target"] + @pytest.mark.asyncio + async def test_agent_runtime_dispatch_signs_off_the_event_loop(self, monkeypatch): + """Regression for issue #40165: the agent-runtime pass-through signed on the loop, so botocore's + blocking credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe: Final = EventLoopProbe() + release: Final = asyncio.create_task(probe.release_refresh_from_the_loop()) + + with self._patched_dispatch(MappingProxyType({}), credentials=probe.credentials()) as (create_route, forwarder): + result: Final = await bedrock_proxy_route( + endpoint=self.AGENT_RUNTIME_ENDPOINT, + request=self._mock_request(), + fastapi_response=Mock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + await release + + assert result == "forwarded" + assert create_route.call_args.kwargs["custom_headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + @pytest.mark.asyncio @pytest.mark.parametrize("value", (True, "true", "True")) async def test_agent_runtime_dispatch_rejected_when_disabled(self, value: bool | str): @@ -4999,7 +5022,11 @@ class TestPassthroughRouterModelBudgetReservation: monkeypatch.setattr(proxy_server, "llm_router", RecordingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) return captured def _assert_metadata_carries_attribution(self, captured: list[dict], user_api_key_dict: UserAPIKeyAuth) -> None: @@ -5098,7 +5125,11 @@ class TestAzureRouterModelStreamingDispatch: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5158,7 +5189,11 @@ class TestAzureRouterModelStreamingKeepalive: monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) - monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) request = MagicMock(spec=Request) request.method = "POST" @@ -5208,6 +5243,97 @@ class TestAzureRouterModelStreamingKeepalive: assert chunks == [b"data: hello\n\n"] +class TestRouterModelRelayUpstreamContract: + def _request(self, content_type: str) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": content_type} + request.query_params = {} + return request + + def _install_router(self, monkeypatch, router, body: dict) -> None: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr( + ep, + "is_passthrough_request_using_router_model", + lambda request_body, llm_router=None: request_body.get("model") in ("gpt-5", "router-model"), + ) + + def _recording_router(self, captured: list[dict]): + class RecordingRouter: + async def allm_passthrough_route(self, **kwargs): + captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + return RecordingRouter() + + @pytest.mark.asyncio + async def test_azure_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_vllm_relay_keeps_the_json_body_when_the_content_type_carries_a_charset(self, monkeypatch): + body = {"model": "router-model", "messages": [{"role": "user", "content": "hi"}]} + captured: list[dict] = [] + self._install_router(monkeypatch, self._recording_router(captured), body) + + await vllm_proxy_route( + endpoint="/chat/completions", + request=self._request("application/json; charset=utf-8"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert captured[0]["json"] == body + + @pytest.mark.asyncio + async def test_azure_relay_returns_the_upstream_status_and_body_when_the_deployment_rejects_the_call( + self, monkeypatch + ): + upstream_body = {"error": {"code": "DeploymentNotFound", "message": "The API deployment does not exist."}} + + class RejectingRouter: + async def allm_passthrough_route(self, **kwargs): + upstream_request = httpx.Request( + "POST", "https://my-azure.openai.azure.com/openai/deployments/gpt-5/chat/completions" + ) + upstream = httpx.Response( + 404, json=upstream_body, headers={"x-ms-request-id": "req-1"}, request=upstream_request + ) + raise httpx.HTTPStatusError("404", request=upstream_request, response=upstream) + + self._install_router(monkeypatch, RejectingRouter(), {"model": "gpt-5", "stream": False}) + + result = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request("application/json"), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert result.status_code == 404 + assert json.loads(result.body) == upstream_body + assert result.headers["x-ms-request-id"] == "req-1" + + @pytest.mark.asyncio async def test_bedrock_count_tokens_error_forwards_provider_headers(): """The count tokens route converts BedrockError into an HTTPException, and dropping the @@ -5240,3 +5366,88 @@ async def test_bedrock_count_tokens_error_forwards_provider_headers(): assert exc_info.value.status_code == 500 assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" + + +class _AzureGroupRouter: + def __init__(self, captured: list[dict]) -> None: + self.captured = captured + + def get_model_names(self, team_id=None): + return ["gpt", "other-group"] + + def get_model_list(self, model_name=None, team_id=None): + rows = [ + {"model_name": "gpt", "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_key": "k"}}, + {"model_name": "other-group", "litellm_params": {"model": "azure/gpt-5.4", "api_key": "k"}}, + ] + return [row for row in rows if model_name is None or row["model_name"] == model_name] + + async def allm_passthrough_route(self, **kwargs): + self.captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + +class TestAzureRelayDeploymentSegment: + """A key allowed one model group must not reach another deployment by naming it in the + ``openai/deployments/`` segment while the group segment picks the credential.""" + + def test_models_served_by_group_resolves_each_deployment_to_its_model_name(self): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _models_served_by_group + + assert _models_served_by_group(_AzureGroupRouter([]), "gpt") == frozenset({"gpt-5.4-mini"}) + assert _models_served_by_group(_AzureGroupRouter([]), "missing-group") == frozenset() + + def _install(self, monkeypatch, body: dict) -> list[dict]: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + captured: list[dict] = [] + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", _AzureGroupRouter(captured)) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + return captured + + def _request(self) -> Request: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + return request + + @pytest.mark.asyncio + async def test_azure_relay_rejects_a_deployment_the_group_does_not_serve(self, monkeypatch): + from fastapi import HTTPException + + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + with pytest.raises(HTTPException) as exc_info: + await azure_proxy_route( + endpoint="gpt/openai/deployments/gpt-5.4/chat/completions", + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert exc_info.value.status_code == 400 + assert "gpt-5.4" in exc_info.value.detail["error"] + assert captured == [] + + @pytest.mark.asyncio + async def test_azure_relay_dispatches_the_group_and_its_own_deployment_name(self, monkeypatch): + captured = self._install(monkeypatch, {"model": "gpt", "messages": []}) + + for endpoint in ( + "gpt/openai/deployments/gpt/chat/completions", + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + ): + await azure_proxy_route( + endpoint=endpoint, + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token", models=["gpt"]), + ) + + assert [call["model"] for call in captured] == ["gpt", "gpt"] diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 61880a8c6f6..0a2641082dc 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -675,8 +675,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch): ], ) - monkeypatch.setattr(litellm, "callbacks", []) - result = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, @@ -1065,7 +1063,7 @@ class _TextReturningGuardrail(CustomGuardrail): class _TextTranslation: - delivers_ended_stream_text_rewrites = False + delivers_ended_stream_rewrites = False def __init__(self): self.seen_guardrail_names = [] @@ -1087,7 +1085,7 @@ class _WritingTranslation: """Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the chat/Responses/Messages handlers do on an ended stream.""" - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True async def process_output_streaming_response( self, @@ -1106,12 +1104,13 @@ class _WritingTranslation: logging_obj=litellm_logging_obj, ) responses_so_far[0]["text"] = outputs["texts"][0] - responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] + if len(outputs["tool_calls"]) == 1: + responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] return responses_so_far class _RefusingTranslation: - delivers_ended_stream_text_rewrites = True + delivers_ended_stream_rewrites = True async def process_output_streaming_response( self, @@ -1148,6 +1147,7 @@ def _assert_passed_with_discard_warning(result, caplog): assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) @pytest.mark.asyncio @@ -1229,13 +1229,47 @@ async def test_streaming_step_delivers_text_rewrite_through_writing_translation( @pytest.mark.asyncio -async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog): +async def test_streaming_step_delivers_tool_call_rewrite_through_writing_translation(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) chunks = [_chunk()] with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_WritingTranslation(), chunks) + assert result.terminal_action == "allow" + assert chunks[0]["text"] == "hello [MASKED]" + assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "[MASKED]"}' + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +class _ToolCallDroppingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool_call(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_ToolCallDroppingGuardrail()]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation(), chunks) + _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] @@ -1294,3 +1328,350 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] + + +class _LegacyHookGuardrail(CustomGuardrail): + """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" + + def __init__(self, replacement=None, raises=None, guardrail_name="masker", rewrite_in_place=None): + super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True) + self.replacement = replacement + self.raises = raises + self.rewrite_in_place = rewrite_in_place + self.calls = [] + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response}) + if self.raises is not None: + raise self.raises + if self.rewrite_in_place is not None: + response["text"] = self.rewrite_in_place + return self.replacement + + +class _NativeHooksGuardrail(_LegacyHookGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + +class _LegacyScanningTranslation: + """Stores the assembled response under request_data["response"] before scanning, like the + chat, Responses, and Messages handlers, hands hooks a route-native shape, and re-extracts one + text per entry of a replacement's "texts".""" + + delivers_ended_stream_rewrites = True + + def post_call_hook_response(self, response): + return {"native": True, "text": response["text"], "tool_calls": response["tool_calls"]} + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + request_data.setdefault( + "response", {"text": responses_so_far[0]["text"], "tool_calls": [dict(responses_so_far[0]["tool_call"])]} + ) + outputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + responses_so_far[0]["text"] = outputs["texts"][0] + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + inputs = {"texts": [response["text"]] if "text" in response else list(response["texts"])} + if response.get("tool_calls"): + inputs["tool_calls"] = list(response["tool_calls"]) + await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +class _ToolOnlyLegacyScanningTranslation(_LegacyScanningTranslation): + """Like the Messages handler on a tool-only message: the ended-stream scan omits "texts" from + the inputs, while the non-streaming scan of the same response sends an empty list.""" + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + request_data.setdefault("response", {"text": "", "tool_calls": [dict(responses_so_far[0]["tool_call"])]}) + await guardrail_to_apply.apply_guardrail( + inputs={"tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": [], "tool_calls": list(response.get("tool_calls") or [])}, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +def _tool_only_chunk(): + return {"text": "", "tool_call": _chunk()["tool_call"]} + + +def _native(text): + return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]} + + +def _legacy_replacement(*texts, tool_calls=None): + return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} + + +async def _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, on_fail="block", on_error="next", translation=None +): + return await _run_legacy_streaming_steps( + monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error, translation=translation + ) + + +async def _run_legacy_streaming_steps( + monkeypatch, guardrails, chunks, on_fail="block", on_error="next", translation=None +): + monkeypatch.setattr(litellm, "callbacks", list(guardrails)) + return await PipelineExecutor.execute_steps( + steps=[ + PipelineStep( + guardrail=guardrail.guardrail_name, + on_pass="next" if position + 1 < len(guardrails) else "allow", + on_fail=on_fail, + on_error=on_error, + ) + for position, guardrail in enumerate(guardrails) + ], + mode="post_call", + data={"model": "m"}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="p", + streaming_chunks=chunks, + endpoint_translation=_LegacyScanningTranslation() if translation is None else translation, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail]) +async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class): + guardrail = guardrail_class(replacement=_legacy_replacement("[REWRITTEN] hello world")) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in guardrail.calls] == [_native("hello world")] + assert guardrail.calls[0]["data"]["model"] == "m" + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_delivers_a_legacy_rewrite_made_in_place(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(rewrite_in_place="[REWRITTEN] hello world") + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert len(guardrail.calls) == 1 + assert chunks == [_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_streaming_step_blocks_with_the_legacy_hook_exception(monkeypatch): + exc = HTTPException(status_code=400, detail={"error": "output blocked"}) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, _LegacyHookGuardrail(raises=exc), chunks) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["fail"] + assert result.original_exception is exc + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatch): + chunks = [_chunk()] + + result = await _run_legacy_streaming_step( + monkeypatch, _LegacyHookGuardrail(raises=ValueError("boom")), chunks, on_error="block" + ) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["error"] + assert result.step_results[0].error_detail == "boom" + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("split", "in two")) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail( + replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call]) + ) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[])) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alone(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert chunks == [_tool_only_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call])) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_tool_only_chunk()] + + +class _NoHooksGuardrail(CustomGuardrail): + pass + + +class _IteratorAndLegacyHookGuardrail(_LegacyHookGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for item in response: + yield item + + +class _UnscannableRewriteTranslation(_LegacyScanningTranslation): + """Like the chat handler on a response whose choices are plain dicts: the non-streaming scan + never hands anything to the guardrail.""" + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + return response + + +def test_streaming_execution_runs_legacy_hooks_only_when_that_hook_is_their_only_streaming_path(): + assert PipelineExecutor.supports_streaming_execution(_LegacyHookGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_NativeHooksGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_IteratorAndLegacyHookGuardrail()) is False + assert PipelineExecutor.supports_streaming_execution(_NoHooksGuardrail(guardrail_name="neither")) is False + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_rescan(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("hello [MASKED]")) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation()) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): + masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) + auditor = _LegacyHookGuardrail(replacement=None, guardrail_name="auditor") + chunks = [_chunk()] + + result = await _run_legacy_streaming_steps(monkeypatch, [masker, auditor], chunks, on_fail="next") + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass", "pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in masker.calls] == [_native("hello world")] + assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")] diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py new file mode 100644 index 00000000000..37849101b3a --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -0,0 +1,262 @@ +import logging +from collections.abc import Iterator, Mapping + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.router import Deployment, LiteLLM_Params + +GOVERNED_MODEL_GROUP = "gpt-5.4-mini" +GOVERNED_MODEL_ID = "deployment-governed" +UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini" +UNGOVERNED_MODEL_ID = "deployment-ungoverned" +WILDCARD_MODEL_GROUP = "openai/*" +WILDCARD_MODEL_ID = "deployment-wildcard" + + +class FakeRouter: + def __init__(self, deployments: dict[str, Deployment], model_group_alias: dict[str, object] | None = None): + self._deployments = deployments + self.model_group_alias = model_group_alias or {} + + def get_deployment(self, model_id: str) -> Deployment | None: + return self._deployments.get(model_id) + + +def _deployment(model_group: str, model_id: str) -> Deployment: + return Deployment( + model_name=model_group, + litellm_params=LiteLLM_Params(model=f"openai/{model_group}"), + model_info={"id": model_id}, + ) + + +def _router(model_group_alias: dict[str, object] | None = None) -> FakeRouter: + return FakeRouter( + { + GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID), + UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID), + WILDCARD_MODEL_ID: _deployment(WILDCARD_MODEL_GROUP, WILDCARD_MODEL_ID), + }, + model_group_alias, + ) + + +def _encoded_response_id(model_id: str) -> str: + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=model_id, response_id="resp_upstream" + ) + + +def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, object]: + return { + "guardrails": {"add": [guardrail]}, + "pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": "block"}]}, + } + + +@pytest.fixture +def policy_engine() -> Iterator[None]: + policy_registry = get_policy_registry() + attachment_registry = get_attachment_registry() + policy_registry.load_policies( + { + "response-governance": _pipeline_policy("output-word-filter"), + "input-governance": _pipeline_policy("input-word-filter", mode="pre_call"), + "team-governance": _pipeline_policy("team-word-filter"), + "tag-governance": _pipeline_policy("tag-word-filter"), + } + ) + attachment_registry.load_attachments( + [ + {"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]}, + {"policy": "team-governance", "teams": ["governed-team"]}, + {"policy": "tag-governance", "tags": ["governed"]}, + ] + ) + yield + policy_registry.clear() + attachment_registry.clear() + + +def _retrieval_data(model_id: str) -> dict[str, object]: + return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}} + + +def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + bucket = data["litellm_metadata"] + assert isinstance(bucket, dict) + return tuple( + (policy_name, ",".join(step.guardrail for step in pipeline.steps)) + for policy_name, pipeline in bucket["_guardrail_pipelines"] + ) + + +def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine: None) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["_pipeline_managed_guardrails"] == frozenset({"output-word-filter"}) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["litellm_metadata"]["applied_guardrails"] == ["output-word-filter"] + assert data["litellm_metadata"]["policy_sources"] == {"response-governance": "model:gpt-5.4-mini"} + assert "model" not in data + assert "guardrails" not in data["litellm_metadata"] + + +def test_key_and_team_context_also_governs_retrieval(policy_engine: None) -> None: + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval( + data=data, user_api_key_dict=UserAPIKeyAuth(team_alias="governed-team"), llm_router=_router() + ) + + assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),) + + +def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine: None) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + + +def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine: None) -> None: + data: dict[str, object] = { + "response_id": _encoded_response_id(UNGOVERNED_MODEL_ID), + "litellm_metadata": {"tags": ["governed"]}, + } + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert _attached_pipelines(data) == (("tag-governance", "tag-word-filter"),) + assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"} + + +def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine: None) -> None: + data = _retrieval_data(UNGOVERNED_MODEL_ID) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(UNGOVERNED_MODEL_ID) + + +def test_already_attached_policy_is_not_attached_twice(policy_engine: None) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router() + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + + +def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "the model name it was submitted as" in record.getMessage() + ] + + +def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: + data = _retrieval_data(WILDCARD_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(WILDCARD_MODEL_ID) + assert [ + "as model group openai/* (a wildcard deployment)" in message + for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_aliased_model_group_still_attaches_its_own_policies_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}}) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert [ + "(the target of model_group_alias gpt-mini, gpt-hidden)" in message + for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model( + policy_engine: None, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval( + data=_retrieval_data(GOVERNED_MODEL_ID), + user_api_key_dict=UserAPIKeyAuth(), + llm_router=_router({"other-alias": UNGOVERNED_MODEL_GROUP}), + ) + + assert _hidden_submit_model_warnings(caplog) == [] + + +def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + and "retrieved without its post_call policy pipelines" in record.getMessage() + ] + + +@pytest.mark.parametrize( + ("response_id", "reason"), + [ + ("resp_plain_upstream_id", "response id names no deployment"), + (_encoded_response_id("deployment-missing-from-router"), "deployment no longer in the router"), + (None, "response id names no deployment"), + ], +) +def test_unresolvable_response_id_attaches_nothing_and_warns( + policy_engine: None, caplog: pytest.LogCaptureFixture, response_id: str, reason: str +) -> None: + data = {"response_id": response_id, "litellm_metadata": {}} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == {"response_id": response_id, "litellm_metadata": {}} + assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True] + + +def test_without_a_router_attaches_nothing_and_warns(policy_engine: None, caplog: pytest.LogCaptureFixture) -> None: + data = _retrieval_data(GOVERNED_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True] + + +def test_without_policy_engine_attaches_nothing_quietly(caplog: pytest.LogCaptureFixture) -> None: + get_policy_registry().clear() + data = _retrieval_data(GOVERNED_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None) + + assert data == _retrieval_data(GOVERNED_MODEL_ID) + assert _ungoverned_retrieval_warnings(caplog) == [] diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index dcb63b8ca82..2d9c1bd8b46 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -13,9 +13,12 @@ Routes covered: from __future__ import annotations +import asyncio import json from unittest.mock import AsyncMock, MagicMock +import pytest + from .conftest import VOLATILE_KEYS, normalize @@ -229,10 +232,7 @@ def test_config_update_no_db_error(client, auth_as, monkeypatch): json={"general_settings": {"alerting": ["slack"]}}, ) assert response.status_code != 200 - assert ( - "db" in str(response.json()).lower() - or "connect" in str(response.json()).lower() - ) + assert "db" in str(response.json()).lower() or "connect" in str(response.json()).lower() # --------------------------------------------------------------------------- @@ -273,9 +273,7 @@ def test_config_field_update_happy_admin(client, auth_as, mock_prisma, monkeypat } -def test_config_field_update_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_update_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin cannot update config fields — returns 400 with not-allowed detail (handler uses 400 for the auth gate, not 403).""" from litellm.proxy import proxy_server as ps @@ -335,9 +333,7 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 200 assert normalize(response.json()) == { "field_name": "max_parallel_requests", @@ -345,9 +341,7 @@ def test_config_field_info_happy_admin(client, auth_as, mock_prisma, monkeypatch } -def test_config_field_info_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin (INTERNAL_USER) is denied — admin-view gate fires.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -356,9 +350,7 @@ def test_config_field_info_non_admin_rejected( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 400 assert "error" in response.json().get("detail", {}) @@ -375,16 +367,12 @@ def test_config_field_info_field_not_in_db(client, auth_as, mock_prisma, monkeyp monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + response = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert response.status_code == 400 assert "not in DB" in response.json().get("detail", {}).get("error", "") -def test_config_field_info_redacts_nested_secret_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_redacts_nested_secret_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): """A view-only admin reading a structured field must not receive nested credentials. database_args carries aws_web_identity_token (a DynamoDB role-assumption credential); it must come back redacted while non-secret @@ -405,9 +393,7 @@ def test_config_field_info_redacts_nested_secret_for_view_only_admin( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - response = client.get( - "/config/field/info", params={"field_name": "database_args"} - ) + response = client.get("/config/field/info", params={"field_name": "database_args"}) assert response.status_code == 200 value = response.json()["field_value"] assert value["aws_web_identity_token"] == "REDACTED" @@ -415,9 +401,7 @@ def test_config_field_info_redacts_nested_secret_for_view_only_admin( assert value["user_table_name"] == "LiteLLM_UserTable" -def test_config_field_info_full_admin_sees_nested_secret( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_full_admin_sees_nested_secret(client, auth_as, mock_prisma, monkeypatch): """The redaction must not over-redact for a full PROXY_ADMIN, who needs the real nested value to populate the edit form.""" from litellm.proxy import proxy_server as ps @@ -435,18 +419,14 @@ def test_config_field_info_full_admin_sees_nested_secret( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/field/info", params={"field_name": "database_args"} - ) + response = client.get("/config/field/info", params={"field_name": "database_args"}) assert response.status_code == 200 value = response.json()["field_value"] assert value["aws_web_identity_token"] == "sk-super-secret-token" assert value["region_name"] == "us-east-1" -def test_config_field_info_redacts_top_level_scalar_for_view_only( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_info_redacts_top_level_scalar_for_view_only(client, auth_as, mock_prisma, monkeypatch): """The top-level scalar branch must also redact for a view-only admin. database_url carries DB credentials and is not caught by the name masker, so it is in the explicit secret set.""" @@ -460,9 +440,7 @@ def test_config_field_info_redacts_top_level_scalar_for_view_only( monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - response = client.get( - "/config/field/info", params={"field_name": "database_url"} - ) + response = client.get("/config/field/info", params={"field_name": "database_url"}) assert response.status_code == 200 assert response.json()["field_value"] == "REDACTED" @@ -476,17 +454,12 @@ def test_redact_general_setting_value_recurses_list_of_dicts(): {"path": "/foo", "headers": {"Authorization": "Bearer sk-x"}}, {"path": "/bar", "client_secret": "sk-y"}, ] - redacted = ps._redact_general_setting_value( - "some_list_field", value, is_full_admin=False - ) + redacted = ps._redact_general_setting_value("some_list_field", value, is_full_admin=False) assert redacted[0]["headers"]["Authorization"] == "REDACTED" assert redacted[0]["path"] == "/foo" assert redacted[1]["client_secret"] == "REDACTED" assert redacted[1]["path"] == "/bar" - assert ( - ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) - == value - ) + assert ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) == value def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): @@ -504,22 +477,16 @@ def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): for _ in range(ps._REDACT_SECRET_MAX_DEPTH + 2): nested = {"wrap": nested} - out = ps._redact_general_setting_value( - "some_struct_field", nested, is_full_admin=False - ) + out = ps._redact_general_setting_value("some_struct_field", nested, is_full_admin=False) # the secret must not survive anywhere in the returned tree assert "sk-leak-bottom" not in repr(out) # full admin is unaffected by the cap — the value comes back untouched - admin_out = ps._redact_general_setting_value( - "some_struct_field", nested, is_full_admin=True - ) + admin_out = ps._redact_general_setting_value("some_struct_field", nested, is_full_admin=True) assert admin_out is nested -def test_config_list_redacts_pass_through_secret_for_view_only( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_list_redacts_pass_through_secret_for_view_only(client, auth_as, mock_prisma, monkeypatch): """/config/list must not leak pass_through_endpoints upstream credentials to a view-only admin. pass_through_endpoints is a known secret-bearing field, so a non-admin gets it redacted; a full admin still sees it.""" @@ -546,24 +513,16 @@ def test_config_list_redacts_pass_through_secret_for_view_only( ) def _pass_through_value(body): - return next( - entry["field_value"] - for entry in body - if entry["field_name"] == "pass_through_endpoints" - ) + return next(entry["field_value"] for entry in body if entry["field_name"] == "pass_through_endpoints") with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): - view_resp = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + view_resp = client.get("/config/list", params={"config_type": "general_settings"}) assert view_resp.status_code == 200 assert "sk-UPSTREAM-SECRET" not in view_resp.text assert _pass_through_value(view_resp.json()) == "REDACTED" with auth_as(LitellmUserRoles.PROXY_ADMIN): - admin_resp = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + admin_resp = client.get("/config/list", params={"config_type": "general_settings"}) assert admin_resp.status_code == 200 admin_value = _pass_through_value(admin_resp.json()) assert admin_value[0]["headers"]["Authorization"] == "Bearer sk-UPSTREAM-SECRET" @@ -587,9 +546,7 @@ def test_config_list_happy_admin(client, auth_as, mock_prisma, monkeypatch): monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 200 body = response.json() assert isinstance(body, list) @@ -695,9 +652,7 @@ def test_config_list_non_admin_rejected(client, auth_as, mock_prisma, monkeypatc monkeypatch.setattr(ps, "prisma_client", mock_prisma) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 400 assert "role" in response.json().get("detail", {}).get("error", "").lower() @@ -710,9 +665,7 @@ def test_config_list_no_db_error(client, auth_as, monkeypatch): monkeypatch.setattr(ps, "prisma_client", None) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get( - "/config/list", params={"config_type": "general_settings"} - ) + response = client.get("/config/list", params={"config_type": "general_settings"}) assert response.status_code == 400 assert "error" in response.json().get("detail", {}) @@ -756,9 +709,7 @@ def test_config_field_delete_happy_admin(client, auth_as, mock_prisma, monkeypat } -def test_config_field_delete_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_delete_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin caller hits the 400 not-allowed branch with role in detail.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -778,9 +729,7 @@ def test_config_field_delete_non_admin_rejected( assert "role" in response.json().get("detail", {}).get("error", "").lower() -def test_config_field_delete_field_not_in_config( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_field_delete_field_not_in_config(client, auth_as, mock_prisma, monkeypatch): """If there is no general_settings row at all, returns 400 'not in config'.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -825,9 +774,7 @@ def test_config_callback_delete_happy_admin(client, auth_as, mock_prisma, monkey monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) assert response.status_code == 200 # `deleted_at` is an ISO timestamp generated at request time — extend # the volatile set just for this assertion so dict-equality still works. @@ -840,9 +787,7 @@ def test_config_callback_delete_happy_admin(client, auth_as, mock_prisma, monkey } -def test_config_callback_delete_non_admin_rejected( - client, auth_as, mock_prisma, monkeypatch -): +def test_config_callback_delete_non_admin_rejected(client, auth_as, mock_prisma, monkeypatch): """Non-admin caller is rejected with 400 not-allowed.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -852,9 +797,7 @@ def test_config_callback_delete_non_admin_rejected( monkeypatch.setattr(ps, "store_model_in_db", True) with auth_as(LitellmUserRoles.INTERNAL_USER): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) assert response.status_code == 400 assert "role" in response.json().get("detail", {}).get("error", "").lower() @@ -869,22 +812,15 @@ def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypa monkeypatch.setattr(ps, "store_model_in_db", True) fake_proxy_config = MagicMock() - fake_proxy_config.get_config = AsyncMock( - return_value={"litellm_settings": {"success_callback": ["slack"]}} - ) + fake_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {"success_callback": ["slack"]}}) monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.post( - "/config/callback/delete", json={"callback_name": "langfuse"} - ) + response = client.post("/config/callback/delete", json={"callback_name": "langfuse"}) # The handler re-raises HTTPException(404) verbatim (only generic # `Exception` becomes a 500 ProxyException), so pin 404 strictly. assert response.status_code == 404 - assert ( - "langfuse" in str(response.json()).lower() - or "not found" in str(response.json()).lower() - ) + assert "langfuse" in str(response.json()).lower() or "not found" in str(response.json()).lower() # --------------------------------------------------------------------------- @@ -948,10 +884,7 @@ def test_get_config_callbacks_internal_error(client, auth_as, mock_prisma, monke with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/get/config/callbacks") assert response.status_code >= 400 - assert ( - "boom" in str(response.json()).lower() - or "error" in str(response.json()).lower() - ) + assert "boom" in str(response.json()).lower() or "error" in str(response.json()).lower() _CALLBACK_ENV_FIXTURE = { @@ -985,14 +918,10 @@ def _install_callbacks_config(monkeypatch, mock_prisma): def _callback_variables(body: dict, name: str) -> dict: - return next( - cb["variables"] for cb in body["callbacks"] if cb["name"] == name - ) + return next(cb["variables"] for cb in body["callbacks"] if cb["name"] == name) -def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy._types import LitellmUserRoles _install_callbacks_config(monkeypatch, mock_prisma) @@ -1024,9 +953,7 @@ def test_get_config_callbacks_redacts_secret_env_vars_for_view_only_admin( assert otel_vars["OTEL_ENDPOINT"] == _CALLBACK_ENV_FIXTURE["OTEL_ENDPOINT"] -def test_get_config_callbacks_full_admin_still_sees_secret_env_vars( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_full_admin_still_sees_secret_env_vars(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy._types import LitellmUserRoles _install_callbacks_config(monkeypatch, mock_prisma) @@ -1047,9 +974,7 @@ def test_get_config_callbacks_full_admin_still_sees_secret_env_vars( assert otel_vars["OTEL_HEADERS"] == _CALLBACK_ENV_FIXTURE["OTEL_HEADERS"] -def test_get_config_callbacks_redacts_slack_webhook_urls_for_view_only_admin( - client, auth_as, mock_prisma, monkeypatch -): +def test_get_config_callbacks_redacts_slack_webhook_urls_for_view_only_admin(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -1170,6 +1095,413 @@ def test_get_config_callbacks_redacts_email_alerting_vars_for_view_only_admin( assert admin_email["SMTP_HOST"] == "smtp.resend.com" +def test_get_config_callbacks_appends_runtime_only_callbacks(client, auth_as, mock_prisma, monkeypatch): + """LIT-5281: a YAML callback that the DB callback list replaced in the merged config still runs, so it must + show up as a read_only row next to the editable DB-configured one.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": ["langfuse"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.opentelemetry import OpenTelemetry + + monkeypatch.setattr(litellm, "success_callback", ["langfuse", LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", [OpenTelemetry()]) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("langsmith", "success", True), + ("otel", "success_and_failure", True), + ] + + +def test_get_config_callbacks_accepts_scalar_and_null_yaml_callbacks(client, auth_as, mock_prisma, monkeypatch): + """`success_callback: langfuse` (a YAML scalar) is one configured callback, not eight single-letter ones, and a + `callbacks: null` key contributes nothing.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": "langfuse", "failure_callback": None, "callbacks": None}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.langsmith import LangsmithLogger + + monkeypatch.setattr(litellm, "success_callback", ["langfuse", LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("langsmith", "success", True), + ] + + +def test_get_config_callbacks_deduplicates_configured_and_runtime(client, auth_as, mock_prisma, monkeypatch): + """A configured callback shows once as editable, whether the runtime holds its string or an initialized instance + (arize initializes an ArizeLogger, logfire a bare OpenTelemetry that only its class identifies).""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": ["langfuse", "arize", "logfire"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.arize.arize import ArizeLogger + from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + + arize_logger = ArizeLogger(config=OpenTelemetryConfig(exporter="console"), callback_name="arize") + monkeypatch.setattr(litellm, "success_callback", ["langfuse", arize_logger, OpenTelemetry()]) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse", "success", False), + ("arize", "success", False), + ("logfire", "success", False), + ] + + +def test_get_config_callbacks_keeps_yaml_otel_family_callbacks_next_to_configured_one( + client, auth_as, mock_prisma, monkeypatch +): + """LIT-5281: arize, weave_otel and langfuse_otel all initialize OpenTelemetry subclasses. Saving one of them + from the dashboard replaces the YAML `callbacks` list, so the YAML siblings keep running and must stay listed + under their own names instead of being hidden as duplicates of the configured OTel callback.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"callbacks": ["langfuse_otel"]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + from litellm.integrations.arize.arize import ArizeLogger + from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.opentelemetry import OpenTelemetryConfig + from litellm.integrations.weave.weave_otel import WeaveOtelLogger + + console_config = OpenTelemetryConfig(exporter="console") + monkeypatch.setattr(litellm, "success_callback", [LangsmithLogger()]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr( + litellm, + "callbacks", + [ + ArizeLogger(config=console_config, callback_name="arize"), + WeaveOtelLogger(config=console_config), + LangfuseOtelLogger(config=console_config, callback_name="langfuse_otel"), + ], + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + + assert [(cb["name"], cb["type"], cb.get("read_only", False)) for cb in response.json()["callbacks"]] == [ + ("langfuse_otel", "success_and_failure", False), + ("arize", "success_and_failure", True), + ("langsmith", "success", True), + ("weave_otel", "success_and_failure", True), + ] + + +def _dotted_path_test_function(*args, **kwargs): + pass + + +@pytest.mark.parametrize("handler_kind", ["instance", "function"]) +@pytest.mark.parametrize( + "config_key,expected_type", + [ + ("success_callback", "success"), + ("failure_callback", "failure"), + ("callbacks", "success_and_failure"), + ], +) +def test_get_config_callbacks_deduplicates_dotted_path_callback( + client, auth_as, mock_prisma, monkeypatch, config_key, expected_type, handler_kind +): + """A dotted-path callback stays a single editable row instead of duplicating under its class or function name.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class _DottedPathTestHandler(CustomLogger): + pass + + dotted_handler = _DottedPathTestHandler() if handler_kind == "instance" else _dotted_path_test_function + dotted_path = f"{__name__}.dotted_handler" + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {config_key: [dotted_path]}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + monkeypatch.setattr(litellm, "callbacks", [dotted_handler]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + callbacks = response.json()["callbacks"] + assert [(callback["name"], callback["type"], callback.get("read_only", False)) for callback in callbacks] == [ + (dotted_path, expected_type, False) + ] + + +def test_get_config_callbacks_lists_dict_shaped_config_callbacks(client, auth_as, mock_prisma, monkeypatch): + """Dict-shaped success_callback config values list their keys as editable rows.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": {"langsmith": {"batch_size": 1}}}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_callbacks_by_type", + MagicMock(return_value={"success": ["langsmith"], "failure": [], "success_and_failure": []}), + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + callbacks = response.json()["callbacks"] + assert [(callback["name"], callback.get("read_only", False)) for callback in callbacks] == [("langsmith", False)] + + +def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_as, mock_prisma, monkeypatch): + """Proxy infrastructure callbacks are excluded from callback inventory.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": []}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + import litellm + from litellm._service_logger import ServiceLogging + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_logger import CustomLogger + from litellm.integrations.langsmith import LangsmithLogger + from litellm.integrations.s3_v2 import S3Logger + from litellm.integrations.sqs import SQSLogger + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook + from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + from litellm.router import Router + + class _InventoryTestGuardrail(CustomGuardrail): + pass + + class _UserCodeLogger(CustomLogger): + pass + + def user_code_function(*args, **kwargs): + pass + + async def build_aws_loggers() -> tuple[S3Logger, SQSLogger]: + return S3Logger(s3_bucket_name="inventory-bucket"), SQSLogger(sqs_queue_url="https://sqs.example/inventory") + + s3_logger, sqs_logger = asyncio.run(build_aws_loggers()) + router = Router(model_list=[]) + monkeypatch.setattr(litellm, "input_callback", []) + monkeypatch.setattr( + litellm, "success_callback", [LangsmithLogger(), s3_logger, router.sync_deployment_callback_on_success] + ) + monkeypatch.setattr(litellm, "_async_success_callback", [sqs_logger, router.deployment_callback_on_success]) + monkeypatch.setattr(litellm, "failure_callback", [user_code_function]) + monkeypatch.setattr(litellm, "_async_failure_callback", [router.async_deployment_callback_on_failure]) + monkeypatch.setattr( + litellm, + "callbacks", + [ + _PROXY_MaxBudgetLimiter(), + _PROXY_LiteLLMManagedFiles(internal_usage_cache=MagicMock(), prisma_client=MagicMock()), + ServiceLogging(), + VectorStorePreCallHook(), + _InventoryTestGuardrail(guardrail_name="inventory-test-guardrail"), + _UserCodeLogger(), + ], + ) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + assert "cache" in litellm.success_callback and "cache" in litellm._async_success_callback + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/get/config/callbacks") + + assert response.status_code == 200 + assert [ + (callback["name"], callback["type"], callback["read_only"]) for callback in response.json()["callbacks"] + ] == [ + ("_UserCodeLogger", "success_and_failure", True), + ("langsmith", "success", True), + ("s3", "success", True), + ("sqs", "success", True), + ("user_code_function", "failure", True), + ] + + +def test_get_config_callbacks_redacts_runtime_only_row_secrets_for_view_only_admin( + client, auth_as, mock_prisma, monkeypatch +): + """Runtime-only callback rows are subject to the same redaction gate as configured.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "llm_router", None) + + fake_proxy_config = MagicMock() + fake_proxy_config.get_config = AsyncMock( + return_value={ + "litellm_settings": {"success_callback": []}, + "general_settings": {}, + "environment_variables": dict(_CALLBACK_ENV_FIXTURE), + } + ) + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + import litellm + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", ["otel"]) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + response = client.get("/get/config/callbacks") + assert response.status_code == 200 + body = response.json() + + callbacks = body["callbacks"] + otel_cb = next((cb for cb in callbacks if cb["name"] == "otel"), None) + assert otel_cb is not None + assert otel_cb["type"] == "success_and_failure" + assert otel_cb["read_only"] is True + assert otel_cb["variables"]["OTEL_HEADERS"] == "REDACTED" + assert otel_cb["variables"]["OTEL_ENDPOINT"] == _CALLBACK_ENV_FIXTURE["OTEL_ENDPOINT"] + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + admin_response = client.get("/get/config/callbacks") + assert admin_response.status_code == 200 + admin_body = admin_response.json() + admin_otel = next((cb for cb in admin_body["callbacks"] if cb["name"] == "otel"), None) + assert admin_otel is not None + assert admin_otel["variables"]["OTEL_HEADERS"] == _CALLBACK_ENV_FIXTURE["OTEL_HEADERS"] + + # --------------------------------------------------------------------------- # GET /config/yaml # --------------------------------------------------------------------------- @@ -1183,9 +1515,7 @@ def test_config_yaml_returns_demo_payload(client, auth_as): response = client.request("GET", "/config/yaml", json={}) shape = { "status": response.status_code, - "media_type_yaml": response.headers.get("content-type", "").startswith( - "application/json" - ), + "media_type_yaml": response.headers.get("content-type", "").startswith("application/json"), "has_body": len(response.content) > 0, } assert shape == { diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index bc6106a06f8..e3d6b46e2d2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -343,3 +343,53 @@ def test_anthropic_format_returns_public_team_model_name( assert response.status_code == 200 assert [m["id"] for m in response.json()["data"]] == ["gpt-4-team"] assert internal_name not in response.text + + +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +@pytest.mark.parametrize( + "caller_headers", + [ + {"anthropic-version": "2023-06-01", "user-agent": "claude-code/2.1.267"}, + {"anthropic-version": "2023-06-01", "user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, + {"anthropic-version": "2023-06-01", "x-gateway-client": "claude-code"}, + ], +) +def test_anthropic_format_lists_claude_code_view_ids_for_claude_code( + client, auth_as, patched_models, monkeypatch, path, caller_headers +): + """Claude Code drops every id without claude/anthropic in it and reads [1m] as its 1M marker, so for Claude + Code (its discovery fetch's own user agent, its SDK's, or the gateway-client header a launcher sends) every + group is listed under a Claude-shaped id with the marker where the window reaches 1M; the display name stays + the served name.""" + + def _create_model_info_response(model_id, provider="openai", **kwargs): + if model_id != "claude-sonnet": + return _stub_model_info_response(model_id=model_id, provider=provider) + return {**_stub_model_info_response(model_id=model_id, provider=provider), "max_input_tokens": 1000000} + + patched_models.model_group_alias = {} + patched_models.has_model_id.return_value = False + patched_models.get_candidate_model_ids_for_route.side_effect = lambda name, team_id=None: frozenset({name}) if name in ("gpt-4", "claude-sonnet") else frozenset() + monkeypatch.setattr(proxy_utils, "create_model_info_response", _create_model_info_response) + + with auth_as(): + response = client.get(path, headers=caller_headers) + + assert response.status_code == 200 + body = response.json() + assert [(m["id"], m["display_name"]) for m in body["data"]] == [ + ("claude-router-6770742d34", "gpt-4"), + ("claude-sonnet[1m]", "claude-sonnet"), + ] + assert (body["first_id"], body["last_id"]) == ("claude-router-6770742d34", "claude-sonnet[1m]") + assert [row["source_model"] for row in body["data"]] == ["gpt-4", "claude-sonnet"] + + +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +def test_anthropic_format_keeps_served_ids_for_other_anthropic_clients(client, auth_as, patched_models, path): + """An Anthropic SDK asking for the vendor shape gets the served ids: the view is Claude Code's alone.""" + with auth_as(): + response = client.get(path, headers={"anthropic-version": "2023-06-01", "user-agent": "anthropic-sdk-python/0.40"}) + + assert response.status_code == 200 + assert [m["id"] for m in response.json()["data"]] == ["gpt-4", "claude-sonnet"] diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 86e97a334df..c343652efd9 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -4,11 +4,12 @@ Pins covered: - ``get_current_spend`` - ``increment_spend_counters`` - ``_reconcile_budget_reservation_for_counter_update`` -- ``_increment_end_user_and_tag_spend_counters`` -- ``_increment_org_spend_counter`` -- ``_init_and_increment_unreserved_spend_counter`` -- ``_init_and_increment_spend_counter`` -- ``_init_and_increment_window_spend_counter`` +- ``_prepare_end_user_and_tag_spend_increments`` +- ``_prepare_org_spend_increment`` +- ``_prepare_unreserved_spend_counter_increment`` +- ``_prepare_spend_counter_increment`` +- ``_prepare_window_spend_counter_increment`` +- ``_apply_spend_counter_increments`` - ``_ensure_spend_counter_initialized`` - ``_get_source_cache_base_spend`` - ``_ensure_window_spend_counter_initialized`` @@ -48,9 +49,7 @@ def _make_spend_counter_cache( cache.in_memory_cache.delete_cache = MagicMock() if with_redis: cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock( - return_value=redis_get_value, side_effect=redis_get_side_effect - ) + cache.redis_cache.async_get_cache = AsyncMock(return_value=redis_get_value, side_effect=redis_get_side_effect) cache.redis_cache.async_increment = AsyncMock( return_value=redis_increment_value, side_effect=redis_increment_side_effect, @@ -58,6 +57,8 @@ def _make_spend_counter_cache( cache.redis_cache.async_delete_cache = AsyncMock() cache.redis_cache.async_set_cache = AsyncMock() cache.redis_cache.async_set_max = AsyncMock() + cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None) + cache.redis_cache.get_ttl = MagicMock(return_value=None) else: cache.redis_cache = None cache.async_increment_cache = AsyncMock(return_value=redis_increment_value) @@ -70,9 +71,7 @@ def _make_spend_counter_cache( def _make_user_api_key_cache(get_value=None, get_side_effect=None): cache = MagicMock() - cache.async_get_cache = AsyncMock( - return_value=get_value, side_effect=get_side_effect - ) + cache.async_get_cache = AsyncMock(return_value=get_value, side_effect=get_side_effect) cache.async_set_cache_pipeline = AsyncMock() return cache @@ -109,9 +108,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch ) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=99.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=99.0) assert result == 17.0 @@ -136,9 +133,7 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # the stale counter is repaired up to the authoritative DB value via a # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key="spend:key:abc", value=12.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0) @pytest.mark.asyncio @@ -169,9 +164,7 @@ async def test_get_current_spend_no_floor_without_max_budget(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0) assert result == 2.0 assert from_db.await_count == 0 @@ -210,12 +203,8 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - first = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) - second = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) + first = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) + second = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) assert first == 12.0 assert second == 12.0 @@ -336,9 +325,7 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) def _make_window_spend_prisma(row=None, spend_logs_total=0.0): @@ -379,9 +366,7 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc assert result == 15.0 fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) @pytest.mark.asyncio @@ -393,9 +378,7 @@ async def test_get_current_spend_floors_window_against_logs_when_row_stale(monke window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) fake_prisma = _make_window_spend_prisma( - row=SimpleNamespace( - window_start=window_start - timedelta(days=7), spend=999.0 - ), + row=SimpleNamespace(window_start=window_start - timedelta(days=7), spend=999.0), spend_logs_total=15.0, ) fake_cache = _make_spend_counter_cache(redis_get_value=2.0) @@ -423,21 +406,13 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat rather than admitted on an unverifiable budget.""" from fastapi import HTTPException - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) with pytest.raises(HTTPException) as exc: - await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert exc.value.status_code == 503 @@ -445,18 +420,12 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkeypatch): """Default (flag off): an unverifiable read keeps the existing behavior and admits using the cached fallback — no new rejection.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "general_settings", {}) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -466,13 +435,9 @@ async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypa authoritative, so an under-budget request is admitted normally.""" fake_cache = _make_spend_counter_cache(redis_get_value=1.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -481,16 +446,10 @@ async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monke """End-user/tag callers pass fallback_authoritative=True (their spend is loaded fresh from the DB in auth), so fail-closed does not reject them even when the counter path is unreadable.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) result = await ps.get_current_spend( counter_key="spend:end_user:e1", @@ -508,9 +467,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa re-checks the authoritative DB and enforces against it.""" fake_cache = _make_spend_counter_cache(redis_get_value=0.00001) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) from_db = AsyncMock(return_value=0.5) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -532,9 +489,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa @pytest.mark.asyncio async def test_increment_spend_counters_increments_all_buckets(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -543,9 +498,7 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): async def _fake_coalesced(**kwargs): return None - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) await ps.increment_spend_counters( token="hashed-tok", @@ -554,25 +507,36 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): response_cost=5.0, ) + pipeline = fake_cache.redis_cache.async_increment_pipeline + pipeline.assert_awaited_once() + increment_list = pipeline.await_args.kwargs["increment_list"] + assert {op["key"] for op in increment_list} == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + } + assert all(op["increment_value"] == 5.0 for op in increment_list) observed = { "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, + "pipeline_calls": pipeline.await_count, "user_cache_used": fake_user_cache.async_get_cache.called, } assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 4, + "redis_increment_called": False, + "pipeline_calls": 1, "user_cache_used": True, } class _ConcurrencyProbe: - """Stand-in for redis_cache.async_increment that pins concurrency. + """Stand-in for redis_cache.async_get_cache that pins concurrency. - Each call registers itself as in-flight and blocks on ``release`` until the - test lets it proceed. ``all_arrived`` fires once ``expected`` distinct scope - increments are simultaneously suspended here, which can only happen if the - per-scope increments are gathered rather than awaited one after another. + Each warm-check read registers itself as in-flight and blocks on ``release`` + until the test lets it proceed. ``all_arrived`` fires once ``expected`` + distinct scope warm-checks are simultaneously suspended here, which can only + happen if the per-scope prepares are gathered rather than awaited one after + another. """ def __init__(self, expected_concurrency: int): @@ -581,36 +545,45 @@ class _ConcurrencyProbe: self.max_in_flight = 0 self.all_arrived = asyncio.Event() self.release = asyncio.Event() - self.values: dict[str, float] = {} + self.keys: list[str] = [] - async def async_increment(self, *, key, value, refresh_ttl=True): + async def async_get_cache(self, *, key, **kwargs): self.in_flight += 1 self.max_in_flight = max(self.max_in_flight, self.in_flight) + self.keys.append(key) if self.in_flight >= self.expected: self.all_arrived.set() if not self.release.is_set(): await self.release.wait() self.in_flight -= 1 - self.values[key] = self.values.get(key, 0.0) + value - return self.values[key] + return 1.0 @pytest.mark.asyncio async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): """The six independent scopes (key, team, team_member, user, end_user+tags, - org) must be incremented concurrently. The probe only fires once all six are - suspended in async_increment at the same time, which is impossible if the + org) must prepare their increments concurrently. The probe only fires once + all eight warm-check reads (one per counter: 6 scopes + 2 tags) are + suspended in async_get_cache at the same time, which is impossible if the awaits are chained sequentially.""" - probe = _ConcurrencyProbe(expected_concurrency=6) - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = probe.async_increment + probe = _ConcurrencyProbe(expected_concurrency=8) + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = probe.async_get_cache + recorded: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results + + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) task = asyncio.create_task( ps.increment_spend_counters( @@ -630,16 +603,16 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): probe.release.set() await task pytest.fail( - "scope increments did not run concurrently; sequential awaits " - f"detected (peak in-flight was {probe.max_in_flight}, expected 6)" + "scope prepares did not run concurrently; sequential awaits " + f"detected (peak in-flight was {probe.max_in_flight}, expected 8)" ) - assert probe.in_flight == 6 - assert probe.max_in_flight == 6 + assert probe.in_flight == 8 + assert probe.max_in_flight == 8 probe.release.set() await task - assert probe.values == { + assert recorded == { "spend:key:hashed-tok": 5.0, "spend:team:t1": 5.0, "spend:team_member:u1:t1": 5.0, @@ -659,26 +632,25 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) import litellm.proxy.spend_tracking.budget_reservation as br reserved = {"spend:key:hashed-tok", "spend:org:org1"} - monkeypatch.setattr( - br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved)) - ) + monkeypatch.setattr(br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved))) monkeypatch.setattr(br, "reconcile_budget_reservation", AsyncMock()) recorded: dict[str, float] = {} - async def _record_increment(*, key, value, refresh_ttl=True): - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] + results.append(recorded[op["key"]]) + return results fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _record_increment + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) reservation = {"finalized": False} await ps.increment_spend_counters( @@ -708,27 +680,46 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ): """A failure in one scope must propagate to the caller (so it can invalidate reserved counters) while every other scope still settles rather than being - left as an orphaned background task, and the reservation is not finalized.""" - recorded: dict[str, float] = {} + left as an orphaned background task, and the reservation is not finalized. + The surviving scopes' increments are still applied in the single pipeline: + dropping them would under-count spend, the unsafe direction for budget + enforcement.""" + warmed_keys: list[str] = [] - async def _increment(*, key, value, refresh_ttl=True): + async def _warm_check(*, key, **kwargs): + warmed_keys.append(key) if key == "spend:team:t1": - raise RuntimeError("redis increment failed") - recorded[key] = recorded.get(key, 0.0) + value - return recorded[key] + raise RuntimeError("redis get failed") + return 1.0 - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment = _increment + async def _reseed_fails(*, counter_key, **kwargs): + if counter_key == "spend:team:t1": + raise RuntimeError("reseed failed") + + applied: dict[str, float] = {} + + async def _record_pipeline(increment_list, **_): + results = [] + for op in increment_list: + applied[op["key"]] = op["increment_value"] + results.append(op["increment_value"]) + return results + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_warm_check) + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) + ps.SpendCounterReseed, + "coalesced", + AsyncMock(side_effect=_reseed_fails), ) reservation = {"finalized": False} - with pytest.raises(RuntimeError, match="redis increment failed"): + with pytest.raises(RuntimeError, match="reseed failed"): await ps.increment_spend_counters( token="hashed-tok", team_id="t1", @@ -741,7 +732,19 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ ) assert reservation["finalized"] is False - assert recorded == { + # every sibling scope settled (its warm-check ran) before the error propagated + assert set(warmed_keys) == { + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + "spend:end_user:eu1", + "spend:tag:a", + "spend:org:org1", + } + # the surviving scopes' increments were still applied, in one pipeline call + fake_cache.redis_cache.async_increment_pipeline.assert_awaited_once() + assert applied == { "spend:key:hashed-tok": 5.0, "spend:team_member:u1:t1": 5.0, "spend:user:u1": 5.0, @@ -749,6 +752,7 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ "spend:tag:a": 5.0, "spend:org:org1": 5.0, } + fake_cache.redis_cache.async_increment.assert_not_awaited() @pytest.mark.asyncio @@ -772,6 +776,108 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( assert reservation == {"finalized": True} assert fake_cache.redis_cache.async_increment.called is False + fake_cache.redis_cache.async_increment_pipeline.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( + monkeypatch, +): + """Every scope's increment must go out in a single async_increment_pipeline + call, not one INCRBYFLOAT round-trip per scope.""" + counter_cache = ps.DualCache() + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + + async def _pipeline(increment_list, **_): + return [1.5] * len(increment_list) + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=_pipeline) + fake_redis.async_increment = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + fake_redis.async_increment_pipeline.assert_awaited_once() + assert fake_redis.async_increment.await_count == 0 + increment_list = fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"] + expected_keys = { + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + } + assert {op["key"] for op in increment_list} == expected_keys + assert all(op["increment_value"] == 0.5 for op in increment_list) + for key in expected_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( + monkeypatch, +): + """A failing pipeline must invalidate every pending counter so the next + request reseeds from the DB (which already holds this request's cost) + instead of trusting a value the write may have partially applied.""" + from redis.exceptions import MaxConnectionsError + + counter_cache = ps.DualCache() + pending_keys = ( + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + ) + for key in pending_keys: + counter_cache.in_memory_cache.set_cache(key=key, value=1.0) + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm + fake_redis.async_increment_pipeline = AsyncMock(side_effect=MaxConnectionsError()) + fake_redis.async_increment = AsyncMock() + fake_redis.async_delete_cache = AsyncMock() + fake_redis.get_ttl = MagicMock(return_value=None) + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", None) + + with pytest.raises(MaxConnectionsError): + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + assert fake_redis.async_increment.await_count == 0 + deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} + assert deleted_keys == set(pending_keys) + for key in pending_keys: + assert counter_cache.in_memory_cache.get_cache(key=key) is None # --------------------------------------------------------------------------- @@ -781,9 +887,7 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( @pytest.mark.asyncio async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set_when_none(): - result = await ps._reconcile_budget_reservation_for_counter_update( - budget_reservation=None, response_cost=1.0 - ) + result = await ps._reconcile_budget_reservation_for_counter_update(budget_reservation=None, response_cost=1.0) assert result == set() @@ -818,179 +922,151 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat # --------------------------------------------------------------------------- -# _increment_end_user_and_tag_spend_counters +# _prepare_end_user_and_tag_spend_increments # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_increments_each_unique_tag( +async def test_prepare_end_user_and_tag_spend_increments_returns_each_unique_tag( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=3.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id="eu1", tags=["a", "b", "a", "", None], response_cost=3.0, reserved_counter_keys=set(), ) - observed = { - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - "called": fake_cache.redis_cache.async_increment.called, - } - assert normalize(observed) == { - "increment_calls": 3, - "in_memory_set_calls": 3, - "called": True, + assert {item.counter_key for item in pending} == { + "spend:end_user:eu1", + "spend:tag:a", + "spend:tag:b", } + assert all(item.increment == 3.0 for item in pending) @pytest.mark.asyncio -async def test_increment_end_user_and_tag_spend_counters_no_end_user_no_tags_invalid_input_noop( +async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_invalid_input_noop( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_end_user_and_tag_spend_counters( + pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id=None, tags=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _increment_org_spend_counter +# _prepare_org_spend_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_increment_org_spend_counter_increments_when_org_present(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=10.0 - ) +async def test_prepare_org_spend_increment_returns_pending_when_org_present(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=1.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id="org-1", response_cost=10.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "counter_key_arg": fake_cache.redis_cache.async_increment.call_args.kwargs[ - "key" - ], - } - assert normalize(observed) == { - "increment_called": True, - "increment_calls": 1, - "counter_key_arg": "spend:org:org-1", - } + assert len(pending) == 1 + assert pending[0].counter_key == "spend:org:org-1" + assert pending[0].increment == 10.0 @pytest.mark.asyncio -async def test_increment_org_spend_counter_no_org_is_noop_invalid_id(monkeypatch): +async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._increment_org_spend_counter( + pending = await ps._prepare_org_spend_increment( org_id=None, response_cost=1.0, reserved_counter_keys=set(), ) + assert pending == () assert fake_cache.redis_cache.async_increment.called is False # --------------------------------------------------------------------------- -# _init_and_increment_unreserved_spend_counter +# _prepare_unreserved_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_skips_reserved_keys( +async def test_prepare_unreserved_spend_counter_increment_skips_reserved_keys( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:x", source_cache_key="tag:x", increment=1.0, reserved_counter_keys={"spend:tag:x"}, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False @pytest.mark.asyncio -async def test_init_and_increment_unreserved_spend_counter_proceeds_when_not_reserved( +async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_reserved( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=2.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None) fake_user_cache = _make_user_api_key_cache() + reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_unreserved_spend_counter( + pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:y", source_cache_key="tag:y", increment=2.0, reserved_counter_keys=set(), ) - observed = { - "increment_called": fake_cache.redis_cache.async_increment.called, - "redis_get_called": fake_cache.redis_cache.async_get_cache.called, - "reseed_consulted": True, - } - assert observed == { - "increment_called": True, - "redis_get_called": True, - "reseed_consulted": True, - } + assert pending is not None + assert pending.counter_key == "spend:tag:y" + assert pending.increment == 2.0 + assert fake_cache.redis_cache.async_get_cache.called is True + assert reseed.called is True # --------------------------------------------------------------------------- -# _init_and_increment_spend_counter +# _prepare_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=11.0, redis_increment_value=14.0 - ) +async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypatch): + fake_cache = _make_spend_counter_cache(redis_get_value=11.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -998,12 +1074,14 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) - await ps._init_and_increment_spend_counter( + pending = await ps._prepare_spend_counter_increment( counter_key="spend:key:k", source_cache_key="k", increment=3.0, ) + assert pending.counter_key == "spend:key:k" + assert pending.increment == 3.0 observed = { "reseed_called": reseed.called, "increment_called": fake_cache.redis_cache.async_increment.called, @@ -1011,23 +1089,21 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa } assert normalize(observed) == { "reseed_called": False, - "increment_called": True, + "increment_called": False, "in_memory_seeded_from_redis": True, } # --------------------------------------------------------------------------- -# _init_and_increment_window_spend_counter +# _prepare_window_spend_counter_increment # --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_increments_when_initialized( +async def test_prepare_window_spend_counter_increment_returns_pending_when_initialized( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=0.0, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=0.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( @@ -1036,7 +1112,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ AsyncMock(return_value=0.0), ) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1045,26 +1121,19 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ increment=5.0, ) - observed = { - "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "in_memory_set_calls": fake_cache.in_memory_cache.set_cache.call_count, - } - assert normalize(observed) == { - "redis_increment_called": True, - "increment_calls": 1, - "in_memory_set_calls": 2, - } + assert pending is not None + assert pending.counter_key == "spend:key:k:window:1d" + assert pending.increment == 5.0 @pytest.mark.asyncio -async def test_init_and_increment_window_spend_counter_missing_window_start_invalid_skips( +async def test_prepare_window_spend_counter_increment_missing_window_start_invalid_skips( monkeypatch, ): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._init_and_increment_window_spend_counter( + pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", @@ -1073,6 +1142,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva increment=5.0, ) + assert pending is None assert fake_cache.redis_cache.async_increment.called is False @@ -1114,16 +1184,12 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=7.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0) fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0}) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._ensure_spend_counter_initialized( counter_key="spend:user:u", @@ -1163,9 +1229,7 @@ async def test_get_source_cache_base_spend_reads_first_hit_from_list(monkeypatch fake_user_cache.async_get_cache = AsyncMock(side_effect=_get) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) - result = await ps._get_source_cache_base_spend( - source_cache_key=["miss", "hit-obj", "miss2"] - ) + result = await ps._get_source_cache_base_spend(source_cache_key=["miss", "hit-obj", "miss2"]) observed = { "result": result, @@ -1294,9 +1358,7 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey fake_cache = _make_spend_counter_cache(redis_increment_value=44.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=4.0 - ) + result = await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=4.0) observed = { "result": result, @@ -1314,15 +1376,11 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_increment_side_effect=RuntimeError("incr fail") - ) + fake_cache = _make_spend_counter_cache(redis_increment_side_effect=RuntimeError("incr fail")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) with pytest.raises(RuntimeError): - await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=1.0 - ) + await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=1.0) assert fake_cache.in_memory_cache.delete_cache.called is True assert fake_cache.redis_cache.async_delete_cache.called is True @@ -1343,9 +1401,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) observed = { "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called, "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called, - "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs[ - "key" - ], + "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"], } assert normalize(observed) == { "in_memory_delete_called": True, @@ -1357,9 +1413,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) @pytest.mark.asyncio async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch): fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_delete_cache = AsyncMock( - side_effect=RuntimeError("redis down") - ) + fake_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) await ps._invalidate_spend_counter(counter_key="spend:key:k") diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py index c123eeeed36..6165af4920d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py @@ -41,6 +41,15 @@ class _FlakyRedisCache: self._store[key] = float(value) return True + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs): + return None + @pytest.mark.asyncio async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failure( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6dab054d8ea..40ebc03781c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2220,6 +2220,15 @@ class _ExpiringRedisCache: async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: self.store.pop(key, None) + async def async_increment_pipeline(self, increment_list, **kwargs): + results = [] + for op in increment_list: + results.append(await self.async_increment(op["key"], op["increment_value"])) + return results + + def get_ttl(self, **kwargs) -> None: + return None + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index be666607823..50b26577e5c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ import copy import datetime import json from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Optional +from typing import AsyncGenerator, Callable, Final, Iterator, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -417,7 +417,7 @@ class TestProxyBaseLLMRequestProcessing: ) fake_llm_router = MagicMock() - fake_llm_router.get_model_list.return_value = [ + fake_llm_router.deployments_for_request.return_value = [ { "model_name": "smart-router", "litellm_params": { @@ -8307,3 +8307,116 @@ async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status assert exc_info.value.headers is not None assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500" + + +class TestBackgroundResponseRetrievalGovernance: + """LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines.""" + + GOVERNED_MODEL_GROUP = "gpt-5.4-mini" + GOVERNED_MODEL_ID = "deployment-governed" + + @pytest.fixture + def policy_engine(self) -> Iterator[None]: + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + get_policy_registry().load_policies( + { + "response-governance": { + "guardrails": {"add": ["output-word-filter"]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": "output-word-filter", "on_pass": "allow", "on_fail": "block"}], + }, + } + } + ) + get_attachment_registry().load_attachments( + [{"policy": "response-governance", "models": [self.GOVERNED_MODEL_GROUP]}] + ) + yield + get_policy_registry().clear() + get_attachment_registry().clear() + + def _router(self) -> MagicMock: + from litellm.types.router import Deployment, LiteLLM_Params + + router = MagicMock() + router.get_deployment.side_effect = lambda model_id: ( + Deployment( + model_name=self.GOVERNED_MODEL_GROUP, + litellm_params=LiteLLM_Params(model=f"openai/{self.GOVERNED_MODEL_GROUP}"), + model_info={"id": model_id}, + ) + if model_id == self.GOVERNED_MODEL_ID + else None + ) + return router + + async def _pre_call(self, route_type: str, monkeypatch: pytest.MonkeyPatch) -> dict[str, object]: + from litellm.responses.utils import ResponsesAPIRequestUtils + + client_facing_response_id = "resp_opaque-client-facing-id" + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=self.GOVERNED_MODEL_ID, response_id="resp_upstream" + ) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"response_id": client_facing_response_id, "litellm_metadata": {}} + ) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def passthrough_add_litellm_data_to_request( + data: dict[str, object], **kwargs: object + ) -> dict[str, object]: + return data + + async def decrypting_pre_call_hook( + user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> dict[str, object]: + if data.get("response_id") == client_facing_response_id: + data["response_id"] = encoded_response_id + return data + + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + passthrough_add_litellm_data_to_request, + ) + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=decrypting_pre_call_hook) + proxy_config = MagicMock(spec=ProxyConfig) + proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(), + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type=route_type, + llm_router=self._router(), + ) + return returned_data + + @pytest.mark.asyncio + async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline( + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + data = await self._pre_call("aget_responses", monkeypatch) + + assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f") + pipelines = data["litellm_metadata"]["_guardrail_pipelines"] + assert [(policy_name, [step.guardrail for step in pipeline.steps]) for policy_name, pipeline in pipelines] == [ + ("response-governance", ["output-word-filter"]) + ] + assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] + assert data["model"] is None + + @pytest.mark.asyncio + async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id( + self, policy_engine: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + data = await self._pre_call("aresponses", monkeypatch) + + assert "_guardrail_pipelines" not in data["litellm_metadata"] + assert "applied_policies" not in data["litellm_metadata"] diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 892fa484ab4..94aaa32519e 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -35,6 +35,7 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.litellm_core_utils.redact_messages import _get_turn_off_message_logging_from_dynamic_params from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -7678,6 +7679,107 @@ async def test_missing_session_id_omit_keeps_client_supplied_session_id(): assert _spend_log_session_id(updated) == "client-session-1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_body", + [ + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1"}, + {"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1", "metadata": {"trace_id": "trace-1"}}, + ], +) +async def test_missing_session_id_omit_keeps_body_litellm_session_id( + monkeypatch: pytest.MonkeyPatch, client_body: dict[str, object] +): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + updated = await add_litellm_data_to_request( + data=client_body, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + callback_session_id = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=SimpleNamespace(litellm_session_id=""), + litellm_params=get_litellm_params(litellm_session_id="cust-sess-1", metadata=updated["metadata"]), + ) + assert callback_session_id == "cust-sess-1" + assert updated["metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated) == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_body_litellm_session_id_does_not_override_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "messages": [], + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated) == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_metadata_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "input": "hi", + "litellm_session_id": "cust-sess-1", + "metadata": {"session_id": "meta-sess-1"}, + }, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "meta-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "meta-sess-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +async def test_missing_session_id_omit_keeps_body_litellm_session_id_on_litellm_metadata_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "input": "hi", "litellm_session_id": "cust-sess-1"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "cust-sess-1" + assert _spend_log_session_id(updated, "litellm_metadata") == "cust-sess-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_ignores_empty_body_litellm_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "litellm_session_id": ""}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert "session_id" not in updated["metadata"] + assert _spend_log_session_id(updated) is None + + @pytest.mark.asyncio async def test_missing_session_id_generate_reuses_traceparent_trace_id(): """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" @@ -7808,3 +7910,36 @@ async def test_client_supplied_omit_marker_never_reaches_the_spend_log( if general_settings.get("missing_session_id") == "generate" else "per-call-random-trace-id" ) + + +def test_default_team_settings_bool_turn_off_message_logging_redacts(): + from litellm.proxy.proxy_server import ProxyConfig + + pc = ProxyConfig() + pc.config = { + "litellm_settings": { + "default_team_settings": [ + { + "team_id": "team-redact", + "success_callback": ["gcs_bucket"], + "failure_callback": ["gcs_bucket"], + "turn_off_message_logging": True, + } + ] + } + } + + callback_metadata = LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config( + team_id="team-redact", + proxy_config=pc, + ) + + assert callback_metadata is not None + assert callback_metadata.success_callback == ["gcs_bucket"] + assert callback_metadata.callback_vars == {"turn_off_message_logging": "True"} + assert ( + _get_turn_off_message_logging_from_dynamic_params( + {"standard_callback_dynamic_params": dict(callback_metadata.callback_vars)} + ) + is True + ) diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 22930a26974..28ff4571b44 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -671,6 +671,95 @@ async def test_deferred_stream_guardrails_run_native_hook_when_opted_out(monkeyp assert routed.native_hooks_ran == [] +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monkeypatch): + """A post_call pipeline step already ran the opted-out guardrail's own hook against + the buffered stream, so the deferred audit must not run it a second time.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == [] + + +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch): + """A pipeline step with neither streaming interface keeps the whole pipeline off the + stream, so the deferred audit is the only place the opted-out guardrail's own hook + still runs, the way it did before pipelines ran on streams.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + class NeitherHookGuardrail(CustomGuardrail): + pass + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither]) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="keeps_native", on_fail="next"), + PipelineStep(guardrail="gr-neither", on_fail="block"), + ], + ) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch): + """A route with no endpoint guardrail translation cannot gate the stream through its + pipelines, so the deferred audit still owes the opted-out guardrail its own hook.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + @pytest.mark.asyncio async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch): """The realtime path calls apply_guardrail directly, so the opt-out has to be diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 54579e6cb7c..e058a4f6396 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6,6 +6,7 @@ import os import re import socket import subprocess +import time import types from datetime import datetime, timedelta, timezone from pathlib import Path @@ -28,7 +29,7 @@ from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map @@ -7749,7 +7750,7 @@ async def test_increment_spend_counters_team_and_member(): @pytest.mark.asyncio -async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss(): +async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss(): """When the Redis counter is missing, the reseed path reads the authoritative spend from the DB (not a stale cache), so the next increment continues from the correct base value.""" @@ -7762,8 +7763,17 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( recorded_increments.append({"key": key, "value": value, "ttl": ttl}) return value + async def record_pipeline(increment_list, **kwargs): + results = [] + for op in increment_list: + await record_increment(key=op["key"], value=op["increment_value"], ttl=op["ttl"]) + results.append(op["increment_value"]) + return results + fake_redis = AsyncMock() fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_increment_pipeline = AsyncMock(side_effect=record_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing fake_redis.async_set_cache = AsyncMock(return_value=True) # SET NX wins counter_cache.redis_cache = fake_redis @@ -7782,7 +7792,10 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( stale_cache.in_memory_cache.set_cache(key="team_id:team-9", value=stale_team) import litellm.proxy.proxy_server as ps - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) orig_user, orig_counter, orig_prisma = ( ps.user_api_key_cache, @@ -7793,11 +7806,12 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key="spend:team:team-9", source_cache_key="team_id:team-9", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. @@ -7976,7 +7990,10 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): @pytest.mark.asyncio async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) @@ -7992,7 +8009,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", @@ -8000,6 +8017,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -8015,7 +8033,10 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): @pytest.mark.asyncio async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:team:team-stale-local" @@ -8037,6 +8058,15 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis db_row = MagicMock() @@ -8055,11 +8085,12 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.prisma_client = fake_prisma ps.user_api_key_cache = DualCache() try: - await _init_and_increment_spend_counter( + pending = await _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="team_id:team-stale-local", increment=1.5, ) + await _apply_spend_counter_increments(pending=(pending,)) fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. @@ -8074,7 +8105,10 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-stale-local:window:1h" @@ -8097,6 +8131,15 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8111,7 +8154,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", @@ -8119,6 +8162,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( by=["api_key"], @@ -8138,7 +8182,10 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import ( + _apply_spend_counter_increments, + _prepare_window_spend_counter_increment, + ) counter_cache = DualCache() counter_key = "spend:key:key-window-concurrent-seed:window:1h" @@ -8161,6 +8208,15 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) fake_redis.async_set_cache = AsyncMock(return_value=False) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + + async def redis_increment_pipeline(increment_list, **_): + results = [] + for op in increment_list: + results.append(await redis_increment(key=op["key"], value=op["increment_value"])) + return results + + fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) + fake_redis.get_ttl = MagicMock(return_value=None) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -8175,7 +8231,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", @@ -8183,6 +8239,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() window_start=window_start, increment=0.5, ) + await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) fake_redis.async_set_cache.assert_awaited_once_with( key=counter_key, @@ -8199,7 +8256,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() @pytest.mark.asyncio async def test_window_spend_counter_skips_invalid_window_start(): from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + from litellm.proxy.proxy_server import _prepare_window_spend_counter_increment counter_cache = DualCache() @@ -8208,7 +8265,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): orig_counter = ps.spend_counter_cache ps.spend_counter_cache = counter_cache try: - await _init_and_increment_window_spend_counter( + pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", @@ -8216,6 +8273,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): window_start=None, increment=0.5, ) + assert pending is None assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-invalid-window:window:not-a-duration") is None finally: @@ -8279,6 +8337,9 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): async def assert_reservation_not_finalized_yet(**kwargs): assert budget_reservation["finalized"] is False incremented_counters.append(kwargs["counter_key"]) + return ps._PendingSpendIncrement( + counter_key=kwargs["counter_key"], increment=kwargs["increment"] + ) import litellm.proxy.proxy_server as ps @@ -8287,7 +8348,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): ps.user_api_key_cache = DualCache() try: with patch( - "litellm.proxy.proxy_server._init_and_increment_spend_counter", + "litellm.proxy.proxy_server._prepare_spend_counter_increment", new=AsyncMock(side_effect=assert_reservation_not_finalized_yet), ): await increment_spend_counters( @@ -8620,7 +8681,7 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): async def test_concurrent_read_and_write_paths_share_one_db_query(): """ The read path (`get_current_spend`) and the write path - (`_init_and_increment_spend_counter`) both reseed cold counters from + (`_prepare_spend_counter_increment`) both reseed cold counters from the DB. They must share the per-counter lock so a concurrent pre-call enforcement read and post-call increment for the same counter collapse to one DB query, not two. @@ -8629,7 +8690,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import ( - _init_and_increment_spend_counter, + _prepare_spend_counter_increment, get_current_spend, ) @@ -8683,7 +8744,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): try: results = await _asyncio.gather( get_current_spend(counter_key=counter_key, fallback_spend=0.0), - _init_and_increment_spend_counter( + _prepare_spend_counter_increment( counter_key=counter_key, source_cache_key="ignored", increment=1.5, @@ -12894,3 +12955,59 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() import litellm.proxy.proxy_server as ps assert ps.general_settings["enable_openai_websocket_passthrough"] is False + + +async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + warm_tokenizer("claude-fable-5") + + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=text * 100)) + ) + + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) + + +async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch): + from tokenizers import Tokenizer + + from litellm import Router + from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags + + claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + + class SlowHubTokenizer: + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + time.sleep(0.3) + return claude_tokenizer + + monkeypatch.setattr(litellm.utils, "Tokenizer", SlowHubTokenizer) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + { + "model_name": "self-hosted", + "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, + "model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}}, + } + ] + ), + ) + + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop")) + ) + + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9462f2c8eb0..b78ec7dcff6 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1823,6 +1823,44 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): assert lifted["standard_logging_object"] == {"id": "log-1"} +@pytest.mark.asyncio +async def test_a_dispatched_failure_is_counted_off_the_event_loop(): + from unittest.mock import AsyncMock, patch + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("claude-fable-5") + request_data = { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "call_type": "acompletion", + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + ), + "metadata": {}, + } + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + _, took, lags = await timed_with_loop_lags( + lambda: proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + ) + + assert request_data["combined_usage_object"].prompt_tokens > 0 + assert_loop_stayed_free(took, lags) + + @pytest.mark.asyncio async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch): """Regression for LIT-6043: an expected 4xx must not format a traceback for diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 5cb595840fc..dfe106a3f52 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,7 +11,9 @@ from __future__ import annotations import asyncio import json +from copy import deepcopy import logging +from collections.abc import Iterator from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -25,11 +27,13 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header -from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines +from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -154,9 +158,7 @@ async def test_execute_guardrail_hook_unknown_hook_type_raises(proxy_logging, ma @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_routes_through_router( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_routes_through_router(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": cb}) @@ -172,9 +174,7 @@ async def test_execute_guardrail_with_load_balancing_routes_through_router( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_router_none_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_router_none_raises(proxy_logging, make_user_api_key_auth): with patch("litellm.proxy.proxy_server.llm_router", None): with pytest.raises(ValueError, match="Router not initialized"): await proxy_logging._execute_guardrail_with_load_balancing( @@ -187,9 +187,7 @@ async def test_execute_guardrail_with_load_balancing_router_none_raises( @pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_no_callback_raises( - proxy_logging, make_user_api_key_auth -): +async def test_execute_guardrail_with_load_balancing_no_callback_raises(proxy_logging, make_user_api_key_auth): router = MagicMock() router.get_available_guardrail = MagicMock(return_value={"callback": None}) with patch("litellm.proxy.proxy_server.llm_router", router): @@ -209,9 +207,7 @@ async def test_execute_guardrail_with_load_balancing_no_callback_raises( @pytest.mark.asyncio -async def test_process_guardrail_callback_skipped_when_should_run_false( - proxy_logging, make_user_api_key_auth -): +async def test_process_guardrail_callback_skipped_when_should_run_false(proxy_logging, make_user_api_key_auth): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=False) out = await proxy_logging._process_guardrail_callback( @@ -225,9 +221,7 @@ async def test_process_guardrail_callback_skipped_when_should_run_false( @pytest.mark.asyncio -async def test_process_guardrail_callback_returns_data_on_success( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_process_guardrail_callback_returns_data_on_success(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=True) proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False) @@ -342,14 +336,14 @@ async def test_maybe_execute_pipelines_no_pipelines_returns_data(proxy_logging, @pytest.mark.asyncio -async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode( + proxy_logging, make_user_api_key_auth, monkeypatch +): pipeline = MagicMock() pipeline.mode = "post_call" # not pre_call data = {"metadata": {"_guardrail_pipelines": [("p1", pipeline)]}, "model": "m", "messages": []} executed = MagicMock() - monkeypatch.setattr( - "litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed - ) + monkeypatch.setattr("litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed) out, replacement = await proxy_logging._maybe_execute_pipelines( data=data, user_api_key_dict=make_user_api_key_auth(), @@ -537,9 +531,7 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode(): litellm.callbacks = [cb] try: with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result( - result=result, data={"model": "m"}, policy_name="p" - ) + ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") finally: litellm.callbacks = saved @@ -651,9 +643,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch monkeypatch.setattr(litellm, "callbacks", [prom]) with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_with_metrics( - callback=cb, coro=task(), hook_type="post_call" - ) + await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call") assert detail["guardrail_name"] == "presidio" recorded = prom._record_guardrail_metrics.call_args.kwargs @@ -681,9 +671,7 @@ def _moderation_guardrail() -> MagicMock: @pytest.mark.asyncio -async def test_during_call_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -702,9 +690,7 @@ async def test_during_call_hook_records_latency_metric( @pytest.mark.asyncio -async def test_post_call_success_hook_records_latency_metric( - proxy_logging, make_user_api_key_auth, monkeypatch -): +async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail() prom = _prometheus_callback() monkeypatch.setattr(litellm, "callbacks", [prom, cb]) @@ -732,9 +718,7 @@ async def test_post_call_success_hook_records_latency_metric( async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, monkeypatch): from litellm.proxy.prompts import prompt_registry - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None) data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} await proxy_logging._process_prompt_template( data=data, @@ -759,9 +743,7 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging, "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -809,9 +791,7 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt")) with pytest.raises(RuntimeError): @@ -912,9 +892,7 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec - ) + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock( @@ -965,6 +943,7 @@ def _post_call_pipeline_data( "metadata": { "_guardrail_pipelines": [("response-governance", pipeline)], "_pipeline_managed_guardrails": {guardrail}, + "policy_sources": {"response-governance": "model:m"}, }, **extra, } @@ -1123,9 +1102,7 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe }, } - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" - ) + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") assert seen["count"] == 1 @@ -1288,11 +1265,7 @@ async def test_post_call_pipeline_block_keeps_guardrail_metadata_writes( monkeypatch.setattr( litellm, "callbacks", - [ - BlockingWriterGuardrail( - guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False - ) - ], + [BlockingWriterGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data() @@ -1378,9 +1351,7 @@ async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once( }, } - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" - ) + await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion") assert seen["count"] == 1 @@ -1419,10 +1390,295 @@ async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_ver assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog)) +def _background_response(status: str, text: str = "") -> ResponsesAPIResponse: + output = ( + [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}] if text else [] + ) + return ResponsesAPIResponse(id="resp_bg", created_at=0, output=output, status=status) + + +def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail]: + class OutputBlockingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + return [ + OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] + + @pytest.mark.asyncio -async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog +@pytest.mark.parametrize("pending_status", ["queued", "in_progress"]) +async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + pending_status: str, +) -> None: + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(background=True) + response = _background_response(pending_status) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert out is response + assert "response" not in seen + assert not _warnings(caplog) + assert any( + "response-governance" in record.getMessage() and pending_status in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_status", ["completed", "incomplete"]) +async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + final_status: str, +) -> None: + seen: dict[str, object] = {} + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen)) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = _background_response(final_status, text="kumquat") + + with pytest.raises(HTTPException) as info: + await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert info.value.detail["error"] == "output blocked" + assert seen["response"] is response + + +def _output_passing_callbacks() -> list[CustomGuardrail]: + class OutputPassingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + return [ + OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + ] + + +def _claimed_post_call_pipeline_data( + *policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str | None = "model:m" ): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"} + get_policy_registry().load_policies( + { + policy_name: { + "guardrails": {"add": ["gr-post", *(extra_guardrails or {}).get(policy_name, [])]}, + "pipeline": {"mode": "post_call", "steps": [step]}, + } + for policy_name in policy_names + } + ) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(**step)]) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [(policy_name, pipeline) for policy_name in policy_names], + "_pipeline_managed_guardrails": {"gr-post"}, + "applied_policies": list(policy_names), + "applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)], + "policy_sources": {policy_name: policy_source for policy_name in policy_names if policy_source is not None}, + }, + } + + +@pytest.fixture +def clear_policy_registry() -> Iterator[None]: + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + yield + get_policy_registry().clear() + + +@pytest.mark.asyncio +async def test_pending_background_response_withdraws_the_deferred_policy_claims( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "applied_policies" not in data["metadata"] + assert "policy_sources" not in data["metadata"] + assert "applied_guardrails" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert [message for message in _warnings(caplog) if "through a request tag" in message] == [ + "Policy engine: background response resp_bg matched post_call policies through a request tag at submit; " + "retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body " + "does not govern the completed response: response-governance" + ] + + +@pytest.mark.asyncio +async def test_pending_background_response_warns_when_the_deferred_policy_came_from_the_request_body( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("body-governance", policy_source=None) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert out.status == "queued" + assert "policy_sources" not in data["metadata"] + assert _warnings(caplog) == [ + "Policy engine: background response resp_bg matched post_call policies through the request body's policies " + "list at submit; retrieval carries no request body, so those policies do not govern the completed " + "response: body-governance" + ] + + +@pytest.mark.asyncio +async def test_pending_background_response_matched_through_its_model_does_not_warn( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await proxy_logging.post_call_success_hook( + data=_claimed_post_call_pipeline_data("response-governance"), + response=_background_response("queued"), + user_api_key_dict=make_user_api_key_auth(), + ) + + assert _warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({})) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data( + "input-and-output-governance", + "response-governance", + extra_guardrails={"input-and-output-governance": ["gr-pre"]}, + ) + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("in_progress"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["input-and-output-governance"] + assert data["metadata"]["applied_guardrails"] == ["gr-pre"] + assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"} + + +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + class DualStageGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-post", event_hook=["pre_call", "post_call"], default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert "applied_policies" not in data["metadata"] + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + +@pytest.mark.asyncio +async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + clear_policy_registry: None, +) -> None: + monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("completed", text="fine"), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_policies"] == ["response-governance"] + assert data["metadata"]["policy_sources"] == {"response-governance": "model:m"} + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: monkeypatch.setattr(litellm, "callbacks", []) data = _post_call_pipeline_data(background=True) @@ -1436,36 +1692,7 @@ async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline( assert out is not None assert out.get("background") is True - assert any("response-governance" in message and "background" in message for message in _warnings(caplog)) - - -@pytest.mark.asyncio -async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch, caplog -): - seen: Dict[str, Any] = {} - monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) - pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) - data = { - "model": "m", - "messages": [{"role": "user", "content": "hi"}], - "background": True, - "metadata": { - "_guardrail_pipelines": [("request-governance", pre_call)], - "_pipeline_managed_guardrails": {"gr-post"}, - }, - } - - with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), - data=data, - call_type="aresponses", - guardrails_only=True, - ) - - assert out is not None - assert not any("background" in message for message in _warnings(caplog)) + assert not _warnings(caplog) # --------------------------------------------------------------------------- @@ -1497,29 +1724,131 @@ async def _async_chunk_iter(chunks: List[Any]): yield chunk -def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported( +def _legacy_hook_stream_guardrail( + seen: Dict[str, Any], + rewrite: Callable[[Any], Any] | None = None, + raises: Exception | None = None, + native_lifecycle: bool = False, +) -> CustomGuardrail: + class LegacyHookGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = native_lifecycle + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + seen["data"] = data + seen["user_api_key_dict"] = user_api_key_dict + seen["response"] = deepcopy(response) + if raises is not None: + raise raises + return None if rewrite is None else rewrite(response) + + if native_lifecycle: + + class NativeLifecycleGuardrail(LegacyHookGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + return NativeLifecycleGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) + return LegacyHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["count"] = seen.get("count", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + +def _iterator_and_legacy_hook_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorAndLegacyHookGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["success_hook_calls"] = seen.get("success_hook_calls", 0) + 1 + return None + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["iterator_hook_calls"] = seen.get("iterator_hook_calls", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorAndLegacyHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + +def _rewritten_model_response(response: Any) -> litellm.ModelResponse: + payload = response.model_dump() + payload["choices"][0]["message"]["content"] = "[REWRITTEN] " + payload["choices"][0]["message"]["content"] + return litellm.ModelResponse(**payload) + + +def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only( make_user_api_key_auth, monkeypatch, caplog ): - class NativeOnlyGuardrail(CustomGuardrail): - pass - supported = _unified_stream_guardrail({}) - native_only = NativeOnlyGuardrail(guardrail_name="gr-native", event_hook=GuardrailEventHooks.post_call) - monkeypatch.setattr(litellm, "callbacks", [supported, native_only]) - governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + legacy = _legacy_hook_stream_guardrail({}) + legacy.guardrail_name = "gr-legacy" + iterator_only = _iterator_hook_only_guardrail("gr-iterator", {}) + monkeypatch.setattr(litellm, "callbacks", [supported, legacy, iterator_only]) + governed = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-legacy", on_fail="block")], + ) ungoverned = GuardrailPipeline( mode="post_call", - steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")], + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-iterator", on_fail="block")], ) - pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")]) - data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}} + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-iterator", on_fail="block")]) + data = { + "metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]} + } with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions")) assert streamable == (("governed", governed),) - assert any("'ungoverned'" in message and "gr-native" in message for message in _warnings(caplog)) - assert not any("'governed'" in message for message in _warnings(caplog)) + assert any("'ungoverned'" in message and "gr-iterator" in message for message in _warnings(caplog)) + assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog)) + + +@pytest.mark.parametrize( + "request_route", + ["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], +) +def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response( + make_user_api_key_auth, monkeypatch, caplog, request_route +): + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail({})]) + legacy = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("legacy-governance", legacy)]}} + auth = make_user_api_key_auth(request_route=request_route) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'legacy-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_keeps_guardrails_with_their_own_iterator_hook_on_their_own_path( + make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", {})]) + both_hooks = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("both-hooks", both_hooks)]}} + auth = make_user_api_key_auth(request_route="/v1/chat/completions") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'both-hooks'" in message and "gr-post" in message for message in _warnings(caplog)) def test_streamable_post_call_pipelines_is_empty_on_route_without_translation( @@ -1569,55 +1898,123 @@ async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_u @pytest.mark.asyncio @pytest.mark.parametrize("native_lifecycle", [False, True]) -async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_lacks_unified_support( +async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite( proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog ): seen: Dict[str, Any] = {} - if native_lifecycle: - - class NativeOnlyGuardrail(CustomGuardrail): - use_native_lifecycle_hooks = True - - async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): - seen["count"] = seen.get("count", 0) + 1 - return inputs - - else: - - class NativeOnlyGuardrail(CustomGuardrail): - async def async_post_call_success_hook(self, data, user_api_key_dict, response): - seen["count"] = seen.get("count", 0) + 1 - return response - - monkeypatch.setattr( - litellm, - "callbacks", - [NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], - ) + guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_model_response, native_lifecycle=native_lifecycle) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) chunks = _stream_chunks() - delivered: List[Any] = [] + auth = make_user_api_key_auth(request_route="/v1/chat/completions") with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), - data=data, - call_type="completion", - guardrails_only=True, + user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True ) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert out is not None and out.get("stream") is True + assert seen["count"] == 1 + assert isinstance(seen["response"], litellm.ModelResponse) + assert seen["response"].choices[0].message.content == "hello world" + assert seen["data"]["messages"] == data["messages"] + assert seen["user_api_key_dict"] is auth + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert delivered[0].choices[0].delta.content == "[REWRITTEN] hello world" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + assert _warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item async for item in proxy_logging.async_post_call_streaming_iterator_hook( user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), response=_async_chunk_iter(chunks), request_data=data, + ) + ] + + assert seen["count"] == 1 + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert [item.choices[0].delta.content for item in delivered] == ["hello ", "world"] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_ends_stream_with_legacy_hook_exception( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + blocked = HTTPException(status_code=400, detail={"error": "output blocked"}) + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, raises=blocked)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, ): delivered.append(item) - assert out is not None - assert out.get("stream") is True - assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] - assert len(delivered) == 2 - assert seen.get("count") is None - assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + with pytest.raises(HTTPException) as info: + await _drain() + + assert seen["count"] == 1 + assert delivered == [] + assert info.value is blocked + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + def rewrite(response: Any) -> Dict[str, Any]: + return {**response, "content": [{"type": "text", "text": "[REWRITTEN] " + response["content"][0]["text"]}]} + + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=rewrite)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_sse_chunks()), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert seen["response"]["content"][0]["text"] == "hello world" + assert seen["response"]["role"] == "assistant" + raw = b"".join(delivered).decode() + assert "[REWRITTEN] hello world" in raw + assert raw.count("event: content_block_delta") == 1 + for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + assert f"event: {expected_event}" in raw @pytest.mark.asyncio @@ -1625,19 +2022,7 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} - - class IteratorHookGuardrail(CustomGuardrail): - async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): - seen["count"] = seen.get("count", 0) + 1 - async for item in response: - item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" - yield item - - monkeypatch.setattr( - litellm, - "callbacks", - [IteratorHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], - ) + monkeypatch.setattr(litellm, "callbacks", [_iterator_hook_only_guardrail("gr-post", seen)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) @@ -1656,6 +2041,30 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_runs_the_iterator_hook_of_a_guardrail_that_also_has_a_post_call_hook( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ) + ] + + assert seen == {"iterator_hook_calls": 1} + assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"] + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "rewrite_attribute, value", @@ -1812,7 +2221,9 @@ def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str, async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): return {**inputs, **transform(inputs)} - return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + return RewritingStreamGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) def _tool_call_stream_chunks() -> List[Any]: @@ -1836,7 +2247,7 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: @pytest.mark.asyncio @pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) -async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_tool_call_rewrite( +async def test_streaming_iterator_hook_pipeline_delivers_runtime_tool_call_rewrite( proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog ): transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731 @@ -1855,9 +2266,12 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_to delivered.append(item) assert len(delivered) == 2 - assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + delivered_tool_call = delivered[0].choices[0].delta.tool_calls[0] + assert delivered_tool_call.function.arguments == '{"ssn": "[MASKED]"}' + assert delivered_tool_call.function.name == "lookup" + assert delivered_tool_call.id == "call_1" assert delivered[1].choices[0].finish_reason == "tool_calls" - assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + assert not any("discarded" in message for message in _warnings(caplog)) @pytest.mark.asyncio @@ -1965,37 +2379,97 @@ async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_anothe @pytest.mark.asyncio -async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvable_response_shape( +async def test_streaming_iterator_hook_skips_pipeline_and_warns_without_request_route( proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) - chunks = [object(), object()] - delivered: List[Any] = [] + chunks = _stream_chunks() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - async for item in proxy_logging.async_post_call_streaming_iterator_hook( - user_api_key_dict=make_user_api_key_auth(), - response=_async_chunk_iter(chunks), - request_data=data, - ): - delivered.append(item) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] assert len(delivered) == 2 assert seen.get("count") is None - assert any("response-governance" in message and "shape" in message for message in _warnings(caplog)) + assert any("response-governance" in message and "route None" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_runs_pipeline_managed_guardrail_without_request_route( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class UnifiedRecordingGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 + return None + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + monkeypatch.setattr( + litellm, + "callbacks", + [UnifiedRecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(), + ) + + assert result is not None + assert seen["gr-post"] == 1 def _anthropic_sse_chunks() -> List[bytes]: events = [ - ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), - ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}), + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "m", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}, + ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), - ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + ), ("message_stop", {"type": "message_stop"}), ] return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] @@ -2062,7 +2536,13 @@ async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthrop assert "hello [MASKED]" in raw assert "hello world" not in raw assert raw.count("event: content_block_delta") == 1 - for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + for expected_event in ( + "message_start", + "content_block_start", + "content_block_stop", + "message_delta", + "message_stop", + ): assert f"event: {expected_event}" in raw @@ -2135,9 +2615,7 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( managed = UnifiedRecordingGuardrail( guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True ) - free = RecordingGuardrail( - guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True - ) + free = RecordingGuardrail(guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True) monkeypatch.setattr(litellm, "callbacks", [managed, free]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) @@ -2182,3 +2660,177 @@ async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_str assert result is not None assert seen["count"] == 1 assert seen["response"] == "hello " + + +def _mask_tool_call_arguments(inputs: Dict[str, Any]) -> Dict[str, Any]: + return { + "tool_calls": [ + { + "id": stream_item_field(tool_call, "id"), + "type": "function", + "function": { + "name": stream_item_field(stream_item_field(tool_call, "function"), "name"), + "arguments": '{"fruit": "[MASKED]"}', + }, + } + for tool_call in inputs.get("tool_calls", []) + ] + } + + +def _anthropic_tool_use_sse_chunks() -> List[bytes]: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_tool_use_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert '{\\"fruit\\": \\"[MASKED]\\"}' in raw + assert "persim" not in raw + assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw + assert '"stop_reason": "tool_use"' in raw + assert raw.count("event: content_block_delta") == 2 + + +def _responses_function_call_events() -> List[Dict[str, Any]]: + def item(arguments: str, status: str) -> Dict[str, Any]: + return { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup_fruit", + "arguments": arguments, + "status": status, + } + + return [ + {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": '{"fruit":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": ' "persimmon"}'}, + {"type": "response.function_call_arguments.done", "item_id": "fc_1", "output_index": 0, "arguments": '{"fruit": "persimmon"}'}, + {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.completed", + "response": {"id": "resp_1", "created_at": 1, "model": "m", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed"}, + }, + ] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert [event["type"] for event in delivered] == [event["type"] for event in _responses_function_call_events()] + assert [event["delta"] for event in delivered if event["type"] == "response.function_call_arguments.delta"] == ['{"fruit": "[MASKED]"}', ""] + assert delivered[3]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}' + assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}' + assert "persimmon" not in json.dumps(delivered) + + +def _drop_tool_calls(inputs: Dict[str, Any]) -> Dict[str, Any]: + return {"tool_calls": []} + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_chat_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_tool_call_stream_chunks()), + request_data=data, + ) + ] + + assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + assert delivered[1].choices[0].finish_reason == "tool_calls" + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()), + request_data=data, + ) + ] + + assert delivered == _anthropic_tool_use_sse_chunks() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_responses_events( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"), + response=_async_chunk_iter(_responses_function_call_events()), + request_data=data, + ) + ] + + assert delivered == _responses_function_call_events() + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2068f10ea2d..46249e50572 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1,4 +1,5 @@ import json +from typing import Final import pytest @@ -1421,6 +1422,88 @@ class TestToolChoiceTransformation: ) assert result == "required" + @pytest.mark.parametrize( + "request_tool_choice,expected", + [ + ({"type": "function", "name": "run_command"}, {"type": "function", "name": "run_command"}), + ({"type": "function", "function": {"name": "run_command"}}, {"type": "function", "name": "run_command"}), + ({"type": "custom", "name": "ApplyPatch"}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "custom", "custom": {"name": "ApplyPatch"}}, {"type": "custom", "name": "ApplyPatch"}), + ({"type": "function"}, "required"), + ({"type": "tool"}, "required"), + ({"type": "auto"}, "auto"), + ("required", "required"), + ("none", "none"), + (None, "auto"), + ("any", "auto"), + ("run_command", "auto"), + ({"name": "run_command"}, "auto"), + ], + ) + def test_transform_tool_choice_for_responses_api_response( + self, request_tool_choice: object, expected: str | dict[str, str] + ) -> None: + result: Final = LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response( + request_tool_choice + ) + assert result == expected + + def test_non_streamed_response_echoes_named_tool_choice_in_responses_api_shape(self) -> None: + chat_completion_response: Final = ModelResponse( + id="chatcmpl-named-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_pwd", + type="function", + function=Function(name="run_command", arguments='{"command":"pwd"}'), + ) + ], + ), + ) + ], + ) + + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": {"type": "function", "name": "run_command"}}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == {"type": "function", "name": "run_command"} + + def test_non_streamed_response_with_unrecognized_tool_choice_echoes_auto(self) -> None: + chat_completion_response: Final = ModelResponse( + id="chatcmpl-unrecognized-tool-choice", + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message(role="assistant", content="/Users/dev"), + ) + ], + ) + + responses_api_response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Run the command pwd.", + responses_api_request={"tool_choice": "any"}, + chat_completion_response=chat_completion_response, + ) + + assert responses_api_response.tool_choice == "auto" + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 719d51c11e3..850ee7ba623 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -11,6 +11,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -628,3 +629,79 @@ def test_streamed_anthropic_tool_call_events_correlate_on_normalized_item_id(): assert item_dones[0].item.call_id == "toolu_01AbCdEf" for evt in deltas + dones: assert evt.item_id == added[0].item.id + + +def _tool_call_chunk(finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_pwd", + "type": "function", + "function": {"name": "run_command", "arguments": '{"command":"pwd"}'}, + "index": 0, + } + ], + ), + finish_reason=finish_reason, + ) + ], + ) + + +def test_streamed_named_tool_choice_is_echoed_in_responses_api_shape() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": {"type": "function", "name": "run_command"}, + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + events: Final = list(iterator) + + response_events: Final = [event for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] + assert [event.type for event in response_events] == [ + "response.created", + "response.in_progress", + "response.completed", + ] + assert [event.response.tool_choice for event in response_events] == [ + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + {"type": "function", "name": "run_command"}, + ] + assert any(getattr(event, "type", None) == "response.output_item.done" for event in events) + + +def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: + iterator: Final = LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper([_tool_call_chunk(finish_reason="tool_calls")]), + request_input="Run the command pwd.", + responses_api_request={ + "tools": [{"type": "function", "name": "run_command", "parameters": {"type": "object"}}], + "tool_choice": "any", + }, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + response_events: Final = [ + event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + + assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c226c0b4d09..5e0e794d93e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -18,6 +18,7 @@ from litellm.responses.streaming_iterator import ( SyncResponsesAPIStreamingIterator, ) from litellm.types.llms.openai import ( + ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, @@ -329,8 +330,6 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): def _responses_api_response_with_usage() -> ResponsesAPIResponse: - from litellm.types.llms.openai import ResponseAPIUsage - return ResponsesAPIResponse( id="resp_lit6427", created_at=int(datetime(2025, 1, 1).timestamp()), @@ -368,6 +367,53 @@ def test_stamp_responses_usage_cost_keeps_provider_reported_cost(): logging_obj._response_cost_calculator.assert_not_called() +def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: + return ResponsesAPIResponse.model_construct( + id="resp_lit7391", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="perplexity/deepseek-v4-flash-0731", + object="response", + output=[], + truncation="", + usage=usage, + ) + + +def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + response = _unvalidated_response_with_dict_usage( + { + "input_tokens": 29, + "output_tokens": 120, + "output_tokens_details": {"reasoning_tokens": 117}, + "total_tokens": 149, + "cost": {"currency": "USD", "input_cost": 0, "output_cost": 3e-05, "total_cost": 3e-05}, + } + ) + logging_obj = Mock(spec=LiteLLMLoggingObj) + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(3e-05) + assert response.usage.output_tokens_details.reasoning_tokens == 117 + logging_obj._response_cost_calculator.assert_not_called() + + +def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.return_value = 0.000704 + + _stamp_responses_usage_cost(response, logging_obj) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.cost == pytest.approx(0.000704) + logging_obj._response_cost_calculator.assert_called_once_with(result=response) + + def test_stamp_responses_usage_cost_survives_calculator_failure(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost @@ -535,5 +581,50 @@ async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): iterator._log_completed_response(is_async=True) - assert logged == [iterator.completed_response] + assert len(logged) == 1 + assert logged[0] is not iterator.completed_response + assert logged[0].response is not iterator.completed_response.response + assert logged[0].response._hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" assert iterator.completed_response.response._hidden_params == {} + + +def _unvalidated_completed_config() -> Mock: + """Config whose completed event carries a Perplexity-style response that fails validation + (``truncation: ""``) and already holds the stamped ``ResponseAPIUsage``.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + response = _unvalidated_response_with_dict_usage( + ResponseAPIUsage(input_tokens=29, output_tokens=373, total_tokens=402, cost={"total_cost": 0.0001}) + ) + return ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_validation(): + """LIT-7391: the logging copy cannot round-trip a response that fails validation, and logging + rewrites the assembled response's usage to chat shape in place, so the event handed to logging + must never be the one the caller receives.""" + logging_obj = _logging_obj_stub() + logging_obj.stream = True + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator(headers={}, config=_unvalidated_completed_config(), logging_obj=logging_obj) + events = [event async for event in iterator] + + assert len(logged) == 1 + now = datetime.now() + LiteLLMLoggingObj._get_assembled_streaming_response( + logging_obj, logged[0], start_time=now, end_time=now, is_async=True, streaming_chunks=[] + ) + assert logged[0].response.usage["prompt_tokens"] == 29 + + client_usage = events[-1].response.usage + assert isinstance(client_usage, ResponseAPIUsage) + assert client_usage.input_tokens == 29 + assert client_usage.cost == pytest.approx(0.0001) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 9f925499f1a..51103297c58 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3846,11 +3846,11 @@ class TestRouterPreRoutingSharedAliasName: def test_forwardable_alias_marker_params_reads_the_marker_entry_only(self): router = Router(model_list=[self._plain_entry(), self._marker_entry(), self._tier_entry()]) - forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=())) + forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=(), request_kwargs={})) assert forwarded["drop_params"] is True assert "api_key" not in forwarded and "api_base" not in forwarded - assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=()) == () + assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=(), request_kwargs={}) == () @staticmethod def _region_marker_entry() -> dict: diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index 812d7bbff32..1a8614e3fca 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -8,11 +8,12 @@ import json from datetime import datetime, timedelta import pytest - +from pydantic import ValidationError import litellm from litellm.caching.caching import DualCache -from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler +from litellm.router import Router +from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler, RoutingArgs DEPLOYMENT_ID = "9876" KWARGS = { @@ -58,9 +59,9 @@ def test_sync_embedding_latency_is_json_serializable(): latencies = _recorded_latencies(cache) assert latencies, "expected a latency entry to be recorded" - assert all( - not isinstance(value, timedelta) for value in latencies - ), f"raw timedelta leaked into latency list: {latencies}" + assert all(not isinstance(value, timedelta) for value in latencies), ( + f"raw timedelta leaked into latency list: {latencies}" + ) assert latencies[-1] == pytest.approx(2.0) # the exact failure mode from production: redis cache sync json.dumps json.dumps({"latency": latencies}) @@ -84,9 +85,9 @@ async def test_async_embedding_latency_is_json_serializable(): latencies = _recorded_latencies(cache) assert latencies, "expected a latency entry to be recorded" - assert all( - not isinstance(value, timedelta) for value in latencies - ), f"raw timedelta leaked into latency list: {latencies}" + assert all(not isinstance(value, timedelta) for value in latencies), ( + f"raw timedelta leaked into latency list: {latencies}" + ) assert latencies[-1] == pytest.approx(3.0) json.dumps({"latency": latencies}) @@ -292,6 +293,85 @@ async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_worke assert picked["model_info"]["id"] == FAST_TTFT_ID +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("ttft_percentile", "first_samples", "second_samples", "expected_id"), + [ + (None, [0.1, 0.1, 1.0], [0.3, 0.3, 0.3], SLOW_TTFT_ID), + (0.5, [0.1, 0.1, 1.0], [0.3, 0.3, 0.3], FAST_TTFT_ID), + (0.9, [0.1, 0.1, 0.1, 0.1, 1.5], [0.3, 0.3, 0.3, 0.3, 0.3], SLOW_TTFT_ID), + ], + ids=["default_average", "p50", "p90"], +) +async def test_streaming_ttft_ranking_percentile( + sync_mode: bool, + ttft_percentile: float | None, + first_samples: list[float], + second_samples: list[float], + expected_id: str, +): + cache = DualCache() + routing_args = {} if ttft_percentile is None else {"ttft_percentile": ttft_percentile} + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args=routing_args) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token_seconds": first_samples}, + SLOW_TTFT_ID: {"time_to_first_token_seconds": second_samples}, + }, + ) + + if sync_mode: + picked = handler.get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + else: + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == expected_id + + +@pytest.mark.parametrize("ttft_percentile", [0, -0.1, 1.1]) +def test_ttft_percentile_validation(ttft_percentile: float): + with pytest.raises(ValidationError): + RoutingArgs(ttft_percentile=ttft_percentile) + + +@pytest.mark.parametrize("ttft_percentile", [0.5, 0.9, 0.95, 1.0]) +def test_ttft_percentile_accepts_valid_values(ttft_percentile: float): + assert RoutingArgs(ttft_percentile=ttft_percentile).ttft_percentile == ttft_percentile + + +@pytest.mark.asyncio +async def test_ttft_percentile_does_not_change_non_streaming_routing(): + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"ttft_percentile": 0.9}) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"latency": [1.0], "time_to_first_token_seconds": [0.1]}, + SLOW_TTFT_ID: {"latency": [0.2], "time_to_first_token_seconds": [1.5]}, + }, + ) + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": False, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == SLOW_TTFT_ID + + @pytest.mark.asyncio @pytest.mark.parametrize( "cached_entry", @@ -318,3 +398,81 @@ async def test_async_get_available_deployments_treats_missing_samples_as_zero_la assert picked is not None assert picked["model_info"]["id"] == DEPLOYMENT_ID + + +def _latency_router(routing_strategy_args: dict) -> Router: + return Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": {"model": f"openai/{MODEL_GROUP}", "api_key": "sk-fake"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in (FAST_TTFT_ID, SLOW_TTFT_ID) + ], + routing_strategy="latency-based-routing", + routing_strategy_args=routing_strategy_args, + ) + + +def _seed_streaming_ttft(router: Router) -> None: + router.cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token_seconds": [0.1, 0.1, 1.0]}, + SLOW_TTFT_ID: {"time_to_first_token_seconds": [0.3, 0.3, 0.3]}, + }, + ) + + +async def _pick_streaming(router: Router) -> str: + picked = await router.async_get_available_deployment( + model=MODEL_GROUP, + request_kwargs={"stream": True, "metadata": {}}, + ) + return picked["model_info"]["id"] + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_applies_ttft_percentile(): + """A config reload that adds ttft_percentile must reach the live selector, + not sit unused until the proxy restarts.""" + router = _latency_router({"max_latency_list_size": 50}) + _seed_streaming_ttft(router) + + assert await _pick_streaming(router) == SLOW_TTFT_ID + + router.update_settings(routing_strategy_args={"max_latency_list_size": 50, "ttft_percentile": 0.5}) + + assert await _pick_streaming(router) == FAST_TTFT_ID + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_keeps_previous_args_when_invalid(): + router = _latency_router({"ttft_percentile": 0.5}) + _seed_streaming_ttft(router) + + router.update_settings(routing_strategy_args={"ttft_percentile": 5}) + + assert await _pick_streaming(router) == FAST_TTFT_ID + + +@pytest.mark.asyncio +async def test_runtime_routing_strategy_args_update_is_a_noop_without_a_selector(): + """simple-shuffle has no selector to re-link, so an args update must leave + the router alone instead of blowing up on a missing selector attribute.""" + router = Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": {"model": f"openai/{MODEL_GROUP}", "api_key": "sk-fake"}, + "model_info": {"id": FAST_TTFT_ID}, + } + ], + routing_strategy="simple-shuffle", + ) + + router.update_settings(routing_strategy_args={"ttl": 5}) + + assert router.routing_strategy_args == {"ttl": 5} + assert await _pick_streaming(router) == FAST_TTFT_ID diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index dac991a41c4..ea8e2eacaa6 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -16,12 +16,10 @@ The mechanism works without any cache and supports two encoding strategies: """ import time -from typing import List, Optional from unittest.mock import AsyncMock, patch import pytest - import litellm from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -68,9 +66,7 @@ class TestEncryptedItemIdCodec: def test_roundtrip(self): model_id = "deployment-1" original_item_id = "rs_abc123def456" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) assert encoded.startswith("encitem_") decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) assert decoded is not None @@ -81,9 +77,7 @@ class TestEncryptedItemIdCodec: """Decoding must succeed even if base64 padding (=) was stripped in transit.""" model_id = "gpt-5.1-codex-openai-2" original_item_id = "rs_0efb96cb222403210069a01d5d52588196a9dc394ffdb89d00" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) # Strip any trailing '=' to simulate what happens in transit stripped = encoded.rstrip("=") decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(stripped) @@ -100,9 +94,7 @@ class TestEncryptedItemIdCodec: """item_id values containing ';' must survive the roundtrip.""" model_id = "deployment-1" original_item_id = "rs_part1;part2;part3" - encoded = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_item_id - ) + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) assert decoded is not None assert decoded["item_id"] == original_item_id @@ -118,11 +110,7 @@ class TestUpdateEncryptedContentItemIds: {"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}, ], } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, model_id - ) - ) + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, model_id) # Plain message item untouched assert result["output"][0]["id"] == "msg_abc" # Reasoning item with encrypted_content gets encoded @@ -133,16 +121,8 @@ class TestUpdateEncryptedContentItemIds: assert decoded["item_id"] == "rs_xyz" def test_no_op_when_model_id_is_none(self): - response = { - "output": [ - {"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"} - ] - } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, None - ) - ) + response = {"output": [{"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}]} + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, None) assert result["output"][0]["id"] == "rs_xyz" @@ -151,9 +131,7 @@ class TestEncryptedContentWrapping: """Test wrapping encrypted_content with model_id metadata.""" model_id = "deployment-1" original_content = "gAAAAABpnW_yEYmSNEyOG_original_encrypted_data" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) assert wrapped.startswith("litellm_enc:") assert wrapped != original_content @@ -170,9 +148,7 @@ class TestEncryptedContentWrapping: ( model_id, content, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - plain_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(plain_content) assert model_id is None assert content == plain_content @@ -189,11 +165,7 @@ class TestEncryptedContentWrapping: }, ], } - result = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response, model_id - ) - ) + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, model_id) assert result["output"][0].get("encrypted_content") is None wrapped = result["output"][1]["encrypted_content"] assert wrapped.startswith("litellm_enc:") @@ -210,19 +182,13 @@ class TestRestoreEncryptedContentItemIds: def test_restores_encoded_ids(self): model_id = "deployment-1" original_id = "rs_encrypted_item_456" - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, original_id - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_id) request_input = [ {"type": "message", "id": "msg_abc123", "role": "assistant"}, {"type": "reasoning", "id": encoded_id, "encrypted_content": "secret"}, ] - restored = ( - ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) - ) + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert restored[0]["id"] == "msg_abc123" assert restored[1]["id"] == original_id @@ -230,33 +196,21 @@ class TestRestoreEncryptedContentItemIds: """Test that wrapped encrypted_content is unwrapped before forwarding.""" model_id = "deployment-1" original_content = "gAAAAABpnW_yEYmSNEyOG_original" - wrapped_content = ( - ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) - ) + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) request_input = [ {"type": "reasoning", "encrypted_content": wrapped_content}, ] - restored = ( - ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) - ) + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert restored[0]["encrypted_content"] == original_content def test_no_op_for_plain_string_input(self): - result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - "Hello world" - ) + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input("Hello world") assert result == "Hello world" def test_no_op_for_unencoded_ids(self): request_input = [{"type": "message", "id": "msg_plain"}] - result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - request_input - ) + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(request_input) assert result[0]["id"] == "msg_plain" @@ -283,9 +237,7 @@ async def test_encrypted_content_affinity_tracks_and_routes(): "id": "msg_abc123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], }, { "type": "reasoning", @@ -347,9 +299,9 @@ async def test_encrypted_content_affinity_tracks_and_routes(): # The response must have rewritten the encrypted item's ID to encoded form encoded_item_id = _extract_encoded_item_id(first_response) - assert encoded_item_id.startswith( - "encitem_" - ), f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}" + assert encoded_item_id.startswith("encitem_"), ( + f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}" + ) # Verify the encoded ID decodes back to the correct deployment + original ID decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded_item_id) @@ -371,9 +323,9 @@ async def test_encrypted_content_affinity_tracks_and_routes(): ) second_model_id = second_response._hidden_params["model_id"] - assert ( - second_model_id == first_model_id - ), f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + ) @pytest.mark.asyncio @@ -478,9 +430,7 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): # Extract encoded item ID from the first response output encoded_item_id = _extract_encoded_item_id(first_response) - assert encoded_item_id.startswith( - "encitem_" - ), f"Expected encitem_... but got {encoded_item_id!r}" + assert encoded_item_id.startswith("encitem_"), f"Expected encitem_... but got {encoded_item_id!r}" # Follow-up with the encoded item ID — should pin to same deployment second_response = await router.aresponses( @@ -628,17 +578,13 @@ async def test_encrypted_content_affinity_with_wrapped_content_no_id(): if hasattr(first_item, "encrypted_content") else first_item.get("encrypted_content") ) - assert wrapped_content.startswith( - "litellm_enc:" - ), f"Expected wrapped content but got {wrapped_content[:50]}..." + assert wrapped_content.startswith("litellm_enc:"), f"Expected wrapped content but got {wrapped_content[:50]}..." # Verify we can extract model_id from wrapped content ( extracted_model_id, _, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - wrapped_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped_content) assert extracted_model_id == first_model_id # Second request: use wrapped encrypted_content WITHOUT an ID (Codex behavior) @@ -653,9 +599,9 @@ async def test_encrypted_content_affinity_with_wrapped_content_no_id(): ) second_model_id = second_response._hidden_params["model_id"] - assert ( - second_model_id == first_model_id - ), f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + ) def test_encrypted_content_wrapping_preserves_original_content(): @@ -664,13 +610,9 @@ def test_encrypted_content_wrapping_preserves_original_content(): This is critical for streaming responses where content must round-trip correctly. """ model_id = "test-deployment-1" - original_encrypted_content = ( - "gAAAAABpnW_yEYmSNEyOG_streaming_test_content_with_special_chars==+/" - ) + original_encrypted_content = "gAAAAABpnW_yEYmSNEyOG_streaming_test_content_with_special_chars==+/" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_encrypted_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_encrypted_content, model_id) assert wrapped.startswith("litellm_enc:") assert wrapped != original_encrypted_content @@ -691,9 +633,7 @@ def test_encrypted_content_wrapping_with_multiple_semicolons(): model_id = "deployment-with-semicolons" original_content = "gAAAAAB;some;content;with;semicolons" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) ( extracted_model_id, @@ -764,9 +704,7 @@ async def test_encrypted_content_affinity_preserves_litellm_metadata_for_respons request_kwargs=request_kwargs, ) - assert ( - request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True - ) + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True assert request_kwargs["litellm_metadata"]["model_info"] == {"id": "dep-1"} @@ -777,9 +715,7 @@ def test_encrypted_content_wrapping_empty_string(): model_id = "test-deployment" original_content = "" - wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - original_content, model_id - ) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(original_content, model_id) assert wrapped.startswith("litellm_enc:") @@ -1132,9 +1068,7 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): "api_key": "fake-azure-resource-key-a", } - pydantic_key = EncryptedContentAffinityCheck._encryption_boundary_key( - pydantic_params - ) + pydantic_key = EncryptedContentAffinityCheck._encryption_boundary_key(pydantic_params) plain_key = EncryptedContentAffinityCheck._encryption_boundary_key(plain_params) assert pydantic_key is not None @@ -1161,18 +1095,8 @@ def test_boundary_key_rejects_non_dict_like_inputs(): for bad in (None, [], "not a dict", 42, object()): assert EncryptedContentAffinityCheck._encryption_boundary_key(bad) is None - assert ( - EncryptedContentAffinityCheck._encryption_boundary_key( - {"api_base": "", "api_key": "k"} - ) - is None - ) - assert ( - EncryptedContentAffinityCheck._encryption_boundary_key( - {"api_base": "https://x"} - ) - is None - ) + assert EncryptedContentAffinityCheck._encryption_boundary_key({"api_base": "", "api_key": "k"}) is None + assert EncryptedContentAffinityCheck._encryption_boundary_key({"api_base": "https://x"}) is None # --------------------------------------------------------------------------- @@ -1180,10 +1104,11 @@ def test_boundary_key_rejects_non_dict_like_inputs(): # --------------------------------------------------------------------------- -def _make_originating_mock(api_base: str, api_key: str): +def _make_originating_mock(api_base: str, api_key: str, model_name: str = "gpt-5.4"): from unittest.mock import MagicMock originating = MagicMock() + originating.model_name = model_name originating.litellm_params.model_dump.return_value = { "api_base": api_base, "api_key": api_key, @@ -1192,19 +1117,23 @@ def _make_originating_mock(api_base: str, api_key: str): def _make_router_mock_with_cooldown( - originating, cooldown_entries: Optional[List[tuple]] = None + originating, + cooldown_entries: list[tuple] | None = None, + routed_group_model_ids: list[str] | None = None, ): """ Build a MagicMock router whose ``cooldown_cache.async_get_active_cooldowns`` - returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown). + returns ``cooldown_entries`` (defaulting to ``[]`` — no active cooldown), and + whose ``get_candidate_model_ids_for_route`` returns ``routed_group_model_ids`` + (the deployment ids the router resolves for the routed model; defaulting to ``[]`` + — origin absent from the routed group, i.e. a tier change). """ from unittest.mock import AsyncMock, MagicMock mock_router = MagicMock() mock_router.get_deployment.return_value = originating - mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock( - return_value=list(cooldown_entries or []) - ) + mock_router.cooldown_cache.async_get_active_cooldowns = AsyncMock(return_value=list(cooldown_entries or [])) + mock_router.get_candidate_model_ids_for_route.return_value = frozenset(routed_group_model_ids or []) return mock_router @@ -1235,15 +1164,15 @@ async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_42 }, ) ], + routed_group_model_ids=["deployment-a-cooled", "deployment-b"], ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-cooled", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1297,15 +1226,15 @@ async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_fo }, ) ], + routed_group_model_ids=["deployment-a-cooled-429", "deployment-b"], ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-cooled-429", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-cooled-429", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1345,15 +1274,16 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ ) originating = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") - mock_router = _make_router_mock_with_cooldown(originating, cooldown_entries=[]) + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-a-filtered", "deployment-b"] + ) check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a-filtered", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a-filtered", "rs_test") healthy_only_b = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1377,15 +1307,18 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ @pytest.mark.asyncio -async def test_affinity_raises_bad_request_when_origin_removed(): +async def test_affinity_strips_and_dispatches_when_origin_is_unknown_or_removed(): """ - Originating deployment was removed from the router config and no boundary - peer is available. This is permanent (the stale encrypted_content cannot - be honored), so surface a 400 with actionable text. + A removed deployment, or a forged/unknown affinity marker, resolves to no + originating deployment. It is handled like a cross-group origin: the encrypted + reasoning is stripped and the request dispatches with its readable history, + rather than returning a distinguishable error. That uniform handling denies an + authenticated caller a deployment-id existence oracle, an existing cross-group id + and a nonexistent id both strip and proceed, so responses cannot be told apart. + The membership lookup is skipped entirely when the origin is unknown. """ from unittest.mock import MagicMock - from litellm.exceptions import BadRequestError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1394,12 +1327,11 @@ async def test_affinity_raises_bad_request_when_origin_removed(): mock_router.get_deployment.return_value = None check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-removed", "rs_test" - ) - healthy_only_b = [ + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-removed") + routed_pool = [ { "model_info": {"id": "deployment-b"}, + "model_name": "gpt-5.4", "litellm_params": { "api_base": "https://account-b.openai.azure.com/", "api_key": "key-b", @@ -1408,18 +1340,28 @@ async def test_affinity_raises_bad_request_when_origin_removed(): } ] request_kwargs = { - "input": [{"id": encoded_id, "type": "reasoning"}], + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], } - with pytest.raises(BadRequestError) as excinfo: - await check.async_filter_deployments( - model="gpt-5.4", - healthy_deployments=healthy_only_b, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=routed_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert "deployment-removed" not in str(excinfo.value) + assert result is routed_pool + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in request_kwargs["input"]) + mock_router.get_candidate_model_ids_for_route.assert_not_called() @pytest.mark.asyncio @@ -1444,9 +1386,7 @@ async def test_affinity_does_not_raise_when_boundary_peer_available(): mock_router.get_deployment.return_value = originating check = EncryptedContentAffinityCheck(router=mock_router) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-a", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a", "rs_test") peer = { "model_info": {"id": "deployment-a-peer"}, "litellm_params": { @@ -1490,9 +1430,7 @@ async def test_model_group_affinity_config_enables_encrypted_content_affinity(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {}, @@ -1536,9 +1474,7 @@ async def test_model_group_affinity_config_does_not_disable_global_encrypted_con }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {}, @@ -1600,15 +1536,9 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen try: callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) assert encrypted_content_callback.enable_global_affinity is False cache_key = DeploymentAffinityCheck.get_affinity_cache_key( @@ -1620,9 +1550,7 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [ { @@ -1643,16 +1571,384 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen ) assert after_deployment_affinity == [deployment_a, deployment_b] - after_encrypted_content_affinity = ( - await encrypted_content_callback.async_filter_deployments( - model=model_group, - healthy_deployments=after_deployment_affinity, - messages=None, - request_kwargs=request_kwargs, - ) + after_encrypted_content_affinity = await encrypted_content_callback.async_filter_deployments( + model=model_group, + healthy_deployments=after_deployment_affinity, + messages=None, + request_kwargs=request_kwargs, ) assert after_encrypted_content_affinity == [deployment_b] assert request_kwargs.get("_encrypted_content_affinity_pinned") is True finally: router.discard() + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_through_the_bridge(): + """ + Claude Code behind /v1/messages replays the encrypted reasoning the bridge packed + into a thinking block's signature (or a redacted block's data). The pin has to be + read from those blocks because the bridge builds the Responses `input` only after + the router has picked a deployment. + """ + check = EncryptedContentAffinityCheck() + deployments = [ + {"model_info": {"id": "openai-org-a"}, "litellm_params": {"model": "openai/gpt-5.1"}}, + {"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5.1"}}, + ] + request_kwargs = {"model": "gpt-5.1"} + + pinned = await check.async_filter_deployments( + model="gpt-5.1", + healthy_deployments=deployments, + messages=_bridge_replayed_anthropic_messages(minted_by="openai-org-b"), + request_kwargs=request_kwargs, + ) + + assert [d["model_info"]["id"] for d in pinned] == ["openai-org-b"] + assert request_kwargs["_encrypted_content_affinity_pinned"] is True + + +def _bridge_replayed_anthropic_messages(minted_by: str) -> list: + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", minted_by) + return [ + {"role": "user", "content": "Solve the zebra puzzle"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + { + "type": "thinking", + "thinking": "The bridge packed this one", + "signature": f"litellm_encrypted_reasoning:{wrapped}", + }, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ], + }, + {"role": "user", "content": "And who drinks water?"}, + ] + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_routed_to_another_group(): + """ + The /v1/messages twin of the tier-change case: the routed group holds no deployment + of the org that minted the reasoning, so the bridge-tagged blocks are dropped whole + and the request dispatches to the routed pool. No unsigned thinking block may be left + behind: Anthropic and Bedrock reject a thinking block with a missing signature the + same way they reject a foreign one. + """ + originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["openai-org-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + routed_pool = [{"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5-nano"}}] + messages = _bridge_replayed_anthropic_messages(minted_by="openai-org-a") + assistant_content = messages[1]["content"] + request_kwargs = {"model": "gpt-5.1"} + + result = await check.async_filter_deployments( + model="gpt-simple-tier", + healthy_deployments=routed_pool, + messages=messages, + request_kwargs=request_kwargs, + ) + + assert result is routed_pool + assert "_encrypted_content_affinity_pinned" not in request_kwargs + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ] + assert all(block["signature"] for block in assistant_content if block["type"] == "thinking") + + +class TestStripEncryptedReasoningFromInput: + def test_keeps_summary_and_drops_encrypted_content_and_id(self): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-a", "rs_1") + request_input = [ + {"role": "user", "content": "first turn"}, + { + "type": "reasoning", + "id": encoded_id, + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "thought about it"}], + }, + {"type": "reasoning", "id": encoded_id, "encrypted_content": wrapped}, + {"type": "reasoning", "encrypted_content": wrapped, "summary": []}, + {"type": "message", "id": "msg_1", "role": "assistant", "content": "hi"}, + {"role": "user", "content": "second turn"}, + ] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == [ + {"role": "user", "content": "first turn"}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "thought about it"}], + }, + {"type": "message", "id": "msg_1", "role": "assistant", "content": "hi"}, + {"role": "user", "content": "second turn"}, + ] + + def test_keeps_string_form_summary_when_stripping(self): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + request_input = [ + {"type": "reasoning", "encrypted_content": wrapped, "summary": "plain string thought"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "content": [{"type": "output_text", "text": "in content"}], + }, + {"type": "reasoning", "encrypted_content": wrapped, "summary": "", "content": []}, + ] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == [ + {"type": "reasoning", "summary": "plain string thought"}, + {"type": "reasoning", "content": [{"type": "output_text", "text": "in content"}]}, + ] + + def test_leaves_input_untouched_when_no_encrypted_reasoning(self): + request_input = [ + {"role": "user", "content": "first turn"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "no blob here"}]}, + {"role": "user", "content": "second turn"}, + ] + before = [dict(item) for item in request_input] + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + assert request_input == before + + +def _cross_group_request_kwargs(): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + return { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "ZEBRA: why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"type": "message", "role": "assistant", "content": "Rayleigh scattering."}, + {"role": "user", "content": "KIWI: and sunsets?"}, + ], + } + + +@pytest.mark.asyncio +async def test_affinity_strips_encrypted_reasoning_when_routed_to_another_model_group(): + """ + An auto-router tier change (or a model switch with no boundary peer): the + routed pool holds no deployment of the origin's model group. The origin is + healthy, so a 503 would be wrong; the follow-up dispatches to the routed + pool with the origin's encrypted reasoning stripped and its summary kept. + """ + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + routed_pool = [ + { + "model_info": {"id": "deployment-b"}, + "model_name": "gpt-simple-tier", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5-nano", + }, + } + ] + request_kwargs = _cross_group_request_kwargs() + original_input = request_kwargs["input"] + + result = await check.async_filter_deployments( + model="gpt-simple-tier", + healthy_deployments=routed_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert result is routed_pool + assert "_encrypted_content_affinity_pinned" not in request_kwargs + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] is True + assert request_kwargs["input"] is original_input + assert [item.get("type") or item["role"] for item in original_input] == [ + "user", + "reasoning", + "message", + "user", + ] + assert original_input[1] == { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "scattering"}], + } + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in original_input) + + +@pytest.mark.asyncio +async def test_affinity_fails_fast_within_the_origins_own_group(): + """ + Negative class for the tier-change discriminator: the routed group IS the + origin's group (a same-group cooldown, not a tier change), so even with a + healthy non-origin sibling that cannot decrypt the content, the request + still fails fast and the encrypted reasoning is left intact rather than + stripped. Preserves the LIT-3051 cooldown contract. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock( + "https://account-a.openai.azure.com/", "key-a", model_name="gpt-reasoning-tier" + ) + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-a", "deployment-a2"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + sibling_pool = [ + { + "model_info": {"id": "deployment-a2"}, + "model_name": "gpt-reasoning-tier", + "litellm_params": { + "api_base": "https://account-a2.openai.azure.com/", + "api_key": "key-a2", + "model": "azure/gpt-5.4", + }, + } + ] + request_kwargs = _cross_group_request_kwargs() + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="gpt-reasoning-tier", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + + +@pytest.mark.asyncio +async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_same_by_id(): + """ + The discriminator must key on deployment-id membership, not on the model-group + name string. Here the origin's configured group is spelled ``openai/gpt-5.4-mini`` + while the routed group is the canonical ``gpt-5.4-mini``: same group, different + spelling. A name compare (``originating.model_name != model``) would read this as + a tier change and strip the reasoning it did not have to. Because the origin's id + is a member of the routed group, this is a same-group cooldown instead: the request + fails fast and the encrypted reasoning is left intact. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="openai/gpt-5.4-mini") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-mini-a", "deployment-mini-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-mini-a") + sibling_pool = [ + { + "model_info": {"id": "deployment-mini-b"}, + "model_name": "gpt-5.4-mini", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5.4-mini", + }, + } + ] + request_kwargs = { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + } + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="gpt-5.4-mini", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + + +@pytest.mark.asyncio +async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes(): + """ + The exact `model_name` index does not include team-public or pattern routes, so a + same-group cooldown reached only through one of those would be misread as a tier change + and stripped. The check asks the router for the candidate ids it resolves for the route + (`get_candidate_model_ids_for_route`), which covers those paths, rather than the bare + index. Here that set marks the origin as a candidate, so the request fails fast with its + reasoning intact, and the routed group and team are passed through to the router. + """ + from litellm.exceptions import ServiceUnavailableError + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock(None, "key-a", model_name="model_name_teamA_uuid") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["deployment-team-a", "deployment-team-b"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-team-a") + sibling_pool = [ + { + "model_info": {"id": "deployment-team-b"}, + "model_name": "team-public-model", + "litellm_params": { + "api_base": "https://gateway.example/v1", + "api_key": "key-b", + "model": "openai/gpt-5.4-mini", + }, + } + ] + request_kwargs = { + "litellm_metadata": {"user_api_key_team_id": "teamA"}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + } + + with pytest.raises(ServiceUnavailableError): + await check.async_filter_deployments( + model="team-public-model", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) + + assert request_kwargs["input"][1].get("encrypted_content") + mock_router.get_candidate_model_ids_for_route.assert_called_once_with(model="team-public-model", team_id="teamA") diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 030bdfe03e9..333e7b2ff31 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -477,3 +477,65 @@ async def test_wildcard_route_resolves_underlying_model_minimum(local_model_cost assert deployments[0]["litellm_params"]["model"] == "anthropic/claude-opus-4-6" assert _get_min_token_count_for_deployments(deployments) == 4096 + + +@pytest.mark.asyncio +async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + check = PromptCachingDeploymentCheck(cache=DualCache()) + deployments = _deployments("anthropic/claude-fable-5") + messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}]) + + result, took, lags = await timed_with_loop_lags( + lambda: check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages + ) + ) + + assert result == deployments + assert_loop_stayed_free(took, lags) + + +@pytest.mark.asyncio +async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}], + ) + standard_logging_object = { + "call_type": "acompletion", + "model": "anthropic/claude-fable-5", + "messages": messages, + "model_id": "dep-1", + } + + _, took, lags = await timed_with_loop_lags( + lambda: check.async_log_success_event( + kwargs={"standard_logging_object": standard_logging_object}, + response_obj=None, + start_time=None, + end_time=None, + ) + ) + + assert await PromptCachingCache(cache=cache).async_get_model_id(messages=messages, tools=None) == { + "model_id": "dep-1" + } + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 30f658d7ea2..ac18b4889dd 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -12,6 +12,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, PROVIDER_SCOPED_CREDENTIAL_PARAMS, @@ -756,3 +757,20 @@ class TestWarnOnProviderCredentialMismatch: ) is None ) + + +@pytest.mark.parametrize( + ("litellm_params", "expected"), + [ + ({"model": "azure_ai/gpt-5.4-mini", "custom_llm_provider": "azure"}, "azure"), + ({"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.openai.azure.com"}, "azure_ai"), + ({"model": "cohere/command-r"}, "cohere"), + ({"model": "gpt-5.4-mini"}, "openai"), + ({"model": "no-provider-knows-this-model"}, None), + ({"api_base": "https://my-resource.openai.azure.com"}, None), + ], + ids=["declared_wins", "prefix_beats_host_flip", "prefix_beats_cohere_chat_flip", "unprefixed_inferred", "unknown", "no_model"], +) +def test_provider_for_generic_call(litellm_params, expected, monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") + assert provider_for_generic_call(litellm_params) == expected diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f8fa2231597..8f8a7640c08 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3736,6 +3736,112 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(100 * 3e-08) +def test_completion_cost_logs_the_rates_it_billed_at(monkeypatch): + """A caller reporting the cost lines beside their per-token rates reads both off this one call. + completion_cost infers the provider, and xai's inclusive tier thresholds put a request sitting + exactly on 200k at the tier rate, which a lookup made without that inferred provider would miss. + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setitem( + litellm.model_cost, + "xai/tiered-model", + { + "input_cost_per_token": 3e-6, + "output_cost_per_token": 15e-6, + "cache_read_input_token_cost": 3e-7, + "input_cost_per_token_above_200k_tokens": 6e-6, + "output_cost_per_token_above_200k_tokens": 3e-5, + "cache_read_input_token_cost_above_200k_tokens": 6e-7, + "litellm_provider": "xai", + "mode": "chat", + }, + ) + logging_obj = Logging( + model="xai/tiered-model", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="billed-rates", + function_id="f", + ) + usage = Usage( + prompt_tokens=200_000, + completion_tokens=1_000, + total_tokens=201_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100_000), + ) + + litellm.completion_cost( + completion_response=ModelResponse(model="xai/tiered-model", usage=usage), + model="xai/tiered-model", + custom_llm_provider=None, + litellm_logging_obj=logging_obj, + ) + + rates = logging_obj.billed_token_rates + assert rates is not None + assert rates.input_cost_per_token == pytest.approx(6e-6) + assert rates.cache_read_input_token_cost == pytest.approx(6e-7) + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx( + 100_000 * rates.cache_read_input_token_cost + ) + assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(1_000 * rates.output_cost_per_token) + + +def test_completion_cost_logs_cache_and_reasoning_breakdown_for_custom_pricing(): + """ + A custom-priced deployment bills cache tokens at its custom cache rates, but the + breakdown stored for the spend logs carried no cache or reasoning lines for it. + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CompletionTokensDetailsWrapper, CostPerToken + + logging_obj = Logging( + model="openai/onprem-model", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="custom-pricing-breakdown", + function_id="f", + ) + response = ModelResponse( + model="openai/onprem-model", + usage=Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800, cache_creation_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=200), + ), + ) + + total = completion_cost( + completion_response=response, + model="openai/onprem-model", + custom_llm_provider="openai", + custom_cost_per_token=CostPerToken( + input_cost_per_token=1e-6, + output_cost_per_token=2e-6, + cache_read_input_token_cost=1e-7, + cache_creation_input_token_cost=1.25e-6, + ), + litellm_logging_obj=logging_obj, + ) + + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(800 * 1e-7) + assert logging_obj.cost_breakdown["cache_creation_cost"] == pytest.approx(100 * 1.25e-6) + assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(200 * 2e-6) + assert total == pytest.approx(100 * 1e-6 + 800 * 1e-7 + 100 * 1.25e-6 + 500 * 2e-6) + + def test_cost_per_token_per_second_pricing(monkeypatch): """ Models priced by duration (input/output_cost_per_second) with no per-token rates diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 038df3656fe..5cb52295f94 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -120,7 +120,7 @@ def test_completion_missing_role(openai_api_response): print(f"openai_api_response: {openai_api_response}") with patch.object( - client.chat.completions.with_raw_response, "create", mock_raw_response + client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) ) as mock_create: litellm.completion( model="gpt-4o-mini", @@ -1367,6 +1367,78 @@ def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( } +@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) +def test_responses_bridge_preserves_reasoning_effort_with_drop_params( + reasoning_effort, + restore_model_registry, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + response_body: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "test-responses-bridge", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) + model: Final = "perplexity/test-responses-bridge" + litellm.register_model( + { + model: { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_reasoning": False, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) + + request_body: Final = json.loads(response_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "high"} + + @pytest.mark.parametrize( "model, model_info, expected_model_param, expected_base_model_param", [ diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bc79c5f6589..d59d72a094e 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5049,6 +5049,41 @@ def test_get_deployment_model_info_base_model_merge_priority(): print("✓ Base model merge priority test passed!") +@pytest.mark.parametrize( + "model, litellm_params, endpoint, expected", + [ + ( + "gpt", + {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, + "gpt/openai/deployments/gpt-5.4-mini/chat/completions", + "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", + ), + ( + "aws/anthropic/bedrock-claude", + {"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"}, + "/model/aws/anthropic/bedrock-claude/invoke", + "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", + ), + ( + "my-gemini", + {"model": "gemini/gemini-3.1-pro-preview", "api_key": "key"}, + "v1beta/models/my-gemini:streamGenerateContent", + "v1beta/models/gemini-3.1-pro-preview:streamGenerateContent", + ), + ], +) +def test_add_deployment_model_to_endpoint_rewrites_the_model_group_only_as_whole_path_segments( + model, litellm_params, endpoint, expected +): + router = litellm.Router(model_list=[{"model_name": model, "litellm_params": litellm_params}]) + + result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( + kwargs={"endpoint": endpoint}, model=model, model_name=litellm_params["model"] + ) + + assert result["endpoint"] == expected + + def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): """ Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix @@ -10003,13 +10038,6 @@ class TestTaggedAutoRouterOnSharedModelName: def test_deployment_without_litellm_params_mapping_is_not_a_marker(self): assert litellm.Router._is_strategy_marker_deployment({"model_name": "gpt4o"}) is False - def test_model_name_has_plain_deployments_reflects_the_pool(self): - mixed = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True) - marker_only = self._router(marker_tags=["route"], include_plain_sibling=False, enable_tag_filtering=True) - - assert mixed._model_name_has_plain_deployments("gpt4o") is True - assert marker_only._model_name_has_plain_deployments("gpt4o") is False - class TestAutoRouterSharedModelNameConnectionParams: """A plain deployment sharing its model_name with an `auto_router/` marker must not have @@ -10618,6 +10646,300 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestTeamPublicNameReachesPreRoutingStrategies: + """A team-scoped strategy router is stored under an internal `model_name_{team}_{uuid}` with the + caller-facing name in `model_info.team_public_model_name`, and the four registries key on that + internal name. A team key asks for the public name, so the hook has to resolve it to the team's + marker through the same team-first resolution the deployment path uses, and a resolution that + yields only markers is not callable on any path (LIT-7363).""" + + MARKER_TIMEOUT = 42.0 + REGISTRY_NAMES = ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers") + TEAM = "team-a" + OTHER_TEAM = "team-b" + PUBLIC_NAME = "smart-route" + INTERNAL_NAME = "model_name_team-a_0b3c" + SIBLING_INTERNAL_NAME = "model_name_team-a_9e1d" + + class _RewriteStrategy: + def __init__(self, rewrite_to: str = "gemini-flash"): + self.rewrite_to = rewrite_to + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + return PreRoutingHookResponse(model=self.rewrite_to, messages=messages) + + @classmethod + def _team_marker(cls, internal_name: str, tags: list[str] | None = None) -> dict: + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + return { + "model_name": internal_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + "timeout": cls.MARKER_TIMEOUT, + **({"tags": tags} if tags else {}), + }, + "model_info": {"team_id": cls.TEAM, "team_public_model_name": cls.PUBLIC_NAME}, + } + + @classmethod + def _router( + cls, + registrations: dict[str, "TestTeamPublicNameReachesPreRoutingStrategies._RewriteStrategy"], + registry_name: str = "complexity_routers", + extra_deployments: tuple[dict, ...] = (), + markers: tuple[dict, ...] | None = None, + enable_tag_filtering: bool = False, + ) -> "litellm.Router": + from litellm.types.router import TaggedPreRoutingStrategy + + markers = markers if markers is not None else (cls._team_marker(cls.INTERNAL_NAME),) + tier = { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + } + router = litellm.Router( + model_list=[*markers, tier, *extra_deployments], + enable_tag_filtering=enable_tag_filtering, + ) + tags_by_name = {m["model_name"]: tuple(m["litellm_params"].get("tags") or ()) for m in markers} + for name in cls.REGISTRY_NAMES: + setattr(router, name, {}) + setattr( + router, + registry_name, + { + name: [TaggedPreRoutingStrategy(tags=tags_by_name[name], strategy=strategy)] + for name, strategy in registrations.items() + }, + ) + return router + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + @classmethod + def _team_request(cls, team_id: str | None = "team-a", tags: list[str] | None = None) -> dict: + metadata = {**({"user_api_key_team_id": team_id} if team_id else {}), **({"tags": tags} if tags else {})} + return {"metadata": metadata} + + @pytest.mark.parametrize("registry_name", REGISTRY_NAMES) + @pytest.mark.asyncio + async def test_team_key_dispatches_to_the_strategy_registered_under_the_internal_name(self, registry_name): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}, registry_name=registry_name) + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + + @pytest.mark.asyncio + async def test_team_key_deployment_selection_lands_on_the_tier_and_forwards_the_marker_params(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = self._team_request() + + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert deployment["litellm_params"]["model"] == "gemini/gemini-3.6-flash" + assert request_kwargs["timeout"] == self.MARKER_TIMEOUT + + @pytest.mark.asyncio + async def test_another_team_never_reaches_the_strategy_or_the_marker(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + + assert ( + await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + is None + ) + with pytest.raises(litellm.BadRequestError): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_sibling_team_markers_select_by_request_tag_then_default(self): + router = self._router( + { + self.INTERNAL_NAME: self._RewriteStrategy("cn-model"), + self.SIBLING_INTERNAL_NAME: self._RewriteStrategy("us-model"), + }, + markers=( + self._team_marker(self.INTERNAL_NAME, tags=["cn"]), + self._team_marker(self.SIBLING_INTERNAL_NAME, tags=["us", "default"]), + ), + ) + + async def routed(tags: list[str] | None) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=tags), messages=self._messages() + ) + return response.model if response else None + + assert await routed(["cn"]) == "cn-model" + assert await routed(["us"]) == "us-model" + assert await routed(None) == "us-model" + + @pytest.mark.asyncio + async def test_team_public_name_shadows_a_global_model_for_that_team_only(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": self.PUBLIC_NAME, "litellm_params": {"model": "openai/gpt-4o"}},), + ) + + async def routed(request_kwargs: dict) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return response.model if response else None + + async def selected(request_kwargs: dict) -> str: + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return deployment["litellm_params"]["model"] + + assert await routed(self._team_request()) == "gemini-flash" + assert await selected(self._team_request()) == "gemini/gemini-3.6-flash" + for request_kwargs in (self._team_request(None), self._team_request(self.OTHER_TEAM)): + assert await routed(request_kwargs) is None + assert await selected(request_kwargs) == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_tag_filtering_hands_untagged_team_requests_to_the_team_plain_sibling(self): + plain_sibling = { + "model_name": self.SIBLING_INTERNAL_NAME, + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"team_id": self.TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + markers=(self._team_marker(self.INTERNAL_NAME, tags=["route"]),), + extra_deployments=(plain_sibling,), + enable_tag_filtering=True, + ) + + tagged = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=["route"]), messages=self._messages() + ) + assert tagged is not None and tagged.model == "gemini-flash" + for _ in range(20): + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + assert deployment["litellm_params"]["model"] == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_marker_only_team_resolution_is_rejected_as_uncallable(self): + import re + + from litellm.types.router import RouterErrors + + router = self._router({}) + + with pytest.raises( + litellm.BadRequestError, match=re.escape(RouterErrors.only_strategy_marker_deployments.value) + ): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_proxy_admin_without_a_team_reaches_the_team_strategy_by_public_name(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + @pytest.mark.asyncio + async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": "shared-name", "litellm_params": {"model": "openai/gpt-4o"}},), + ) + principals = { + "team": self._team_request(), + "other-team": self._team_request(self.OTHER_TEAM), + "no-team": self._team_request(None), + "admin": {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}}, + } + for principal, request_kwargs in principals.items(): + for model in (self.PUBLIC_NAME, "shared-name", "gemini-flash", "missing"): + resolved = [d["model_name"] for d in router.deployments_for_request(model, request_kwargs)] + callable_names = [ + name + for name, deployment in zip(resolved, router.deployments_for_request(model, request_kwargs)) + if not router._is_strategy_marker_deployment(deployment) + ] + if resolved and not callable_names: + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + elif not resolved: + with pytest.raises(litellm.BadRequestError): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + else: + _, deployments = router._common_checks_available_deployment( + model=model, request_kwargs=request_kwargs + ) + assert [d["model_name"] for d in deployments] == callable_names, (principal, model) + + def test_drop_strategy_markers_keeps_plain_deployments_and_rejects_marker_only_sets(self): + router = self._router({}) + marker = router.model_list[0] + plain = {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o"}} + + assert router._drop_strategy_markers("x", [marker, plain]) == [plain] + assert router._drop_strategy_markers("x", [plain]) == [plain] + assert router._drop_strategy_markers("x", []) == [] + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._drop_strategy_markers("x", [marker]) + + def test_team_deployments_across_teams_unions_one_team_and_rejects_two(self): + other_team_marker = { + **self._team_marker(self.SIBLING_INTERNAL_NAME), + "model_info": {"team_id": self.OTHER_TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + one_team = self._router({}) + two_teams = self._router({}, markers=(self._team_marker(self.INTERNAL_NAME), other_team_marker)) + + assert [d["model_name"] for d in one_team._team_deployments_across_teams(self.PUBLIC_NAME)] == [ + self.INTERNAL_NAME + ] + assert one_team._team_deployments_across_teams("missing") == [] + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + two_teams._team_deployments_across_teams(self.PUBLIC_NAME) + + + def test_compression_policy_follows_the_same_resolution_for_every_principal(self): + from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model + + marker = self._team_marker(self.INTERNAL_NAME) + marker["litellm_params"]["auto_router_routing_compression"] = "headroom-team" + router = self._router({}, markers=(marker,)) + admin = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + expected = AutoRouterCompressionPolicy(routing="headroom-team", model=None) + + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(), ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, admin, ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(self.OTHER_TEAM), ()) is None + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(None), ()) is None + + class TestAutoRouterCompressionDecoupling: """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` decouple what the routing decision sees from what the model call sees. The one @@ -13850,6 +14172,49 @@ async def test_router_retry_policy_controls_upstream_attempt_count( assert upstream.call_count == expected_upstream_calls +@pytest.mark.asyncio +async def test_generic_call_keeps_the_deployment_name_of_an_azure_ai_model_on_an_azure_openai_host(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "aoai-gpt", + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://my-resource.openai.azure.com", + "api_key": "deployment-key", + }, + } + ] + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream = respx_mock.post(host="my-resource.openai.azure.com", path__regex=r"^/openai/.*responses$").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + }, + ) + ) + await router.aresponses(model="aoai-gpt", input="hi") + + assert json.loads(upstream.calls.last.request.content)["model"] == "gpt-5.4-mini" + + @pytest.mark.asyncio @pytest.mark.parametrize( "retry_policy,upstream_error", @@ -14728,3 +15093,50 @@ def test_router_stays_quiet_when_a_deployment_drop_params_is_a_flag(value, caplo ) assert "is not a flag value" not in caplog.text + + +def test_get_candidate_model_ids_for_route_covers_model_name_and_pattern(): + """ + get_candidate_model_ids_for_route resolves a route the way the router does, so a + pre-call check can tell a genuine cross-group route from same-group unavailability. + A concrete model group returns its member ids; a wildcard/pattern deployment is + included for a concrete model it matches, which the bare model_name index misses. + The unprefixed-name case must resolve through get_deployments_by_pattern (which retries + the provider-qualified form), not a bare pattern_router.route that only sees the literal + name. Regression guard for the LIT-7195 tier-change discriminator's team/pattern gaps. + """ + router = Router( + model_list=[ + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-a", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-b", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-b"}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "sk-c", "api_base": "https://x.invalid"}, + "model_info": {"id": "dep-wild"}, + }, + ] + ) + + assert router.get_candidate_model_ids_for_route(model="grp") == frozenset({"dep-a", "dep-b"}) + assert "dep-wild" in router.get_candidate_model_ids_for_route(model="openai/gpt-4o-some-new-model") + # unprefixed name whose provider resolves to openai: only get_deployments_by_pattern's + # provider-qualified retry matches "openai/*"; a bare route() on the literal name misses it + assert "dep-wild" in router.get_candidate_model_ids_for_route(model="gpt-5") + + +def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_id(): + deployments = ( + {"model_info": {"id": "a"}}, + {"model_info": {"id": 2}}, + {"model_info": {}}, + {"no_model_info": True}, + ) + assert Router._deployment_ids(deployments) == frozenset({"a", "2"}) diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py index a5a68271111..3cef1c7bb63 100644 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ b/tests/test_litellm/test_router/test_io_token_rate_limits.py @@ -1039,3 +1039,31 @@ class TestContextSlotRetention: assert deployment is not None router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) assert get_io_token_rate_limit_request_kwargs() is kwargs + + +@pytest.mark.asyncio +async def test_the_deployment_itpm_reservation_counts_the_request_off_the_event_loop(): + from litellm.utils import get_utc_datetime + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + warm_tokenizer("anthropic/claude-fable-5") + deployment = { + "litellm_params": {"model": "anthropic/claude-fable-5", "itpm": 10_000_000}, + "model_info": {"id": "io-loop-id"}, + "model_name": "claude", + } + set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": text * 100}], "metadata": {}}) + + _, took, lags = await timed_with_loop_lags(lambda: check.async_pre_call_check(deployment)) + + minute = get_utc_datetime().strftime("%H-%M") + reserved = await dual_cache.async_get_cache(key=f"global_router:io-loop-id:anthropic/claude-fable-5:itpm:{minute}") + assert reserved > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index fee5e3a2e4c..8b186be43e5 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2,10 +2,12 @@ import asyncio import json import logging import os +import threading from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from jsonschema import validate @@ -54,6 +56,12 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_cloudflare_model_info_includes_rpm(local_model_cost_map: None) -> None: + assert litellm.get_model_info("cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8")["rpm"] == 300 + assert litellm.get_model_info("cloudflare/@cf/moonshotai/kimi-k2.6")["rpm"] == 20 + assert litellm.get_model_info("cloudflare/@cf/openai/whisper-large-v3-turbo")["rpm"] == 720 + + def test_get_utc_datetime_returns_current_aware_utc_time() -> None: before: Final = datetime.now(timezone.utc) result: Final = litellm.utils.get_utc_datetime() @@ -808,6 +816,7 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_second", "output_cost_per_second", "output_cost_per_second_480p", + "output_cost_per_second_720p", "output_cost_per_second_1080p", "output_cost_per_second_4k", "input_cost_per_query", @@ -1031,6 +1040,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_pixel": {"type": "number"}, "output_cost_per_second": {"type": "number"}, "output_cost_per_second_480p": {"type": "number"}, + "output_cost_per_second_720p": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, "output_cost_per_second_4k": {"type": "number"}, "output_cost_per_token": {"type": "number"}, @@ -1403,23 +1413,35 @@ def test_supports_tool_choice_simple_tests(): is True ) - assert ( - litellm.utils.supports_tool_choice(model="us.amazon.nova-micro-v1:0") is False - ) - assert ( - litellm.utils.supports_tool_choice(model="bedrock/us.amazon.nova-micro-v1:0") - is False - ) - assert ( - litellm.utils.supports_tool_choice( - model="us.amazon.nova-micro-v1:0", custom_llm_provider="bedrock_converse" - ) - is False - ) - assert litellm.utils.supports_tool_choice(model="perplexity/sonar") is False +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "model", + [ + "amazon.nova-lite-v1:0", + "amazon.nova-micro-v1:0", + "amazon.nova-pro-v1:0", + "apac.amazon.nova-lite-v1:0", + "apac.amazon.nova-micro-v1:0", + "apac.amazon.nova-pro-v1:0", + "bedrock/us-gov-east-1/amazon.nova-pro-v1:0", + "bedrock/us-gov-west-1/amazon.nova-lite-v1:0", + "bedrock/us-gov-west-1/amazon.nova-micro-v1:0", + "bedrock/us-gov-west-1/amazon.nova-pro-v1:0", + "eu.amazon.nova-lite-v1:0", + "eu.amazon.nova-micro-v1:0", + "eu.amazon.nova-pro-v1:0", + "us.amazon.nova-lite-v1:0", + "us.amazon.nova-micro-v1:0", + "us.amazon.nova-pro-v1:0", + ], +) +def test_amazon_nova_v1_understanding_models_support_tool_choice(model: str) -> None: + assert litellm.utils.supports_tool_choice(model=model) is True + + def test_check_provider_match(): """ Test the _check_provider_match function for various provider scenarios @@ -2382,6 +2404,28 @@ def test_register_model_with_scientific_notation(): _invalidate_model_cost_lowercase_map() +@respx.mock +def test_register_model_url_fetch_uses_single_attempt(monkeypatch): + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + before = dict(litellm.model_cost) + threads_before = {thread.name for thread in threading.enumerate()} + route = respx.get("https://example.invalid/custom_pricing.json").mock( + return_value=httpx.Response(503) + ) + + litellm.register_model(model_cost="https://example.invalid/custom_pricing.json") + + threads_after = {thread.name for thread in threading.enumerate()} + assert route.call_count == 1 + assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"} + assert not any( + thread.name == "litellm-model-cost-map-retry" and thread.is_alive() + for thread in threading.enumerate() + ) + assert litellm.model_cost.keys() >= before.keys() + + def test_register_model_openrouter_without_slash(): """ Test that register_model handles openrouter models without '/' in the name. diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 2a60ff9c4b5..f3cd4618078 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -532,6 +532,32 @@ class TestVideoGeneration: assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001 assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001 + def test_completion_cost_xai_imagine_video_720p_tier_from_cost_map(self, monkeypatch): + """720p xAI Imagine Video requests bill the published 720p rate, not the 480p base rate.""" + from litellm.cost_calculator import completion_cost + + local_map_path = os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + with open(local_map_path, "r") as f: + monkeypatch.setattr(litellm, "model_cost", json.load(f)) + + def cost_for(model: str, resolution: str, duration: float) -> float: + mock_response = MagicMock() + mock_response.usage = {"duration_seconds": duration, "video_resolution": resolution} + type(mock_response)._hidden_params = {} + return completion_cost( + completion_response=mock_response, + model=model, + call_type="create_video", + custom_llm_provider="xai", + ) + + assert abs(cost_for("xai/grok-imagine-video", "720p", 10.0) - 0.7) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "720p", 10.0) - 1.4) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "480p", 10.0) - 0.8) < 0.001 + assert abs(cost_for("xai/grok-imagine-video-1.5", "1080p", 10.0) - 2.5) < 0.001 + def test_completion_cost_veo_31_tiers_pin_published_rates(self, monkeypatch): """The gemini and vertex_ai veo 3.1 entries bill Google's published per-second tier rates.""" from litellm.cost_calculator import completion_cost diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 44360e94392..1f459bc50ea 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -11402,9 +11402,9 @@ } }, "node_modules/smol-toml": { - "version": "1.6.1", - "resolved": "https://registry.npmjs.org/smol-toml/-/smol-toml-1.6.1.tgz", - "integrity": "sha512-dWUG8F5sIIARXih1DTaQAX4SsiTXhInKf1buxdY9DIg4ZYPZK5nGM1VRIYmEbDbsHt7USo99xSLFu5Q1IqTmsg==", + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/smol-toml/-/smol-toml-1.8.0.tgz", + "integrity": "sha512-kCZr2V3ch9i00x8zXRhjUNVcjG9ijES5dDudkXvUVCT5QlJNQWElSJdZqyPemffHoLNUYwOcou0Fy+ojN0uHSQ==", "dev": true, "license": "BSD-3-Clause", "engines": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index 64363da9933..1f73671caae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -1,3 +1,4 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; @@ -14,7 +15,9 @@ vi.mock("./useShadowEval", () => ({ })); const authorizedRoleMock = vi.fn(() => ({ accessToken: "token", isViewOnly: false })); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => authorizedRoleMock() })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ userId: "test-user-id", userRole: "Admin", ...authorizedRoleMock() }), +})); vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useInfiniteKeys: vi.fn(() => ({ @@ -68,27 +71,33 @@ vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({ })), })); -vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ +vi.mock("@/app/(dashboard)/hooks/models/useModels", async (importOriginal) => ({ + ...(await importOriginal()), useAutoRouters: vi.fn(() => ({ data: [ { model_name: "claude-auto", litellm_params: { model: "auto_router/claude-auto" } }, { model_name: "gpt-auto", litellm_params: { model: "auto_router/gpt-auto" } }, ], })), - usePlainModelGroups: vi.fn(() => new Set(["prod-claude"])), + usePlainModelGroups: vi.fn(() => new Set(["prod-claude", "prod-judge"])), + usePlainChatModelGroups: vi.fn(() => new Set(["prod-claude", "prod-judge"])), + usePlainChatModelDeployments: vi.fn(() => [ + { + model_name: "prod-judge", + litellm_params: { model: "anthropic/claude-sonnet-5" }, + model_info: { mode: "chat" }, + }, + ]), })); -vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ - useModelCostMap: vi.fn(() => ({ - data: { - "claude-sonnet-5": { litellm_provider: "anthropic", mode: "chat" }, - "gpt-4o": { litellm_provider: "openai", mode: "chat" }, - "gemini/gemini-2.5-pro": { litellm_provider: "gemini", mode: "chat" }, - "text-embedding-3-large": { litellm_provider: "openai", mode: "embedding" }, - }, - })), +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + modelInfoCall: vi.fn(), })); +import { usePlainChatModelGroups, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; +import { modelInfoCall } from "@/components/networking"; + import ShadowEvalSection, { shadowedTargetLabel } from "./ShadowEvalSection"; import { useShadowEvalJob, @@ -107,7 +116,7 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ models: [], direction: "forward", baseline_model: null, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", shadow_percentage: 10, targets: [ { @@ -249,6 +258,85 @@ describe("ShadowEvalSection", () => { if (defaultKeysImpl) vi.mocked(useInfiniteKeys).mockImplementation(defaultKeysImpl); }); + it("labels only configured judge recommendations", async () => { + const user = userEvent.setup(); + mockHooks({}); + render(); + + await user.click(screen.getByPlaceholderText("Select a judge model")); + expect(screen.getByRole("option", { name: /prod-judge.*Recommended/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /openai\/gpt-4o/ })).not.toBeInTheDocument(); + + await user.keyboard("{Escape}"); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); + await user.click(screen.getByPlaceholderText("Select a baseline model")); + expect(screen.getByRole("option", { name: "prod-judge", exact: true })).toBeInTheDocument(); + expect(screen.queryByText("Recommended")).not.toBeInTheDocument(); + }); + + it("keeps custom models selectable through the real model hooks without widening chat choices to traffic filters", async () => { + const hooks = await vi.importActual( + "@/app/(dashboard)/hooks/models/useModels", + ); + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const deployments = [ + { model_name: "custom-chat", litellm_params: { model: "openai/private-chat" } }, + { model_name: "custom-judge", litellm_params: { model: "openai/private-judge" }, model_info: { mode: null } }, + { + model_name: "embedding", + litellm_params: { model: "openai/private-embedding" }, + model_info: { mode: "embedding" }, + }, + { + model_name: "responses-only", + litellm_params: { model: "openai/private-responses" }, + model_info: { mode: "responses" }, + }, + { model_name: "auto-router", litellm_params: { model: "auto_router/complexity_router" } }, + ]; + vi.mocked(modelInfoCall).mockResolvedValue({ data: deployments, total_pages: 1 }); + const user = userEvent.setup(); + const { start } = mockHooks({}); + await vi.mocked(usePlainModelGroups).withImplementation(hooks.usePlainModelGroups, async () => { + await vi.mocked(usePlainChatModelGroups).withImplementation(hooks.usePlainChatModelGroups, async () => { + render( + + + , + ); + await chooseSelectOption(user, screen.getByPlaceholderText("Every model the targets use"), "responses-only"); + await chooseSelectOption(user, screen.getByPlaceholderText("Every model the targets use"), "custom-chat"); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); + await user.click(screen.getByPlaceholderText("Search keys by alias")); + await user.click(within(await screen.findByTestId("paginated-multi-select-list")).getByText("prod-alpha")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); + await user.click(screen.getByPlaceholderText("Select a judge model")); + expect(screen.getAllByRole("option")).toHaveLength(2); + expect(screen.getByRole("option", { name: "custom-chat", exact: true })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "custom-judge", exact: true })).toBeInTheDocument(); + await user.click(screen.getByRole("option", { name: "custom-judge", exact: true })); + await user.click(screen.getByPlaceholderText("Select a baseline model")); + expect(screen.getAllByRole("option")).toHaveLength(2); + expect(screen.getByRole("option", { name: "custom-chat", exact: true })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "custom-judge", exact: true })).toBeInTheDocument(); + await user.click(screen.getByRole("option", { name: "custom-chat", exact: true })); + await user.click(screen.getByText("Start shadow eval")); + expect(start.mutate).toHaveBeenCalledWith( + expect.objectContaining({ judge_model: "custom-judge", baseline_model: "custom-chat", models: [] }), + ); + }); + }); + client.clear(); + }); + it("offers the start form while the list is still loading", () => { mockHooks({ isPending: true }); render(); @@ -444,7 +532,8 @@ describe("ShadowEvalSection", () => { expect(screen.getByText("Start shadow eval")).toBeDisabled(); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + expect(screen.queryByRole("option", { name: /openai\/gpt-4o/ })).not.toBeInTheDocument(); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); const expectedBody = { @@ -457,7 +546,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -474,7 +563,7 @@ describe("ShadowEvalSection", () => { await user.click(within(teamList).getByText("engineering")); await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); const expectedBody = { @@ -487,7 +576,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -503,7 +592,7 @@ describe("ShadowEvalSection", () => { await chooseSelectOption(user, screen.getByPlaceholderText("Every model the targets use"), "prod-claude"); await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); expect(start.mutate).toHaveBeenCalledWith( @@ -524,20 +613,23 @@ describe("ShadowEvalSection", () => { expect(screen.queryByPlaceholderText("Select a baseline model")).not.toBeInTheDocument(); expect(screen.getByPlaceholderText("Every model the targets use")).toBeInTheDocument(); - await user.click(screen.getByText("Adoption check: key's traffic vs the router")); - await user.click(await screen.findByText("Regression check: router's picks vs a baseline")); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); expect(screen.queryByPlaceholderText("Every model the targets use")).not.toBeInTheDocument(); await user.click(screen.getByPlaceholderText("Search keys by alias")); const keyList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(keyList).getByText("prod-alpha")); await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); expect(screen.getByText("Start shadow eval")).toBeDisabled(); await user.click(screen.getByPlaceholderText("Select a baseline model")); - expect(await screen.findByRole("option", { name: /openai\/gpt-4o/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /openai\/gpt-4o/ })).not.toBeInTheDocument(); await user.click(screen.getByRole("option", { name: /prod-claude/ })); await user.click(screen.getByText("Start shadow eval")); @@ -552,7 +644,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -574,7 +666,7 @@ describe("ShadowEvalSection", () => { screen.getByText("Every router sees the same sampled requests, judged against the same live responses"), ).toBeInTheDocument(); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByText("Start shadow eval")); const expectedBody = { @@ -587,7 +679,7 @@ describe("ShadowEvalSection", () => { shadow_percentage: 10, duration_days: 7, max_budget: 10, - judge_model: "anthropic/claude-sonnet-5", + judge_model: "prod-judge", }; expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); @@ -605,10 +697,13 @@ describe("ShadowEvalSection", () => { await user.click(await screen.findByText("gpt-auto")); await user.click(routerInput); await user.click(await screen.findByText("claude-auto")); - await user.click(screen.getByText("Adoption check: key's traffic vs the router")); - await user.click(await screen.findByText("Regression check: router's picks vs a baseline")); + await chooseSelectOption( + user, + screen.getByText("Adoption check: key's traffic vs the router"), + "Regression check: router's picks vs a baseline", + ); await user.click(screen.getByPlaceholderText("Select a judge model")); - await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(await screen.findByRole("option", { name: /prod-judge/ })); await user.click(screen.getByPlaceholderText("Select a baseline model")); await user.click(screen.getByRole("option", { name: /prod-claude/ })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx index 2eb5fa9c945..d85d26a21a8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx @@ -5,8 +5,13 @@ import React, { useMemo, useState } from "react"; import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; -import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; +import { + useAutoRouters, + usePlainChatModelDeployments, + usePlainChatModelGroups, + usePlainModelGroups, +} from "@/app/(dashboard)/hooks/models/useModels"; +import { buildModelAvailability, deploymentRefsFromModelInfo, resolveAvailableModels } from "@/lib/autorouter_presets"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; @@ -24,53 +29,8 @@ type ShadowEvalDirection = ShadowEvalJob["direction"]; const MAX_ROUTERS = 4; const MAX_MODELS = 100; - const RECOMMENDED_JUDGE_MODELS = ["anthropic/claude-sonnet-5", "openai/gpt-4o", "gemini/gemini-2.5-pro"] as const; -interface CostMapEntry { - litellm_provider?: string; - mode?: string; -} - -const useChatModelNames = (): string[] => { - const { data: costMap } = useModelCostMap(); - return useMemo(() => { - if (!costMap) return []; - const chatModels = Object.entries(costMap as Record) - .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) - .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); - return [...new Set(chatModels)].toSorted((a, b) => a.localeCompare(b)); - }, [costMap]); -}; - -const useJudgeModelOptions = (): SearchSelectOption[] => { - const chatModels = useChatModelNames(); - return useMemo(() => { - const pinned: SearchSelectOption[] = RECOMMENDED_JUDGE_MODELS.map((model) => ({ - label: model, - value: model, - sublabel: "Recommended", - })); - const pinnedNames = new Set(RECOMMENDED_JUDGE_MODELS); - const rest = chatModels.filter((model) => !pinnedNames.has(model)).map((model) => ({ label: model, value: model })); - return [...pinned, ...rest]; - }, [chatModels]); -}; - -const useBaselineModelOptions = (): SearchSelectOption[] => { - const configuredGroups = usePlainModelGroups(); - const chatModels = useChatModelNames(); - return useMemo(() => { - const configured = [...configuredGroups] - .toSorted((a, b) => a.localeCompare(b)) - .map((model) => ({ label: model, value: model, sublabel: "Configured on this gateway" })); - const rest = chatModels - .filter((model) => !configuredGroups.has(model)) - .map((model) => ({ label: model, value: model })); - return [...configured, ...rest]; - }, [configuredGroups, chatModels]); -}; - const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[] = [ { value: "forward", label: "Adoption check: key's traffic vs the router" }, { value: "reverse", label: "Regression check: router's picks vs a baseline" }, @@ -276,13 +236,32 @@ export const StartForm: React.FC = () => { const [judgeModel, setJudgeModel] = useState(""); const [maxBudget, setMaxBudget] = useState("10"); const { data: autoRouters } = useAutoRouters(); - const judgeModelOptions = useJudgeModelOptions(); - const baselineModelOptions = useBaselineModelOptions(); const configuredGroups = usePlainModelGroups(); + const chatGroups = usePlainChatModelGroups(); + const chatDeployments = usePlainChatModelDeployments(); const modelOptions = useMemo( () => [...configuredGroups].toSorted((a, b) => a.localeCompare(b)).map((name) => ({ label: name, value: name })), [configuredGroups], ); + const chatOptions = useMemo( + () => modelOptions.filter((option) => chatGroups.has(option.value)), + [modelOptions, chatGroups], + ); + const chatAvailability = useMemo( + () => buildModelAvailability(chatGroups, deploymentRefsFromModelInfo(chatDeployments)), + [chatDeployments, chatGroups], + ); + const recommendedJudgeModels = useMemo( + () => new Set(RECOMMENDED_JUDGE_MODELS.flatMap((model) => resolveAvailableModels(model, chatAvailability))), + [chatAvailability], + ); + const judgeOptions = useMemo( + () => + chatOptions.map((option) => + recommendedJudgeModels.has(option.value) ? { ...option, sublabel: "Recommended" } : option, + ), + [chatOptions, recommendedJudgeModels], + ); const start = useStartShadowEval(); const routerOptions = useMemo(() => { @@ -434,7 +413,7 @@ export const StartForm: React.FC = () => { {direction === "reverse" && ( { )} => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index cfafe82ee30..6489bc2171d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -5,17 +5,19 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { isAutoRouterDeployment, selectAutoRouterModelGroups, - selectPlainModelGroups, + selectPlainChatModelGroups, useAllProxyModels, useAutoRouterModelGroups, useAutoRouters, useInfiniteModelInfo, useModelHub, useModelsInfo, + usePlainChatModelGroups, useSelectedTeamModels, useUserModels, type AllProxyModelsResponse, type AutoRouterCandidateDeployment, + type AutoRouterDeployment, type PaginatedModelInfoResponse, type ProxyModel, } from "./useModels"; @@ -984,29 +986,45 @@ describe("selectAutoRouterModelGroups", () => { }); }); -describe("selectPlainModelGroups", () => { - it("keeps only non-auto-router model groups", () => { - const deployments: AutoRouterCandidateDeployment[] = [ - { model_name: "smart-router", litellm_params: { model: "auto_router/complexity_router" } }, - { model_name: "claude-haiku", litellm_params: { model: "anthropic/claude-haiku-4-5" } }, - { model_name: "claude-sonnet", litellm_params: { model: "anthropic/claude-sonnet-4-5" } }, - { model_name: "cheap-router", litellm_params: { model: "auto_router/adaptive_router" } }, +describe("selectPlainChatModelGroups", () => { + it("keeps chat-capable groups when mode metadata is absent or any sibling is compatible", () => { + const deployments: AutoRouterDeployment[] = [ + { model_name: "no-info" }, + { model_name: "null-info", model_info: null }, + { model_name: "empty-info", model_info: {} }, + { model_name: "missing-mode", model_info: { db_model: false } }, + { model_name: "null-mode", model_info: { mode: null } }, + { model_name: "empty-mode", model_info: { mode: "" } }, + { model_name: "chat", model_info: { mode: "chat", db_model: true } }, + { model_name: "completion", model_info: { mode: "completion" } }, + { model_name: "chat-and-missing", model_info: { mode: "chat" } }, + { model_name: "chat-and-missing" }, + { model_name: "chat-then-embedding", model_info: { mode: "chat" } }, + { model_name: "chat-then-embedding", model_info: { mode: "embedding" } }, + { model_name: "embedding-then-chat", model_info: { mode: "embedding" } }, + { model_name: "embedding-then-chat", model_info: { mode: "chat" } }, + { model_name: "embedding-only", model_info: { mode: "embedding" } }, + { model_name: "speech-only", model_info: { mode: "speech" } }, + { model_name: "shared-router", litellm_params: { model: "openai/gpt-4o" } }, + { model_name: "shared-router", litellm_params: { model: "auto_router/complexity_router" } }, + { model_name: "", model_info: { mode: "chat" } }, ]; - expect(selectPlainModelGroups(deployments)).toEqual(new Set(["claude-haiku", "claude-sonnet"])); - }); - - it("drops a group name that also fronts an auto-router deployment", () => { - const deployments: AutoRouterCandidateDeployment[] = [ - { model_name: "shared-name", litellm_params: { model: "auto_router/complexity_router" } }, - { model_name: "shared-name", litellm_params: { model: "anthropic/claude-sonnet-4-5" } }, - ]; - - expect(selectPlainModelGroups(deployments)).toEqual(new Set()); - }); - - it("drops deployments that have no public model_name", () => { - expect(selectPlainModelGroups([{ model_name: "", litellm_params: { model: "openai/gpt-4o" } }])).toEqual(new Set()); + expect(selectPlainChatModelGroups(deployments)).toEqual( + new Set([ + "no-info", + "null-info", + "empty-info", + "missing-mode", + "null-mode", + "empty-mode", + "chat", + "completion", + "chat-and-missing", + "chat-then-embedding", + "embedding-then-chat", + ]), + ); }); }); @@ -1103,6 +1121,47 @@ describe("useAutoRouterModelGroups", () => { expect(modelInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", 3, 1000); }); + it("uses every page for configured chat groups and keeps custom deployments without mode metadata", async () => { + (modelInfoCall as any).mockImplementation((_t: string, _u: string, _r: string, page: number) => + Promise.resolve( + page === 1 + ? { + data: [ + { model_name: "configured-chat", model_info: { mode: "chat" } }, + { model_name: "embedding-only", model_info: { mode: "embedding" } }, + ], + total_pages: 2, + } + : { + data: [ + { model_name: "custom-no-mode", model_info: { db_model: true } }, + { model_name: "speech-only", model_info: { mode: "speech" } }, + ], + total_pages: 2, + }, + ), + ); + + const { result } = renderHook(() => usePlainChatModelGroups(), { wrapper }); + + await waitFor(() => expect(result.current.size).toBe(2)); + expect(result.current).toEqual(new Set(["configured-chat", "custom-no-mode"])); + expect(modelInfoCall).toHaveBeenCalledTimes(2); + }); + + it("returns an empty chat group set while loading and after failure", async () => { + (modelInfoCall as any).mockReturnValueOnce(new Promise(() => {})); + const loading = renderHook(() => usePlainChatModelGroups(), { wrapper }); + expect(loading.result.current).toEqual(new Set()); + loading.unmount(); + + queryClient.clear(); + (modelInfoCall as any).mockRejectedValueOnce(new Error("boom")); + const failed = renderHook(() => usePlainChatModelGroups(), { wrapper }); + await waitFor(() => expect(modelInfoCall).toHaveBeenCalledTimes(2)); + expect(failed.result.current).toEqual(new Set()); + }); + it("returns an empty set before the model list resolves", () => { (modelInfoCall as any).mockReturnValue(new Promise(() => {})); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index b3a783a71dc..579ee7ff81a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -2,6 +2,7 @@ import { useQuery, useInfiniteQuery, useQueryClient, UseQueryResult } from "@tan import { createQueryKeys } from "../common/queryKeysFactory"; import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "../useAuthorized"; +import { EndpointType, isModeCompatibleWithEndpoint } from "@/components/chat_ui/mode_endpoint_mapping"; export interface ProxyModel { id: string; @@ -87,6 +88,7 @@ export const useModelsInfo = ( const AUTO_ROUTER_MODEL_PREFIX = "auto_router/"; const AUTO_ROUTER_LOOKUP_PAGE_SIZE = 1000; const NO_AUTO_ROUTERS: ReadonlySet = new Set(); +const NO_DEPLOYMENTS: AutoRouterDeployment[] = []; export interface AutoRouterCandidateDeployment { model_name?: string | null; @@ -96,6 +98,7 @@ export interface AutoRouterCandidateDeployment { export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { litellm_params?: { model?: string | null; + base_model?: string | null; complexity_router_config?: unknown; complexity_router_default_model?: string | null; auto_router_config?: unknown; @@ -111,6 +114,7 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { /** False for config.yaml-defined deployments, which the update and delete routes refuse. */ db_model?: boolean | null; base_model?: string | null; + mode?: string | null; created_at?: string | null; updated_at?: string | null; team_id?: string | null; @@ -142,6 +146,22 @@ export const selectPlainModelGroups = (deployments: AutoRouterCandidateDeploymen ); }; +export const selectPlainChatModelDeployments = (deployments: AutoRouterDeployment[]): AutoRouterDeployment[] => { + const plainGroups = selectPlainModelGroups(deployments); + return deployments.filter( + (deployment) => + plainGroups.has(deployment.model_name ?? "") && + isModeCompatibleWithEndpoint(deployment.model_info?.mode, EndpointType.CHAT), + ); +}; + +export const selectPlainChatModelGroups = (deployments: AutoRouterDeployment[]): ReadonlySet => + new Set( + selectPlainChatModelDeployments(deployments) + .map((deployment) => deployment.model_name) + .filter((name): name is string => Boolean(name)), + ); + export const fetchAllModelDeployments = async ( accessToken: string, userId: string, @@ -180,37 +200,32 @@ export const autoRouterListKey = (userId: string | null, userRole: string | null }, }); -export const useAutoRouterModelGroups = (): ReadonlySet => { +const useDeployments = ( + select: (deployments: AutoRouterDeployment[]) => TSelected, +): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); - const { data } = useQuery>({ + return useQuery({ queryKey: autoRouterListKey(userId, userRole), queryFn: async () => await fetchAllModelDeployments(accessToken!, userId!, userRole!), enabled: Boolean(accessToken && userId && userRole), - select: selectAutoRouterModelGroups, + select, }); - return data ?? NO_AUTO_ROUTERS; }; -export const usePlainModelGroups = (): ReadonlySet => { - const { accessToken, userId, userRole } = useAuthorized(); - const { data } = useQuery>({ - queryKey: autoRouterListKey(userId, userRole), - queryFn: async () => await fetchAllModelDeployments(accessToken!, userId!, userRole!), - enabled: Boolean(accessToken && userId && userRole), - select: selectPlainModelGroups, - }); - return data ?? NO_AUTO_ROUTERS; -}; +export const useAutoRouterModelGroups = (): ReadonlySet => + useDeployments(selectAutoRouterModelGroups).data ?? NO_AUTO_ROUTERS; -export const useAutoRouters = (): UseQueryResult => { - const { accessToken, userId, userRole } = useAuthorized(); - return useQuery({ - queryKey: autoRouterListKey(userId, userRole), - queryFn: async () => await fetchAllModelDeployments(accessToken!, userId!, userRole!), - enabled: Boolean(accessToken && userId && userRole), - select: selectAutoRouterDeployments, - }); -}; +export const usePlainModelGroups = (): ReadonlySet => + useDeployments(selectPlainModelGroups).data ?? NO_AUTO_ROUTERS; + +export const usePlainChatModelGroups = (): ReadonlySet => + useDeployments(selectPlainChatModelGroups).data ?? NO_AUTO_ROUTERS; + +export const usePlainChatModelDeployments = (): AutoRouterDeployment[] => + useDeployments(selectPlainChatModelDeployments).data ?? NO_DEPLOYMENTS; + +export const useAutoRouters = (): UseQueryResult => + useDeployments(selectAutoRouterDeployments); export const useInvalidateAutoRouters = (): (() => Promise) => { const queryClient = useQueryClient(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 7f1f4cc4bd5..3fe34610260 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -29,6 +29,10 @@ vi.mock("@/components/NoRedisWarningBanner", () => ({ NoRedisWarningBanner: () => null, })); +vi.mock("@/components/EnvCredentialLoginWarningBanner", () => ({ + EnvCredentialLoginWarningBanner: () => null, +})); + vi.mock("@/components/LicenseExpiryBanner", () => ({ LicenseExpiryBanner: () => null, })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 98f2a36d6f3..fa6df7f176a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -10,6 +10,7 @@ import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import { useRouter, useSearchParams } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner"; +import { EnvCredentialLoginWarningBanner } from "@/components/EnvCredentialLoginWarningBanner"; import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner"; import { UserBanner } from "@/components/UserBanner"; import { uiHref } from "@/utils/uiHref"; @@ -113,6 +114,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
@@ -132,6 +134,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
{children}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts new file mode 100644 index 00000000000..c7147466d43 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts @@ -0,0 +1,105 @@ +import { describe, expect, it } from "vitest"; +import { getEditToolPreview } from "./editToolPreview"; + +const saved = { + url: "https://example.com/mcp", + transport: "http", + auth_type: "basic", + static_headers: [{ header: "X-Tenant", value: "original" }], +}; + +describe("getEditToolPreview", () => { + it("keeps saved discovery for unchanged settings and unrelated edits", () => { + expect(getEditToolPreview({ ...saved, server_name: "renamed" }, saved)).toEqual({ kind: "saved" }); + }); + + it("previews URL changes with the existing server credential left for server-side inheritance", () => { + expect(getEditToolPreview({ ...saved, url: "https://example.com/corrected-mcp" }, saved)).toEqual({ + kind: "preview", + config: { + url: "https://example.com/corrected-mcp", + transport: "http", + auth_type: "basic", + static_headers: { "X-Tenant": "original" }, + credentials: undefined, + }, + }); + }); + + it.each(["https://other.example/mcp", "http://example.com/mcp", "https://example.com:8443/mcp"])( + "requires explicit credentials for a changed origin: %s", + (url) => { + expect(getEditToolPreview({ ...saved, url, static_headers: [] }, saved)).toEqual({ + kind: "incomplete", + message: expect.stringContaining("origin changed"), + }); + const explicit = { ...saved, url, static_headers: [], credentials: { auth_value: "new:secret" } }; + expect(getEditToolPreview(explicit, saved).kind).toBe("preview"); + }, + ); + + it("does not automatically send saved static headers to a new origin", () => { + expect(getEditToolPreview({ ...saved, url: "https://other.example/mcp", auth_type: "none" }, saved).kind).toBe( + "incomplete", + ); + }); + + it("uses edited static headers and only the static auth value", () => { + expect( + getEditToolPreview( + { + ...saved, + static_headers: [{ header: "X-Tenant", value: "corrected" }], + credentials: { auth_value: "user:password", access_token: "old-oauth-token", client_secret: "old-client" }, + }, + saved, + ), + ).toEqual({ + kind: "preview", + config: { + url: saved.url, + transport: "http", + auth_type: "basic", + static_headers: { "X-Tenant": "corrected" }, + credentials: { auth_value: "user:password" }, + }, + }); + }); + + it.each(["", "https://", "file:///tmp/server"])("does not connect to an incomplete or unsupported URL: %s", (url) => { + expect(getEditToolPreview({ ...saved, url }, saved)).toEqual({ kind: "incomplete" }); + }); + + it("waits for a static header value before connecting", () => { + const values = { ...saved, static_headers: [{ header: "X-Tenant", value: "" }] }; + expect(getEditToolPreview(values, saved)).toEqual({ kind: "incomplete" }); + }); + + it("waits for credentials when switching from None to Basic Auth", () => { + expect(getEditToolPreview(saved, { ...saved, auth_type: "none" })).toEqual({ kind: "incomplete" }); + }); + + it("does not forward old credentials when switching to None", () => { + const result = getEditToolPreview( + { ...saved, auth_type: "none", credentials: { auth_value: "old-secret" } }, + saved, + ); + expect(result.kind).toBe("preview"); + if (result.kind === "preview") expect(result.config.credentials).toBeUndefined(); + }); + + it.each(["oauth2", "true_passthrough", "oauth_delegate", "oauth2_token_exchange", "oauth2_id_jag", "aws_sigv4"])( + "preserves the existing discovery path for %s", + (auth_type) => { + expect(getEditToolPreview({ ...saved, auth_type, url: "https://changed.example/mcp" }, saved)).toEqual({ + kind: "saved", + }); + }, + ); + + it("keeps stdio and OpenAPI on their existing discovery path", () => { + for (const transport of ["stdio", "openapi"]) { + expect(getEditToolPreview({ ...saved, transport }, saved)).toEqual({ kind: "saved" }); + } + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts new file mode 100644 index 00000000000..6dea7d7d18e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts @@ -0,0 +1,67 @@ +import { AUTH_TYPE, TRANSPORT } from "@/components/mcp_tools/types"; +import { AUTH_TYPES_REQUIRING_AUTH_VALUE, reduceStaticHeaders } from "./createServerPayload"; + +const connectionConfig = (values: Readonly>) => { + const credentials = values.credentials; + const authValue = + credentials && typeof credentials === "object" && "auth_value" in credentials ? credentials.auth_value : undefined; + const needsAuthValue = + typeof values.auth_type === "string" && AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(values.auth_type); + return { + url: typeof values.url === "string" ? values.url : "", + transport: typeof values.transport === "string" ? values.transport : "", + auth_type: typeof values.auth_type === "string" ? values.auth_type : "", + static_headers: Object.fromEntries( + Object.entries(reduceStaticHeaders(values.static_headers)).sort(([a], [b]) => a.localeCompare(b)), + ), + credentials: + needsAuthValue && typeof authValue === "string" && authValue.trim() ? { auth_value: authValue } : undefined, + }; +}; + +type EditToolPreview = + | { readonly kind: "saved" } + | { readonly kind: "incomplete"; readonly message?: string } + | { readonly kind: "preview"; readonly config: ReturnType }; + +export const getEditToolPreview = ( + values: Readonly>, + initialValues: Readonly>, +): EditToolPreview => { + const staticAuth = + values.auth_type === AUTH_TYPE.NONE || + (typeof values.auth_type === "string" && AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(values.auth_type)); + if (!staticAuth || ![TRANSPORT.HTTP, TRANSPORT.SSE].includes(String(values.transport))) { + return { kind: "saved" }; + } + + const config = connectionConfig(values); + if (JSON.stringify(config) === JSON.stringify(connectionConfig(initialValues))) { + return { kind: "saved" }; + } + + const missingNewCredential = + config.auth_type !== initialValues.auth_type && + AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) && + config.credentials === undefined; + const validUrl = URL.canParse(config.url) && ["http:", "https:"].includes(new URL(config.url).protocol); + const incompleteHeaders = Object.values(config.static_headers).some((value) => !value.trim()); + if (!validUrl || missingNewCredential || incompleteHeaders) { + return { kind: "incomplete" }; + } + const savedConfig = connectionConfig(initialValues); + const changedOrigin = + !URL.canParse(savedConfig.url) || new URL(config.url).origin !== new URL(savedConfig.url).origin; + const reusesHeader = Object.entries(config.static_headers).some( + ([key, value]) => savedConfig.static_headers[key] === value, + ); + const needsSavedCredential = AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) && !config.credentials; + if (changedOrigin && (needsSavedCredential || reusesHeader)) { + return { + kind: "incomplete", + message: + "The server origin changed. Enter credentials and replace or remove saved static headers to preview tools.", + }; + } + return { kind: "preview", config }; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx index a9664191f3f..f3cd37cc580 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.integration.test.tsx @@ -1,6 +1,9 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; -import { render, screen, waitFor, act } from "@testing-library/react"; +import { render, screen, waitFor, act, fireEvent } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; + +import { selectOption } from "./testUtils"; import MCPServerEdit from "./mcp_server_edit"; import * as networking from "@/components/networking"; @@ -27,10 +30,6 @@ vi.mock("./mcp_server_cost_config", () => ({ default: () =>
, })); -vi.mock("./mcp_tool_configuration", () => ({ - default: () =>
, -})); - const BASE: MCPServer = { server_id: "srv_1", server_name: "srv", @@ -369,3 +368,105 @@ describe("mcp_server_edit save payload contract", () => { } }); }); + +describe("MCPServerEdit live tool preview", () => { + beforeEach(() => { + vi.resetAllMocks(); + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: [], + error: "connection_error", + message: "Saved credentials rejected", + }); + vi.mocked(networking.testMCPToolsListRequest).mockResolvedValue({ + tools: [ + { name: "echo", description: "Echo the supplied message", inputSchema: { type: "object", properties: {} } }, + ], + }); + }); + + const renderEditor = (server: MCPServer = BASE) => + render( + , + ); + + it("replaces the saved connection failure with tools after correcting Basic Auth without saving", async () => { + renderEditor(); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + await selectOption("Authentication", "Basic Auth"); + fireEvent.change(screen.getByLabelText("Authentication Value"), { target: { value: "preview:correct" } }); + expect(screen.queryByText("Saved credentials rejected")).not.toBeInTheDocument(); + expect(screen.getByText("Loading tools...")).toBeInTheDocument(); + expect(networking.testMCPToolsListRequest).not.toHaveBeenCalled(); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + const expectedConfig = { + server_id: BASE.server_id, + url: BASE.url, + auth_type: "basic", + credentials: { auth_value: "preview:correct" }, + }; + expect(networking.testMCPToolsListRequest).toHaveBeenCalledExactlyOnceWith( + "access-token", + expect.objectContaining(expectedConfig), + ); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + }); + + it("refreshes tools when a static header is corrected", async () => { + renderEditor({ ...BASE, static_headers: { "X-Preview-Key": "wrong" } }); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + fireEvent.change(screen.getByPlaceholderText("Header value"), { target: { value: "correct" } }); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(networking.testMCPToolsListRequest).toHaveBeenCalledExactlyOnceWith( + "access-token", + expect.objectContaining({ static_headers: { "X-Preview-Key": "correct" } }), + ); + }); + + it("coalesces URL edits and ignores an older failed preview after the latest preview succeeds", async () => { + const user = userEvent.setup(); + const older = Promise.withResolvers<{ tools: never[]; error: string; message: string }>(); + vi.mocked(networking.testMCPToolsListRequest).mockImplementationOnce(() => older.promise); + renderEditor(); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("MCP Server URL"), { target: { value: "https://first.example/mcp" } }); + await waitFor(() => expect(networking.testMCPToolsListRequest).toHaveBeenCalledTimes(1)); + await user.clear(screen.getByLabelText("MCP Server URL")); + await user.type(screen.getByLabelText("MCP Server URL"), "https://latest.example/mcp"); + expect(networking.testMCPToolsListRequest).toHaveBeenCalledTimes(1); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(networking.testMCPToolsListRequest).toHaveBeenCalledTimes(2); + expect(networking.testMCPToolsListRequest).toHaveBeenLastCalledWith( + "access-token", + expect.objectContaining({ url: "https://latest.example/mcp" }), + ); + await act(async () => older.resolve({ tools: [], error: "connection_error", message: "Older request failed" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(screen.queryByText("Older request failed")).not.toBeInTheDocument(); + }); + + it("ignores a saved-record response after editing and restores saved discovery when changes are reverted", async () => { + const saved = Promise.withResolvers<{ tools: never[]; error: string; message: string }>(); + vi.mocked(networking.listMCPTools).mockImplementationOnce(() => saved.promise); + renderEditor(); + await waitFor(() => expect(networking.listMCPTools).toHaveBeenCalledTimes(1)); + fireEvent.change(screen.getByLabelText("MCP Server URL"), { target: { value: "https://correct.example/mcp" } }); + fireEvent.click(await screen.findByRole("button", { name: "Flat List" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + await act(async () => saved.resolve({ tools: [], error: "connection_error", message: "Stale saved response" })); + expect(screen.getByText("echo")).toBeInTheDocument(); + expect(screen.queryByText("Stale saved response")).not.toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("MCP Server URL"), { target: { value: BASE.url } }); + expect(await screen.findByText("Saved credentials rejected")).toBeInTheDocument(); + expect(networking.listMCPTools).toHaveBeenCalledTimes(2); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 8793c45371a..2a37029a2c4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -52,6 +52,7 @@ import EnvVarsSection from "./EnvVarsSection"; import { validateMCPServerUrl, validateMCPServerName, normalizeToolOverrideMap } from "./utils"; import { EditServerFormValues, buildEditServerPayload, editPayloadErrorMessage } from "./editServerPayload"; import { toast } from "@/lib/toast"; +import { getEditToolPreview } from "./editToolPreview"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { MountedFormField, @@ -449,14 +450,30 @@ const MCPServerEdit: React.FC = ({ } }, [mcpServer]); - // Fetch tools when component mounts for a saved server + const toolPreview = getEditToolPreview(allFieldsValue(form), initialValues); + const toolPreviewKey = JSON.stringify(toolPreview); + useEffect(() => { - if (!mcpServer.server_id || mcpServer.server_id.trim() === "") { + const controller = new AbortController(); + setTools([]); + setToolsError(null); + setIsLoadingTools(false); + if (!accessToken || !mcpServer.server_id) return; + if (toolPreview.kind === "incomplete") { + setToolsError(toolPreview.message ?? "Complete the URL, authentication, and header settings to load tools."); return; } - fetchTools(); + setIsLoadingTools(true); + const timer = setTimeout( + () => fetchTools(() => !controller.signal.aborted), + toolPreview.kind === "preview" ? 500 : 0, + ); + return () => { + controller.abort(); + clearTimeout(timer); + }; // eslint-disable-next-line react-hooks/exhaustive-deps - }, [mcpServer, accessToken, userID, oauthTokenResponse?.access_token]); + }, [mcpServer, accessToken, userID, oauthTokenResponse?.access_token, toolPreviewKey]); // Invalidate a token authorized in this edit session once any mint-relevant field diverges from the // identity it was minted against (url, auth_type, oauth_flow_type, client creds/scopes, or the @@ -519,6 +536,7 @@ const MCPServerEdit: React.FC = ({ const previewWithStagedInteractiveToken = async ( isPassthrough: boolean, isBrowserHeldTokenMode: boolean, + isCurrent: () => boolean, ): Promise => { const stagedToken = !isPassthrough && !isBrowserHeldTokenMode && getEffectiveAuthType() === AUTH_TYPE.OAUTH2 @@ -550,6 +568,7 @@ const MCPServerEdit: React.FC = ({ registration_url: values.registration_url, }; const toolsResponse = await testMCPToolsListRequest(accessToken, previewConfig, stagedToken); + if (!isCurrent()) return true; if (toolsResponse.tools && !toolsResponse.error) { setTools(toolsResponse.tools); } else { @@ -557,15 +576,16 @@ const MCPServerEdit: React.FC = ({ setToolsError(toolsResponse.message || "Failed to load tools"); } } catch (error) { + if (!isCurrent()) return true; setTools([]); setToolsError(error instanceof Error ? error.message : "Failed to load tools"); } finally { - setIsLoadingTools(false); + if (isCurrent()) setIsLoadingTools(false); } return true; }; - const fetchTools = async () => { + const fetchTools = async (isCurrent: () => boolean) => { if (!accessToken || !mcpServer.server_id) return; // OBO/M2M/static auth is attached server-side from the stored credential, so @@ -574,6 +594,7 @@ const MCPServerEdit: React.FC = ({ // same way the Tools playground does. let customHeaders: Record | undefined; const isPassthrough = + toolPreview.kind === "saved" && getMcpOAuthMode({ auth_type: mcpServer.auth_type, oauth2_flow: mcpServer.oauth2_flow, @@ -581,9 +602,10 @@ const MCPServerEdit: React.FC = ({ }) === "passthrough"; const isBrowserHeldTokenMode = isClientForwardedTokenMode(getEffectiveAuthType()); - if (await previewWithStagedInteractiveToken(isPassthrough, isBrowserHeldTokenMode)) { + if (await previewWithStagedInteractiveToken(isPassthrough, isBrowserHeldTokenMode, isCurrent)) { return; } + if (!isCurrent()) return; if (isPassthrough || isBrowserHeldTokenMode) { const token = oauthTokenResponse?.access_token ?? @@ -591,6 +613,7 @@ const MCPServerEdit: React.FC = ({ ? getToken(mcpServer.server_id, userID)?.access_token ?? null : null); if (!token) { + setIsLoadingTools(false); setTools([]); setToolsError( isBrowserHeldTokenMode @@ -608,7 +631,15 @@ const MCPServerEdit: React.FC = ({ try { // include_disabled_tools: configuring the allowlist needs the full server // catalog, so tools toggled off still render (as unchecked) instead of vanishing. - const toolsResponse = await listMCPTools(accessToken, mcpServer.server_id, customHeaders, true); + const toolsResponse = + toolPreview.kind === "preview" + ? await testMCPToolsListRequest(accessToken, { + ...toolPreview.config, + server_id: mcpServer.server_id, + server_name: mcpServer.server_name || mcpServer.alias, + }) + : await listMCPTools(accessToken, mcpServer.server_id, customHeaders, true); + if (!isCurrent()) return; if (toolsResponse.tools && !toolsResponse.error) { setTools(toolsResponse.tools); @@ -617,10 +648,11 @@ const MCPServerEdit: React.FC = ({ setToolsError(toolsResponse.message || "Failed to load tools"); } } catch (error) { + if (!isCurrent()) return; setTools([]); setToolsError(error instanceof Error ? error.message : "Failed to load tools"); } finally { - setIsLoadingTools(false); + if (isCurrent()) setIsLoadingTools(false); } }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx index 274bdf63e32..a3d62de9941 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_tool_configuration.tsx @@ -433,7 +433,7 @@ const MCPToolConfiguration: React.FC = ({ {isLoadingTools && (
-

Loading tools from spec...

+

Loading tools...

)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index f07b66efdf8..bf6b092a2b0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import ChatUI from "./ChatUI"; import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages"; vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), @@ -14,6 +15,10 @@ vi.mock("@/components/llm_calls/chat_completion", () => ({ makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined), })); +vi.mock("../../llm_calls/anthropic_messages", () => ({ + makeAnthropicMessagesRequest: vi.fn().mockResolvedValue(undefined), +})); + vi.mock("@/components/networking", () => ({ tagListCall: vi.fn().mockResolvedValue({}), vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), @@ -32,6 +37,8 @@ beforeEach(() => { const CHAT_REQUEST_ARG_COUNT = 26; const STREAMING_ENABLED_ARG_INDEX = 25; +const MESSAGES_REQUEST_ARG_COUNT = 19; +const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18; async function openComboboxByPlaceholder(placeholder: string) { const user = userEvent.setup(); @@ -378,6 +385,52 @@ describe("ChatUI", () => { expect(requestArgs[STREAMING_ENABLED_ARG_INDEX]).toBe(false); }); + it("should send the /v1/messages request non-streaming after Stream responses is unchecked", async () => { + const user = userEvent.setup(); + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/messages"); + await selectComboboxOption("Select a Model", "Model 1"); + + await user.click(await screen.findByTestId("model-settings-button")); + + const streamingCheckbox = await screen.findByRole("checkbox", { name: /Stream responses/i }); + expect(streamingCheckbox).toBeChecked(); + await user.click(streamingCheckbox); + + await waitFor(() => { + expect(screen.getByRole("checkbox", { name: /Stream responses/i })).not.toBeChecked(); + }); + + const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)"); + await act(async () => { + fireEvent.change(messageInput, { target: { value: "hello" } }); + }); + await act(async () => { + fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" }); + }); + + await waitFor(() => { + expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1); + }); + + const requestArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0]; + expect(requestArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT); + expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false); + }); + it("should force streaming in simplified mode even when the playground setting is off", async () => { sessionStorage.setItem("streamingEnabled", "false"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 35378d3d4e7..eca9e2323ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -1025,6 +1025,7 @@ const ChatUI: React.FC = ({ mcpServers, mcpServerToolRestrictions, mcpToolsets, + streamingEnabled, ); } else if (endpointType === EndpointType.EMBEDDINGS) { await makeOpenAIEmbeddingsRequest( @@ -1174,7 +1175,10 @@ const ChatUI: React.FC = ({ return !model.mode || model.mode === "chat"; }; - const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES; + const supportsStreamingToggle = + endpointType === EndpointType.CHAT || + endpointType === EndpointType.RESPONSES || + endpointType === EndpointType.ANTHROPIC_MESSAGES; const modelsForEndpoint = useMemo( () => filterModelsForEndpoint(modelInfo, endpointType as EndpointType), [modelInfo, endpointType], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx index 8fd4a57dbfd..87d569f6490 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx @@ -1,7 +1,9 @@ import { ModelGroup } from "@/components/llm_calls/fetch_models"; -import { EndpointType, getEndpointType, ModelMode } from "@/components/chat_ui/mode_endpoint_mapping"; - -const KNOWN_MODEL_MODES = new Set(Object.values(ModelMode)); +import { + EndpointType, + getEndpointType, + isModeCompatibleWithEndpoint, +} from "@/components/chat_ui/mode_endpoint_mapping"; export const determineEndpointType = (selectedModel: string, modelInfo: ModelGroup[]): EndpointType => { const selectedModelInfo = modelInfo.find((option) => option.model_group === selectedModel); @@ -13,31 +15,8 @@ export const determineEndpointType = (selectedModel: string, modelInfo: ModelGro return EndpointType.CHAT; }; -export const isModelCompatibleWithEndpoint = (model: ModelGroup, endpointType: EndpointType): boolean => { - if (!model.mode) { - return true; - } - - if (!KNOWN_MODEL_MODES.has(model.mode)) { - return false; - } - - const optionEndpoint = getEndpointType(model.mode); - - if ( - endpointType === EndpointType.RESPONSES || - endpointType === EndpointType.ANTHROPIC_MESSAGES || - endpointType === EndpointType.INTERACTIONS - ) { - return optionEndpoint === endpointType || optionEndpoint === EndpointType.CHAT; - } - - if (endpointType === EndpointType.IMAGE_EDITS) { - return optionEndpoint === endpointType || optionEndpoint === EndpointType.IMAGE; - } - - return optionEndpoint === endpointType; -}; +export const isModelCompatibleWithEndpoint = (model: ModelGroup, endpointType: EndpointType): boolean => + isModeCompatibleWithEndpoint(model.mode, endpointType); export const filterModelsForEndpoint = (models: ModelGroup[], endpointType: EndpointType): ModelGroup[] => models.filter((model) => isModelCompatibleWithEndpoint(model, endpointType)); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx index 96ace129f87..9f030d8b2d4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx @@ -7,13 +7,27 @@ vi.mock("@/components/networking", () => ({ })); const mockMessagesStream = vi.fn(); +const mockMessagesCreate = vi.fn(); vi.mock("@anthropic-ai/sdk", () => ({ default: vi.fn(function () { - return { messages: { stream: mockMessagesStream } }; + return { messages: { stream: mockMessagesStream, create: mockMessagesCreate } }; }), })); +const NON_STREAMING_ARGS = [ + undefined, // traceId + undefined, // vector_store_ids + undefined, // guardrails + undefined, // policies + undefined, // selectedMCPServers + undefined, // customBaseUrl + undefined, // mcpServers + undefined, // mcpServerToolRestrictions + undefined, // mcpToolsets + false, // streamingEnabled +] as const; + describe("anthropic_messages prompt cache usage", () => { const captureUsage = async (usage: Record): Promise => { async function* mockStream() { @@ -59,3 +73,53 @@ describe("anthropic_messages prompt cache usage", () => { expect(usageData.promptTokens).toBe(5000); }); }); + +describe("anthropic_messages non-streaming", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("sends stream:false through messages.create and renders the full reply at once", async () => { + mockMessagesCreate.mockResolvedValue({ + content: [ + { type: "thinking", thinking: "considering" }, + { type: "text", text: "OK" }, + ], + usage: { input_tokens: 12, output_tokens: 3, cache_read_input_tokens: 7 }, + }); + const updateTextUI = vi.fn(); + const onReasoningContent = vi.fn(); + const onUsageData = vi.fn(); + + await makeAnthropicMessagesRequest( + [{ role: "user", content: "Hello" }], + updateTextUI, + "claude-haiku-4-5", + "test-token", + undefined, + undefined, + onReasoningContent, + undefined, + onUsageData, + ...NON_STREAMING_ARGS, + ); + + expect(mockMessagesStream).not.toHaveBeenCalled(); + expect(mockMessagesCreate).toHaveBeenCalledTimes(1); + expect(mockMessagesCreate.mock.calls[0][0]).toMatchObject({ model: "claude-haiku-4-5", stream: false }); + expect(updateTextUI).toHaveBeenCalledWith("assistant", "OK", "claude-haiku-4-5"); + expect(onReasoningContent).toHaveBeenCalledWith("considering"); + const expectedUsage: TokenUsage = { completionTokens: 3, promptTokens: 12, totalTokens: 15, cacheReadTokens: 7 }; + expect(onUsageData).toHaveBeenCalledWith(expectedUsage); + }); + + it("keeps streaming as the default when the flag is omitted", async () => { + async function* emptyStream() {} + mockMessagesStream.mockReturnValue(emptyStream()); + + await makeAnthropicMessagesRequest([{ role: "user", content: "Hello" }], vi.fn(), "claude-haiku-4-5", "test-token"); + + expect(mockMessagesCreate).not.toHaveBeenCalled(); + expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx index 14afa013768..9dd2f675c44 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx @@ -7,6 +7,13 @@ import { getProxyBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +const toTokenUsage = (usage: Anthropic.Usage): TokenUsage => ({ + completionTokens: usage.output_tokens, + promptTokens: usage.input_tokens, + totalTokens: usage.input_tokens + usage.output_tokens, + ...extractPromptCacheTokens(usage), +}); + export async function makeAnthropicMessagesRequest( messages: MessageType[], updateTextUI: (role: string, delta: string, model?: string) => void, @@ -26,6 +33,7 @@ export async function makeAnthropicMessagesRequest( mcpServers?: MCPServer[], mcpServerToolRestrictions?: Record, mcpToolsets?: MCPToolset[], + streamingEnabled: boolean = true, ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -58,7 +66,7 @@ export async function makeAnthropicMessagesRequest( const requestBody: any = { model: selectedModel, messages: messages.map((m) => ({ role: m.role, content: m.content })), - stream: true, + stream: streamingEnabled, max_tokens: 1024, // @ts-ignore - litellm specific parameter litellm_trace_id: traceId, @@ -74,6 +82,20 @@ export async function makeAnthropicMessagesRequest( if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids; if (guardrails) requestBody.guardrails = guardrails; if (policies) requestBody.policies = policies; + + if (!streamingEnabled) { + const message: Anthropic.Message = await client.messages.create({ ...requestBody, stream: false }, { signal }); + for (const block of message.content) { + if (block.type === "text") { + updateTextUI("assistant", block.text, selectedModel); + } else if (block.type === "thinking" && onReasoningContent) { + onReasoningContent(block.thinking); + } + } + onUsageData?.(toTokenUsage(message.usage)); + return; + } + // Use the streaming helper method for cleaner async iteration // @ts-ignore - The SDK types might not include all litellm-specific parameters const stream = client.messages.stream(requestBody, { signal }); @@ -105,14 +127,7 @@ export async function makeAnthropicMessagesRequest( // Process usage data from message_delta events if (messageStreamEvent.type === "message_delta" && (messageStreamEvent as any).usage && onUsageData) { - const usage = (messageStreamEvent as any).usage; - const usageData: TokenUsage = { - completionTokens: usage.output_tokens, - promptTokens: usage.input_tokens, - totalTokens: usage.input_tokens + usage.output_tokens, - ...extractPromptCacheTokens(usage), - }; - onUsageData(usageData); + onUsageData(toTokenUsage((messageStreamEvent as any).usage)); } } } catch (error) { diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx new file mode 100644 index 00000000000..535694f14da --- /dev/null +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx @@ -0,0 +1,76 @@ +import { renderWithProviders, screen } from "../../tests/test-utils"; +import { vi } from "vitest"; +import { EnvCredentialLoginWarningBanner } from "./EnvCredentialLoginWarningBanner"; +import type { HealthReadinessDetailsResponse } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import type { UseQueryResult } from "@tanstack/react-query"; + +vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails", () => ({ + useHealthReadinessDetails: vi.fn(), +})); +vi.mock("@/contexts/AuthContext", () => ({ + useAuth: vi.fn(), +})); + +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useAuth } from "@/contexts/AuthContext"; + +const mockDetails = (data: Partial | undefined) => { + vi.mocked(useHealthReadinessDetails).mockReturnValue({ data } as UseQueryResult); +}; + +const mockRole = (userRole: string) => { + vi.mocked(useAuth).mockReturnValue({ userRole } as ReturnType); +}; + +describe("EnvCredentialLoginWarningBanner", () => { + it("should warn an admin when env-credential login is enabled", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(screen.getByText("Environment-credential login is enabled")).toBeInTheDocument(); + }); + + it("should tell the admin to create a regular admin account before disabling", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByText(/First create a regular admin account/i)).toBeInTheDocument(); + expect(screen.getByText("general_settings.disable_env_credential_login: true")).toBeInTheDocument(); + }); + + it("should warn an admin viewer too", () => { + mockRole("Admin Viewer"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + renderWithProviders(); + expect(screen.getByRole("alert")).toBeInTheDocument(); + }); + + it("should render nothing for a non-admin even when the proxy reports the warning", () => { + mockRole("Internal User"); + mockDetails({ status: "healthy", show_env_credential_login_warning: true }); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render nothing when env-credential login is disabled", () => { + mockRole("Admin"); + mockDetails({ status: "healthy", show_env_credential_login_warning: false }); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render nothing when readiness details are unavailable", () => { + mockRole("Admin"); + mockDetails(undefined); + const { container } = renderWithProviders(); + expect(container).toBeEmptyDOMElement(); + }); + + it("should pass the access token to the readiness hook", () => { + mockRole("Admin"); + mockDetails(undefined); + renderWithProviders(); + expect(useHealthReadinessDetails).toHaveBeenCalledWith("my-token"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx new file mode 100644 index 00000000000..3a9d50011f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx @@ -0,0 +1,35 @@ +"use client"; + +import React from "react"; +import { TriangleAlert } from "lucide-react"; +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useAuth } from "@/contexts/AuthContext"; +import { isAdminRole } from "@/utils/roles"; + +export const EnvCredentialLoginWarningBanner: React.FC<{ accessToken: string | null }> = ({ accessToken }) => { + const { userRole } = useAuth(); + const { data: healthData } = useHealthReadinessDetails(accessToken); + + if (!isAdminRole(userRole) || !healthData?.show_env_credential_login_warning) { + return null; + } + + return ( +
+
+ ); +}; diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx index 040fa463f9c..af3f98b746b 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.test.tsx @@ -89,6 +89,65 @@ describe("PassThroughEndpointsTable", () => { expect(onDeleteClick).toHaveBeenCalledWith("ep-1"); }); + it("should disable edit and delete for config-defined endpoints", async () => { + const user = userEvent.setup(); + const onEndpointClick = vi.fn(); + const onDeleteClick = vi.fn(); + const configEndpoint: passThroughItem = { + id: "ep-config", + path: "/from-config", + target: "https://config.example.com", + headers: {}, + is_from_config: true, + }; + render( + , + ); + + await user.click(screen.getByTestId("endpoint-actions-ep-config")); + const editItem = await screen.findByTestId("endpoint-action-edit"); + const deleteItem = await screen.findByTestId("endpoint-action-delete"); + + expect(editItem).toHaveAttribute("data-disabled"); + expect(deleteItem).toHaveAttribute("data-disabled"); + expect(screen.getByTestId("endpoint-config-hint")).toHaveTextContent( + "This endpoint is defined in the config file and cannot be edited or deleted on the dashboard.", + ); + + await user.click(editItem); + await user.click(deleteItem); + + expect(onEndpointClick).not.toHaveBeenCalled(); + expect(onDeleteClick).not.toHaveBeenCalled(); + }); + + it("should not show the config hint for DB endpoints", async () => { + const user = userEvent.setup(); + render(); + + await user.click(screen.getByTestId("endpoint-actions-ep-1")); + await screen.findByTestId("endpoint-action-delete"); + expect(screen.queryByTestId("endpoint-config-hint")).not.toBeInTheDocument(); + }); + + it("should label endpoint source as Config or DB", () => { + const configEndpoint: passThroughItem = { + id: "ep-config", + path: "/from-config", + target: "https://config.example.com", + headers: {}, + is_from_config: true, + }; + render(); + expect(screen.getByText("Config")).toBeInTheDocument(); + expect(screen.getAllByText("DB")).toHaveLength(2); + }); + it("should disable edit and delete for endpoints without an id", async () => { const user = userEvent.setup(); const onEndpointClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx index d22b274861a..5b18685a140 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTableColumns.tsx @@ -18,6 +18,9 @@ import { cn } from "@/lib/cva.config"; import type { passThroughItem } from "./PassThroughSettings"; +const CONFIG_ENDPOINT_HINT = + "This endpoint is defined in the config file and cannot be edited or deleted on the dashboard."; + function HeaderWithTooltip({ title, tooltip }: { title: string; tooltip: string }) { return (
@@ -73,6 +76,7 @@ interface EndpointRowActionsProps { function EndpointRowActions({ endpoint, onEndpointClick, onDeleteClick }: EndpointRowActionsProps) { const endpointId = endpoint.id; + const isFromConfig = endpoint.is_from_config ?? false; return ( endpointId && onEndpointClick(endpointId)} + disabled={isFromConfig || !endpointId} + onClick={() => !isFromConfig && endpointId && onEndpointClick(endpointId)} > Edit @@ -95,12 +99,17 @@ function EndpointRowActions({ endpoint, onEndpointClick, onDeleteClick }: Endpoi endpointId && onDeleteClick(endpointId)} + disabled={isFromConfig || !endpointId} + onClick={() => !isFromConfig && endpointId && onDeleteClick(endpointId)} > Delete + {isFromConfig && ( +
+ {CONFIG_ENDPOINT_HINT} +
+ )}
); @@ -124,7 +133,9 @@ export const getPassThroughEndpointsTableColumns = ({ enableSorting: false, cell: ({ row }) => { const endpointId = row.original.id; - if (!endpointId) return ; + if (!endpointId || row.original.is_from_config) { + return ; + } return ( { + const isFromConfig = row.original.is_from_config ?? false; + return ; + }, + }, { id: "path", accessorKey: "path", diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx index 2dc6fdbd32c..8ef2766d412 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughSettings.tsx @@ -1,4 +1,13 @@ import React, { useState, useEffect } from "react"; +import { + AlertDialog, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { deletePassThroughEndpointsCall, getPassThroughEndpointsCall } from "../networking"; import AddPassThroughEndpoint from "../add_pass_through"; @@ -25,6 +34,7 @@ export interface passThroughItem { methods?: string[]; guardrails?: Record; default_query_params?: Record; + is_from_config?: boolean; } const PassThroughSettings: React.FC = ({ accessToken, userRole, userID, premiumUser }) => { @@ -133,42 +143,22 @@ const PassThroughSettings: React.FC = ({ accessToken, onDeleteClick={handleDelete} /> - {isDeleteModalOpen && ( -
-
- - - - -
-
-
-
-

Delete Pass-Through Endpoint

-
-

- Are you sure you want to delete this pass-through endpoint? This action cannot be undone. -

-
-
-
-
-
- - -
-
-
-
- )} + !open && cancelDelete()}> + + + Delete Pass-Through Endpoint + + Are you sure you want to delete this pass-through endpoint? This action cannot be undone. + + + + Cancel + + + +
); }; diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx index 3ef7dee01b0..a4d76e2cc06 100644 --- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx @@ -97,6 +97,24 @@ describe("LoggingCallbacksTable", () => { expect(onDelete).toHaveBeenCalledWith(callback); }); + it("shows a read-only label instead of the actions menu for runtime-only callback rows", () => { + render( + , + ); + expect(screen.getByTestId("callback-actions-langfuse-success")).toBeInTheDocument(); + expect(screen.queryByTestId("callback-actions-datadog-success")).not.toBeInTheDocument(); + expect(screen.getAllByText("Read only")).toHaveLength(1); + }); + // Regression: `/get_callbacks` returns the same `name` twice when a // callback is registered for both success and failure (e.g. `generic_api` // → POST to spend-log on both 200 and 4xx/5xx). The UI used to ignore diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx index 2263fe03b3d..950ef5de6e4 100644 --- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTableColumns.tsx @@ -50,6 +50,16 @@ interface CallbackRowActionsProps { } function CallbackRowActions({ callback, onTest, onEdit, onDelete }: CallbackRowActionsProps) { + if (callback.read_only) { + return ( + + Read only + + ); + } return ( { await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]).toMatchObject({ complexity_router_config: { - tier_model_configs: { - REASONING: [{ model_name: "claude-opus-5", litellm_params: { reasoning_effort: "high" } }], - }, + tier_model_configs: ANTHROPIC_PRESET.complexity_router_config.tier_model_configs, }, }); }); diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx index f31e1839739..4ba6d8d50b4 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx @@ -51,4 +51,23 @@ describe("ResponseMetrics prompt cache chips", () => { expect(screen.queryByText(/Response Cache/)).not.toBeInTheDocument(); }); + + it("does not render the Cost chip when a persisted cost is null", () => { + render(); + + expect(screen.queryByText(/Cost:/)).not.toBeInTheDocument(); + expect(screen.getByText("In: 1")).toBeInTheDocument(); + }); + + it("does not render the Cost chip for NaN", () => { + render(); + + expect(screen.queryByText(/Cost:/)).not.toBeInTheDocument(); + }); + + it("renders the Cost chip for a finite cost", () => { + render(); + + expect(screen.getByText("Cost: $0.000063")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx index ec62d0618d7..bb7debc7aee 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx @@ -159,7 +159,7 @@ const ResponseMetrics: React.FC = ({ timeToFirstToken, tot /> )} - {usage?.cost !== undefined && ( + {typeof usage?.cost === "number" && Number.isFinite(usage.cost) && ( { // else default to chat return EndpointType.CHAT; }; + +export const isModeCompatibleWithEndpoint = (mode: string | null | undefined, endpointType: EndpointType): boolean => { + if (!mode) return true; + if (!Object.values(ModelMode).includes(mode as ModelMode)) return false; + const optionEndpoint = getEndpointType(mode); + if ( + endpointType === EndpointType.RESPONSES || + endpointType === EndpointType.ANTHROPIC_MESSAGES || + endpointType === EndpointType.INTERACTIONS + ) { + return optionEndpoint === endpointType || optionEndpoint === EndpointType.CHAT; + } + if (endpointType === EndpointType.IMAGE_EDITS) { + return optionEndpoint === endpointType || optionEndpoint === EndpointType.IMAGE; + } + return optionEndpoint === endpointType; +}; diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx index 08b02cacb21..ec477441586 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx @@ -468,6 +468,18 @@ describe("chat_completion prompt cache usage", () => { expect(usageData).not.toHaveProperty("cacheReadTokens"); expect(usageData).not.toHaveProperty("cacheCreationTokens"); }); + + it("omits cost when the provider reports a non-numeric value", async () => { + const usageData = await captureUsage({ cost: "not-a-number" }); + + expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); + }); + + it("omits cost when the provider reports a blank value", async () => { + const usageData = await captureUsage({ cost: " " }); + + expect(usageData).toEqual(expect.not.objectContaining({ cost: expect.anything() })); + }); }); describe("chat_completion response cache", () => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index cd0852bc06e..ffa2877fbd9 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -5,6 +5,7 @@ import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; import { MCPServer, MCPToolset, type MCPEvent } from "@/components/mcp_tools/types"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +import { parseUsageCost } from "./usage_cost"; const completionAsSingleChunk = (completion: ChatCompletion): ChatCompletionChunk => ({ @@ -243,9 +244,9 @@ export async function makeOpenAIChatCompletionRequest( usageData.reasoningTokens = chunkWithUsage.usage.completion_tokens_details.reasoning_tokens; } - // Extract cost from usage object if available - if (chunkWithUsage.usage.cost !== undefined && chunkWithUsage.usage.cost !== null) { - usageData.cost = parseFloat(chunkWithUsage.usage.cost); + const parsedCost = parseUsageCost(chunkWithUsage.usage.cost); + if (parsedCost !== undefined) { + usageData.cost = parsedCost; } onUsageData(usageData); diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index ae5e224cfbf..290c8b0b619 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -233,6 +233,64 @@ describe("responses_api", () => { expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); }); + it("should omit cost when the proxy reports a non-numeric cost", async () => { + async function* streamWithNonNumericCost() { + yield { + type: "response.completed", + response: { + id: "resp_non_numeric_cost", + usage: { output_tokens: 12, input_tokens: 12, total_tokens: 24, cost: "not-a-number" }, + }, + }; + } + mockResponsesCreate.mockResolvedValueOnce(streamWithNonNumericCost()); + + const onUsageData = vi.fn(); + + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + onUsageData, + ); + + expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); + }); + + it("should omit cost when the proxy reports a blank cost", async () => { + async function* streamWithBlankCost() { + yield { + type: "response.completed", + response: { + id: "resp_blank_cost", + usage: { output_tokens: 12, input_tokens: 12, total_tokens: 24, cost: " " }, + }, + }; + } + mockResponsesCreate.mockResolvedValueOnce(streamWithBlankCost()); + + const onUsageData = vi.fn(); + + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gpt-4", + "test-token", + undefined, + undefined, + undefined, + undefined, + onUsageData, + ); + + expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ cost: expect.anything() }), ""); + }); + it("should replay MCP output items as events for a non-streaming response", async () => { mockResponsesCreate.mockReturnValueOnce( nonStreamingResponse({ diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 7ab76488504..ce54c7c6b40 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -4,6 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; +import { parseUsageCost } from "./usage_cost"; import type { MCPEvent } from "@/components/mcp_tools/types"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { @@ -311,8 +312,9 @@ export async function makeOpenAIResponsesRequest( usageData.reasoningTokens = reasoningTokens; } - if (usage.cost !== undefined && usage.cost !== null) { - usageData.cost = Number(usage.cost); + const parsedCost = parseUsageCost(usage.cost); + if (parsedCost !== undefined) { + usageData.cost = parsedCost; } onUsageData(usageData, mcpToolUsed); diff --git a/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts new file mode 100644 index 00000000000..ee1021c0629 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it } from "vitest"; +import { parseUsageCost } from "./usage_cost"; + +describe("parseUsageCost", () => { + it("keeps finite numbers, including zero", () => { + expect(parseUsageCost(0)).toBe(0); + expect(parseUsageCost(0.000063)).toBe(0.000063); + }); + + it("keeps numeric strings", () => { + expect(parseUsageCost("0.00019")).toBe(0.00019); + expect(parseUsageCost(" 0.00019 ")).toBe(0.00019); + }); + + it("drops blank strings instead of fabricating a zero cost", () => { + expect(parseUsageCost("")).toBeUndefined(); + expect(parseUsageCost(" ")).toBeUndefined(); + expect(parseUsageCost("\t\n")).toBeUndefined(); + }); + + it("drops strings with a numeric prefix instead of truncating them", () => { + expect(parseUsageCost("1oops")).toBeUndefined(); + expect(parseUsageCost("0.5 USD")).toBeUndefined(); + }); + + it("drops non-finite numbers", () => { + expect(parseUsageCost(Number.NaN)).toBeUndefined(); + expect(parseUsageCost(Number.POSITIVE_INFINITY)).toBeUndefined(); + }); + + it("drops values that are not numbers or strings", () => { + expect(parseUsageCost(null)).toBeUndefined(); + expect(parseUsageCost(undefined)).toBeUndefined(); + expect(parseUsageCost(true)).toBeUndefined(); + expect(parseUsageCost([])).toBeUndefined(); + expect(parseUsageCost(["0.5"])).toBeUndefined(); + expect(parseUsageCost({ total_cost: 0.5 })).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts new file mode 100644 index 00000000000..79f56dcb2df --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/usage_cost.ts @@ -0,0 +1,23 @@ +/** + * Providers and upstream gateways report `usage.cost` unvalidated: it arrives as a number, a numeric + * string, an empty string, or something non-numeric. A cost that does not resolve to a finite number + * must be dropped rather than coerced, because `NaN` survives `JSON.stringify` as `null` and crashes + * the metrics row on the next load. + */ +export function parseUsageCost(rawCost: unknown): number | undefined { + if (typeof rawCost === "number") { + return Number.isFinite(rawCost) ? rawCost : undefined; + } + + if (typeof rawCost !== "string") { + return undefined; + } + + const trimmed = rawCost.trim(); + if (trimmed === "") { + return undefined; + } + + const parsed = Number(trimmed); + return Number.isFinite(parsed) ? parsed : undefined; +} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx similarity index 61% rename from ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.test.tsx rename to ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx index d6e614da5eb..b1174d1d37d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.integration.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { render, screen } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect } from "vitest"; import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "./MCPToolArgumentsForm"; @@ -26,6 +26,118 @@ const submitError = async (ref: React.RefObject) }; describe("MCPToolArgumentsForm", () => { + it("keeps dotted arguments separate from a same-prefix object and converts their values", async () => { + const ref = renderForm({ + type: "object", + properties: { + "filter.category": { type: "string" }, + filter: { type: "object" }, + "page.limit": { type: "integer" }, + query: { type: "string" }, + }, + required: ["filter.category"], + }); + + fireEvent.change(screen.getByRole("textbox", { name: "filter.category *" }), { + target: { value: "invoices" }, + }); + fireEvent.change(screen.getByRole("textbox", { name: "filter" }), { + target: { value: '{"category":"receipts","metadata":{"region":"eu"}}' }, + }); + fireEvent.change(screen.getByRole("spinbutton", { name: "page.limit" }), { target: { value: "7" } }); + fireEvent.change(screen.getByRole("textbox", { name: "query" }), { target: { value: "September" } }); + + const expected = { + "filter.category": "invoices", + filter: { category: "receipts", metadata: { region: "eu" } }, + "page.limit": 7, + query: "September", + }; + await expect(submit(ref)).resolves.toEqual(expected); + }); + + it("shows required validation on the literal dotted field and accepts a correction", async () => { + const ref = renderForm({ + type: "object", + properties: { "filter.category": { type: "string" } }, + required: ["filter.category"], + }); + + expect(await submitError(ref)).toEqual({ + errorFields: [{ name: ["filter.category"], errors: ["Please enter filter.category"] }], + }); + expect(await screen.findByText("Please enter filter.category")).toBeInTheDocument(); + expect(screen.getByRole("textbox", { name: "filter.category *" })).toHaveAttribute("aria-invalid", "true"); + + fireEvent.change(screen.getByRole("textbox", { name: "filter.category *" }), { + target: { value: "invoices" }, + }); + await expect(submit(ref)).resolves.toEqual({ "filter.category": "invoices" }); + }); + + it("validates JSON for dotted arguments inside params and preserves their literal names", async () => { + const ref = renderForm({ + type: "object", + properties: { + params: { + type: "object", + properties: { "filter.options": { type: "object" } }, + required: ["filter.options"], + }, + }, + required: [], + }); + const field = screen.getByRole("textbox", { name: "filter.options *" }); + fireEvent.change(field, { target: { value: "invalid" } }); + + expect(await submitError(ref)).toEqual({ + errorFields: [{ name: ["filter.options"], errors: ["Invalid JSON"] }], + }); + expect(await screen.findByText("Invalid JSON")).toBeInTheDocument(); + + fireEvent.change(field, { target: { value: '{"region":"eu"}' } }); + await expect(submit(ref)).resolves.toEqual({ params: { "filter.options": { region: "eu" } } }); + }); + + it("resets dotted defaults and positional values when the selected tool changes", async () => { + const ref = React.createRef(); + const { rerender } = render( + , + ); + expect(screen.getByRole("textbox", { name: "filter.category" })).toHaveValue("invoices"); + await expect(submit(ref)).resolves.toEqual({ "filter.category": "invoices" }); + fireEvent.change(screen.getByRole("textbox", { name: "filter.category" }), { + target: { value: "edited" }, + }); + await expect(submit(ref)).resolves.toEqual({ "filter.category": "edited" }); + + rerender( + , + ); + expect(screen.getByRole("textbox", { name: "filter.category" })).toHaveValue("receipts"); + await expect(submit(ref)).resolves.toEqual({ query: "new tool", "filter.category": "receipts" }); + }); + it("returns typed values for a string, integer, number and boolean field", async () => { const user = userEvent.setup(); const ref = renderForm({ diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx index ca8c5697e6e..ab3213d9b37 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx @@ -1,6 +1,6 @@ import React, { forwardRef, useImperativeHandle, useMemo } from "react"; import { CircleHelp } from "lucide-react"; -import { useForm, type Resolver } from "react-hook-form"; +import { useForm, type Resolver, type ResolverResult } from "react-hook-form"; import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Input } from "@/components/ui/input"; @@ -9,7 +9,10 @@ import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { MCPTool, InputSchema, InputSchemaProperty } from "./types"; -type ToolFormValues = Record; +type ToolFormValues = { args: unknown[] }; + +const argumentValues = (schema: InputSchema, values: ToolFormValues): Record => + Object.fromEntries(Object.keys(schema.properties ?? {}).map((key, index) => [key, values.args[index]])); const STRING_SCHEMA_MESSAGES: Readonly> = { input: "Please enter input for this tool" }; @@ -38,7 +41,7 @@ type FieldError = { type: string; message: string }; const collectErrors = ( actualSchema: InputSchema, requiredMessages: Readonly>, - values: ToolFormValues, + values: Record, ): Record => { const entries = Object.entries(actualSchema.properties ?? {}).flatMap<[string, FieldError]>(([key, prop]) => { const value = values[key]; @@ -56,9 +59,19 @@ const collectErrors = ( const buildResolver = (actualSchema: InputSchema, requiredMessages: Readonly> = {}): Resolver => - (values) => { - const errors = collectErrors(actualSchema, requiredMessages, values); - return Object.keys(errors).length > 0 ? { values: {}, errors } : { values, errors: {} }; + (values): ResolverResult => { + const errors = collectErrors(actualSchema, requiredMessages, argumentValues(actualSchema, values)); + if (Object.keys(errors).length === 0) return { values, errors: {} }; + return { + values: {}, + errors: { + args: Object.fromEntries( + Object.keys(actualSchema.properties ?? {}).flatMap((key, index) => + Object.hasOwn(errors, key) ? [[index, errors[key]]] : [], + ), + ), + }, + }; }; const labelFor = (key: string, prop: InputSchemaProperty, required: boolean): React.ReactNode => ( @@ -238,10 +251,7 @@ const MCPToolArgumentsForm = forwardRef( - () => - Object.fromEntries( - Object.entries(actualSchema.properties ?? {}).map(([key, prop]) => [key, getInitialValueForField(prop)]), - ), + () => ({ args: Object.values(actualSchema.properties ?? {}).map(getInitialValueForField) }), [actualSchema], ); @@ -255,7 +265,7 @@ const MCPToolArgumentsForm = forwardRef ({ getSubmitValues: async () => { - const values = form.getValues(); + const values = argumentValues(actualSchema, form.getValues()); const errors = collectErrors(actualSchema, requiredMessages, values); if (Object.keys(errors).length > 0) { await form.trigger(); @@ -286,14 +296,16 @@ const MCPToolArgumentsForm = forwardRef Input * } > - {(field) => } + {(field) => ( + + )} @@ -318,13 +330,13 @@ const MCPToolArgumentsForm = forwardRef - {Object.entries(actualSchema.properties).map(([key, prop]) => { + {Object.entries(actualSchema.properties).map(([key, prop], index) => { const required = actualSchema.required?.includes(key) ?? false; return ( {(field) => { @@ -375,7 +387,7 @@ const MCPToolArgumentsForm = forwardRef ); @@ -385,7 +397,7 @@ const MCPToolArgumentsForm = forwardRef ); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 8a5f83adbdf..fed11454c23 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -13,6 +13,7 @@ import { buildModelAvailability, deploymentRefsFromModelInfo, normalizeModelName, + resolveAvailableModels, } from "./autorouter_presets"; import { DEFAULT_MATCH_THRESHOLD } from "@/components/add_model/SemanticKeywordMatching"; import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKeywords"; @@ -380,6 +381,18 @@ describe("autorouter_presets", () => { expect(availability.underlyingIndex.size).toBe(0); }); + it("returns every configured group serving the same underlying model", () => { + const availability = buildModelAvailability( + ["z-group", "a-group"], + [ + { modelGroup: "z-group", underlyingModels: ["anthropic/claude-sonnet-5"] }, + { modelGroup: "a-group", underlyingModels: ["bedrock/us.anthropic.claude-sonnet-5-v1:0"] }, + ], + ); + + expect(resolveAvailableModels("anthropic/claude-sonnet-5", availability)).toEqual(["a-group", "z-group"]); + }); + it("breaks ties between groups serving the same model deterministically, alphabetically", () => { const availability = buildModelAvailability( ["z-group", "a-group"], diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index f2df55fc310..02096cada41 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -159,16 +159,19 @@ export const deploymentRefsFromModelInfo = ( return row.model_name && underlyingModels.length > 0 ? [{ modelGroup: row.model_name, underlyingModels }] : []; }); -export const resolveAvailableModel = (requiredModel: string, availability: ModelAvailability): string | undefined => { +export const resolveAvailableModels = (requiredModel: string, availability: ModelAvailability): readonly string[] => { const { modelGroups, underlyingIndex } = availability; - if (modelGroups.has(requiredModel)) return requiredModel; + if (modelGroups.has(requiredModel)) return [requiredModel]; const normalized = normalizeModelName(requiredModel); - const groupMatch = Array.from(modelGroups).find((available) => normalizeModelName(available) === normalized); - if (groupMatch !== undefined) return groupMatch; + const groupMatches = Array.from(modelGroups).filter((available) => normalizeModelName(available) === normalized); + if (groupMatches.length > 0) return groupMatches; const key = normalizeUnderlyingModel(requiredModel); - return key === null ? undefined : underlyingIndex.get(key)?.[0]; + return key === null ? [] : underlyingIndex.get(key) ?? []; }; +export const resolveAvailableModel = (requiredModel: string, availability: ModelAvailability): string | undefined => + resolveAvailableModels(requiredModel, availability)[0]; + export const getMissingModels = ( config: Parameters[0], availability: ModelAvailability, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..6d62ce2b675 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -3318,11 +3318,14 @@ export interface paths { * - model: Model name (e.g., "gpt-4", "claude-3-opus") * - input_tokens: Expected input tokens per request * - output_tokens: Expected output tokens per request + * - cache_read_input_tokens: Cache-read tokens per request, counted within input_tokens (optional) + * - cache_creation_input_tokens: Cache-write tokens per request, counted within input_tokens (optional) + * - reasoning_tokens: Reasoning tokens per request, counted within output_tokens (optional) * - num_requests_per_day: Number of requests per day (optional) * - num_requests_per_month: Number of requests per month (optional) * * Returns cost breakdown including: - * - Per-request costs (input, output, margin) + * - Per-request costs (input, output, margin, plus the cache-read, cache-write and reasoning shares) * - Daily costs (if num_requests_per_day provided) * - Monthly costs (if num_requests_per_month provided) * @@ -3331,7 +3334,9 @@ export interface paths { * { * "model": "gpt-4", * "input_tokens": 1000, + * "cache_read_input_tokens": 800, * "output_tokens": 500, + * "reasoning_tokens": 200, * "num_requests_per_day": 100, * "num_requests_per_month": 3000 * } @@ -25865,6 +25870,11 @@ export interface components { * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). An INFO notice is logged once per worker at config load while this flag is active as a reminder that hard enforcement is relaxed. */ disable_budget_reservation?: boolean | null; + /** + * Disable Env Credential Login + * @description If True, disables signing in to the Admin UI with the environment credentials: UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback means env-credential login is always live by default). Database users with passwords are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password before enabling, or nobody can sign in to the UI. A locked-out admin can still administer the proxy over the API with the master key, and can unset this setting and restart the proxy to restore env-credential login. Default is False. + */ + disable_env_credential_login?: boolean | null; /** * Disable Password Login When Sso Enabled * @description If True and SSO is configured (MICROSOFT_CLIENT_ID, GOOGLE_CLIENT_ID, GENERIC_CLIENT_ID, or SAML_IDP_METADATA_URL/XML), disables username/password login on /login, /v2/login, and /v3/login so SSO is the only way to reach the Admin UI. An admin locked out of the UI can still administer the proxy over the API with the master key; unset this setting and restart the proxy to restore UI username/password login. Default is False. @@ -26468,6 +26478,18 @@ export interface components { * @description Request body for /cost/estimate endpoint. */ CostEstimateRequest: { + /** + * Cache Creation Input Tokens + * @description Input tokens written to the prompt cache; counted within input_tokens + * @default 0 + */ + cache_creation_input_tokens: number; + /** + * Cache Read Input Tokens + * @description Input tokens read from the prompt cache; counted within input_tokens + * @default 0 + */ + cache_read_input_tokens: number; /** * Input Tokens * @description Expected input tokens per request @@ -26493,17 +26515,65 @@ export interface components { * @description Expected output tokens per request */ output_tokens: number; + /** + * Reasoning Tokens + * @description Reasoning tokens the model emits; counted within output_tokens + * @default 0 + */ + reasoning_tokens: number; }; /** * CostEstimateResponse * @description Response body for /cost/estimate endpoint. */ CostEstimateResponse: { + /** + * Cache Creation Cost Per Request + * @description Cache-write share of input_cost_per_request + * @default 0 + */ + cache_creation_cost_per_request: number; + /** + * Cache Creation Input Token Cost + * @description Rate billed per cache-write token + */ + cache_creation_input_token_cost?: number | null; + /** + * Cache Creation Input Tokens + * @default 0 + */ + cache_creation_input_tokens: number; + /** + * Cache Read Cost Per Request + * @description Cache-read share of input_cost_per_request + * @default 0 + */ + cache_read_cost_per_request: number; + /** + * Cache Read Input Token Cost + * @description Rate billed per cache-read token + */ + cache_read_input_token_cost?: number | null; + /** + * Cache Read Input Tokens + * @default 0 + */ + cache_read_input_tokens: number; /** * Cost Per Request * @description Total cost per request (includes margin) */ cost_per_request: number; + /** + * Daily Cache Creation Cost + * @description Cache-write share of daily_input_cost + */ + daily_cache_creation_cost?: number | null; + /** + * Daily Cache Read Cost + * @description Cache-read share of daily_input_cost + */ + daily_cache_read_cost?: number | null; /** * Daily Cost * @description Total daily cost (includes margin) @@ -26524,12 +26594,20 @@ export interface components { * @description Daily output token cost */ daily_output_cost?: number | null; + /** + * Daily Reasoning Cost + * @description Reasoning share of daily_output_cost + */ + daily_reasoning_cost?: number | null; /** * Input Cost Per Request * @description Input token cost per request (before margin) */ input_cost_per_request: number; - /** Input Cost Per Token */ + /** + * Input Cost Per Token + * @description Rate billed per input token + */ input_cost_per_token?: number | null; /** Input Tokens */ input_tokens: number; @@ -26541,6 +26619,16 @@ export interface components { margin_cost_per_request: number; /** Model */ model: string; + /** + * Monthly Cache Creation Cost + * @description Cache-write share of monthly_input_cost + */ + monthly_cache_creation_cost?: number | null; + /** + * Monthly Cache Read Cost + * @description Cache-read share of monthly_input_cost + */ + monthly_cache_read_cost?: number | null; /** * Monthly Cost * @description Total monthly cost (includes margin) @@ -26561,21 +26649,45 @@ export interface components { * @description Monthly output token cost */ monthly_output_cost?: number | null; + /** + * Monthly Reasoning Cost + * @description Reasoning share of monthly_output_cost + */ + monthly_reasoning_cost?: number | null; /** Num Requests Per Day */ num_requests_per_day?: number | null; /** Num Requests Per Month */ num_requests_per_month?: number | null; + /** + * Output Cost Per Reasoning Token + * @description Rate billed per reasoning token + */ + output_cost_per_reasoning_token?: number | null; /** * Output Cost Per Request * @description Output token cost per request (before margin) */ output_cost_per_request: number; - /** Output Cost Per Token */ + /** + * Output Cost Per Token + * @description Rate billed per output token + */ output_cost_per_token?: number | null; /** Output Tokens */ output_tokens: number; /** Provider */ provider?: string | null; + /** + * Reasoning Cost Per Request + * @description Reasoning share of output_cost_per_request + * @default 0 + */ + reasoning_cost_per_request: number; + /** + * Reasoning Tokens + * @default 0 + */ + reasoning_tokens: number; }; /** CreateCredentialItem */ CreateCredentialItem: { @@ -29629,6 +29741,8 @@ export interface components { output_cost_per_second_480p?: number | null; /** Output Cost Per Second 4K */ output_cost_per_second_4k?: number | null; + /** Output Cost Per Second 720P */ + output_cost_per_second_720p?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; /** Output Cost Per Token Above 128K Tokens */ @@ -39820,6 +39934,8 @@ export interface components { output_cost_per_second_480p?: number | null; /** Output Cost Per Second 4K */ output_cost_per_second_4k?: number | null; + /** Output Cost Per Second 720P */ + output_cost_per_second_720p?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; /** Output Cost Per Token Above 128K Tokens */ diff --git a/whitelisted_bedrock_models.txt b/whitelisted_bedrock_models.txt index 578edddec8d..6124cb41044 100644 --- a/whitelisted_bedrock_models.txt +++ b/whitelisted_bedrock_models.txt @@ -6,6 +6,7 @@ ai21.jamba-instruct-v1:0 twelvelabs.pegasus-1-2-v1:0 us.twelvelabs.pegasus-1-2-v1:0 eu.twelvelabs.pegasus-1-2-v1:0 +global.twelvelabs.pegasus-1-2-v1:0 amazon.titan-text-express-v1 amazon.titan-text-lite-v1 amazon.titan-text-premier-v1:0