diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d693c251df8..45af8e1d7c9 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 15626 + "limit": 15468 }, "reportArgumentType": { - "limit": 2220 + "limit": 2218 }, "reportAssignmentType": { "limit": 319 @@ -18,13 +18,13 @@ "limit": 40 }, "reportDeprecated": { - "limit": 212 + "limit": 211 }, "reportDuplicateImport": { "limit": 19 }, "reportExplicitAny": { - "limit": 4879 + "limit": 4846 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5610 + "limit": 5609 }, "reportMissingTypeArgument": { - "limit": 15338 + "limit": 15330 }, "reportMissingTypeStubs": { "limit": 40 @@ -93,25 +93,25 @@ "limit": 181 }, "reportTypedDictNotRequiredAccess": { - "limit": 25 + "limit": 24 }, "reportUndefinedVariable": { "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44389 + "limit": 44368 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38449 + "limit": 38418 }, "reportUnknownParameterType": { - "limit": 19657 + "limit": 19649 }, "reportUnknownVariableType": { - "limit": 30012 + "limit": 29987 }, "reportUnnecessaryCast": { "limit": 111 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 829 + "limit": 828 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 2864663df93..c45c6bd98c0 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,8 +1,9 @@ import json -from collections.abc import AsyncIterator, Iterator, Sequence -from typing import Any, Final, TypedDict, cast +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from types import MappingProxyType +from typing import Any, Final, TypeAlias, cast -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, TypedDict from litellm import verbose_logger from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema @@ -11,7 +12,6 @@ from litellm.types.llms.openai import ( ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, ChatCompletionImageObject, - ChatCompletionRequest, ChatCompletionSystemMessage, ChatCompletionTextObject, ChatCompletionToolCallFunctionChunk, @@ -24,38 +24,62 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( AdapterCompletionStreamWrapper, ChatCompletionDeltaCustomToolCall, + ChatCompletionDeltaToolCall, ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, Choices, Delta, + Function, Message, ModelResponse, ModelResponseStream, StreamingChoices, - Usage, ) - -class _GenAITextPart(TypedDict, total=False): - text: ReadOnly[str] +_JsonDict: TypeAlias = dict[str, object] +_JsonDictList: TypeAlias = list[_JsonDict] -class _GenAISystemInstruction(TypedDict, total=False): - parts: ReadOnly[list[_GenAITextPart]] +class _ToolCallAccumulator(TypedDict): + name: ReadOnly[str] + arguments: ReadOnly[str] + + +class _GenAIFunctionCall(TypedDict): + name: ReadOnly[str] + args: ReadOnly[Mapping[str, object]] class _GenAIPart(TypedDict, total=False): text: ReadOnly[str] - functionCall: ReadOnly[dict[str, object]] + functionCall: ReadOnly[_GenAIFunctionCall] + + +class _GenAIFunctionResponse(TypedDict, total=False): + name: ReadOnly[str] + response: ReadOnly[object] + + +class _GenAIRequestFunctionCall(TypedDict, total=False): + name: ReadOnly[str] + args: ReadOnly[Mapping[str, object]] + + +class _GenAIContentPart(TypedDict, total=False): + text: ReadOnly[str] + inline_data: ReadOnly[Mapping[str, str]] + functionResponse: ReadOnly[_GenAIFunctionResponse] + functionCall: ReadOnly[_GenAIRequestFunctionCall] class _GenAIFunctionDeclaration(TypedDict, total=False): name: ReadOnly[str] description: ReadOnly[str] - parametersJsonSchema: ReadOnly[dict[str, object]] + parametersJsonSchema: ReadOnly[object] class _GenAITool(TypedDict, total=False): - functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]] + functionDeclarations: ReadOnly[Sequence[_GenAIFunctionDeclaration]] class _GenAIFunctionCallingConfig(TypedDict, total=False): @@ -66,9 +90,11 @@ class _GenAIToolConfig(TypedDict, total=False): functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig] -def _decode_tool_call_arguments(raw_arguments: str) -> object: - """Decode a tool call's JSON-encoded arguments into the value Google GenAI expects.""" - return json.loads(raw_arguments) +class _GenAISystemInstruction(TypedDict, total=False): + parts: ReadOnly[Sequence[Mapping[str, str]]] + + +_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({}) class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): @@ -78,12 +104,12 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ sent_first_chunk: bool = False - # State tracking for accumulating partial tool calls - accumulated_tool_calls: dict[int, dict[str, str]] + _parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads) def __init__(self, completion_stream: object): self.sent_first_chunk = False - self.accumulated_tool_calls = {} + # State tracking for accumulating partial tool calls + self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]() self._returned_response = False super().__init__(completion_stream) @@ -128,7 +154,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): # After the stream is exhausted, check for any remaining accumulated tool calls if self.accumulated_tool_calls: try: - parts: Final[list[_GenAIPart]] = [] + parts: Final = list[_GenAIPart]() for ( tool_call_index, tool_call_data, @@ -136,7 +162,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): try: # For tool calls with no arguments, accumulated_args will be "", which is not valid JSON. # We default to an empty JSON object in this case. - parsed_args = _decode_tool_call_arguments(tool_call_data["arguments"] or "{}") + parsed_args: Mapping[str, object] = self._parse_accumulated_args( + tool_call_data["arguments"] or "{}" + ) function_call_part: _GenAIPart = { "functionCall": { "name": tool_call_data["name"] or "undefined_tool_name", @@ -153,7 +181,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): tool_call_data["arguments"], ) if parts: - final_chunk: Final[dict[str, object]] = { + final_chunk: Final = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -215,14 +243,16 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): class GoogleGenAIAdapter: """Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format""" + _parse_tool_call_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads) + def __init__(self) -> None: pass def translate_generate_content_to_completion( self, model: str, - contents: list[dict[str, Any]] | dict[str, Any], - config: dict[str, Any] | None = None, + contents: _JsonDictList | _JsonDict, + config: Mapping[str, object] | None = None, litellm_params: GenericLiteLLMParams | None = None, **kwargs, ) -> dict[str, Any]: @@ -254,7 +284,7 @@ class GoogleGenAIAdapter: messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction) # Create base request as dict (which is compatible with ChatCompletionRequest) - completion_request: Final[ChatCompletionRequest] = { + completion_request: Final[_JsonDict] = { "model": model, "messages": messages, } @@ -316,9 +346,9 @@ class GoogleGenAIAdapter: def _add_generic_litellm_params_to_request( self, - completion_request_dict: dict[str, object], + completion_request_dict: _JsonDict, litellm_params: GenericLiteLLMParams | None = None, - ) -> dict[str, object]: + ) -> _JsonDict: """Add generic litellm params to request. e.g add api_base, api_key, api_version, etc. Args: @@ -330,7 +360,7 @@ class GoogleGenAIAdapter: """ allowed_fields: Final = GenericLiteLLMParams.model_fields.keys() if litellm_params: - litellm_dict: Final = litellm_params.model_dump(exclude_none=True) + litellm_dict: Final[_JsonDict] = litellm_params.model_dump(exclude_none=True) for key, value in litellm_dict.items(): if key in allowed_fields: completion_request_dict[key] = value @@ -350,12 +380,12 @@ class GoogleGenAIAdapter: tools: Sequence[_GenAITool], ) -> list[ChatCompletionToolParam]: """Transform Google GenAI tools to OpenAI tools format""" - openai_tools: Final[list[dict[str, object]]] = [] + openai_tools: Final = list[_JsonDict]() for tool in tools: if "functionDeclarations" in tool: for func_decl in tool["functionDeclarations"]: - function_chunk: dict[str, object] = { + function_chunk: _JsonDict = { "name": func_decl.get("name", ""), } @@ -364,7 +394,7 @@ class GoogleGenAIAdapter: if "parametersJsonSchema" in func_decl: function_chunk["parameters"] = func_decl["parametersJsonSchema"] - openai_tool: dict[str, object] = {"type": "function", "function": function_chunk} + openai_tool: _JsonDict = {"type": "function", "function": function_chunk} openai_tools.append(openai_tool) # normalize the tool schemas @@ -395,13 +425,13 @@ class GoogleGenAIAdapter: # Handle system instruction if system_instruction: - system_parts: Final = system_instruction.get("parts", []) + system_parts: Final[Sequence[Mapping[str, str]]] = system_instruction.get("parts", []) if system_parts and "text" in system_parts[0]: messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"])) for content in contents: role = content.get("role", "user") - parts = content.get("parts", []) + parts: Sequence[_GenAIContentPart | str | None] = content.get("parts", []) if role == "user": # Handle user messages with potential function responses @@ -504,7 +534,7 @@ class GoogleGenAIAdapter: def translate_completion_to_generate_content( self, response: ModelResponse, - ) -> dict[str, object]: + ) -> _JsonDict: """ Transform litellm completion response to Google GenAI generate_content format @@ -527,13 +557,13 @@ class GoogleGenAIAdapter: parts = self._transform_openai_message_to_google_genai_parts(choice.message) else: # Fallback for generic choice objects - message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get( - "content", "" - ) + message_content: str = getattr(choice, "message", _EMPTY_STR_MAPPING).get("content", "") or getattr( + choice, "delta", _EMPTY_STR_MAPPING + ).get("content", "") parts = [{"text": message_content}] if message_content else [] # Create Google GenAI format response - generate_content_response: Final[dict[str, object]] = { + generate_content_response: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -567,7 +597,7 @@ class GoogleGenAIAdapter: self, response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper, - ) -> dict[str, object] | None: + ) -> Mapping[str, object] | None: """ Transform streaming litellm completion chunk to Google GenAI generate_content format @@ -594,7 +624,7 @@ class GoogleGenAIAdapter: finish_reason: str | None = getattr(choice, "finish_reason", None) else: # Fallback for generic choice objects - message_content: Final = getattr(choice, "delta", {}).get("content", "") + message_content: Final[str] = getattr(choice, "delta", _EMPTY_STR_MAPPING).get("content", "") parts = [{"text": message_content}] if message_content else [] finish_reason = getattr(choice, "finish_reason", None) @@ -603,7 +633,7 @@ class GoogleGenAIAdapter: return None # Create Google GenAI streaming format response - streaming_chunk: Final[dict[str, object]] = { + streaming_chunk: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -640,9 +670,9 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, message: Message, - ) -> list[_GenAIPart]: + ) -> Sequence[_GenAIPart]: """Transform OpenAI message to Google GenAI parts format""" - parts: Final[list[_GenAIPart]] = [] + parts: Final = list[_GenAIPart]() # Add text content if present if hasattr(message, "content") and message.content: @@ -650,24 +680,22 @@ class GoogleGenAIAdapter: # Add tool calls if present if hasattr(message, "tool_calls") and message.tool_calls: - for tool_call in message.tool_calls: - if ( - hasattr(tool_call, "function") - and not isinstance(tool_call, ChatCompletionMessageCustomToolCall) - and tool_call.function - ): + tool_calls: Final[Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]] = ( + message.tool_calls + ) + for tool_call in tool_calls: + function: Function | None = getattr(tool_call, "function", None) + if function: try: - args = ( - _decode_tool_call_arguments(tool_call.function.arguments) - if tool_call.function.arguments - else {} + args: Mapping[str, object] = ( + self._parse_tool_call_args(function.arguments) if function.arguments else {} ) except json.JSONDecodeError: args = {} function_call_part: _GenAIPart = { "functionCall": { - "name": tool_call.function.name or "undefined_tool_name", + "name": function.name or "undefined_tool_name", "args": args, } } @@ -677,20 +705,22 @@ class GoogleGenAIAdapter: def _transform_openai_delta_to_google_genai_parts_with_accumulation( self, delta: Delta, wrapper: GoogleGenAIStreamWrapper - ) -> list[_GenAIPart]: + ) -> Sequence[_GenAIPart]: """Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls.""" # 1. Initialize wrapper state if it doesn't exist if not hasattr(wrapper, "accumulated_tool_calls"): wrapper.accumulated_tool_calls = {} - parts: Final[list[_GenAIPart]] = [] + parts: Final = list[_GenAIPart]() if hasattr(delta, "content") and delta.content: parts.append({"text": delta.content}) # 2. Ensure tool_calls is iterable - tool_calls: Final = delta.tool_calls or [] + tool_calls: Final[Sequence[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = ( + delta.tool_calls or [] + ) for tool_call in tool_calls: if not hasattr(tool_call, "function") or isinstance(tool_call, ChatCompletionDeltaCustomToolCall): @@ -709,19 +739,20 @@ class GoogleGenAIAdapter: } # Accumulate name and arguments - function_name = getattr(tool_call.function, "name", None) - args_chunk = getattr(tool_call.function, "arguments", None) + delta_function: Function | None = getattr(tool_call, "function", None) + function_name: str | None = getattr(delta_function, "name", None) + args_chunk: str | None = getattr(delta_function, "arguments", None) # Optimization: Skip chunks that have no new data if not function_name and not args_chunk: verbose_logger.debug("Skipping empty tool call chunk for index: %s", tool_call_index) continue - if function_name: - wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name - - if args_chunk: - wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += args_chunk + previous_data: _ToolCallAccumulator = wrapper.accumulated_tool_calls[tool_call_index] + wrapper.accumulated_tool_calls[tool_call_index] = _ToolCallAccumulator( + name=function_name or previous_data["name"], + arguments=previous_data["arguments"] + (args_chunk or ""), + ) # Attempt to parse and emit a complete tool call accumulated_data = wrapper.accumulated_tool_calls[tool_call_index] @@ -731,7 +762,7 @@ class GoogleGenAIAdapter: # 5. Attempt to parse arguments even if name hasn't arrived. try: # Attempt to parse the accumulated arguments string - parsed_args = _decode_tool_call_arguments(accumulated_args) + parsed_args: Mapping[str, object] = self._parse_tool_call_args(accumulated_args) # If parsing succeeds, but we don't have a name yet, wait. # The part will be created by a later chunk that brings the name. @@ -765,7 +796,7 @@ class GoogleGenAIAdapter: return mapping.get(finish_reason, "STOP") - def _map_usage(self, usage: Usage | None) -> dict[str, int]: + def _map_usage(self, usage: object) -> Mapping[str, int]: """Map OpenAI usage to Google GenAI usage format""" return { "promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0, diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index a474a11601d..c9e511905a6 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -8,11 +8,12 @@ import uuid from collections import Counter from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload import httpx -from typing_extensions import Never, ReadOnly +from typing_extensions import Never, ReadOnly, Required from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -30,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.utils import ( ChatCompletionMessageToolCall, Function, @@ -52,17 +54,102 @@ _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0 _EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({}) -class _ServiceToolCall(TypedDict): - id: ReadOnly[str] +class _ModerationToolCall(TypedDict, total=False): + id: ReadOnly[Required[str]] -class _ServiceMessage(TypedDict, total=False): +class _ModerationMessage(TypedDict, total=False): + content: ReadOnly[str | None] + tool_calls: ReadOnly[Sequence[_ModerationToolCall] | None] + + +class _ModerationChoice(TypedDict, total=False): + message: ReadOnly[_ModerationMessage | None] + + +class _ModerationResponse(TypedDict, total=False): + choices: ReadOnly[Sequence[_ModerationChoice]] + + +class _LogEventKwargs(TypedDict, total=False): + standard_logging_object: ReadOnly[Required[StandardLoggingPayload]] + litellm_call_id: ReadOnly[str] + + +class _HasCallId(Protocol): + def get(self, key: Literal["litellm_call_id"], /) -> str | None: ... + + +class _HasModelAttr(Protocol): + model: str | None + + +class _ResponseSource(Protocol): + def get(self, key: Literal["response"], /) -> "_HasModelAttr | None": ... + + +class _ModelSource(Protocol): + def get(self, key: Literal["model"], default: str, /) -> str: ... + + +class _FallbackSource(Protocol): + @overload + def get(self, key: Literal["start_time"], /) -> datetime | None: ... + @overload + def get(self, key: str, /) -> object | None: ... + + +class _RequestContextSource(Protocol): + @overload + def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ... + @overload + def get(self, key: str, /) -> object | None: ... + def __contains__(self, key: object, /) -> bool: ... + def __getitem__(self, key: str, /) -> object: ... + + +class _ToolCallLike(Protocol): + id: str | None + type: str | None + function: Function + + +class _ModerationSourceToolCall(TypedDict, total=False): + function: ReadOnly[Mapping[str, object] | None] + + +class _ModerationSourceMessage(TypedDict, total=False): + role: ReadOnly[str] + function_call: ReadOnly[Mapping[str, object] | None] + tool_calls: ReadOnly[Sequence[_ModerationSourceToolCall | None] | None] + + +class _FlattenedModerationMessage(TypedDict): + role: ReadOnly[str | None] content: ReadOnly[str] - tool_calls: ReadOnly[Sequence[_ServiceToolCall]] -class _ServiceChoice(TypedDict, total=False): - message: ReadOnly[_ServiceMessage] +class _CorrelatablePayload(TypedDict): + id: str # writable-ok: _apply_correlation_id overwrites the provider id on a deep-copied payload + + +class _SystemPromptCarrier(TypedDict, total=False): + messages: object # writable-ok: _prepend_system_prompt rebinds messages on the copied payload by design + + +class _BlockFailurePayload(TypedDict, total=False): + id: object # writable-ok: correlation id is pinned after copying the base payload + model: ReadOnly[object] + model_group: ReadOnly[object] + model_id: ReadOnly[str] + model_parameters: ReadOnly[object] + startTime: ReadOnly[float | None] + endTime: ReadOnly[float | None] + completionStartTime: ReadOnly[float | None] + messages: object # writable-ok: passed to _prepend_system_prompt, which rebinds messages + metadata: ReadOnly[StandardLoggingUserAPIKeyMetadata] + response: str # writable-ok: block failure text replaces the copied response + status: ReadOnly[str] class _MalformedToolBlockingResponseError(Exception): @@ -385,7 +472,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _stash_block_context( logging_obj: Optional["LiteLLMLoggingObj"], - request_data: dict, + request_data: dict[str, object], ) -> None: """Stash signals so the deferred success-event skips this request and ``async_post_call_failure_hook`` can build the failure payload. @@ -414,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): request_data["_rubrik_logging_obj"] = logging_obj @staticmethod - def _normalize_tool_calls(tool_calls: Sequence[object]) -> tuple[ChatCompletionMessageToolCall, ...]: + def _normalize_tool_calls( + tool_calls: Sequence[ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike], + ) -> tuple[ChatCompletionMessageToolCall, ...]: """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) @staticmethod - def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall: + def _normalize_tool_call( + tc: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike, + ) -> ChatCompletionMessageToolCall: if isinstance(tc, ChatCompletionMessageToolCall): return tc if isinstance(tc, dict): @@ -460,12 +551,15 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ``content`` is sent so the webhook can moderate the response text; ``None`` when the assistant produced no text (tool-call-only response). """ - message: Final[dict[str, object]] = { + message: Final[Mapping[str, object]] = { "role": "assistant", "content": content or None, + **( + {"tool_calls": tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)} + if tool_calls + else _EMPTY_MAPPING + ), } - if tool_calls: - message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls) return { "id": request_id or f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", @@ -481,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]: + def _flatten_messages_for_moderation( + messages: Sequence[AllMessageValues | None] | None, + ) -> tuple[_FlattenedModerationMessage, ...]: """Collapse each message's content to a plain string for the webhook. litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, @@ -502,7 +598,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ) @staticmethod - def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]: + def _moderation_text_parts(message: _ModerationSourceMessage) -> tuple[str, ...]: """Every attacker-controlled text segment of a message: its content plus the arguments of any tool call or deprecated function call.""" fc: Final = message.get("function_call") @@ -530,16 +626,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ``/v1/messages`` requests too. Optional fields are sent only when present so the payload stays clean. """ - payload: Final[dict[str, object]] = { - "model": inputs.get("model") or request_data.get("model") or "", - "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), - } tools: Final = inputs.get("tools") - if tools is not None: - payload["tools"] = tools user: Final = request_data.get("user") - if user: - payload["user"] = user # Fall back to litellm_call_id, the stable cross-provider join key the # response/tool path uses (see _correlation_id). LiteLLM does not # populate request_data["correlation_key"]; it carries litellm_call_id. @@ -547,14 +635,18 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # when correlation_key is empty, so without this the block fires but no # log is ever written. An explicit correlation_key still wins. correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id") - if correlation_key: - payload["correlation_key"] = correlation_key - return payload + return { + "model": inputs.get("model") or request_data.get("model") or "", + "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), + **({"tools": tools} if tools is not None else _EMPTY_MAPPING), + **({"user": user} if user else _EMPTY_MAPPING), + **({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING), + } @staticmethod def _extract_request_data( - call_details: Mapping[str, Any], - request_data: Mapping[str, object] | None, + call_details: _RequestContextSource, + request_data: _RequestContextSource | None, ) -> Mapping[str, object]: """Extract original request data from model_call_details for the response moderation service envelope. @@ -590,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _sanitize_proxy_server_request(proxy_server_request: object) -> object: + def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object: """Allowlist only routing fields (``url``, ``method``) when forwarding ``proxy_server_request`` to an external webhook, dropping inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw @@ -600,18 +692,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request} @staticmethod - def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, str]) -> str: + def _resolve_model(request_data: _ResponseSource, call_details: _ModelSource) -> str: """Get the model name for the ModifyResponseException.""" response: Final = request_data.get("response") if response and hasattr(response, "model"): - response_model: Final[str | None] = getattr(response, "model", None) - return response_model or "unknown" + return response.model or "unknown" return call_details.get("model", "unknown") # -- Logging hooks --------------------------------------------------------- @staticmethod - def _correlation_id(call_details: Mapping[str, str], request_data: Mapping[str, str] | None = None) -> str | None: + def _correlation_id( + call_details: _HasCallId | _LogEventKwargs, request_data: _HasCallId | None = None + ) -> str | None: """The id that joins a blocked request's two S3 logs by filename: the moderation (``_blocking``) log and the failure (response) log. @@ -625,7 +718,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id") @classmethod - def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, str]) -> None: + def _apply_correlation_id(cls, payload: _CorrelatablePayload, source: _HasCallId | _LogEventKwargs) -> None: """Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log shares its S3 filename id with the moderation (``_blocking``) and failure logs for the same request -- for every provider. @@ -645,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): payload["id"] = correlated @staticmethod - def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None: + def _prepend_system_prompt(payload: _SystemPromptCarrier, source: Mapping[str, object]) -> None: """Prepend ``source["system"]`` onto ``payload["messages"]``. Builds a NEW messages list rather than mutating ``payload["messages"]`` @@ -673,9 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): exc_info=True, ) - async def _prepare_log_payload( - self, kwargs: Mapping[str, object], event_type: str - ) -> StandardLoggingPayload | None: + async def _prepare_log_payload(self, kwargs: _LogEventKwargs, event_type: str) -> StandardLoggingPayload | None: """Shared logic for success logging (sampled).""" if random.random() > self.sampling_rate: verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate) @@ -684,12 +775,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # Deep-copy so mutations don't affect other callbacks sharing this object standard_logging_payload: Final[StandardLoggingPayload] = safe_deep_copy(kwargs["standard_logging_object"]) - self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime + self._apply_correlation_id(standard_logging_payload, kwargs) self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime return standard_logging_payload - async def _append_and_maybe_flush(self, payload) -> None: + async def _append_and_maybe_flush(self, payload: Mapping[str, object]) -> None: self._ensure_periodic_flush_task() self.log_queue.append(payload) self._enforce_max_queue_size() @@ -714,7 +805,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self._dropped_since_warning = 0 self._last_drop_warning_time = now - async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str): + async def _enqueue_log_event(self, kwargs: _LogEventKwargs, event_type: str): try: payload: Final = await self._prepare_log_payload(kwargs, event_type) if payload is None: @@ -835,7 +926,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): logging_obj: "LiteLLMLoggingObj", exception: "ModifyResponseException", user_api_key_dict: "UserAPIKeyAuth", - ) -> StandardLoggingPayload: + ) -> _BlockFailurePayload: """Build a failure-style payload using the exception text as response. Blocked-tool events are security-relevant and **bypass sampling**: @@ -877,9 +968,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): call_details: Final = logging_obj.model_call_details exception_text: Final = f"{type(exception).__name__}: {exception.message}" - base: Final = call_details.get("standard_logging_object") + base: Final[StandardLoggingPayload | None] = call_details.get("standard_logging_object") if base is not None: - payload: dict[str, object] = safe_deep_copy(base) + payload: _BlockFailurePayload = self._copy_block_payload_base(base) else: verbose_logger.debug( "Rubrik: standard_logging_object not yet on model_call_details " @@ -901,6 +992,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return payload + @staticmethod + def _copy_block_payload_base(base: StandardLoggingPayload) -> _BlockFailurePayload: + return safe_deep_copy(base) + @staticmethod def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata: """Identify the caller whose request was blocked. @@ -923,9 +1018,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @classmethod def _build_fallback_payload( cls, - call_details: Mapping[str, Any], + call_details: _FallbackSource, user_api_key_dict: "UserAPIKeyAuth", - ) -> dict[str, object]: + ) -> _BlockFailurePayload: # Convert datetime to a Unix float so json.dumps can serialize it. # httpx's json= parameter uses stdlib json.dumps with no custom encoder. _raw_start: Final = call_details.get("start_time") @@ -959,7 +1054,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): response: Final = await self.async_httpx_client.post( url=self.logging_endpoint, json=data, - headers=self._headers, + headers=dict(self._headers), ) response.raise_for_status() except httpx.HTTPStatusError as e: @@ -1013,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Webhook services ------------------------------------------------------ - async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, Any]: + async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> _ModerationResponse: """POST ``payload`` to a Rubrik webhook and return its dict response. Raises: @@ -1023,11 +1118,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): verbose_logger.debug("Sending request to %s: %s", service_name, endpoint) http_response: Final = await self.moderation_client.post( endpoint, - json=payload, - headers=self._headers, + json=dict(payload), + headers=dict(self._headers), ) http_response.raise_for_status() - result: Final[object] = http_response.json() + result: Final[_ModerationResponse | None] = http_response.json() if not isinstance(result, dict): raise TypeError( f"{service_name} returned non-dict JSON " @@ -1040,7 +1135,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self, response_data: Mapping[str, object], request_data: Mapping[str, object], - ) -> Mapping[str, Any]: + ) -> _ModerationResponse: """Post the ``{request, response}`` envelope to the after_completion webhook and return its (possibly rewritten) response. @@ -1056,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "Response moderation service", ) - async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, Any]: + async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> _ModerationResponse: """Post a bare OpenAI request to the before_prompt webhook. Returns ``{}`` (passthrough) or a synthetic chat.completion (block). @@ -1064,14 +1159,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service") @staticmethod - def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None: + def _extract_prompt_refusal(service_response: _ModerationResponse) -> str | None: """Return the refusal text when the prompt was blocked, else None. The before_prompt webhook returns ``{}`` (passthrough) or a synthetic chat.completion whose ``choices[0].message.content`` is the refusal explanation. """ - choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices") + choices: Final = service_response.get("choices") if not choices: return None message: Final = choices[0].get("message") or _EMPTY_MAPPING @@ -1080,7 +1175,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _extract_response_block( - service_response: Mapping[str, Any], + service_response: _ModerationResponse, all_tool_calls: Sequence[ChatCompletionMessageToolCall], sent_content: str, ) -> BlockedResponseResult | None: @@ -1103,7 +1198,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Expects service_response in OpenAI chat completion format: {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} """ - choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or () + choices: Final = service_response.get("choices") or () if not choices: raise _MalformedToolBlockingResponseError("Response moderation service returned empty response") diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index af82cf33796..81310b9ddc3 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,7 +10,7 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, TypedDict, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast from typing_extensions import ReadOnly @@ -46,7 +46,13 @@ from litellm.types.integrations.websearch_interception import ( AnthropicServerToolUseBlock, WebSearchInterceptionConfig, ) -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.anthropic import AnthropicThinkingParam +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAudioParam, + ChatCompletionPredictionContentParam, + OpenAIWebSearchOptions, +) from litellm.types.utils import ( AgenticLoopParams, CallTypes, @@ -56,6 +62,8 @@ from litellm.types.utils import ( from litellm.utils import ProviderConfigManager if TYPE_CHECKING: + from aiohttp import ClientSession + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, @@ -94,23 +102,98 @@ class _WebSearchSettingsView(TypedDict): websearch_interception_params: WebSearchInterceptionConfig +class _SearchToolLitellmParams(TypedDict, total=False): + search_provider: ReadOnly[str | None] + + class _SearchToolConfig(TypedDict, total=False): search_tool_name: str - litellm_params: Mapping[str, object] | None + litellm_params: ReadOnly[_SearchToolLitellmParams | None] -class _DeploymentKwargsView(TypedDict): - """Typed reads of the untyped request kwargs seen by the deployment hook.""" - +class _LitellmParamsProviderView(TypedDict, total=False): custom_llm_provider: ReadOnly[str] - litellm_params: ReadOnly[Mapping[str, object]] + + +class _DeploymentCallKwargsView(TypedDict): + custom_llm_provider: ReadOnly[str] + litellm_params: ReadOnly[_LitellmParamsProviderView] model: ReadOnly[str] -class _UserAuthView(TypedDict): - """Typed read of the optional team attached to the caller's auth object.""" +class _AcreateNamedParams(TypedDict, total=False): + metadata: ReadOnly[Never] + stop_sequences: ReadOnly[Never] + stream: ReadOnly[bool | None] + system: ReadOnly[str | None] + temperature: ReadOnly[float | None] + thinking: ReadOnly[Never] + tool_choice: ReadOnly[Never] + tools: ReadOnly[Never] + top_k: ReadOnly[int | None] + top_p: ReadOnly[float | None] + container: ReadOnly[Never] - team_id: ReadOnly[str | None] + +class _AsearchNamedParams(TypedDict, total=False): + max_results: ReadOnly[int | None] + search_domain_filter: ReadOnly[Never] + max_tokens_per_page: ReadOnly[int | None] + country: ReadOnly[str | None] + api_key: ReadOnly[str | None] + api_base: ReadOnly[str | None] + timeout: ReadOnly[float | None] + extra_headers: ReadOnly[Never] + + +class _AcompletionNamedParams(TypedDict, total=False): + functions: ReadOnly[Never] + function_call: ReadOnly[str | None] + timeout: ReadOnly[float | None] + temperature: ReadOnly[float | None] + top_p: ReadOnly[float | None] + n: ReadOnly[int | None] + stream: ReadOnly[bool | None] + stream_options: ReadOnly[Never] + stop: ReadOnly[Never] + max_tokens: ReadOnly[int | None] + max_completion_tokens: ReadOnly[int | None] + modalities: ReadOnly[Never] + prediction: ReadOnly[ChatCompletionPredictionContentParam | None] + audio: ReadOnly[ChatCompletionAudioParam | None] + presence_penalty: ReadOnly[float | None] + frequency_penalty: ReadOnly[float | None] + logit_bias: ReadOnly[Never] + user: ReadOnly[str | None] + response_format: ReadOnly[Never] + seed: ReadOnly[int | None] + tools: ReadOnly[Never] + tool_choice: ReadOnly[Never] + parallel_tool_calls: ReadOnly[bool | None] + logprobs: ReadOnly[bool | None] + top_logprobs: ReadOnly[int | None] + deployment_id: ReadOnly[str | None] + reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None] + verbosity: ReadOnly[Literal["low", "medium", "high"] | None] + safety_identifier: ReadOnly[str | None] + service_tier: ReadOnly[str | None] + store: ReadOnly[bool | None] + prompt_cache_key: ReadOnly[str | None] + base_url: ReadOnly[str | None] + api_version: ReadOnly[str | None] + api_key: ReadOnly[str | None] + model_list: ReadOnly[Never] + extra_headers: ReadOnly[Never] + thinking: ReadOnly[AnthropicThinkingParam | None] + web_search_options: ReadOnly[OpenAIWebSearchOptions | None] + include_server_side_tool_invocations: ReadOnly[bool | None] + shared_session: ReadOnly["ClientSession | None"] + enable_json_schema_validation: ReadOnly[bool | None] + + +_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {} +_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {} +_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {} class WebSearchInterceptionLogger(CustomLogger): @@ -312,17 +395,17 @@ class WebSearchInterceptionLogger(CustomLogger): """ # Check if this is for an enabled provider # Try top-level kwargs first, then nested litellm_params, then derive from model name - kwargs_view: Final[_DeploymentKwargsView] = { + call_kwargs_view: Final[_DeploymentCallKwargsView] = { "custom_llm_provider": kwargs.get("custom_llm_provider", ""), "litellm_params": kwargs.get("litellm_params", {}), "model": kwargs.get("model", ""), } - custom_llm_provider = kwargs_view["custom_llm_provider"] or kwargs_view["litellm_params"].get( + custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get( "custom_llm_provider", "" ) if not custom_llm_provider: try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs_view["model"]) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"]) except Exception: custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: @@ -1218,10 +1301,10 @@ class WebSearchInterceptionLogger(CustomLogger): messages: list[dict], tool_calls: list[dict], thinking_blocks: list[dict], - anthropic_messages_optional_request_params: dict, + anthropic_messages_optional_request_params: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None", stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], ) -> "AnthropicMessagesResponse | AsyncIterator[object]": """Legacy path: execute search + build patch + run follow-up call.""" request_patch, structured_results = await self._build_anthropic_request_patch( @@ -1229,9 +1312,9 @@ class WebSearchInterceptionLogger(CustomLogger): messages=messages, tool_calls=tool_calls, thinking_blocks=thinking_blocks, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params), logging_obj=logging_obj, - kwargs=kwargs, + kwargs=dict[str, object](kwargs), ) if request_patch.messages is None: raise ValueError("WebSearchInterception: missing follow-up messages") @@ -1246,12 +1329,14 @@ class WebSearchInterceptionLogger(CustomLogger): if max_tokens is None: max_tokens = cast(int, kwargs.get("max_tokens", 1024)) + patch_kwargs: Final = dict[str, object](request_patch.kwargs) response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( max_tokens=max_tokens, messages=request_patch.messages, model=request_patch.model or model, + **_NO_ACREATE_NAMED, **optional_params, - **request_patch.kwargs, + **patch_kwargs, ) # Legacy path: the new path goes through the typed plan + core @@ -1393,12 +1478,13 @@ class WebSearchInterceptionLogger(CustomLogger): search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None - search_litellm_params: dict[str, Any] = {} + search_litellm_params: Mapping[str, object] = {} search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) - search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) - search_provider = search_litellm_params.get("search_provider") + tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {} + search_litellm_params = dict[str, object](tool_params) + search_provider = tool_params.get("search_provider") # Fallback to perplexity if no router or no search tools configured if not search_provider: @@ -1426,12 +1512,15 @@ class WebSearchInterceptionLogger(CustomLogger): if key != "search_provider" and value is not None } result: Final = ( - await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + await litellm.asearch( + query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs + ) if search_metadata is None else await litellm.asearch( query=query, search_provider=search_provider, litellm_metadata=search_metadata, + **_NO_ASEARCH_NAMED, **search_kwargs, ) ) @@ -1471,8 +1560,7 @@ class WebSearchInterceptionLogger(CustomLogger): valid_token=user_api_key_auth, ) - auth_view: Final[_UserAuthView] = {"team_id": getattr(user_api_key_auth, "team_id", None)} - team_id: Final = auth_view["team_id"] + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) if team_id: from litellm.proxy.proxy_server import ( prisma_client, @@ -1587,10 +1675,10 @@ class WebSearchInterceptionLogger(CustomLogger): model: str, messages: list[dict], tool_calls: list[dict], - optional_params: dict, + optional_params: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None", stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], response_format: str = "openai", ) -> "ModelResponse | CustomStreamWrapper": """Legacy path: execute search + build patch + run follow-up call.""" @@ -1598,8 +1686,8 @@ class WebSearchInterceptionLogger(CustomLogger): model=model, messages=messages, tool_calls=tool_calls, - optional_params=optional_params, - kwargs=kwargs, + optional_params=dict[str, object](optional_params), + kwargs=dict[str, object](kwargs), response_format=response_format, ) if request_patch.messages is None: @@ -1607,11 +1695,13 @@ class WebSearchInterceptionLogger(CustomLogger): params: Final = dict(optional_params) params.update(request_patch.optional_params) params.pop("tool_choice", None) + patch_kwargs: Final = dict[str, object](request_patch.kwargs) return await litellm.acompletion( model=request_patch.model or model, messages=request_patch.messages, + **_NO_ACOMPLETION_NAMED, **params, - **request_patch.kwargs, + **patch_kwargs, ) async def _build_chat_completion_request_patch( diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 4f739dc4c8d..6f6b82400f9 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -16,7 +16,7 @@ import json from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable from typing_extensions import ReadOnly, TypedDict, assert_never @@ -58,6 +58,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -98,6 +100,38 @@ InputWriteBackTarget = ( ) +def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: + return value + + +def _content_block_at(blocks: Sequence[object], index: int) -> object: + return blocks[index] + + +@runtime_checkable +class _ModelDumpBlock(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +@runtime_checkable +class _TextAttrBlock(Protocol): + text: str + + +class _WritableMessage(Protocol): + @overload + def get(self, key: str, /) -> object | None: ... + + @overload + def get(self, key: str, default: object, /) -> object: ... + + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _as_writable(value: _WritableMessage) -> _WritableMessage: + return value + + @dataclass(frozen=True, slots=True) class ScannedText: text: str @@ -136,7 +170,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _build_streaming_usage_response( - responses_so_far: list[object], + responses_so_far: Sequence[object], request_data: dict | None, ) -> ModelResponse | None: chunks: Final = tuple(response for response in responses_so_far if isinstance(response, (str, bytes))) @@ -154,7 +188,7 @@ class AnthropicMessagesHandler(BaseTranslation): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: list[object] | None = None, + responses_so_far: Sequence[object] | None = None, ) -> list[bytes]: """ Build an Anthropic SSE sequence delivering the guardrail block message @@ -172,9 +206,22 @@ class AnthropicMessagesHandler(BaseTranslation): would make Anthropic clients reject the stream. """ if stream_started: - return self._block_continuation_chunks(exc, responses_so_far or []) + return list(self._block_continuation_chunks(exc, responses_so_far or [])) return self._standalone_block_chunks(exc) + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames + + message, _ = serialize_http_exception_detail(exc.detail) + return tuple(anthropic_sse_error_frames(message)) + def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: import uuid @@ -197,7 +244,9 @@ class AnthropicMessagesHandler(BaseTranslation): ) return list(FakeAnthropicMessagesStreamIterator(response=block_response)) - def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[object]) -> list[bytes]: + def _block_continuation_chunks( + self, exc: "ModifyResponseException", responses_so_far: Sequence[object] + ) -> Sequence[bytes]: """Continue an already-started message: close the open content block, append the block message as a new text block, then end the message -- without a second message_start.""" @@ -209,7 +258,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _sse(event_type: str, payload: dict) -> bytes: return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() - output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None))["output_tokens"] + output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None)).get("output_tokens", 0) open_index, max_index = self._content_block_state(responses_so_far) new_index: Final = (max_index + 1) if max_index is not None else 0 chunks: list[bytes] = [] @@ -247,7 +296,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _content_block_state( - responses_so_far: list[object], + responses_so_far: Sequence[object], ) -> tuple[int | None, int | None]: """From the SSE chunks already sent to the client, return (open content-block index or None, highest content-block index seen or None). @@ -273,7 +322,20 @@ class AnthropicMessagesHandler(BaseTranslation): return open_index, max_index @staticmethod - def _iter_sse_events(item: object) -> list[dict[str, object]]: + def _parse_sse_data_line(raw_line: str) -> tuple[Mapping[str, object], ...]: + line: Final = raw_line.strip() + if not line.startswith("data:"): + return () + try: + parsed: Final[object] = json.loads(line[len("data:") :].strip()) + except json.JSONDecodeError: + return () + if not isinstance(parsed, dict): + return () + return (_as_str_mapping(parsed),) + + @staticmethod + def _iter_sse_events(item: object) -> Sequence[Mapping[str, object]]: """Yield the event-data dicts in one stream chunk. Handles both formats this stream can carry (see @@ -281,24 +343,15 @@ class AnthropicMessagesHandler(BaseTranslation): several events separated by a blank line -- and an already-parsed event ``dict``.""" if isinstance(item, dict): - return [item] + return (_as_str_mapping(item),) if not isinstance(item, (bytes, bytearray)): - return [] - events: Final[list[dict[str, object]]] = [] - for block in item.decode("utf-8", errors="replace").split("\n\n"): - for line in block.split("\n"): - line = line.strip() - if not line.startswith("data:"): - continue - try: - parsed: str | int | float | bool | None | Sequence[object] | Mapping[str, object] = json.loads( - line[len("data:") :].strip() - ) - except json.JSONDecodeError: - continue - if isinstance(parsed, dict): - events.append(parsed) - return events + return () + return tuple( + event + for block in item.decode("utf-8", errors="replace").split("\n\n") + for line in block.split("\n") + for event in AnthropicMessagesHandler._parse_sse_data_line(line) + ) def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" @@ -331,7 +384,7 @@ class AnthropicMessagesHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - ) -> Any: + ) -> Mapping[str, object]: """ Process input messages by applying guardrails to text content. """ @@ -491,7 +544,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _openai_system_message_to_anthropic( - message: dict[str, object], + message: Mapping[str, object], ) -> dict[str, object] | None: # mutable-ok: API message payload """Convert an OpenAI system message to the client's Anthropic-shaped entry.""" content: Final = message.get("content") @@ -571,7 +624,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _defer_systems_inside_tool_exchanges( - structured_messages: list, # mutable-ok: API message payload + structured_messages: Sequence[Mapping[str, object]], ) -> list: """Hold a system row until the tool exchange around it completes so the call/result pair converts together.""" from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges @@ -765,7 +818,7 @@ class AnthropicMessagesHandler(BaseTranslation): if scan_only_tool_results: return EMPTY_EXTRACTED_INPUT - text_str: Final[str | None] = content_item.get("text", None) + text_str: Final[str | None] = content_item.get("text") return ExtractedInput( scanned=( () if text_str is None else (ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx)),) @@ -815,7 +868,7 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: list[dict[str, object]], + messages: Sequence[_WritableMessage], responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: @@ -941,7 +994,7 @@ class AnthropicMessagesHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> list[Any]: + ) -> Sequence[object]: """ Process output streaming response by applying guardrails to text content. @@ -1037,7 +1090,7 @@ class AnthropicMessagesHandler(BaseTranslation): return request_data @staticmethod - def _get_response_content(response: object) -> list[Any]: + def _get_response_content(response: object) -> Sequence[object]: """Extract content list from a dict or object response.""" if isinstance(response, dict): return response.get("content", []) or [] @@ -1047,7 +1100,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _extract_from_content_blocks( self, - response_content: list[Any], + response_content: Sequence[object], texts_to_check: list[str], images_to_check: list[str], task_mappings: list[tuple[int, int | None]], @@ -1055,21 +1108,10 @@ class AnthropicMessagesHandler(BaseTranslation): ) -> None: """Extract text, images, and tool calls from content blocks.""" for content_idx, content_block in enumerate(response_content): - block_dict: dict[str, object] = {} - if isinstance(content_block, dict): - block_type = content_block.get("type") - block_dict = cast(dict[str, object], content_block) - elif hasattr(content_block, "type"): - block_type = getattr(content_block, "type", None) - if hasattr(content_block, "model_dump"): - block_dict = content_block.model_dump() - else: - block_dict = { - "type": block_type, - "text": getattr(content_block, "text", None), - } - else: + fields = self._output_block_fields(content_block) + if fields is None: continue + block_type, block_dict = fields if block_type in ["text", "tool_use"]: self._extract_output_text_and_images( @@ -1081,6 +1123,21 @@ class AnthropicMessagesHandler(BaseTranslation): tool_calls_to_check=tool_calls_to_check, ) + @staticmethod + def _output_block_fields(content_block: object) -> "tuple[object, Mapping[str, object]] | None": + if isinstance(content_block, dict): + block_dict: Final = _as_str_mapping(content_block) + return block_dict.get("type"), block_dict + if not hasattr(content_block, "type"): + return None + block_type: Final = getattr(content_block, "type", None) + if isinstance(content_block, _ModelDumpBlock): + return block_type, content_block.model_dump() + return block_type, { + "type": block_type, + "text": getattr(content_block, "text", None), + } + @staticmethod def _build_guardrail_inputs( texts_to_check: list[str], @@ -1103,7 +1160,7 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs - def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. @@ -1186,7 +1243,7 @@ class AnthropicMessagesHandler(BaseTranslation): return text - def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if streaming response has ended by looking for non-null stop_reason. @@ -1281,7 +1338,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _extract_output_text_and_images( self, - content_block: dict[str, object], + content_block: Mapping[str, object], content_idx: int, texts_to_check: list[str], images_to_check: list[str], @@ -1304,7 +1361,7 @@ class AnthropicMessagesHandler(BaseTranslation): task_mappings.append((content_idx, None)) # Extract tool calls - elif content_type == "tool_use": + elif content_type == "tool_use" and isinstance(content_block, dict): tool_call: Final = AnthropicConfig.convert_tool_use_to_openai_format( anthropic_tool_content=content_block, index=content_idx, @@ -1329,7 +1386,7 @@ class AnthropicMessagesHandler(BaseTranslation): content_idx = cast(int, mapping[0]) # Handle both dict and object responses - response_content: list[Any] = [] + response_content: Sequence[object] = [] if isinstance(response, dict): response_content = response.get("content", []) or [] elif hasattr(response, "content"): @@ -1345,14 +1402,15 @@ class AnthropicMessagesHandler(BaseTranslation): if content_idx >= len(response_content): continue - content_block = response_content[content_idx] + content_block = _content_block_at(response_content, content_idx) # Verify it's a text block and update the text field # Handle both dict and Pydantic object content blocks if isinstance(content_block, dict): - if content_block.get("type") == "text": - cast(dict[str, object], content_block)["text"] = guardrail_response + block = _as_writable(content_block) + if block.get("type") == "text": + block["text"] = guardrail_response elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text": # Update Pydantic object's text attribute - if hasattr(content_block, "text"): + if isinstance(content_block, _TextAttrBlock): content_block.text = guardrail_response diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index a45bfb93640..050ab67c86c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -13,21 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers: """ import re -from collections.abc import Mapping, Sequence -from typing import ( - TYPE_CHECKING, - Final, - Literal, - NotRequired, - Optional, - Protocol, - TypedDict, - Union, - cast, - runtime_checkable, -) +from collections.abc import Awaitable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, Union, cast -from typing_extensions import ReadOnly +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -36,11 +25,11 @@ from litellm.types.llms.anthropic import ( CompactionBlock, UsageIteration, ) -from litellm.types.llms.openai import AllMessageValues if TYPE_CHECKING: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor, RateLimitResponse from litellm.router import Router from litellm.types.llms.anthropic import ( AllAnthropicPassThroughMessageValues, @@ -96,6 +85,77 @@ _PROPAGATED_METADATA_KEYS: Final = ( _SUMMARY_TAG_RE: Final = re.compile(r"(.*?)", re.IGNORECASE | re.DOTALL) +_MsgT: Final = TypeVar("_MsgT", bound=Mapping[str, object]) + + +def _as_object(value: object) -> object: + return value + + +def _is_tool_result_block(block: object) -> bool: + return isinstance(block, dict) and block.get("type") in ("tool_result",) + + +class _SummaryCallKwargs(TypedDict): + model: ReadOnly[str] + max_tokens: ReadOnly[int] + timeout: ReadOnly[float] + litellm_metadata: ReadOnly[Mapping[str, object]] + user: ReadOnly[NotRequired[str]] + allowed_model_region: ReadOnly[NotRequired[str]] + + +class _SummaryOptionalKwargs(TypedDict, total=False): + user: ReadOnly[str] + allowed_model_region: ReadOnly[str] + + +class _SummaryAcompletion(Protocol): + def __call__( + self, + *, + messages: Sequence[Mapping[str, object]], + **kwargs: Unpack[_SummaryCallKwargs], # kwargs-ok: forwarded verbatim to acompletion, which owns them + ) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ... + + +class _CreateRateLimitDescriptors(Protocol): + def __call__( + self, + *, + user_api_key_dict: "UserAPIKeyAuth", + data: Mapping[str, str], + rpm_limit_type: object, + tpm_limit_type: object, + model_has_failures: bool, + ) -> "Sequence[RateLimitDescriptor]": ... + + +class _AddModelRateLimitDescriptor(Protocol): + def __call__( + self, + *, + user_api_key_dict: "UserAPIKeyAuth", + requested_model: str, + descriptors: "Sequence[RateLimitDescriptor]", + ) -> None: ... + + +class _CreateOrgRateLimitDescriptors(Protocol): + def __call__( + self, user_api_key_dict: "UserAPIKeyAuth", requested_model: str | None = None + ) -> "Sequence[RateLimitDescriptor]": ... + + +class _ShouldRateLimit(Protocol): + def __call__( + self, + *, + descriptors: "Sequence[RateLimitDescriptor]", + parent_otel_span: object, + read_only: bool, + ) -> "Awaitable[RateLimitResponse]": ... + def _read_summary_model_setting() -> str | None: """Look up the configured summarization model from proxy general_settings.""" @@ -438,40 +498,57 @@ async def _check_summary_model_rate_limit( except Exception: return True - limiter: Final = getattr(proxy_logging_obj, "max_parallel_request_limiter", None) + limiter: Final[object] = getattr(proxy_logging_obj, "max_parallel_request_limiter", None) + should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None) + create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr( + limiter, "_create_rate_limit_descriptors", None + ) + add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr( + limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None + ) + add_project_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr( + limiter, "_add_project_model_rate_limit_descriptor_from_metadata", None + ) + create_org_descriptors: Final[_CreateOrgRateLimitDescriptors | None] = getattr( + limiter, "create_organization_rate_limit_descriptor", None + ) if ( limiter is None - or not hasattr(limiter, "should_rate_limit") - or not hasattr(limiter, "_create_rate_limit_descriptors") + or should_rate_limit_check is None + or create_descriptors is None + or add_team_descriptor is None + or add_project_descriptor is None + or create_org_descriptors is None ): return True try: - metadata: Final = getattr(user_api_key_auth, "metadata", None) or {} + metadata: Final[Mapping[str, object]] = getattr(user_api_key_auth, "metadata", None) or {} data: Final = {"model": summary_model} - descriptors: Final = limiter._create_rate_limit_descriptors( + base_descriptors: Final = create_descriptors( user_api_key_dict=user_api_key_auth, data=data, rpm_limit_type=metadata.get("rpm_limit_type"), tpm_limit_type=metadata.get("tpm_limit_type"), model_has_failures=False, ) - limiter._add_team_model_rate_limit_descriptor_from_metadata( + add_team_descriptor( user_api_key_dict=user_api_key_auth, requested_model=summary_model, - descriptors=descriptors, + descriptors=base_descriptors, ) - limiter._add_project_model_rate_limit_descriptor_from_metadata( + add_project_descriptor( user_api_key_dict=user_api_key_auth, requested_model=summary_model, - descriptors=descriptors, + descriptors=base_descriptors, ) - descriptors.extend(limiter.create_organization_rate_limit_descriptor(user_api_key_auth, summary_model)) + descriptors: Final = (*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model)) if not descriptors: return True - response: Final = await limiter.should_rate_limit( + parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None) + response: Final[RateLimitResponse] = await should_rate_limit_check( descriptors=descriptors, - parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None), + parent_otel_span=parent_otel_span, read_only=True, ) except Exception as e: @@ -485,7 +562,7 @@ async def _check_summary_model_rate_limit( def _find_latest_compaction_index( - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], ) -> tuple[int | None, int | None]: """Return (message_index, block_index) of the most recent compaction block. @@ -504,8 +581,8 @@ def _find_latest_compaction_index( def _slice_around_compaction_block( - messages: list[dict[str, object]], -) -> tuple[list[dict[str, object]], dict[str, object] | None]: + messages: Sequence[_MsgT], +) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]: """Apply Anthropic's "drop everything before the compaction block" rule. Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)`` @@ -519,21 +596,22 @@ def _slice_around_compaction_block( return messages, None original_msg: Final = messages[msg_idx] - raw_content: Final = original_msg.get("content") - original_content: Final[list[object]] = raw_content if isinstance(raw_content, list) else [] - compaction_block: Final = cast(dict[str, object], original_content[blk_idx]) + original_content: Final = original_msg["content"] + if not isinstance(original_content, list): + return messages, None + original_blocks: Final = cast("Sequence[dict[str, object]]", original_content) + compaction_block: Final = original_blocks[blk_idx] # Per Anthropic's contract everything before the compaction block is # dropped, including earlier blocks within the same assistant message. - sliced_content: Final = list(original_content[blk_idx:]) + sliced_content: Final = list(original_blocks[blk_idx:]) - sliced_messages: Final[list[dict[str, object]]] = [{**original_msg, "content": sliced_content}] - sliced_messages.extend(messages[msg_idx + 1 :]) + sliced_messages: Final = [{**original_msg, "content": sliced_content}, *messages[msg_idx + 1 :]] return sliced_messages, compaction_block def _strip_compaction_blocks( - messages: list[dict[str, object]], + messages: Sequence[dict[str, object]], ) -> list[dict[str, object]]: """Drop any ``compaction`` content blocks from messages. @@ -640,7 +718,7 @@ def _propagate_metadata( def _count_effective_tokens( model: str, - effective_messages: list[dict[str, object]], + effective_messages: Sequence[dict[str, object]], compaction_block: CompactionBlock | None, tools: list[dict[str, object]] | None, system: str | list[dict[str, object]] | None = None, @@ -719,17 +797,18 @@ def _system_to_text( return "" if isinstance(system, str): return system - parts: Final[list[str]] = [] - for block in system: - if isinstance(block, dict) and block.get("type") == "text": - text = block.get("text") - if isinstance(text, str) and text: - parts.append(text) - return "\n".join(parts) + return "\n".join( + text + for block in system + if isinstance(block, dict) + and block.get("type") == "text" + and isinstance(text := block.get("text"), str) + and text + ) def _select_last_user_question( - messages: list[dict[str, object]], + messages: Sequence[dict[str, object]], ) -> list[dict[str, object]]: """Pick the most recent ``user`` turn that is a real question. @@ -744,16 +823,18 @@ def _select_last_user_question( turns, or contained no user turns at all). The downstream call always needs a non-empty user message. """ + blocks: Sequence[object] for msg in reversed(messages): if msg.get("role") != "user": continue content = msg.get("content") if isinstance(content, list): - filtered = [blk for blk in content if not (isinstance(blk, dict) and blk.get("type") == "tool_result")] + blocks = [*map(_as_object, content)] + filtered = [blk for blk in blocks if not _is_tool_result_block(blk)] if not filtered: # Purely tool_result — skip and look for an earlier turn. continue - if len(filtered) < len(content): + if len(filtered) < len(blocks): return [{**msg, "content": filtered}] return [msg] return [ @@ -796,10 +877,10 @@ def _system_to_openai_message( def _build_summary_messages( - effective_messages: list[dict[str, object]], + effective_messages: Sequence[dict[str, object]], prompt: str, system: str | list[dict[str, object]] | None = None, -) -> list[dict[str, object]]: +) -> Sequence[Mapping[str, object]]: """Build the OpenAI-shape message list for the summary call. The caller's ``system`` prompt is prepended (the default summarization @@ -827,7 +908,7 @@ def _build_summary_messages( ) openai_messages = stripped - summary_messages: Final[list[dict[str, object]]] = [] + summary_messages: Final[list[Mapping[str, object]]] = [] system_message: Final = _system_to_openai_message(system) if system_message is not None: summary_messages.append(system_message) @@ -862,33 +943,15 @@ def _append_text_to_content(content: object, extra_text: str) -> object: if isinstance(content, str): return f"{content}\n\n{extra_text}" if isinstance(content, list): - appended: Final[list[object]] = [*content, {"type": "text", "text": extra_text}] + appended: Final[Sequence[object]] = [*map(_as_object, content), {"type": "text", "text": extra_text}] return appended return [content, {"type": "text", "text": extra_text}] -class _SummaryCallUserKwarg(TypedDict, total=False): - user: ReadOnly[str] - - -class _SummaryCallRegionKwarg(TypedDict, total=False): - allowed_model_region: ReadOnly[str] - - -class _SummaryCallKwargs(TypedDict): - model: ReadOnly[str] - messages: ReadOnly[list[AllMessageValues]] - max_tokens: ReadOnly[int] - timeout: ReadOnly[float] - litellm_metadata: ReadOnly[Mapping[str, object]] - user: NotRequired[ReadOnly[str]] - allowed_model_region: NotRequired[ReadOnly[str]] - - async def _call_summary_model( *, summary_model: str, - summary_messages: list[dict[str, object]], + summary_messages: Sequence[Mapping[str, object]], metadata: Mapping[str, object], llm_router: Optional["Router"], allowed_model_region: str | None = None, @@ -926,39 +989,38 @@ async def _call_summary_model( # than from ``litellm_metadata``, so without it the summary tokens would not # debit the caller's end-user counters. end_user_id: Final = metadata.get("user_api_key_end_user_id") + user_kwargs: Final = ( + _SummaryOptionalKwargs(user=end_user_id) + if isinstance(end_user_id, str) and end_user_id + else _SummaryOptionalKwargs() + ) + region_kwargs: Final = ( + _SummaryOptionalKwargs(allowed_model_region=allowed_model_region) + if allowed_model_region is not None + else _SummaryOptionalKwargs() + ) call_kwargs: Final[_SummaryCallKwargs] = { "model": summary_model, - "messages": cast("list[AllMessageValues]", summary_messages), # cast-ok: built as OpenAI chat messages "max_tokens": max_tokens, "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, "litellm_metadata": metadata, - **( - _SummaryCallUserKwarg(user=end_user_id) - if isinstance(end_user_id, str) and end_user_id - else _SummaryCallUserKwarg() - ), - **( - _SummaryCallRegionKwarg(allowed_model_region=allowed_model_region) - if allowed_model_region is not None - else _SummaryCallRegionKwarg() - ), + **user_kwargs, + **region_kwargs, } - if llm_router is not None and hasattr(llm_router, "acompletion"): - return await llm_router.acompletion(**call_kwargs) - return await litellm.acompletion(**call_kwargs) - - -@runtime_checkable -class _ResponseWithChoices(Protocol): - choices: Sequence[object] + router_acompletion: Final[_SummaryAcompletion | None] = getattr(llm_router, "acompletion", None) + if llm_router is not None and router_acompletion is not None: + return await router_acompletion(messages=summary_messages, **call_kwargs) + return await litellm.acompletion(messages=[*summary_messages], **call_kwargs) def _extract_response_text(response: object) -> str | None: - if not isinstance(response, _ResponseWithChoices) or not response.choices: - return None try: - message: Final[object] = getattr(response.choices[0], "message", None) - content: Final[object] = getattr(message, "content", None) + choices: Final[Sequence[object] | None] = getattr(response, "choices", None) + if choices is None: + return None + choice: Final = choices[0] + message: Final = getattr(choice, "message", None) + content: Final = getattr(message, "content", None) if isinstance(content, str): return content # Some providers return a list of content parts. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..b07a6b986d7 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,8 +1,11 @@ from abc import ABC, abstractmethod +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Final, Optional if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -73,6 +76,31 @@ class BaseTranslation(ABC): return transformed + @staticmethod + def merge_user_api_key_metadata_into_request( + request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place + user_api_key_dict: Optional["UserAPIKeyAuth"], + ) -> None: + """ + Add the prefixed ``user_api_key_*`` metadata to the request's resolved + metadata bucket without overwriting existing keys. + + Writes must go through ``get_or_create_metadata_bucket``: creating a + ``litellm_metadata`` key on a route whose bucket is ``metadata`` (chat + completions) flips the bucket for every later metadata write, and spend + logging never sees those writes (e.g. guardrail_information). + """ + from litellm.litellm_core_utils.core_helpers import ( + get_or_create_metadata_bucket, + ) + + user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if not user_metadata: + return + _, metadata_bucket = get_or_create_metadata_bucket(request_data) + for key, value in user_metadata.items(): + metadata_bucket.setdefault(key, value) + @abstractmethod async def process_input_messages( self, @@ -147,6 +175,26 @@ class BaseTranslation(ABC): """ return None + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + """ + Build the stream items that surface a guardrail HTTPException (a block + with the default exception-on-block config, or a failed scan) after the + response has already started streaming, in this endpoint's wire format. + + Called only once chunks have been sent: the HTTP status is gone, so the + failure must travel as an in-stream error frame. ``responses_so_far`` + holds the chunks the client has already received, for formats whose + error frame continues the stream (e.g. sequence numbers). + + Returns None when the format has no in-stream error frame; the caller + then re-raises ``exc``. Override in endpoint subclasses. + """ + return None + def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None: """ Convert request data to OpenAI-spec structured messages. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 54673c77f80..de15fefe943 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,6 +14,7 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Union, cast import litellm @@ -46,6 +47,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -382,11 +385,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "response" not in request_data: request_data["response"] = response - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -555,11 +554,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "responses" not in request_data: request_data["responses"] = responses_so_far - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -591,6 +586,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + import json + + from litellm.proxy.common_request_processing import sse_error_payload + + _, error_obj = sse_error_payload(exc) + return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),) + @staticmethod def _accumulate_string_content_by_choice_index( responses_so_far: list["ModelResponseStream"], @@ -653,10 +660,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): request_data = {"responses": responses_so_far} elif "responses" not in request_data: request_data["responses"] = responses_so_far - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if responses_so_far and getattr(responses_so_far[0], "model", None): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..24dbd7c08d0 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -48,6 +48,8 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam, + ErrorEvent, + ErrorEventError, OpenAIMcpServerTool, ResponsesAPIStreamEvents, ) @@ -59,6 +61,8 @@ from litellm.types.responses.main import ( from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth @@ -80,6 +84,14 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: + sequence_numbers: Final = ( + item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) + for item in reversed(responses_so_far or ()) + ) + return next((n + 1 for n in sequence_numbers if isinstance(n, int)), 0) + + class OpenAIResponsesHandler(BaseTranslation): """ Handler for processing OpenAI Responses API with guardrails. @@ -620,6 +632,29 @@ class OpenAIResponsesHandler(BaseTranslation): } return responses_so_far[-1].get("type") in terminal_types + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + + message, _ = serialize_http_exception_detail(exc.detail) + return ( + ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=_next_stream_sequence_number(responses_so_far), + error=ErrorEventError( + type="guardrail_error", + code=str(exc.status_code), + message=message, + param=None, + ), + ), + ) + def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: """ Get the string so far from the responses so far. diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 8dd77db08c0..c7289eb47f5 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -117,6 +117,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): def __init__(self): super().__init__() + @staticmethod + def _parse_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse: + return raw_response.json() + def get_supported_openai_params(self, model: str) -> list: """ Get the list of supported OpenAI parameters for video generation. @@ -141,7 +145,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: """ Map OpenAI parameters to RunwayML format. @@ -151,37 +155,44 @@ class RunwayMLVideoConfig(BaseVideoConfig): - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ - mapped_params: Final[dict[str, object]] = {} + supported_openai_params: Final = self.get_supported_openai_params(model) + return { + **self._prompt_image_param(video_create_optional_params), + **self._ratio_param(video_create_optional_params), + **self._duration_param(video_create_optional_params), + # Pass through other parameters that aren't OpenAI-specific + **{key: value for key, value in video_create_optional_params.items() if key not in supported_openai_params}, + } + @staticmethod + def _prompt_image_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, object]: # Handle input_reference parameter - map to promptImage + # RunwayML supports URLs and data URIs directly if "input_reference" in video_create_optional_params: - input_reference: Final = video_create_optional_params["input_reference"] - # RunwayML supports URLs and data URIs directly - mapped_params["promptImage"] = input_reference + return {"promptImage": video_create_optional_params["input_reference"]} + return {} + @staticmethod + def _ratio_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, str]: # Handle size parameter - convert "1280x720" to "1280:720" if "size" in video_create_optional_params: size: Final = video_create_optional_params["size"] if isinstance(size, str) and "x" in size: - mapped_params["ratio"] = size.replace("x", ":") + return {"ratio": size.replace("x", ":")} + return {} + @staticmethod + def _duration_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, int]: # Handle seconds parameter - convert to integer if "seconds" in video_create_optional_params: seconds: Final = video_create_optional_params["seconds"] if seconds is not None: try: - mapped_params["duration"] = int(float(seconds)) if isinstance(seconds, str) else int(seconds) + return {"duration": int(float(seconds)) if isinstance(seconds, str) else int(seconds)} except (ValueError, TypeError): # If conversion fails, use default duration pass - - # Pass through other parameters that aren't OpenAI-specific - supported_openai_params: Final = self.get_supported_openai_params(model) - for key, value in video_create_optional_params.items(): - if key not in supported_openai_params: - mapped_params[key] = value - - return mapped_params + return {} def validate_environment( self, @@ -236,7 +247,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): model: str, prompt: str, api_base: str, - video_create_optional_request_params: dict, + video_create_optional_request_params: dict[str, object], litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[dict, RequestFiles, str]: @@ -406,20 +417,18 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Get task status to retrieve video URL url: Final = f"{api_base}/tasks/{encoded_video_id}" - params: Final[dict[str, str]] = {} + return url, dict[str, str]() - return url, params - - def _extract_video_url_from_response(self, response_data: dict[str, Any]) -> str: + def _extract_video_url_from_response(self, response_data: _RunwayTaskResponse) -> str: """ Helper method to extract video URL from RunwayML response. Shared between sync and async transforms. """ # Extract video URL from the output field video_url = None - if "output" in response_data and response_data["output"]: - output: Final = response_data["output"] - video_url = output[0] if isinstance(output, list) else output + raw_output: Final = response_data.get("output") + if raw_output: + video_url = raw_output if isinstance(raw_output, str) else raw_output[0] if not video_url: # Check if the video generation failed or is still processing @@ -453,7 +462,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_url: Final = self._extract_video_url_from_response(response_data) # Download the video from the CloudFront URL synchronously @@ -482,7 +491,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_url: Final = self._extract_video_url_from_response(response_data) # Download the video from the CloudFront URL asynchronously @@ -564,9 +573,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Construct the URL for task cancellation url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel" - data: Final[dict[str, str]] = {} - - return url, data + return url, dict[str, str]() def transform_video_delete_response( self, @@ -604,9 +611,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): url: Final = f"{api_base}/tasks/{encoded_video_id}" # Empty dict for GET request (no body) - data: Final[dict[str, str]] = {} - - return url, data + return url, dict[str, str]() def transform_video_status_retrieve_response( self, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c6a6639409a..ed66bb65ced 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -502,7 +502,7 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch return logging_obj -def _serialize_http_exception_detail( +def serialize_http_exception_detail( detail: object, ) -> tuple[str, dict | None]: """ @@ -535,7 +535,7 @@ def _serialize_http_exception_detail( def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException: raw_detail: Final = _getattr_object(exc, "detail", str(exc)) - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) return ProxyException( @@ -818,7 +818,7 @@ async def _buffer_first_chunk_honoring_disconnect( raise _ClientDisconnectedBeforeFirstChunk() -def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: +def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: """Build the ProxyException-shaped ``{"error": ...}`` body used in SSE error frames. Matches ``ProxyException.to_dict()`` so streaming and non-streaming error frames @@ -827,7 +827,7 @@ def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: # Preserve status code from HTTPException (e.g. guardrail blocks) error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start") - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) @@ -942,7 +942,7 @@ async def create_response( # Unexpected error consuming first chunk. verbose_proxy_logger.exception("Error consuming first chunk from generator: %s", e) - error_status, error_obj = _sse_error_payload(e) + error_status, error_obj = sse_error_payload(e) async def error_gen_message() -> AsyncGenerator[str, None]: for frame in _sse_error_frames(error_obj): @@ -1119,7 +1119,7 @@ async def open_sse_before_first_byte( # would never fire and the failure would go unaudited. The hook # also gets to sanitize what reaches the client, by returning or # raising a replacement, so its answer decides the frame. - _, error_obj = _sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) + _, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) for frame in _sse_error_frames(error_obj): yield frame.encode() return @@ -2380,53 +2380,14 @@ class ProxyBaseLLMRequestProcessing: if requested_model_from_client: self.data["_litellm_client_requested_model"] = requested_model_from_client - # Streaming: attach a closure that fires after all guardrail - # end-of-stream blocks complete. CSW.__anext__ stores the - # assembled response on logging_obj; the outer consumer - # (ProxyLogging._fire_deferred_stream_logging) fires the - # closure after the full streaming pipeline finishes. - # The closure runs non-apply_guardrail hooks on the - # assembled response, then fires success logging. - # Only for CustomStreamWrapper — raw async generators from - # passthrough routes bypass CSW and would orphan the closure. - from litellm.litellm_core_utils.streaming_handler import ( - CustomStreamWrapper, - ) - - if _post_call_guardrails_active and isinstance(response, CustomStreamWrapper): - # Intentionally a live reference (not a copy) — mirrors - # ProxyLogging.post_call_success_hook which also mutates - # data["guardrail_to_apply"] during iteration. - _captured_data: Final = self.data - _captured_user_api_key_dict: Final = user_api_key_dict - _captured_logging_obj: Final = logging_obj - - async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: - await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( - captured_data=_captured_data, - captured_user_api_key_dict=_captured_user_api_key_dict, - captured_logging_obj=_captured_logging_obj, - assembled_response=assembled_response, - cache_hit=cache_hit, - ) - - logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete - elif ( - _post_call_guardrails_active - and route_type == "anthropic_messages" - and self._is_streaming_response(response) - ): - from litellm.litellm_core_utils.logging_worker import ( - GLOBAL_LOGGING_WORKER, + if _post_call_guardrails_active: + self._arm_deferred_stream_dispatch( + response=response, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, ) - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) - - logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete - if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): @@ -3096,6 +3057,94 @@ class ProxyBaseLLMRequestProcessing: except Exception as e: verbose_proxy_logger.exception("Error firing deferred logging: %s", e) + def _arm_deferred_stream_dispatch( + self, + response: object, + route_type: str, + user_api_key_dict: "UserAPIKeyAuth", + logging_obj: LiteLLMLoggingObj, + ) -> None: + """ + Streaming with post-call guardrails active: attach a closure that + ProxyLogging._fire_deferred_stream_logging fires after all guardrail + end-of-stream blocks complete, so the spend log sees + guardrail_information. + + Three closure shapes, matching who owns logging for the stream: + - CustomStreamWrapper (chat completions) stores + (assembled_response, cache_hit); the closure also runs + non-apply_guardrail post-call hooks via + _run_deferred_stream_guardrails. + - Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares + its inner CustomStreamWrapper's logging_obj, so it stores the same + (assembled_response, cache_hit) shape; the closure only dispatches + success logging, matching the route's pre-existing hook surface. + - Native anthropic_messages/aresponses iterators store a single + ready-made logging coroutine to enqueue. + + Raw async generators from passthrough routes bypass all three and + would orphan the closure, so they are not armed here. + + The router wraps iterators that cannot carry _hidden_params in + HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the + unwrapped inner iterator. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper + + unwrapped: Final = response._inner if isinstance(response, HiddenParamsAsyncIteratorWrapper) else response + + if isinstance(unwrapped, CustomStreamWrapper): + # Intentionally a live reference (not a copy) — mirrors + # ProxyLogging.post_call_success_hook which also mutates + # data["guardrail_to_apply"] during iteration. + _captured_data: Final = self.data + _captured_user_api_key_dict: Final = user_api_key_dict + _captured_logging_obj: Final = logging_obj + + async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=_captured_data, + captured_user_api_key_dict=_captured_user_api_key_dict, + captured_logging_obj=_captured_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + return + + if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): + return + + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + if isinstance(unwrapped, LiteLLMCompletionStreamingIterator): + _captured_bridge_logging_obj: Final = logging_obj + + async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: + await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + + logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete + return + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + async def _on_deferred_native_stream_complete( + logging_coroutine: Coroutine[object, object, object], + ) -> None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + + logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete + @staticmethod async def _run_deferred_stream_guardrails( captured_data: dict, diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 367552e783e..cd0aa75b859 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -6,9 +6,11 @@ Admins use the management endpoints to read and update input_policy / output_pol """ import uuid -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem @@ -27,6 +29,13 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +class _ModelDumpMethod(Protocol): + def __call__(self) -> Mapping: ... + + +_ROW_DICT: Final = TypeAdapter(dict) + + def _tool_table_actions(prisma_client: "PrismaClient") -> "TableActions[prisma_db_models.LiteLLM_ToolTable]": table: Final[TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table return table @@ -41,33 +50,35 @@ def _object_permission_table_actions( return table -def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow: +def _row_to_model(row: object) -> LiteLLM_ToolTableRow: """Convert a Prisma model instance or dict to LiteLLM_ToolTableRow.""" - model_dump: Final = getattr(row, "model_dump", None) + model_dump: Final[_ModelDumpMethod | None] = getattr(row, "model_dump", None) if callable(model_dump): row = model_dump() elif not isinstance(row, dict): - row = { - k: getattr(row, k, None) - for k in ( - "tool_id", - "tool_name", - "origin", - "input_policy", - "output_policy", - "call_count", - "assignments", - "key_hash", - "team_id", - "key_alias", - "user_agent", - "last_used_at", - "created_at", - "updated_at", - "created_by", - "updated_by", - ) - } + row = _ROW_DICT.validate_python( + { + k: getattr(row, k, None) + for k in ( + "tool_id", + "tool_name", + "origin", + "input_policy", + "output_policy", + "call_count", + "assignments", + "key_hash", + "team_id", + "key_alias", + "user_agent", + "last_used_at", + "created_at", + "updated_at", + "created_by", + "updated_by", + ) + } + ) return LiteLLM_ToolTableRow( tool_id=row.get("tool_id", ""), tool_name=row.get("tool_name", ""), @@ -190,7 +201,7 @@ async def update_tool_policy( _updated_by: Final = updated_by or "system" now: Final = datetime.now(timezone.utc) - create_data: Final[dict[str, object]] = { + create_data: Final[Mapping[str, str | datetime]] = { "tool_id": str(uuid.uuid4()), "tool_name": tool_name, "input_policy": input_policy or "untrusted", @@ -200,14 +211,16 @@ async def update_tool_policy( "created_at": now, "updated_at": now, } - update_data: Final[dict[str, object]] = { - "updated_by": _updated_by, - "updated_at": now, + update_data: Final[Mapping[str, str | datetime]] = { + key: value + for key, value in ( + ("updated_by", _updated_by), + ("updated_at", now), + ("input_policy", input_policy), + ("output_policy", output_policy), + ) + if value is not None } - if input_policy is not None: - update_data["input_policy"] = input_policy - if output_policy is not None: - update_data["output_policy"] = output_policy await _tool_table_actions(prisma_client).upsert( where={"tool_name": tool_name}, @@ -338,7 +351,7 @@ class ToolPolicyRegistry: self._blocked_tools_by_op_id = {} for row in perms: op_id = getattr(row, "object_permission_id", None) - blocked = getattr(row, "blocked_tools", None) or [] + blocked: Sequence[str] = getattr(row, "blocked_tools", None) or [] if op_id: self._blocked_tools_by_op_id[op_id] = list(blocked) @@ -370,10 +383,12 @@ class ToolPolicyRegistry: """ if not tool_names: return {} - blocked: Final[set[str]] = set() - for op_id in (object_permission_id, team_object_permission_id): - if op_id and op_id.strip(): - blocked.update(self._blocked_tools_by_op_id.get(op_id.strip(), [])) + blocked: Final[frozenset[str]] = frozenset( + tool + for op_id in (object_permission_id, team_object_permission_id) + if op_id and op_id.strip() + for tool in self._blocked_tools_by_op_id.get(op_id.strip(), []) + ) result: Final[dict[str, str]] = {} for name in tool_names: if name in blocked: @@ -408,13 +423,12 @@ async def add_tool_to_object_permission_blocked( ) if row is None: return False - current: Final = list(getattr(row, "blocked_tools", []) or []) + current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or [] if tool_name in current: return True - current.append(tool_name) await _object_permission_table_actions(prisma_client).update( where={"object_permission_id": object_permission_id}, - data={"blocked_tools": current}, + data={"blocked_tools": [*current, tool_name]}, ) return True except Exception as e: @@ -436,13 +450,12 @@ async def remove_tool_from_object_permission_blocked( ) if row is None: return False - current = list(getattr(row, "blocked_tools", []) or []) + current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or [] if tool_name not in current: return False - current = [t for t in current if t != tool_name] await _object_permission_table_actions(prisma_client).update( where={"object_permission_id": object_permission_id}, - data={"blocked_tools": current}, + data={"blocked_tools": [t for t in current if t != tool_name]}, ) return True except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 4a3c8b0628d..14f2aa88a03 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -18,6 +18,7 @@ import time from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import datetime, timezone from itertools import accumulate, groupby +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast import httpx @@ -42,7 +43,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_request_processing import _serialize_http_exception_detail +from litellm.proxy.common_request_processing import serialize_http_exception_detail from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, @@ -52,7 +53,12 @@ from litellm.proxy.guardrails.anthropic_sse import ( model_response_text, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks +from litellm.types.guardrails import ( + BedrockChecksConfigModel, + BedrockGuardrailStreamingParams, + GuardrailEventHooks, + LitellmParams, +) from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksMessage, @@ -221,9 +227,23 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + streaming_buffer_until_moderated: bool | None = None, + streaming_sampling_rate: int | None = None, + streaming_end_of_stream_only: bool | None = None, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self._set_streaming_params( + BedrockGuardrailStreamingParams.from_extras( + MappingProxyType( + { + "streaming_buffer_until_moderated": streaming_buffer_until_moderated, + "streaming_sampling_rate": streaming_sampling_rate, + "streaming_end_of_stream_only": streaming_end_of_stream_only, + } + ) + ) + ) self.guardrailIdentifier = guardrailIdentifier self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" @@ -278,6 +298,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): list(self.checks.keys()) if self.checks else None, ) + def _set_streaming_params(self, streaming_params: BedrockGuardrailStreamingParams) -> None: + self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated + self.streaming_sampling_rate = streaming_params.streaming_sampling_rate + self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only + + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + super().update_in_memory_litellm_params(litellm_params) + self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra)) + + def _streams_incrementally(self) -> bool: + return not self.streaming_buffer_until_moderated and not self.mask_response_content + @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: return [ @@ -2658,6 +2690,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Collect content from the stream and run the bedrock OUTPUT scan (post_call only validates the response). """ + if self._streams_incrementally(): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + async for streamed_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + guardrail_to_apply=self, + buffer_until_moderated_default=False, + ): + yield streamed_chunk + return + # Import here to avoid circular imports from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.main import stream_chunk_builder @@ -2714,7 +2761,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if not raw_sse or (not is_block and not headers_flushed): raise - block_message, _ = _serialize_http_exception_detail(block_detail) + block_message, _ = serialize_http_exception_detail(block_detail) for error_frame in anthropic_sse_error_frames( block_message if is_block else f"{block_exc.status_code}: {block_message}" ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e95e97bfe74..46b00829b74 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -57,6 +57,9 @@ class _EndpointTranslation(Protocol): @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + @property + def build_stream_error_items(self) -> "Callable[..., Sequence[object] | None]": ... + def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: return translation @@ -408,14 +411,32 @@ class UnifiedLLMGuardrails(CustomLogger): call_type: str | None, responses_so_far: Sequence[object], request_data: dict, + endpoint_translation: _EndpointTranslation | None = None, + stream_started: bool = False, + responses_yielded: Sequence[object] | None = None, ) -> AsyncGenerator[object, None]: - """Surface a mid-stream HTTPException. For A2A call types the response has - already started, so emit an in-stream JSON-RPC error chunk; otherwise - re-raise so the proxy can report it. + """Surface a mid-stream HTTPException (a guardrail block with the default + exception-on-block config, or a failed scan). + + A2A call types emit an in-stream JSON-RPC error chunk. For other call + types, once chunks have already reached the client the HTTP status is + gone, so the failure is delegated to the endpoint translation's + ``build_stream_error_items`` and travels as an in-stream error frame in + that endpoint's wire format. Before the first chunk (or when the format + has no in-stream error frame) the exception is re-raised so the proxy + can report it with a real HTTP status. """ if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data)) return + if stream_started and endpoint_translation is not None: + error_items: Final = endpoint_translation.build_stream_error_items( + exc, responses_so_far=tuple(responses_yielded) if responses_yielded is not None else None + ) + if error_items is not None: + for error_item in error_items: + yield error_item + return raise exc def _build_transform_chunk( @@ -586,7 +607,15 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk raise _StreamTerminated() except HTTPException as e: - async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data): + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): yield error_item raise _StreamTerminated() @@ -1070,11 +1099,17 @@ class UnifiedLLMGuardrails(CustomLogger): return except HTTPException as e: # Response already started (we already yielded chunks); cannot send 400. - # For A2A, yield an in-stream JSON-RPC error so the client sees it. - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - return - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=chunks_yielded, + responses_yielded=responses_yielded, + ): + yield error_item + return chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1133,7 +1168,13 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk return except HTTPException as e: - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - else: - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): + yield error_item diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 47aea62f4c2..76dea1b7784 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -11,6 +11,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): BedrockGuardrail, ) + streaming_params: Final = BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra) _bedrock_callback: Final = BedrockGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, @@ -38,6 +39,9 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only, only_scan_new_messages=litellm_params.only_scan_new_messages or False, + streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated, + streaming_sampling_rate=streaming_params.streaming_sampling_rate, + streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 9a98bdbb6b1..9c4c9948c10 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Awaitable, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -735,10 +735,44 @@ def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKey ) -async def _get_user_info_teams( - prisma_client: Any, +class _UserInfoDataClient(Protocol): + @overload + async def get_data(self, *, user_id: str) -> "prisma_models.LiteLLM_UserTable | None": ... + + @overload + async def get_data( + self, + *, + user_id: str | None, + table_name: Literal["key"], + query_type: Literal["find_all"], + ) -> "Sequence[LiteLLM_VerificationToken] | None": ... + + @overload + async def get_data( + self, + *, + team_id_list: list[str], + table_name: Literal["team"], + query_type: Literal["find_all"], + ) -> "Sequence[TeamListResponseObject] | None": ... + + +async def _get_user_info_keys( + prisma_client: "_UserInfoDataClient", user_id: str | None, - user_info: Any | None, +) -> "Sequence[LiteLLM_VerificationToken] | None": + return await prisma_client.get_data( + user_id=user_id, + table_name="key", + query_type="find_all", + ) + + +async def _get_user_info_teams( + prisma_client: "_UserInfoDataClient", + user_id: str | None, + user_info: "prisma_models.LiteLLM_UserTable", user_api_key_dict: UserAPIKeyAuth, ) -> tuple[list[TeamListResponseObject], list[TeamListResponseObject] | None]: """Fetch and merge teams from membership + user.teams field.""" @@ -759,7 +793,7 @@ async def _get_user_info_teams( team_list = teams_1 team_id_list = [team.team_id for team in teams_1] - teams_2: list[TeamListResponseObject] | None = None + teams_2: Sequence[TeamListResponseObject] | None = None target_team_ids: Final = getattr(user_info, "teams", None) if target_team_ids and isinstance(target_team_ids, list): @@ -769,8 +803,8 @@ async def _get_user_info_teams( query_type="find_all", ) elif user_api_key_dict.user_id is not None and user_id is None: - caller_user_info: Final[object] = await prisma_client.get_data(user_id=user_api_key_dict.user_id) - caller_team_ids: Final = getattr(caller_user_info, "teams", None) + caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id) + caller_team_ids: Final = caller_user_info.teams if caller_user_info is not None else None if caller_team_ids: teams_2 = await prisma_client.get_data( team_id_list=caller_team_ids, @@ -807,7 +841,7 @@ def _redact_scim_enterprise_metadata( def _build_user_info_response( user_id: str | None, user_info: Any | None, - keys: list[LiteLLM_VerificationToken] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, team_list: list[TeamListResponseObject], teams_1: list[TeamListResponseObject] | None, model_max_budget_usage: dict[str, dict[str, object]] | None = None, @@ -894,11 +928,7 @@ async def user_info( ) ## GET ALL KEYS ## - keys: Final = await prisma_client.get_data( - user_id=user_id, - table_name="key", - query_type="find_all", - ) + keys: Final = await _get_user_info_keys(prisma_client, user_id) response_data: Final = _build_user_info_response( user_id=user_id, @@ -1077,6 +1107,12 @@ async def user_info_v2( raise handle_exception_on_proxy(e) +async def _fetch_admin_teams_and_keys_rows( + prisma_client: "PrismaClient", sql_query: str +) -> Sequence[Mapping[str, Sequence[Mapping[str, object]] | None]]: + return await prisma_client.db.query_raw(sql_query) + + async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): """ Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying @@ -1100,22 +1136,25 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - results: Final = await prisma_client.db.query_raw(sql_query) + results: Final = await _fetch_admin_teams_and_keys_rows(prisma_client, sql_query) verbose_proxy_logger.debug("results_keys: %s", results) - _keys_in_db: Final[Sequence[dict[str, object]]] = results[0]["keys"] or [] + _keys_in_db: Final[Sequence[Mapping[str, object]]] = results[0]["keys"] or [] # cast all keys to LiteLLM_VerificationToken keys_in_db: Final = [] for key in _keys_in_db: - if key.get("models") is None: - key["models"] = [] - keys_in_db.append(LiteLLM_VerificationToken.model_validate(key)) + key_payload = dict[str, object](key) + if key_payload.get("models") is None: + key_payload["models"] = [] + keys_in_db.append(LiteLLM_VerificationToken.model_validate(key_payload)) # cast all teams to LiteLLM_TeamTable - _teams_in_db: list[LiteLLM_TeamTable] = results[0]["teams"] or [] - _teams_in_db = [LiteLLM_TeamTable.model_validate(team) for team in _teams_in_db] - _teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "") + _teams_rows: Final[Sequence[Mapping[str, object]]] = results[0]["teams"] or [] + _teams_in_db: Final = sorted( + (LiteLLM_TeamTable.model_validate(team) for team in _teams_rows), + key=lambda x: getattr(x, "team_alias", "") or "", + ) returned_keys: Final = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) # Get admin's own user_id and user_info @@ -1140,7 +1179,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): def _process_keys_for_user_info( - keys: list[LiteLLM_VerificationToken] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, ): from litellm.constants import UI_SESSION_TOKEN_TEAM_ID @@ -1231,7 +1270,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda async def _schedule_user_update_audit_log( - response: dict[str, Any], + response: Mapping[str, object], existing_user_row: BaseModel | None, litellm_changed_by: str | None, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f0f684b0fa1..6c3f2d082ec 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,7 +18,8 @@ import os import re import secrets import traceback -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast @@ -230,6 +231,54 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions: ) +class _CustomKeyHooksModule(Protocol): + user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None + user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None + + +def _custom_key_generate_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_generate + + +def _custom_key_update_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_update + + +class _LegacyDumpable(Protocol): + def dict(self) -> Mapping[str, object]: ... + + +def _legacy_model_dict(row: _LegacyDumpable) -> Mapping[str, object]: + return row.dict() + + +def _as_object_dict(values: Mapping[str, object]) -> Mapping[str, object]: + return values + + +def _model_items(model: BaseModel) -> Iterator[tuple[str, object]]: + return iter(model) + + +class _EnvVarsParam(Protocol): + @property + def param_value(self) -> Mapping[str, str] | None: ... + + +def _env_vars_param_value(param: _EnvVarsParam) -> Mapping[str, str] | None: + return param.param_value + + +def _tx_tables_context( + open_tx: Callable[[], AbstractAsyncContextManager[_TxTables]], +) -> AbstractAsyncContextManager[_TxTables]: + return open_tx() + + async def _check_custom_key_allowed(custom_key_value: str | None) -> None: """Raise 403 if custom API keys are disabled and a custom key was provided.""" if custom_key_value is None: @@ -910,7 +959,7 @@ async def _common_key_generation_helper( # check if user set default key/generate params on config.yaml if litellm.default_key_generate_params is not None: - for elem in data: + for elem in _model_items(data): key, value = elem if ( value is None @@ -1692,11 +1741,11 @@ async def generate_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ try: + from litellm.proxy import proxy_server from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, - user_custom_key_generate, ) if prisma_client is None: @@ -1723,7 +1772,7 @@ async def generate_key_fn( ) custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( - user_custom_key_generate + _custom_key_generate_hook(proxy_server) ) if custom_key_generate_hook is not None: if inspect.iscoroutinefunction(custom_key_generate_hook): @@ -1892,11 +1941,11 @@ async def generate_service_account_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ + from litellm.proxy import proxy_server from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, - user_custom_key_generate, ) if prisma_client is None: @@ -1924,7 +1973,9 @@ async def generate_service_account_key_fn( verbose_proxy_logger.debug("entered /key/generate") - custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_generate_hook( + proxy_server + ) if custom_key_generate_hook is not None: if inspect.iscoroutinefunction(custom_key_generate_hook): result: Final = await custom_key_generate_hook(data) @@ -1998,7 +2049,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: Final = _as_object_dict(data.model_dump(exclude_unset=True, exclude_none=True)) try: for k, v in data_json.items(): @@ -2466,6 +2517,12 @@ async def _validate_mcp_servers_for_key_update( return normalized_object_permission +def _require_prisma_client(prisma_client: PrismaClient | None) -> PrismaClient: + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "Database not connected"}) + return prisma_client + + async def _validate_update_key_data( data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken, @@ -2476,8 +2533,7 @@ async def _validate_update_key_data( user_api_key_cache: UserApiKeyCache, ) -> None: """Validate permissions and constraints for key update.""" - if prisma_client is None: - raise HTTPException(status_code=500, detail={"error": "Database not connected"}) + checked_prisma_client: Final = _require_prisma_client(prisma_client) # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(data.spend) @@ -2516,7 +2572,7 @@ async def _validate_update_key_data( await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( user_api_key_dict=user_api_key_dict, route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, existing_key_row=existing_key_row, user_api_key_cache=user_api_key_cache, ) @@ -2602,7 +2658,7 @@ async def _validate_update_key_data( await _check_key_admin_access( user_api_key_dict=user_api_key_dict, hashed_token=hashed_key, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"), ) @@ -2613,7 +2669,7 @@ async def _validate_update_key_data( if _team_id_to_check is not None: team_obj = await get_team_object( team_id=_team_id_to_check, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, check_db_only=True, ) @@ -2629,7 +2685,7 @@ async def _validate_update_key_data( await _check_team_key_limits( team_table=team_obj, data=data, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( @@ -2644,7 +2700,7 @@ async def _validate_update_key_data( await _check_project_key_limits( project_id=_project_id_to_check, data=data, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, ) @@ -2659,7 +2715,7 @@ async def _validate_update_key_data( await _validate_caller_can_assign_key_org( user_api_key_dict=user_api_key_dict, organization_id=data.organization_id, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) # Check org key limits only when throughput-related fields or organization_id change @@ -2675,7 +2731,7 @@ async def _validate_update_key_data( org_table: Final = await get_org_object( org_id=_org_id_to_check, user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) if org_table is None: raise HTTPException( @@ -2685,7 +2741,7 @@ async def _validate_update_key_data( await _check_org_key_limits( org_table=org_table, data=data, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) # if team change - check if this is possible @@ -2715,7 +2771,7 @@ async def _validate_update_key_data( data=data, team_obj=team_obj, existing_key_row=existing_key_row, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, is_proxy_admin=_is_proxy_admin, ) @@ -2808,13 +2864,13 @@ async def update_key_fn( }' ``` """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) try: @@ -2845,7 +2901,9 @@ async def update_key_fn( ) # Custom key update hook - custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update + custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook( + proxy_server + ) if custom_key_update_hook is not None: if inspect.iscoroutinefunction(custom_key_update_hook): result: Final = await custom_key_update_hook(data) @@ -3007,14 +3065,16 @@ async def bulk_update_keys( }' ``` """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) + custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, @@ -3060,7 +3120,7 @@ async def bulk_update_keys( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, - user_custom_key_update=user_custom_key_update, + user_custom_key_update=custom_key_update_hook, ) successful_updates.append( @@ -3138,7 +3198,7 @@ def _build_failed_team_key_update( if hasattr(existing_key_row, "model_dump"): key_info = existing_key_row.model_dump() elif hasattr(existing_key_row, "dict"): - key_info = existing_key_row.dict() + key_info = dict[str, object](_legacy_model_dict(existing_key_row)) if key_info: key_info.pop("token", None) @@ -3169,14 +3229,16 @@ async def bulk_update_team_keys( Callable by proxy admins, or by team admins with `KEY_UPDATE` permission. """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) + custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + if prisma_client is None: raise HTTPException( status_code=500, @@ -3305,7 +3367,7 @@ async def bulk_update_team_keys( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, - user_custom_key_update=user_custom_key_update, + user_custom_key_update=custom_key_update_hook, existing_key_row=existing_by_token[db_token], ) @@ -4304,7 +4366,7 @@ def _transform_verification_tokens_to_deleted_records( "litellm_changed_by": litellm_changed_by, } ) - record = deleted_record.model_dump() + record = dict[str, object](_as_object_dict(deleted_record.model_dump())) # Map org_id to organization_id (model uses org_id, but schema expects organization_id) org_id_value: object = record.pop("org_id", None) @@ -4440,13 +4502,12 @@ async def _rotate_master_key( should_create_model_in_db=False, ) if new_model: - _dumped = new_model.model_dump(exclude_none=True) + _dumped = dict[str, object](_as_object_dict(new_model.model_dump(exclude_none=True))) _dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"]) _dumped["model_info"] = prisma.Json(_dumped["model_info"]) new_models.append(_dumped) verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx_ctx: - tx: Final[_TxTables] = tx_ctx + async with _tx_tables_context(prisma_client.db.tx) as tx: await tx.litellm_proxymodeltable.delete_many() verbose_proxy_logger.debug("Creating %s models", len(new_models)) await tx.litellm_proxymodeltable.create_many( @@ -4461,14 +4522,14 @@ async def _rotate_master_key( if config: """If environment_variables is found, decrypt it and encrypt it with the new master key""" - environment_variables_dict = {} + environment_variables_dict: Mapping[str, str] | None = {} for c in config: if c.param_name == "environment_variables": - environment_variables_dict = c.param_value + environment_variables_dict = _env_vars_param_value(c) if environment_variables_dict: decrypted_env_vars: Final = proxy_config._decrypt_and_set_db_env_variables( - environment_variables=environment_variables_dict + environment_variables=dict[str, str](environment_variables_dict) ) encrypted_env_vars: Final = proxy_config._encrypt_env_variables( environment_variables=decrypted_env_vars, @@ -4534,7 +4595,7 @@ async def _rotate_master_key( updated_patch=decrypted_cred, new_encryption_key=new_master_key, ) - _cred_data = encrypted_cred.model_dump(exclude_none=True) + _cred_data = dict[str, object](_as_object_dict(encrypted_cred.model_dump(exclude_none=True))) if "credential_values" in _cred_data: _cred_data["credential_values"] = prisma.Json(_cred_data["credential_values"]) if "credential_info" in _cred_data: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 5ad41b00890..022a1ecbac4 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -65,6 +65,19 @@ class PassThroughStreamingHandler: route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler ) raw_bytes: Final[list[bytes]] = [] + + def _build_logging_coroutine() -> Coroutine[None, None, None]: + return resolved_route_streaming_logging( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body or {}, + endpoint_type=endpoint_type, + start_time=start_time, + raw_bytes=raw_bytes, + end_time=datetime.now(), + ) + logging_scheduled = False model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection( request_body=request_body, @@ -114,6 +127,21 @@ class PassThroughStreamingHandler: ) if pending: yield pending + # Stream completed cleanly. When the proxy armed deferred + # dispatch (post-call guardrails active), park the logging + # coroutine on logging_obj instead of enqueueing now, so + # ProxyLogging._fire_deferred_stream_logging fires it after + # guardrail end-of-stream blocks populate guardrail_information. + # Disconnect/exception paths skip this and fall through to the + # immediate enqueue in ``finally`` to keep partial billing + # (LIT-2642). + if ( + getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None + and raw_bytes + and response.status_code < 400 + ): + logging_scheduled = True + litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),) except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) raise @@ -128,18 +156,7 @@ class PassThroughStreamingHandler: if not logging_scheduled and raw_bytes and response.status_code < 400: logging_scheduled = True try: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=resolved_route_streaming_logging( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body or {}, - endpoint_type=endpoint_type, - start_time=start_time, - raw_bytes=raw_bytes, - end_time=datetime.now(), - ) - ) + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 9cfd6959a66..b6cbd2d7889 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -6,7 +6,7 @@ import tempfile from collections.abc import Awaitable, Mapping, Sequence from datetime import datetime from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Final, Protocol, cast from fastapi import ( APIRouter, @@ -1317,7 +1317,7 @@ async def test_prompt( async def convert_prompt_file_to_json( file: UploadFile = File(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> dict[str, Any]: +) -> Mapping[str, object]: """ Convert a .prompt file to JSON format. diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 7f00a24a541..ee9ab4ebc97 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -11,11 +11,11 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json from collections.abc import Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Final, TypedDict +from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import Request, Response from fastapi.responses import StreamingResponse -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -29,49 +29,49 @@ if TYPE_CHECKING: from litellm.router import Router -class _StreamContentPart(TypedDict, total=False): - text: ReadOnly[str] +_JsonDict: TypeAlias = dict[str, object] +_JsonList: TypeAlias = list[object] -class _StreamOutputItem(TypedDict, total=False): +class _OutputItem(TypedDict, total=False): id: ReadOnly[str] - content: ReadOnly[Sequence[_StreamContentPart | None]] + content: ReadOnly[Sequence[object]] -class _StreamResponsePayload(TypedDict, total=False): +class _TerminalResponse(TypedDict, total=False): status: ReadOnly[ResponsesAPIStatus] - error: ReadOnly[dict[str, object] | None] - usage: ReadOnly[dict[str, object] | None] - reasoning: ReadOnly[dict[str, object] | None] + error: ReadOnly[_JsonDict] + usage: ReadOnly[_JsonDict] + reasoning: ReadOnly[_JsonDict] tool_choice: ReadOnly[object] - tools: ReadOnly[list[object] | None] - model: ReadOnly[str | None] - instructions: ReadOnly[str | None] - temperature: ReadOnly[float | None] - top_p: ReadOnly[float | None] - max_output_tokens: ReadOnly[int | None] - previous_response_id: ReadOnly[str | None] - text: ReadOnly[dict[str, object] | None] - truncation: ReadOnly[str | None] - parallel_tool_calls: ReadOnly[bool | None] - user: ReadOnly[str | None] - store: ReadOnly[bool | None] - incomplete_details: ReadOnly[dict[str, object] | None] - output: ReadOnly[Sequence[_StreamOutputItem]] + tools: ReadOnly[_JsonList] + model: ReadOnly[str] + instructions: ReadOnly[str] + temperature: ReadOnly[float] + top_p: ReadOnly[float] + max_output_tokens: ReadOnly[int] + previous_response_id: ReadOnly[str] + text: ReadOnly[_JsonDict] + truncation: ReadOnly[str] + parallel_tool_calls: ReadOnly[bool] + user: ReadOnly[str] + store: ReadOnly[bool] + incomplete_details: ReadOnly[_JsonDict] + output: ReadOnly[Sequence[_OutputItem]] class _StreamEvent(TypedDict, total=False): type: ReadOnly[str] - item: ReadOnly[_StreamOutputItem] + item: ReadOnly[_OutputItem] item_id: ReadOnly[str] - part: ReadOnly[_StreamContentPart] content_index: ReadOnly[int] delta: ReadOnly[str] - response: ReadOnly[_StreamResponsePayload] + part: ReadOnly[object] + response: ReadOnly[_TerminalResponse] -def _parse_stream_event(serialized_event: str) -> _StreamEvent: - return json.loads(serialized_event) +class _StreamEventParser: + parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) async def background_streaming_task( @@ -144,9 +144,10 @@ async def background_streaming_task( # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming - output_items: Final[dict[str, _StreamOutputItem]] = {} # Track output items by ID - # Track accumulated text deltas by (item_id, content_index) - accumulated_text: Final[dict[tuple[str, int], str]] = {} + output_items: Final = dict[str, _OutputItem]() # Track output items by ID + accumulated_text: Final = dict[ + tuple[str, int], str + ]() # Track accumulated text deltas by (item_id, content_index) # ResponsesAPIResponse fields to extract from response.completed usage_data = None @@ -216,7 +217,7 @@ async def background_streaming_task( break try: - event = _parse_stream_event(chunk_data) + event: _StreamEvent = _StreamEventParser.parse(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -235,19 +236,18 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update the output item with new content - current_item = output_items[item_id] - appended_item: _StreamOutputItem = { - **current_item, - "content": (*current_item.get("content", ()), content_part), + added_item = output_items[item_id] + output_items[item_id] = { + **added_item, + "content": (*added_item.get("content", ()), content_part), } - output_items[item_id] = appended_item state_dirty = True elif event_type == "response.output_text.delta": # Text delta - accumulate text content # https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta item_id = event.get("item_id") - content_index: int = event.get("content_index", 0) + content_index = event.get("content_index", 0) delta = event.get("delta", "") if item_id and item_id in output_items: @@ -258,24 +258,14 @@ async def background_streaming_task( accumulated_text[key] += delta # Update the content in output_items - current_item = output_items[item_id] - content_list: Sequence[_StreamContentPart | None] = current_item.get("content", ()) - if content_index < len(content_list): - # Update existing content part with accumulated text - content_entry = content_list[content_index] - if isinstance(content_entry, dict): - delta_part: _StreamContentPart = { - **content_entry, - "text": accumulated_text[key], - } - delta_item: _StreamOutputItem = { - **current_item, - "content": tuple( - delta_part if index == content_index else entry - for index, entry in enumerate(content_list) - ), - } - output_items[item_id] = delta_item + delta_item = output_items[item_id] + if "content" in delta_item: + content_list = delta_item["content"] + if content_index < len(content_list): + # Update existing content part with accumulated text + content_entry = content_list[content_index] + if isinstance(content_entry, dict): + content_entry["text"] = accumulated_text[key] state_dirty = True elif event_type == "response.content_part.done": @@ -286,17 +276,17 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update with final content from event - current_item = output_items[item_id] - content_list = current_item.get("content", ()) - if content_index < len(content_list): - finalized_item: _StreamOutputItem = { - **current_item, - "content": tuple( - content_part if index == content_index else entry - for index, entry in enumerate(content_list) - ), - } - output_items[item_id] = finalized_item + done_item = output_items[item_id] + if "content" in done_item: + content_list = done_item["content"] + if content_index < len(content_list): + output_items[item_id] = { + **done_item, + "content": tuple( + content_part if part_index == content_index else existing_part + for part_index, existing_part in enumerate(content_list) + ), + } state_dirty = True elif event_type == "response.output_item.done": diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 278b2bd6ec6..de802b95086 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,11 +5,11 @@ import json import time import traceback import uuid -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable import httpx from openai._streaming import SSEDecoder @@ -42,27 +42,14 @@ from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook if TYPE_CHECKING: + from litellm.caching.caching_handler import LLMCachingHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.types.responses.streaming_websocket import ( PresidioGuardrailCallback, ResponsesBackendWebSocket, ResponsesClientWebSocket, ) - - class _StreamCachingHandler(Protocol): - """The ``_llm_caching_handler`` attached to a logging object, as this module uses it.""" - - original_function: Callable[..., object] - - def _should_store_result_in_cache( - self, original_function: Callable[..., object], kwargs: Mapping[str, object] - ) -> bool: ... - - class PiiUnmaskingGuardrailCallback(PresidioGuardrailCallback, Protocol): - """Guardrail callback that can also reverse its own masking, selected by - ``llm_http_handler`` on exactly this attribute.""" - - def _unmask_pii_text(self, text: str, pii_tokens: Mapping[str, str]) -> str: ... + from litellm.types.router import LiteLLM_Params class ProjectQuotaCallback(Protocol): @@ -104,6 +91,52 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verif return _is_json_object(value) and all(isinstance(item, str) for item in value.values()) +class _MutableJsonObject(Protocol): + @overload + def get(self, key: str, /) -> object | None: ... + @overload + def get(self, key: str, default: object, /) -> object: ... + def __getitem__(self, key: str, /) -> object: ... + def __setitem__(self, key: str, value: object, /) -> None: ... + def __contains__(self, key: object, /) -> bool: ... + def items(self) -> Iterable[tuple[str, object]]: ... + + +class _GetsLitellmParams(Protocol): + def __call__(self, key: str, default: Mapping[str, object], /) -> LiteLLM_Params: ... + + +class _UnmasksPiiText(Protocol): + def __call__(self, text: str, pii_tokens: Mapping[str, str]) -> str: ... + + +class _ShouldStoreResultInCache(Protocol): + def __call__(self, *, original_function: Callable[..., object] | None, kwargs: Mapping[str, object]) -> bool: ... + + +class _PostStreamingDeploymentHook(Protocol): + def __call__( + self, + *, + request_data: Mapping[str, object], + response_chunk: ResponsesAPIStreamingResponse, + call_type: CallTypes | None, + ) -> Awaitable[ResponsesAPIStreamingResponse | None]: ... + + +@runtime_checkable +class _HasPostStreamingDeploymentHook(Protocol): + async_post_call_streaming_deployment_hook: _PostStreamingDeploymentHook + + +def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams: + return fn + + +_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache" +_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text" + + def _load_json_object(payload: str | bytes) -> dict[str, object]: """Parse a JSON payload that the caller consumes as an object.""" return json.loads(payload) @@ -230,7 +263,7 @@ class BaseResponsesAPIStreamingIterator: # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py _api_base: Final = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=_typed_gets_litellm_params(self.logging_obj.model_call_details.get)("litellm_params", {}), ) self._hidden_params: dict[str, object] = { "model_id": _model_id_from_metadata(litellm_metadata), @@ -432,15 +465,20 @@ class BaseResponsesAPIStreamingIterator: end_time: Final = datetime.now() if is_async: - asyncio.create_task( - self.logging_obj.dispatch_success_handlers( - logging_response, - start_time=self.start_time, - end_time=end_time, - cache_hit=self._completed_response_cache_hit, - prefer_async_handlers=True, - ) + logging_coroutine: Final = self.logging_obj.dispatch_success_handlers( + logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + prefer_async_handlers=True, ) + deferred_dispatch_armed: Final = getattr(self.logging_obj, "_on_deferred_stream_complete", None) is not None + if deferred_dispatch_armed: + # End-of-stream guardrail scans write guardrail_information after + # the terminal event; dispatching now would snapshot metadata early. + self.logging_obj._deferred_stream_complete_args = (logging_coroutine,) + else: + asyncio.create_task(logging_coroutine) else: run_async_function( async_function=self.logging_obj.async_success_handler, @@ -559,7 +597,7 @@ class BaseResponsesAPIStreamingIterator: if response_obj is None: return - caching_handler: Final[_StreamCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None) + caching_handler: Final[LLMCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None) if caching_handler is None: return @@ -577,8 +615,11 @@ class BaseResponsesAPIStreamingIterator: if preset_cache_key is not None: request_kwargs["cache_key"] = preset_cache_key - if not caching_handler._should_store_result_in_cache( # pyright: ignore[reportPrivateUsage] # no public API - original_function=caching_handler.original_function, + should_store_result_in_cache: Final[_ShouldStoreResultInCache] = getattr( + caching_handler, _SHOULD_STORE_RESULT_IN_CACHE_ATTR + ) + if not should_store_result_in_cache( + original_function=getattr(caching_handler, "original_function", None), kwargs=request_kwargs, ): return @@ -634,12 +675,15 @@ class BaseResponsesAPIStreamingIterator: typed_call_type = None request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {}) - callbacks: Final = getattr(litellm, "callbacks", None) or [] + callbacks: Final[Sequence[object]] = getattr(litellm, "callbacks", None) or [] hooks_ran = False for callback in callbacks: - if hasattr(callback, "async_post_call_streaming_deployment_hook"): + if isinstance(callback, _HasPostStreamingDeploymentHook): hooks_ran = True - result = await callback.async_post_call_streaming_deployment_hook( + post_streaming_hook: _PostStreamingDeploymentHook = ( + callback.async_post_call_streaming_deployment_hook + ) + result = await post_streaming_hook( request_data=request_data, response_chunk=chunk, call_type=typed_call_type, @@ -1093,7 +1137,7 @@ class _HasModelDumpJson(Protocol): def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... -def _dump_response_object(obj: object) -> dict[str, object]: +def _dump_response_object(obj: object) -> Mapping[str, object]: if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): @@ -1123,21 +1167,20 @@ def _build_content_part_done_event( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], + part_payload: Mapping[str, object], ) -> ResponsesAPIStreamingResponse | None: openai_types: Final = _get_openai_response_types() part_type: Final = part_payload.get("type") part: PART_UNION_TYPES if part_type == "output_text": - annotations: Final = [ - openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) - for annotation in part_payload.get("annotations", []) or [] - ] - part = openai_types.ContentPartDonePartOutputText( - type="output_text", - text=str(part_payload.get("text") or ""), - annotations=annotations, - logprobs=part_payload.get("logprobs"), + raw_annotations: Final[object] = part_payload.get("annotations", []) or [] + part = openai_types.ContentPartDonePartOutputText.model_validate( + { + "type": "output_text", + "text": str(part_payload.get("text") or ""), + "annotations": raw_annotations, + "logprobs": part_payload.get("logprobs"), + } ) elif part_type == "refusal": part = openai_types.ContentPartDonePartRefusal( @@ -1167,7 +1210,7 @@ def _add_text_like_part_events( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], + part_payload: Mapping[str, object], chunk_size: int, ) -> None: openai_types: Final = _get_openai_response_types() @@ -1184,16 +1227,19 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - annotations_payload: Final[Sequence[dict[str, object]]] = part_payload.get("annotations", []) or [] - for annotation_index, annotation in enumerate(annotations_payload): + raw_annotation_items: Final = part_payload.get("annotations") + annotation_items: Final[Sequence[object]] = raw_annotation_items if _is_json_array(raw_annotation_items) else [] + for annotation_index, annotation in enumerate(annotation_items): events.append( - openai_types.OutputTextAnnotationAddedEvent( - type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, - item_id=item_id, - output_index=output_index, - content_index=content_index, - annotation_index=annotation_index, - annotation=annotation, + openai_types.OutputTextAnnotationAddedEvent.model_validate( + { + "type": openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, + "item_id": item_id, + "output_index": output_index, + "content_index": content_index, + "annotation_index": annotation_index, + "annotation": annotation, + } ) ) events.append( @@ -1486,7 +1532,7 @@ class ResponsesWebSocketStreaming: user_api_key_dict: UserAPIKeyAuth | None = None, request_data: dict[str, object] | None = None, first_message: str | None = None, - guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, @@ -1496,17 +1542,17 @@ class ResponsesWebSocketStreaming: self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.request_data: dict[str, object] = request_data or {} - self.messages: list[dict[str, object]] = [] + self.messages: list[_MutableJsonObject] = [] self.input_messages: list[dict[str, object]] = [] self.first_message = first_message - self.guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] = guardrail_callbacks or [] + self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model - def _should_store_event(self, event_obj: Mapping[str, object]) -> bool: + def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES def _store_event(self, event: str | bytes | dict[str, object]) -> None: @@ -1620,7 +1666,7 @@ class ResponsesWebSocketStreaming: finally: await self._log_messages() - def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool: + def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool: """ Overwrite any ``model`` field in a ``response.create`` frame with the connection-authorized model to prevent deployment-substitution attacks. @@ -1689,7 +1735,7 @@ class ResponsesWebSocketStreaming: # forwarded unmasked regardless of where the client places it. nested_candidate = msg_obj.get("response") nested_response = nested_candidate if _is_json_object(nested_candidate) else None - text_containers: list[tuple[dict[str, object], str]] = [] + text_containers: list[tuple[_MutableJsonObject, str]] = [] for container in (msg_obj, nested_response): if container is None: continue @@ -1796,6 +1842,7 @@ class ResponsesWebSocketStreaming: return response_str cb: Final = self.guardrail_callbacks[0] + unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR) event_type: Final = evt_obj.get("type") if event_type == "response.completed": @@ -1815,9 +1862,7 @@ class ResponsesWebSocketStreaming: continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker - text, pii_tokens - ) + unmasked = unmask_pii_text(text, pii_tokens) if unmasked != text: content_block["text"] = unmasked modified = True @@ -1826,9 +1871,7 @@ class ResponsesWebSocketStreaming: if event_type in self._DELTA_EVENT_TYPES: delta: Final = evt_obj.get("delta") if isinstance(delta, str): - unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker - delta, pii_tokens - ) + unmasked = unmask_pii_text(delta, pii_tokens) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) @@ -2030,7 +2073,7 @@ class ManagedResponsesWebSocketHandler: model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_metadata: dict[str, object] | None = None, + litellm_metadata: Mapping[str, object] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, @@ -2043,10 +2086,11 @@ class ManagedResponsesWebSocketHandler: self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.litellm_metadata: dict[str, object] = litellm_metadata or {} - self.model_group: str | None = _optional_str( - self.litellm_metadata.get("model_group") or self.litellm_metadata.get("deployment_model_name") + self.litellm_metadata: Mapping[str, object] = litellm_metadata or {} + raw_model_group: Final = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( + "deployment_model_name" ) + self.model_group: str | None = raw_model_group if isinstance(raw_model_group, str) else None self.api_key = api_key self.api_base = api_base self.timeout = timeout @@ -2110,7 +2154,7 @@ class ManagedResponsesWebSocketHandler: self._session_history[response_id] = messages @staticmethod - def _extract_response_id(completed_event: dict[str, object]) -> str | None: + def _extract_response_id(completed_event: _MutableJsonObject) -> str | None: """ Pull the raw (decoded) response ID out of a ``response.completed`` event. Returns *None* if the event doesn't contain a usable ID. @@ -2125,7 +2169,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, object], + completed_event: _MutableJsonObject, ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into @@ -2182,7 +2226,7 @@ class ManagedResponsesWebSocketHandler: # _process_response_create sub-methods # ------------------------------------------------------------------ - async def _parse_message(self, raw_message: str) -> dict[str, object] | None: + async def _parse_message(self, raw_message: str) -> _MutableJsonObject | None: """Parse raw WS text; return the message dict or None (JSON error / ignored type).""" try: msg_obj: Final = _load_json_object(raw_message) @@ -2195,7 +2239,7 @@ class ManagedResponsesWebSocketHandler: return msg_obj @staticmethod - def _is_warmup_frame(msg_obj: dict[str, object]) -> bool: + def _is_warmup_frame(msg_obj: _MutableJsonObject) -> bool: """Return True for a response.create whose generate flag is false.""" nested: Final = msg_obj.get("response") source: Final = nested if _is_json_object(nested) and nested else msg_obj @@ -2211,13 +2255,13 @@ class ManagedResponsesWebSocketHandler: return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX) @staticmethod - def _warmup_source_params(msg_obj: dict[str, object]) -> dict[str, object]: + def _warmup_source_params(msg_obj: _MutableJsonObject) -> dict[str, object]: nested: Final = msg_obj.get("response") if _is_json_object(nested) and nested: return nested return {k: v for k, v in msg_obj.items() if k != "type"} - def _build_warmup_response(self, msg_obj: dict[str, object]) -> dict[str, object]: + def _build_warmup_response(self, msg_obj: _MutableJsonObject) -> dict[str, object]: """Build a minimal completed Responses API object for a warmup ack.""" source: Final = self._warmup_source_params(msg_obj) wire_model: Final = source.get("model") or self.model_group or self.model @@ -2235,7 +2279,7 @@ class ManagedResponsesWebSocketHandler: }, } - async def _send_warmup_ack(self, msg_obj: dict[str, object]) -> None: + async def _send_warmup_ack(self, msg_obj: _MutableJsonObject) -> None: """ Acknowledge a generate=false prewarm without calling the provider. @@ -2258,7 +2302,7 @@ class ManagedResponsesWebSocketHandler: await self.websocket.send_text(serialized) @staticmethod - def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, object]: + def _build_base_call_kwargs(msg_obj: _MutableJsonObject) -> dict[str, Any]: """ Extract Responses API params from the event, handling both wire formats: Nested: {"type": "response.create", "response": {"input": [...], ...}} @@ -2367,7 +2411,7 @@ class ManagedResponsesWebSocketHandler: call_kwargs.setdefault("litellm_params", {}) call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request - async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None: + async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None: """ Stream ``litellm.aresponses`` and forward every chunk over the WebSocket. @@ -2375,7 +2419,7 @@ class ManagedResponsesWebSocketHandler: directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the completed event dict, or ``None``. """ - completed_event: dict[str, object] | None = ( + completed_event: _MutableJsonObject | None = ( None # rebind-ok: captures the completed event once the stream yields it ) stream_response: Final = await litellm.aresponses(model=model, **call_kwargs) @@ -2401,7 +2445,7 @@ class ManagedResponsesWebSocketHandler: def _save_turn_history( self, - completed_event: dict[str, object] | None, + completed_event: _MutableJsonObject | None, prior_history: list[dict[str, object]], current_messages: list[dict[str, object]], ) -> None: diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index e4ac45df4d5..eabd9278cf6 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -8,16 +8,14 @@ Use this to route requests between Teams """ import re -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict - -from typing_extensions import ReadOnly +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload from litellm._logging import verbose_logger from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs -from litellm.types.router import ConsumedRequestTagsStamp, RouterErrors +from litellm.types.router import ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors if TYPE_CHECKING: from litellm.router import Router as _Router @@ -27,34 +25,63 @@ else: LitellmRouter = Any -class _TagRoutingLitellmParams(TypedDict, total=False): - tags: ReadOnly[Sequence[str] | None] - tag_regex: ReadOnly[Sequence[str] | None] +class _TagLitellmParamsLike(Protocol): + @overload + def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str]: ... + @overload + def get(self, key: Literal["tag_regex"], /) -> Sequence[str] | None: ... -class _TagRoutingDeployment(TypedDict, total=False): - model_name: ReadOnly[str] - litellm_params: ReadOnly[_TagRoutingLitellmParams] - model_info: ReadOnly[Mapping[str, object] | None] +class _ModelInfoLike(Protocol): + @overload + def get(self, key: Literal["allow_fail_open"], /) -> bool | None: ... + @overload + def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ... -class _TagRoutingMatchStamp(TypedDict): - matched_deployment: ReadOnly[str | None] - matched_via: ReadOnly[str] - matched_value: ReadOnly[str] - request_tags: ReadOnly[Sequence[str]] - user_agent: ReadOnly[str] +class _DeploymentLike(Protocol): + @overload + def get(self, key: Literal["litellm_params"], default: Mapping[str, object], /) -> _TagLitellmParamsLike: ... + @overload + def get(self, key: Literal["model_info"], /) -> _ModelInfoLike | None: ... + @overload + def get(self, key: Literal["model_name"], /) -> object: ... -class _TagRoutingMetadata(TypedDict, total=False): - tags: ReadOnly[Sequence[str] | None] - inherited_tags: ReadOnly[Sequence[str] | None] - user_agent: ReadOnly[str] - tag_routing: ReadOnly[_TagRoutingMatchStamp] - _consumed_request_tags: ReadOnly[object] +class _MetadataLike(Protocol): + @overload + def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["user_agent"], default: str, /) -> str: ... + @overload + def get(self, key: Literal["inherited_tags"], /) -> object: ... + def __contains__(self, key: object, /) -> bool: ... + def __setitem__(self, key: Literal["tag_routing"], value: Mapping[str, object], /) -> None: ... -_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) +class _NestedLitellmParamsLike(Protocol): + def get( + self, key: Literal["metadata", "litellm_metadata"], default: Mapping[str, object], / + ) -> _MetadataLike | None: ... + + +class _RequestKwargsLike(Protocol): + @overload + def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ... + @overload + def get(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike | None: ... + def __contains__(self, key: object, /) -> bool: ... + @overload + def __getitem__(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike: ... + @overload + def __getitem__(self, key: Literal["litellm_params"], /) -> _NestedLitellmParamsLike: ... + + +_DeploymentPool = Sequence[_DeploymentLike] | Mapping[_DeploymentLike, object] def _is_valid_deployment_tag_regex( @@ -109,11 +136,11 @@ def is_valid_deployment_tag( def _match_deployment( - deployment: _TagRoutingDeployment, - request_tags: Sequence[str] | None, - header_strings: Sequence[str], + deployment: _DeploymentLike, + request_tags: list[str] | None, + header_strings: list[str], match_any: bool, -) -> Mapping[str, str] | None: +) -> dict[str, str] | None: """ Determine whether *deployment* matches the current request. @@ -198,38 +225,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[ def _exclude_deployments( - deployments: Iterable[_TagRoutingDeployment], + deployments: _DeploymentPool, excluded_set: frozenset[str], -) -> list[_TagRoutingDeployment]: +) -> Sequence[_DeploymentLike]: if not excluded_set: return list(deployments) return [d for d in deployments if not excluded_set.intersection(d.get("litellm_params", {}).get("tags") or [])] def _require_all_tags( - deployments: Iterable[_TagRoutingDeployment], + deployments: _DeploymentPool, required_set: frozenset[str], -) -> tuple[_TagRoutingDeployment, ...]: +) -> tuple[_DeploymentLike, ...]: if not required_set: return tuple(deployments) return tuple(d for d in deployments if required_set.issubset(d.get("litellm_params", {}).get("tags") or [])) def _default_tagged_pool( - deployments: Iterable[_TagRoutingDeployment], -) -> tuple[_TagRoutingDeployment, ...]: + deployments: _DeploymentPool, +) -> tuple[_DeploymentLike, ...]: defaults: Final = tuple(d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or [])) return defaults if defaults else tuple(deployments) -def _known_tag_values(deployments: Iterable[_TagRoutingDeployment]) -> frozenset[str]: +def _known_tag_values(deployments: _DeploymentPool) -> frozenset[str]: return frozenset( - tag for d in deployments for tag in (d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ()) + tag for d in deployments for tag in (d.get("litellm_params", MappingProxyType({})).get("tags") or ()) ) def _unknown_required_tag_hides_an_answer( - healthy_deployments: Iterable[_TagRoutingDeployment], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -253,23 +280,23 @@ def _unknown_required_tag_hides_an_answer( def _chain_allows_fail_open( - healthy_deployments: Iterable[_TagRoutingDeployment], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], ) -> bool: if _unknown_required_tag_hides_an_answer(healthy_deployments, excluded_set, required_set, routing_confirmed): return False - return any((d.get("model_info") or _EMPTY_MODEL_INFO).get("allow_fail_open") is True for d in healthy_deployments) + return any((d.get("model_info") or {}).get("allow_fail_open") is True for d in healthy_deployments) def _trusted_only_pool( - healthy_deployments: Iterable[_TagRoutingDeployment], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, inherited_required_set: frozenset[str] | None, -) -> tuple[_TagRoutingDeployment, ...]: +) -> tuple[_DeploymentLike, ...]: # inherited_*_set is None only when this request carries no origin information # at all (e.g. direct SDK Router usage, bypassing the proxy layer that # populates metadata.inherited_tags) -- treat every constraint as @@ -296,8 +323,8 @@ def _trusted_only_pool( def _resolve_or_fail_open( - pool: Sequence[_TagRoutingDeployment], - healthy_deployments: Iterable[_TagRoutingDeployment], + pool: Sequence[_DeploymentLike], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -305,7 +332,7 @@ def _resolve_or_fail_open( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[_TagRoutingDeployment, ...]: +) -> tuple[_DeploymentLike, ...]: if pool: return tuple(pool) if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed): @@ -325,7 +352,7 @@ def _resolve_or_fail_open( def _resolve_constraint_only_pool( - healthy_deployments: Iterable[_TagRoutingDeployment], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -333,7 +360,7 @@ def _resolve_constraint_only_pool( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[_TagRoutingDeployment, ...]: +) -> tuple[_DeploymentLike, ...]: pool: Final = ( _require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set) if required_set @@ -355,8 +382,8 @@ def _resolve_constraint_only_pool( def _all_deployments_or_fallback( llm_router_instance: LitellmRouter, model: str, - fallback: Iterable[_TagRoutingDeployment], -) -> Iterable[_TagRoutingDeployment]: + fallback: _DeploymentPool, +) -> Sequence[_DeploymentLike | DeploymentTypedDict] | Mapping[_DeploymentLike, object]: try: return llm_router_instance._get_all_deployments(model_name=model) except Exception: # noqa: BLE001 # fail safe toward today's healthy-only behavior on lookup errors @@ -366,8 +393,8 @@ def _all_deployments_or_fallback( def _chain_tag_filtering_override( llm_router_instance: LitellmRouter, model: str, - healthy_deployments: Iterable[_TagRoutingDeployment], -) -> object: + healthy_deployments: _DeploymentPool, +) -> bool | None: # Resolved from every deployment configured for this model group, not just the # ones that survived cooldown/health filtering (async_get_healthy_deployments # filters cooldowns before calling get_deployments_for_tag) -- otherwise the @@ -379,14 +406,14 @@ def _chain_tag_filtering_override( # than crashing the request. all_deployments: Final = _all_deployments_or_fallback(llm_router_instance, model, healthy_deployments) for d in all_deployments: - value = (d.get("model_info") or _EMPTY_MODEL_INFO).get("enable_tag_filtering") + value = (d.get("model_info") or MappingProxyType({})).get("enable_tag_filtering") if value is not None: return value return None def _inherited_constraint_sets( - inherited_tags: Sequence[str] | None, routing_prefix: str + inherited_tags: object, routing_prefix: str ) -> tuple[frozenset[str] | None, frozenset[str] | None]: # None means no origin information is available at all (e.g. this request # bypassed the proxy layer that populates metadata.inherited_tags, as direct @@ -417,43 +444,42 @@ def _tag_known_to_group( if tag_set & routing_confirmed: return True try: - all_deployments: Final[Sequence[_TagRoutingDeployment]] = llm_router_instance._get_all_deployments( - model_name=model - ) + all_deployments: Final = llm_router_instance._get_all_deployments(model_name=model) except Exception: # noqa: BLE001 # fail safe toward "unrecognized" so lookup errors preserve the existing silent-fallback behavior return False return any( - tag_set.intersection(d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ()) - for d in all_deployments + tag_set.intersection(d.get("litellm_params", MappingProxyType({})).get("tags") or ()) for d in all_deployments ) -def _request_tags_after_router_consumption(metadata: _TagRoutingMetadata, model: str) -> Sequence[str] | None: +def _request_tags_after_router_consumption(metadata: object, model: str) -> Sequence[str] | None: # The pre-routing hook stamps which tags selected the router it rewrote the request # to: those tags already did their job and must not also constrain deployment choice # inside the routed group. The request's other tags still apply there, on top of the # inherited_tags snapshot that keeps key/team policy applying. Every other model # group keeps the full list. - stamp: Final = metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY) + if not isinstance(metadata, Mapping): + return None + typed_metadata: Final[Mapping[str, object]] = metadata + request_tags: Final = _tags_in_metadata(typed_metadata) + stamp: Final = typed_metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY) if not isinstance(stamp, ConsumedRequestTagsStamp) or stamp.model_group != model: - return metadata.get("tags") - request_tags: Final = metadata.get("tags") - leftover: Final = tuple( - tag for tag in (request_tags if isinstance(request_tags, (list, tuple)) else ()) if tag not in stamp.tags - ) - inherited_tags: Final = metadata.get("inherited_tags") + return request_tags + leftover: Final = tuple(tag for tag in request_tags if tag not in stamp.tags) + inherited_tags: Final = typed_metadata.get("inherited_tags") if not isinstance(inherited_tags, (list, tuple)): return leftover or None - return tuple(dict.fromkeys((*leftover, *inherited_tags))) + typed_inherited_tags: Final[Sequence[object]] = inherited_tags + return tuple(dict.fromkeys((*leftover, *(tag for tag in typed_inherited_tags if isinstance(tag, str))))) async def get_deployments_for_tag( llm_router_instance: LitellmRouter, model: str, # used to raise the correct error - healthy_deployments: list[Any] | dict[Any, Any], - request_kwargs: dict[Any, Any] | None = None, + healthy_deployments: _DeploymentPool, + request_kwargs: _RequestKwargsLike | None = None, metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", -): +) -> _DeploymentPool: """ Returns a list of deployments that match the requested model and tags in the request. @@ -486,8 +512,7 @@ async def get_deployments_for_tag( verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name)) if metadata_variable_name in request_kwargs: - metadata: Final[_TagRoutingMetadata] = request_kwargs[metadata_variable_name] - stampable_metadata: Final[dict[str, object]] = request_kwargs[metadata_variable_name] + metadata: Final = request_kwargs[metadata_variable_name] request_tags: Final = _request_tags_after_router_consumption(metadata, model) match_any: Final = llm_router_instance.tag_filtering_match_any routing_prefix: Final = llm_router_instance.tag_routing_prefix or "" @@ -532,25 +557,25 @@ async def get_deployments_for_tag( request_tags, ) - new_healthy_deployments: Final[list[_TagRoutingDeployment]] = [] - default_deployments: Final[list[_TagRoutingDeployment]] = [] - if has_positive_filter: verbose_logger.debug( "get_deployments_for_tag routing: request_tags=%s user_agent=%s", request_tags, user_agent, ) - for deployment in candidates: - deployment_tags = deployment.get("litellm_params", {}).get("tags") - - match_result = _match_deployment( - deployment=deployment, - request_tags=positive_tags, - header_strings=header_strings, - match_any=match_any, + deployment_matches: Final = tuple( + ( + deployment, + _match_deployment( + deployment=deployment, + request_tags=positive_tags, + header_strings=header_strings, + match_any=match_any, + ), ) - + for deployment in candidates + ) + for deployment, match_result in deployment_matches: if match_result is not None: verbose_logger.debug( "tag routing match: deployment=%s matched_via=%s matched_value=%s", @@ -559,17 +584,17 @@ async def get_deployments_for_tag( match_result["matched_value"], ) if "tag_routing" not in metadata: - stampable_metadata["tag_routing"] = { + metadata["tag_routing"] = { "matched_deployment": deployment.get("model_name"), "matched_via": match_result["matched_via"], "matched_value": match_result["matched_value"], "request_tags": request_tags or [], "user_agent": user_agent, } - new_healthy_deployments.append(deployment) - - if deployment_tags and "default" in deployment_tags: - default_deployments.append(deployment) + new_healthy_deployments: Final = [d for d, result in deployment_matches if result is not None] + default_deployments: Final = [ + d for d, _ in deployment_matches if "default" in (d.get("litellm_params", {}).get("tags") or ()) + ] if len(new_healthy_deployments) == 0 and len(default_deployments) == 0: return _resolve_or_fail_open( @@ -604,10 +629,11 @@ async def get_deployments_for_tag( return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments # for Untagged requests use default deployments if set - _default_deployments_with_tags: Final[list[_TagRoutingDeployment]] = [] - for deployment in healthy_deployments: - if "default" in deployment.get("litellm_params", {}).get("tags", []): - _default_deployments_with_tags.append(deployment) + _default_deployments_with_tags: Final = [ + deployment + for deployment in healthy_deployments + if "default" in deployment.get("litellm_params", {}).get("tags", []) + ] if len(_default_deployments_with_tags) > 0: return _default_deployments_with_tags diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 9be78757511..4cf4fa62eff 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1,5 +1,7 @@ +from collections.abc import Mapping from datetime import datetime from enum import Enum +from types import MappingProxyType from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -550,6 +552,40 @@ class BedrockGuardrailConfigModel(BaseModel): ) +class BedrockGuardrailStreamingParams(BaseModel): + streaming_buffer_until_moderated: bool = Field( + default=True, + description="If True (default), withhold every streamed chunk until the end-of-stream " + "ApplyGuardrail scan passes, so no flagged content reaches the client before a block. " + "If False, chunks stream through unbuffered, so flagged content can reach the client " + "before the scan finishes; a flagged scan still ends the stream, with a block message " + "when disable_exception_on_block is true and an in-stream error frame otherwise.", + ) + streaming_sampling_rate: int = Field( + default=5, + ge=1, + description="When not buffering and not end-of-stream-only, scan the accumulated response " + "every Nth streamed chunk. Each sampled scan is a full ApplyGuardrail call that delays " + "that chunk, so lower values add latency and AWS text-unit cost.", + ) + streaming_end_of_stream_only: bool = Field( + default=False, + description="When not buffering, skip per-chunk sampling and run one ApplyGuardrail scan " + "on the assembled response at end of stream. Combined with " + "streaming_buffer_until_moderated=false the full response streams live before the scan " + "and the scan result lands in guardrail_information; a flagged response still ends the " + "stream with a block message (disable_exception_on_block=true) or an error frame.", + ) + + @classmethod + def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams": + if not extras: + return cls() + return cls.model_validate( + MappingProxyType({name: extras[name] for name in cls.model_fields if extras.get(name) is not None}) + ) + + class LakeraV2GuardrailConfigModel(BaseModel): """Configuration parameters for the Lakera AI v2 guardrail""" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index be3f13ea7ef..caa8eef252a 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2994 + "limit": 2993 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 809 }, "ANN201": { - "limit": 2003 + "limit": 2002 }, "ANN202": { "limit": 843 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 492 + "limit": 487 }, "ASYNC230": { "limit": 11 @@ -123,7 +123,7 @@ "limit": 12 }, "PERF403": { - "limit": 34 + "limit": 33 }, "PIE804": { "limit": 18 @@ -177,7 +177,7 @@ "limit": 8 }, "RUF019": { - "limit": 32 + "limit": 31 }, "RUF046": { "limit": 4 @@ -198,7 +198,7 @@ "limit": 56 }, "SIM102": { - "limit": 313 + "limit": 312 }, "SIM103": { "limit": 119 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1099 + "limit": 1096 }, "TRY002": { "limit": 524 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 112bc5e6e49..2b43720a126 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -482,23 +482,25 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException when processing streaming harmful content - from fastapi import HTTPException + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - async def _drain(): - result_chunks = [] - async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - result_chunks.append(chunk) + result_chunks = [] + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + result_chunks.append(chunk) - with pytest.raises(HTTPException) as exc_info: - await _drain() - - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + frame = result_chunks[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 914af0e2368..476d443d8d8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -161,19 +161,27 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException - with pytest.raises(HTTPException) as exc_info: - async for ( - _ - ) in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - pass + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + collected = [] + async for ( + chunk + ) in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + collected.append(chunk) + + frame = collected[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio 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 36b356e34d0..235a5c0c09b 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 @@ -5345,3 +5345,248 @@ def test_initialize_bedrock_forwards_aws_external_id(): assert guardrail.optional_params["aws_external_id"] == "external-id-123" finally: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, guardrail) + + +def _chat_chunk(content: str, finish_reason: str | None) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content=content, role="assistant"), + finish_reason=finish_reason, + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + +def _streaming_litellm_params(**extras): + from litellm.types.guardrails import LitellmParams + + return LitellmParams( + guardrail="bedrock", + mode="post_call", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + **extras, + ) + + +def test_initialize_bedrock_wires_streaming_flags(): + from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock + + configured = initialize_bedrock( + _streaming_litellm_params( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=3, + streaming_end_of_stream_only=True, + ), + {"guardrail_name": "bedrock-streaming"}, + ) + defaulted = initialize_bedrock( + _streaming_litellm_params(), + {"guardrail_name": "bedrock-defaults"}, + ) + for registered in (configured, defaulted): + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, registered) + + assert configured.streaming_buffer_until_moderated is False + assert configured.streaming_sampling_rate == 3 + assert configured.streaming_end_of_stream_only is True + assert defaulted.streaming_buffer_until_moderated is True + assert defaulted.streaming_sampling_rate == 5 + assert defaulted.streaming_end_of_stream_only is False + + +def test_initialize_bedrock_rejects_non_positive_sampling_rate(): + from pydantic import ValidationError + + from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock + + with pytest.raises(ValidationError): + initialize_bedrock( + _streaming_litellm_params(streaming_sampling_rate=0), + {"guardrail_name": "bedrock-bad-rate"}, + ) + + +def test_update_in_memory_litellm_params_round_trips_streaming_flags(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-update", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + ) + + guardrail.update_in_memory_litellm_params( + _streaming_litellm_params( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=7, + streaming_end_of_stream_only=True, + ) + ) + assert guardrail.streaming_buffer_until_moderated is False + assert guardrail.streaming_sampling_rate == 7 + assert guardrail.streaming_end_of_stream_only is True + + guardrail.update_in_memory_litellm_params(_streaming_litellm_params()) + assert guardrail.streaming_buffer_until_moderated is True + assert guardrail.streaming_sampling_rate == 5 + assert guardrail.streaming_end_of_stream_only is False + + +async def _run_streaming_hook_recording_order(guardrail: BedrockGuardrail) -> list: + events = [] + minimal = {"action": "NONE", "assessments": [], "outputs": []} + + async def record_scan(*args, **kwargs): + events.append("scan") + return minimal + + async def mock_stream(): + yield _chat_chunk("Hello", None) + yield _chat_chunk(" world", None) + yield _chat_chunk("", "stop") + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + ): + content = chunk.choices[0].delta.content if chunk.choices else None + events.append(("chunk", content)) + return events + + +@pytest.mark.asyncio +async def test_unbuffered_end_of_stream_hook_yields_chunks_before_scan(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-audit-mode", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + streaming_buffer_until_moderated=False, + streaming_end_of_stream_only=True, + ) + + events = await _run_streaming_hook_recording_order(guardrail) + + scan_index = events.index("scan") + chunk_events = [e for e in events if e != "scan"] + assert events.count("scan") == 1 + assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[: scan_index] + assert ("chunk", "Hello") in events[:scan_index] + assert ("chunk", " world") in events[:scan_index] + assert len(chunk_events) == 3 + + +@pytest.mark.asyncio +async def test_buffered_default_hook_scans_before_any_chunk(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-buffered-default", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + + events = await _run_streaming_hook_recording_order(guardrail) + + assert events[0] == "scan" + assert all(e == "scan" or e[0] == "chunk" for e in events) + assert len([e for e in events if e != "scan"]) >= 1 + + +@pytest.mark.asyncio +async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-mask-buffered", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + mask_response_content=True, + streaming_buffer_until_moderated=False, + streaming_end_of_stream_only=True, + ) + + assert guardrail._streams_incrementally() is False + events = await _run_streaming_hook_recording_order(guardrail) + assert events[0] == "scan" + + +@pytest.mark.asyncio +async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_truncating(): + """Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream + scan used to raise after SSE headers were flushed, so the client saw a + silently truncated stream. The unified hook must emit the chat in-stream + error frame instead.""" + from litellm.llms import load_guardrail_translation_mappings + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, + ) + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + streaming_end_of_stream_only=True, + streaming_buffer_until_moderated=False, + guardrail_name="bedrock-eos", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "actionReason": "Guardrail blocked.", + "outputs": [{"text": "Sorry, the model cannot answer this question."}], + "assessments": [ + {"topicPolicy": {"topics": [{"name": "Forbidden topic", "type": "DENY", "action": "BLOCKED"}]}} + ], + } + + def _chunk(content, finish_reason=None): + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta={"content": content, "role": "assistant"}, + finish_reason=finish_reason, + ) + ], + ) + + async def _mock_stream(): + yield _chunk("the forbidden ") + yield _chunk("topic answer", finish_reason="stop") + + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + try: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response) + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"), + response=_mock_stream(), + request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"}, + ): + out.append(item) + finally: + unified_module.endpoint_guardrail_translation_mappings = None + + assert len(out) == 3 + assert isinstance(out[0], ModelResponseStream) + frame = out[-1] + assert isinstance(frame, bytes) + payload = json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8b9ecfbbeee..0b32558a00a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -948,19 +948,24 @@ class TestStreamingTransform: assert streamed == "ABCDEFGHIJ" @pytest.mark.asyncio - async def test_incremental_diff_underflow_raises(self): + async def test_incremental_diff_underflow_emits_error_frame(self): """A transform shorter than what was already streamed cannot retract - bytes: it raises HTTPException(stream_transform_underflow).""" + bytes. Chunks have already been flushed by then, so the underflow + surfaces as the in-stream error frame, not an unraisable HTTPException.""" + import json as _json + # First sample emits "ABCDEF" (6 chars); second sample shrinks to 3. guardrail = _StreamingTextGuardrail(shrink_to="ABC", shrink_after=1) chunks = [_stream_chunk("abcdef"), _stream_chunk("ghij")] - with pytest.raises(unified_module.HTTPException) as exc_info: - await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) - assert exc_info.value.status_code == 400 - assert exc_info.value.detail["error"] == "stream_transform_underflow" + frame = out[-1] + assert isinstance(frame, bytes) + payload = _json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "stream_transform_underflow" + assert payload["error"]["code"] == "400" @pytest.mark.asyncio async def test_incremental_diff_final_chunk_preserves_finish_reason(self): @@ -1747,3 +1752,221 @@ class TestAppliedGuardrailsReflectsExecution: async def test_ordinary_guardrail_is_auto_marked_applied(self): data = await self._run(_AutoLoggingGuardrail()) assert "auto-logging" in _applied_guardrails(data) + + +class _EosHttpBlockingGuardrail(CustomGuardrail): + """Raises the bedrock-shaped block HTTPException at end-of-stream scan time.""" + + def __init__(self): + super().__init__(guardrail_name="eos-http-block") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + raise unified_module.HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "bedrock_guardrail_response": "BLOCKED_TOPIC", + }, + ) + + +def _anthropic_sse_event(event_type, data): + import json as _json + + return f"event: {event_type}\ndata: {_json.dumps(data)}\n\n".encode() + + +def _anthropic_message_chunks(texts): + head = [ + _anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + _anthropic_sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ] + deltas = [ + _anthropic_sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + for text in texts + ] + tail = [ + _anthropic_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + _anthropic_sse_event("message_stop", {"type": "message_stop"}), + ] + return head + deltas + tail + + +class TestStreamingHttpErrorFrames: + """A post-flush end-of-stream guardrail block (HTTPException) must surface as + the endpoint's in-stream error frame instead of an unhandled raise that + silently truncates the SSE stream (PR #38722 defect 1).""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_block_emits_data_error_frame(self): + import json as _json + + guardrail = _EosHttpBlockingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out[:2] == chunks + frame = out[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + + @pytest.mark.asyncio + async def test_messages_eos_block_emits_anthropic_error_event(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages" + ) + + raw = b"".join(c for c in out if isinstance(c, bytes)).decode() + assert "hello " in raw + assert "event: error" in raw + assert "Violated guardrail policy" in raw + assert "guardrail_error" in raw + + @pytest.mark.asyncio + async def test_responses_eos_block_emits_error_event_with_next_sequence(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = [ + {"type": "response.created", "sequence_number": 0}, + {"type": "response.output_text.delta", "sequence_number": 1, "delta": "hello"}, + { + "type": "response.completed", + "sequence_number": 2, + "response": { + "model": "gpt-4", + "output": [{"type": "message", "content": [{"type": "output_text", "text": "hello"}]}], + }, + }, + ] + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses" + ) + + assert chunks[0] in out and chunks[1] in out + assert chunks[2] not in out + error_event = out[-1] + assert error_event.type == "error" + assert error_event.sequence_number == 2 + assert error_event.error.message == "Violated guardrail policy" + assert error_event.error.code == "400" + assert error_event.error.type == "guardrail_error" + + @pytest.mark.asyncio + async def test_pre_flush_block_still_raises_http_exception(self): + guardrail = _EosHttpBlockingGuardrail() + guardrail.streaming_buffer_until_moderated = True + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + with pytest.raises(unified_module.HTTPException) as exc_info: + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Violated guardrail policy" + + +class _AuditRecordingGuardrail(CustomGuardrail): + """Successful scan that records guardrail_information, like a flags-on audit.""" + + def __init__(self): + super().__init__(guardrail_name="audit-recorder") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"action": "NONE"}, + request_data=request_data, + guardrail_status="success", + ) + return inputs + + +class TestStreamingGuardrailInformationBucket: + """guardrail_information written during a chat streaming end-of-stream scan + must land in the request's ``metadata`` bucket that spend logging snapshots. + Regression for PR #38722 defect 2: the chat handler used to plant a + ``litellm_metadata`` key first, flipping the bucket so every later + guardrail_information write was diverted and /spend/logs showed null.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_scan_writes_guardrail_information_to_metadata(self): + guardrail = _AuditRecordingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + async def _mock_stream(): + for chunk in chunks: + yield chunk + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", user_id="user-1", request_route="/v1/chat/completions" + ) + request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}} + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_mock_stream(), + request_data=request_data, + ): + out.append(item) + + assert "litellm_metadata" not in request_data + recorded = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(recorded) == 1 + assert recorded[0]["guardrail_name"] == "audit-recorder" + assert recorded[0]["guardrail_status"] == "success" + assert request_data["metadata"]["user_api_key_user_id"] == "user-1" diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index e70fc61de30..8fde4cc9d5e 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1228,3 +1228,229 @@ class TestFireDeferredStreamLogging: assert info is not None, "guardrail_information should be populated" assert len(info) == 1 assert info[0]["guardrail_name"] == "info-writer" + + +class TestResponsesIteratorDeferredLogging: + """Regression for PR #38722 defect 2 on /v1/responses streams: when the + proxy arms _on_deferred_stream_complete, the responses streaming iterator + must store the logging coroutine for ProxyLogging._fire_deferred_stream_logging + (which runs AFTER end-of-stream guardrail scans write guardrail_information) + instead of dispatching immediately with a premature metadata snapshot.""" + + def _iterator(self, logging_obj): + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + iterator = object.__new__(BaseResponsesAPIStreamingIterator) + iterator.logging_obj = logging_obj + iterator.start_time = None + iterator.completed_response = None + iterator._completed_response_logged = False + iterator._completed_response_cache_hit = None + iterator._persist_completed_response_before_logging = False + return iterator + + def _logging_obj(self): + recorded = {} + + async def dispatch_success_handlers(result=None, **kwargs): + recorded["dispatched"] = True + + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = dispatch_success_handlers + return logging_obj, recorded + + @pytest.mark.asyncio + async def test_armed_iterator_stores_deferred_coroutine(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = MagicMock() + iterator = self._iterator(logging_obj) + + with patch("asyncio.create_task") as mock_create_task: + iterator._log_completed_response(is_async=True) + + mock_create_task.assert_not_called() + args = logging_obj._deferred_stream_complete_args + assert isinstance(args, tuple) and len(args) == 1 + assert "dispatched" not in recorded + await args[0] + assert recorded["dispatched"] is True + + @pytest.mark.asyncio + async def test_unarmed_iterator_dispatches_immediately(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = None + iterator = self._iterator(logging_obj) + + created = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + iterator._log_completed_response(is_async=True) + + assert len(created) == 1 + await created[0] + assert recorded["dispatched"] is True + + +class TestArmDeferredStreamDispatch: + """Regression for PR #38722: the closure shape armed on logging_obj must + match the args the stream's logging owner stores. Bridged /v1/responses + (LiteLLMCompletionStreamingIterator) shares its inner CustomStreamWrapper's + logging_obj, which stores (assembled_response, cache_hit); arming the + single-coroutine native closure there made _fire_deferred_stream_logging + raise TypeError inside the streaming hook, leaking an in-stream 500 error + frame on every streamed /v1/responses request.""" + + def _processor(self): + return ProxyBaseLLMRequestProcessing(data={"model": "gpt-test"}) + + def _dispatch_recording_logging_obj(self): + recorded = {} + + async def dispatch_success_handlers( + result=None, start_time=None, end_time=None, cache_hit=None, prefer_async_handlers=False + ): + recorded["result"] = result + recorded["cache_hit"] = cache_hit + recorded["prefer_async_handlers"] = prefer_async_handlers + + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = dispatch_success_handlers + logging_obj._on_deferred_stream_complete = None + logging_obj._deferred_stream_complete_args = None + return logging_obj, recorded + + @pytest.mark.asyncio + async def test_bridged_responses_iterator_gets_csw_arg_shape(self): + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + logging_obj, recorded = self._dispatch_recording_logging_obj() + bridged = object.__new__(LiteLLMCompletionStreamingIterator) + + self._processor()._arm_deferred_stream_dispatch( + response=bridged, + route_type="aresponses", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + async def test_router_wrapped_bridged_iterator_gets_csw_arg_shape(self): + """The router wraps iterators without _hidden_params in + HiddenParamsAsyncIteratorWrapper before the proxy arms deferral, so + every production streamed /v1/responses reaches arming wrapped; + sniffing the wrapper instead of the inner iterator armed the 1-arg + native closure against the CSW's 2-arg stored shape and leaked a + TypeError 500 frame into the stream.""" + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.router_utils.add_retry_fallback_headers import ( + HiddenParamsAsyncIteratorWrapper, + ) + + logging_obj, recorded = self._dispatch_recording_logging_obj() + wrapped = HiddenParamsAsyncIteratorWrapper(object.__new__(LiteLLMCompletionStreamingIterator)) + + self._processor()._arm_deferred_stream_dispatch( + response=wrapped, + route_type="aresponses", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + async def test_native_stream_closure_enqueues_single_coroutine(self): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, _ = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + closure = logging_obj._on_deferred_stream_complete + assert closure is not None + + async def _logging_coroutine(): + return None + + coro = _logging_coroutine() + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue: + await closure(coro) + mock_enqueue.assert_called_once_with(async_coroutine=coro) + coro.close() + + @pytest.mark.asyncio + async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + logging_obj, recorded = self._dispatch_recording_logging_obj() + csw = object.__new__(CustomStreamWrapper) + processor = self._processor() + + monkeypatch.setattr( # test-quality-ok: empty the process-global callback registry so no ambient guardrail runs + litellm, "callbacks", [] + ) + processor._arm_deferred_stream_dispatch( + response=csw, + route_type="acompletion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + assembled = object() + await logging_obj._on_deferred_stream_complete(assembled, False) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + + def test_non_native_route_generator_not_armed(self): + logging_obj, _ = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="acompletion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assert logging_obj._on_deferred_stream_complete is None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 1d82a5dfc6e..56c89fed79a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -28,12 +28,21 @@ def _make_streaming_response(chunks): return mock +def _unarmed_logging_obj(): + """Real Logging objects only carry _on_deferred_stream_complete when the + proxy arms deferred dispatch; a bare MagicMock's auto-attribute is truthy + and would spuriously trigger the deferral branch.""" + obj = MagicMock() + obj._on_deferred_stream_complete = None + return obj + + @pytest.mark.asyncio async def test_chunk_processor_logs_on_normal_completion(): chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] response = _make_streaming_response(chunks) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -66,7 +75,7 @@ async def test_chunk_processor_logs_on_client_disconnect(): chunks = [b"event-1", b"event-2", b"event-3"] response = _make_streaming_response(chunks) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -104,7 +113,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er response = _make_streaming_response(chunks) response.status_code = 403 - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -134,7 +143,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er async def test_chunk_processor_does_not_schedule_logging_when_no_chunks(): response = _make_streaming_response([]) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -189,7 +198,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker(): async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "claude-3-haiku"}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=EndpointType.GENERIC, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -230,7 +239,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne gen = PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "claude-3-haiku"}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=EndpointType.GENERIC, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -246,7 +255,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne def _logging_obj_with_write_once_cst(): """Build a MagicMock that mirrors the real Logging behavior: _update_completion_start_time latches self.completion_start_time so the write-once guard actually latches.""" - obj = MagicMock() + obj = _unarmed_logging_obj() obj.completion_start_time = None def _update(*, completion_start_time): @@ -301,7 +310,7 @@ async def test_chunk_processor_does_not_reset_completion_start_time_on_later_chu response = _make_streaming_response(chunks) real_first = datetime(2020, 1, 1, 0, 0, 0) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() # Simulate first-chunk stamp having already landed (e.g. under contention or a # prior wrapper that already set it): later chunks must be no-ops. mock_logging_obj.completion_start_time = real_first @@ -387,7 +396,7 @@ async def _collect_openai_passthrough_chunks(chunks, endpoint_type): async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "gpt-4o-mini", "stream": True}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -517,3 +526,109 @@ def test_convert_raw_bytes_survives_truncated_multibyte_sequence(): lines = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) assert any('"type": "message_delta"' in line for line in lines) + + +@pytest.mark.asyncio +async def test_chunk_processor_defers_logging_until_fire_when_armed(): + """Regression for PR #38722: native /v1/messages streams route through + chunk_processor, which enqueued the spend log the moment the stream ended, + racing the guardrail end-of-stream scan and logging + guardrail_information as null. With deferred dispatch armed, the completed + stream must park the logging coroutine on logging_obj and only enqueue it + when ProxyLogging._fire_deferred_stream_logging fires after the scan.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + chunks = [b"event-1", b"event-2"] + response = _make_streaming_response(chunks) + + logging_obj = _unarmed_logging_obj() + logging_obj._deferred_stream_complete_args = None + + enqueued = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + async_coroutine.close() + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, + "ensure_initialized_and_enqueue", + side_effect=_capture, + ) as mock_enqueue: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=gen, + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + received = [] + async for chunk in gen: + received.append(chunk) + await asyncio.sleep(0) + + assert received == chunks + mock_enqueue.assert_not_called() + parked = logging_obj._deferred_stream_complete_args + assert isinstance(parked, tuple) and len(parked) == 1 + assert asyncio.iscoroutine(parked[0]) + + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_called_once() + + +@pytest.mark.asyncio +async def test_chunk_processor_enqueues_immediately_on_disconnect_even_when_armed(): + """Client disconnects never reach _fire_deferred_stream_logging, so parking + the coroutine there would lose the partial-usage spend log (LIT-2642); the + disconnect path must keep enqueueing immediately.""" + chunks = [b"event-1", b"event-2", b"event-3"] + response = _make_streaming_response(chunks) + + logging_obj = _unarmed_logging_obj() + + async def _armed_closure(logging_coroutine): + raise AssertionError("deferred closure must not fire on disconnect") + + logging_obj._on_deferred_stream_complete = _armed_closure + logging_obj._deferred_stream_complete_args = None + + enqueued = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + async_coroutine.close() + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, + "ensure_initialized_and_enqueue", + side_effect=_capture, + ) as mock_enqueue: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ) + await gen.__anext__() + await gen.aclose() + + mock_enqueue.assert_called_once() + assert logging_obj._deferred_stream_complete_args is None diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 833c3754022..629ae6084b8 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1659,29 +1659,29 @@ class TestCommonRequestProcessingHelpers: async def test_serialize_http_exception_detail_helper(self): """Direct unit coverage for the L1 helper across all branches.""" from litellm.proxy.common_request_processing import ( - _serialize_http_exception_detail, + serialize_http_exception_detail, ) import json as _json - assert _serialize_http_exception_detail("plain") == ("plain", None) + assert serialize_http_exception_detail("plain") == ("plain", None) - msg, fields = _serialize_http_exception_detail({"error": "Violated", "extra": "x"}) + msg, fields = serialize_http_exception_detail({"error": "Violated", "extra": "x"}) assert msg == "Violated" assert fields == {"error": "Violated", "extra": "x"} - msg, fields = _serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) + msg, fields = serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) assert msg == "blocked" assert fields == {"error": {"message": "blocked", "code": "x"}} - msg, fields = _serialize_http_exception_detail({"message": "top-level"}) + msg, fields = serialize_http_exception_detail({"message": "top-level"}) assert msg == "top-level" assert fields == {"message": "top-level"} - msg, fields = _serialize_http_exception_detail({"weird": ["a", "b"]}) + msg, fields = serialize_http_exception_detail({"weird": ["a", "b"]}) assert msg == _json.dumps({"weird": ["a", "b"]}) assert fields == {"weird": ["a", "b"]} - assert _serialize_http_exception_detail(42) == ("42", None) + assert serialize_http_exception_detail(42) == ("42", None) async def test_proxy_exception_from_http_exception_helper(self): """The shared HTTPException -> ProxyException conversion keeps a clean 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 542572e1e56..9f1321aec2c 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -346,14 +346,14 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions @pytest.mark.asyncio -async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch): +async def test_unified_guardrail_iterator_accepts_explicit_guardrail(): """ The dispatch passes each guardrail explicitly instead of through a shared request_data key, so chaining two unified-routed guardrails cannot drop - all but the last one. + all but the last one. The block fires after the deltas were already + flushed to the client, so it surfaces as a trailing in-stream error frame + rather than a raised HTTPException. """ - from fastapi import HTTPException - from litellm.proxy.utils import unified_guardrail guardrail = _content_filter_guardrail("BLOCK") @@ -367,14 +367,19 @@ async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch for chunk in _anthropic_stream_chunks(["the", " zebra runs"]): yield chunk - with pytest.raises(HTTPException): - async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), - response=fake_stream(), - request_data=request_data, - guardrail_to_apply=guardrail, - ): - pass + delivered = [] + async for item in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + response=fake_stream(), + request_data=request_data, + guardrail_to_apply=guardrail, + ): + delivered.append(item) + + raw = b"".join(c for c in delivered if isinstance(c, bytes)).decode() + assert "event: error" in raw + assert "guardrail_error" in raw + assert raw.index("guardrail_error") > raw.index(" zebra runs") @pytest.mark.asyncio diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 399d0f6abad..57ee0606dca 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22590 + "limit": 22462 }, "LIT002": { - "limit": 26812 + "limit": 26800 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1041 + "limit": 1039 }, "LIT007": { "limit": 0 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16532 + "limit": 16529 }, "LIT011": { - "limit": 5557 + "limit": 5556 }, "LIT012": { - "limit": 4496 + "limit": 4495 } }