diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index f8ef0ce1732..d60c3e9c0af 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 16389 + "limit": 16171 }, "reportArgumentType": { - "limit": 2229 + "limit": 2226 }, "reportAssignmentType": { "limit": 319 @@ -18,13 +18,13 @@ "limit": 40 }, "reportDeprecated": { - "limit": 212 + "limit": 211 }, "reportDuplicateImport": { "limit": 19 }, "reportExplicitAny": { - "limit": 5242 + "limit": 5199 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5614 + "limit": 5611 }, "reportMissingTypeArgument": { - "limit": 15356 + "limit": 15350 }, "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": 38500 + "limit": 38468 }, "reportUnknownParameterType": { - "limit": 19673 + "limit": 19665 }, "reportUnknownVariableType": { - "limit": 30092 + "limit": 30066 }, "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 7c86ceafd7f..b73a1443330 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, @@ -23,35 +23,63 @@ from litellm.types.llms.openai import ( 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): @@ -62,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): @@ -74,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) @@ -124,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, @@ -132,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", @@ -149,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"}, @@ -211,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]: @@ -250,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, } @@ -312,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: @@ -326,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 @@ -346,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", ""), } @@ -360,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 @@ -391,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 @@ -500,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 @@ -523,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"}, @@ -563,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 @@ -590,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) @@ -599,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"}, @@ -635,10 +669,10 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, - message: Any, - ) -> list[_GenAIPart]: + message: Message, + ) -> 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: @@ -646,20 +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 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, } } @@ -668,21 +704,23 @@ class GoogleGenAIAdapter: return parts if parts else [{"text": ""}] def _transform_openai_delta_to_google_genai_parts_with_accumulation( - self, delta: Any, wrapper: GoogleGenAIStreamWrapper - ) -> list[_GenAIPart]: + self, delta: Delta, wrapper: GoogleGenAIStreamWrapper + ) -> 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"): @@ -701,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] @@ -723,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. @@ -757,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 b9ca18c7843..0d123bc1004 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -16,9 +16,9 @@ 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 assert_never +from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -98,6 +98,48 @@ InputWriteBackTarget = ( ) +class _SSEDelta(TypedDict, total=False): + type: ReadOnly[str] + text: ReadOnly[str] + stop_reason: ReadOnly[str | None] + + +class _SSEEventData(TypedDict, total=False): + delta: ReadOnly[_SSEDelta] + + +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 @@ -126,7 +168,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))) @@ -144,7 +186,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 @@ -162,7 +204,7 @@ 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 _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: @@ -187,7 +229,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.""" @@ -199,7 +243,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] = [] @@ -237,7 +281,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). @@ -263,7 +307,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 @@ -271,24 +328,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.""" @@ -321,7 +369,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. """ @@ -481,7 +529,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") @@ -561,7 +609,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 @@ -755,7 +803,7 @@ class AnthropicMessagesHandler(BaseTranslation): if scan_only_tool_results: return EMPTY_EXTRACTED_INPUT - text_str: Final = 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)),) @@ -805,7 +853,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: @@ -931,7 +979,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. @@ -1027,7 +1075,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 [] @@ -1037,7 +1085,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]], @@ -1045,21 +1093,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( @@ -1071,6 +1108,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], @@ -1093,7 +1145,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. @@ -1164,7 +1216,7 @@ class AnthropicMessagesHandler(BaseTranslation): # Only process content_block_delta events if event_type == "content_block_delta" and data_line: try: - data = json.loads(data_line) + data: _SSEEventData = json.loads(data_line) delta = data.get("delta", {}) if delta.get("type") == "text_delta": text += delta.get("text", "") @@ -1176,7 +1228,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. @@ -1227,7 +1279,7 @@ class AnthropicMessagesHandler(BaseTranslation): # Check for message_delta event with stop_reason if event_type == "message_delta" and data_line: try: - data = json.loads(data_line) + data: _SSEEventData = json.loads(data_line) delta = data.get("delta", {}) stop_reason = delta.get("stop_reason") if stop_reason is not None: @@ -1271,7 +1323,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], @@ -1294,7 +1346,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, @@ -1319,7 +1371,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"): @@ -1335,14 +1387,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 c8cbbba8784..1cf52045a48 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,10 +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, Any, Final, Literal, NotRequired, Optional, TypedDict, Union, cast +from collections.abc import Awaitable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, 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 @@ -29,6 +29,7 @@ from litellm.types.llms.anthropic import ( 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, @@ -84,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.""" @@ -159,11 +231,11 @@ async def _check_summary_model_access( return True key_models: Final = list(getattr(user_api_key_auth, "models", None) or []) - team_id: Final = getattr(user_api_key_auth, "team_id", None) + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None) team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or []) - user_id: Final = getattr(user_api_key_auth, "user_id", None) - project_id: Final = getattr(user_api_key_auth, "project_id", None) + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) + project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None) checks: Final[tuple[tuple[Literal["key", "team"], list[str]], ...]] = ( ("key", key_models), @@ -372,7 +444,7 @@ async def _check_summary_model_budget( return False end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None) - end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None) + end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None) if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: try: await model_max_budget_limiter.is_end_user_within_model_budget( @@ -424,40 +496,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: @@ -471,7 +560,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. @@ -490,8 +579,8 @@ def _find_latest_compaction_index( def _slice_around_compaction_block( - messages: list[dict[str, Any]], -) -> 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)`` @@ -506,19 +595,21 @@ def _slice_around_compaction_block( original_msg: Final = messages[msg_idx] original_content: Final = original_msg["content"] - compaction_block: Final = cast(dict[str, object], original_content[blk_idx]) + 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. @@ -625,7 +716,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, @@ -704,17 +795,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. @@ -729,16 +821,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 [ @@ -761,7 +855,7 @@ def _extract_summary_text(raw: str | None) -> str | None: def _system_to_openai_message( system: str | list[dict[str, Any]] | None, -) -> dict[str, object] | None: +) -> Mapping[str, object] | None: """Translate Anthropic-shaped ``system`` to an OpenAI system message. Accepts a bare string or a list of Anthropic content blocks; returns @@ -772,17 +866,19 @@ def _system_to_openai_message( if isinstance(system, str): return {"role": "system", "content": system} if system else None if isinstance(system, list): - parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"] + parts: Final[tuple[str, ...]] = tuple( + block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text" + ) joined: Final = "\n\n".join(part for part in parts if part) return {"role": "system", "content": joined} if joined else None return None 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 @@ -810,7 +906,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) @@ -845,35 +941,17 @@ 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[object] - - -class _SummaryCallRegionKwarg(TypedDict, total=False): - allowed_model_region: ReadOnly[str] - - -class _SummaryCallKwargs(TypedDict): - model: ReadOnly[str] - messages: ReadOnly[list[dict[str, object]]] - max_tokens: ReadOnly[int] - timeout: ReadOnly[float] - litellm_metadata: ReadOnly[Mapping[str, object]] - user: NotRequired[ReadOnly[object]] - 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: Any, + llm_router: object, allowed_model_region: str | None = None, max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS, ) -> Union["ModelResponse", "CustomStreamWrapper"]: @@ -909,28 +987,37 @@ 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": summary_messages, "max_tokens": max_tokens, "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, "litellm_metadata": metadata, - **(_SummaryCallUserKwarg(user=end_user_id) if 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) + 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: Any) -> str | None: +def _extract_response_text(response: object) -> str | None: try: - choice: Final = response.choices[0] - message: Final = choice.message + 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 @@ -946,7 +1033,7 @@ def _extract_response_text(response: Any) -> str | None: def _extract_usage(response: object) -> tuple[int, int]: - usage: Final = getattr(response, "usage", None) + usage: Final[object] = getattr(response, "usage", None) if usage is None: return 0, 0 return ( 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/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/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 0d12b012c18..52a192f8537 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(): @@ -2805,13 +2856,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: @@ -2842,7 +2893,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) @@ -3004,14 +3057,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, @@ -3057,7 +3112,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( @@ -3135,7 +3190,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) @@ -3166,14 +3221,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, @@ -3302,7 +3359,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], ) @@ -4301,7 +4358,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) @@ -4437,13 +4494,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( @@ -4458,14 +4514,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, @@ -4531,7 +4587,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/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 3c5d889d23f..f755c6d478b 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,12 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Final, TypedDict +from collections.abc import Callable, Mapping, Sequence +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,53 +29,59 @@ 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 _StreamTerminalResponse(TypedDict, total=False): - """Fields of the ``response`` payload carried by a terminal streaming event.""" - +class _TerminalResponse(TypedDict, total=False): status: ReadOnly[ResponsesAPIStatus] + error: ReadOnly[_JsonDict] + usage: ReadOnly[_JsonDict] + reasoning: ReadOnly[_JsonDict] tool_choice: ReadOnly[object] + 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] - output: ReadOnly[Sequence[_StreamOutputItem]] + incomplete_details: ReadOnly[_JsonDict] + output: ReadOnly[Sequence[_OutputItem]] class _StreamEvent(TypedDict, total=False): - """One decoded ``data:`` frame of an OpenAI Responses streaming body.""" - 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[_StreamTerminalResponse] + part: ReadOnly[object] + response: ReadOnly[_TerminalResponse] + + +class _StreamEventParser: + parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) async def background_streaming_task( polling_id: str, - data, + data: dict, polling_handler: ResponsePollingHandler, request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, - general_settings, + general_settings: dict, llm_router: "Router | None", proxy_config: "ProxyConfig", proxy_logging_obj: "ProxyLogging", @@ -138,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 @@ -210,7 +217,7 @@ async def background_streaming_task( break try: - event: _StreamEvent = json.loads(chunk_data) + event: _StreamEvent = _StreamEventParser.parse(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -229,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: @@ -252,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": @@ -280,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 63adf950142..61f624a5d49 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): @@ -94,6 +81,60 @@ 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 _PopsOptionalStr(Protocol): + def __call__(self, key: str, default: None, /) -> str | None: ... + + +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 + + +def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr: + 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) @@ -220,7 +261,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), @@ -549,7 +590,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 @@ -567,8 +608,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 @@ -624,12 +668,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, @@ -1083,7 +1130,7 @@ class _HasModelDumpJson(Protocol): def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... -def _dump_response_object(obj: object) -> dict[str, Any]: +def _dump_response_object(obj: object) -> Mapping[str, object]: if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): @@ -1113,21 +1160,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( @@ -1157,7 +1203,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() @@ -1174,16 +1220,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( @@ -1256,7 +1305,8 @@ def _build_synthetic_response_events( ) if item_type == "message": - content_parts: Sequence[object] = output_item_payload.get("content", []) or [] + raw_content_parts = output_item_payload.get("content") + content_parts: Sequence[object] = raw_content_parts if _is_json_array(raw_content_parts) else [] for content_index, part in enumerate(content_parts): part_payload = _dump_response_object(part) events.append( @@ -1304,8 +1354,9 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - summaries: Sequence[object] = output_item_payload.get("summary", []) or [] - for summary_index, summary in enumerate(summaries): + raw_summary_items = output_item_payload.get("summary") + summary_items: Sequence[object] = raw_summary_items if _is_json_array(raw_summary_items) else [] + for summary_index, summary in enumerate(summary_items): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1476,7 +1527,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, @@ -1486,17 +1537,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: @@ -1610,7 +1661,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. @@ -1679,7 +1730,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 @@ -1786,6 +1837,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": @@ -1805,9 +1857,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 @@ -1816,9 +1866,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) @@ -2020,7 +2068,7 @@ class ManagedResponsesWebSocketHandler: model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, @@ -2033,10 +2081,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, Any] = litellm_metadata or {} - self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( + 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 @@ -2057,7 +2106,7 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ @staticmethod - def _serialize_chunk(chunk: Any) -> str | None: + def _serialize_chunk(chunk: object) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: if isinstance(chunk, _HasModelDumpJson): @@ -2100,7 +2149,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. @@ -2115,7 +2164,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 @@ -2172,7 +2221,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) @@ -2185,7 +2234,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 @@ -2201,13 +2250,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 @@ -2225,7 +2274,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. @@ -2248,7 +2297,7 @@ class ManagedResponsesWebSocketHandler: await self.websocket.send_text(serialized) @staticmethod - def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]: + 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": [...], ...}} @@ -2357,7 +2406,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. @@ -2365,7 +2414,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) @@ -2391,7 +2440,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: @@ -2464,12 +2513,14 @@ class ManagedResponsesWebSocketHandler: # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. - requested_model: Final[str | None] = call_kwargs.pop("model", None) + requested_model: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)("model", None) model: Final[str] = ( self.model if requested_model is None or requested_model == self.model_group else requested_model ) - previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None) + previous_response_id: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)( + "previous_response_id", None + ) current_messages: Final = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history 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/ruff-strict-budget.json b/ruff-strict-budget.json index fce91d57d88..1e855ea1b71 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2998 + "limit": 2995 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 809 }, "ANN201": { - "limit": 2003 + "limit": 2002 }, "ANN202": { "limit": 845 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 597 + "limit": 587 }, "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": 315 + "limit": 314 }, "SIM103": { "limit": 119 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1111 + "limit": 1108 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index bb02618b11e..ab34775c460 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22642 + "limit": 22521 }, "LIT002": { - "limit": 26834 + "limit": 26820 }, "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": 16550 + "limit": 16546 }, "LIT011": { - "limit": 5576 + "limit": 5575 }, "LIT012": { - "limit": 4496 + "limit": 4495 } }