From cb65bf08b8c937196270f14468916de3a627388b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:41:30 -0700 Subject: [PATCH 01/16] chore(typing): clear 1.2k basedpyright Any errors across 16 hotspot files Replace Any-typed seams with real types in the files carrying the highest remaining reportAny/reportExplicitAny density: Literal-keyed structural Protocols for deployment dicts in tag-based routing, typed Prisma table wrappers and row protocols in the key and internal-user management endpoints, TypedDict views for websearch interception kwargs, typed streaming state in the responses iterator and background polling, and concrete request/response types in the google_genai, vertex_ai files, runwayml, rubrik, anthropic context-management, and guardrail translation modules. Mutable annotations introduced along the way were rewritten as read-only views (Mapping/Sequence/tuple) built functionally. No casts, no type: ignore, no noqa, no suppression comments, no new Any annotations, no behavior changes. Whole-tree basedpyright: reportAny 15,496 -> 14,523, reportExplicitAny 5,356 -> 5,102, all rules 145,547 -> 143,989, with no rule increased repo-wide or per-file. Budgets ratcheted: basedpyright -1,545, ruff-strict -73, type-discipline -237. --- basedpyright-code-budget.json | 22 +- .../google_genai/adapters/transformation.py | 186 ++++++++--- litellm/integrations/rubrik.py | 228 +++++++++---- .../websearch_interception/handler.py | 157 +++++++-- .../chat/guardrail_translation/handler.py | 202 +++++++----- .../context_management/editors/compact.py | 213 ++++++++---- .../llms/runwayml/videos/transformation.py | 138 ++++---- .../llms/vertex_ai/files/transformation.py | 114 +++++-- litellm/proxy/db/tool_registry_writer.py | 197 ++++++++---- .../internal_user_endpoints.py | 302 +++++++++++++----- .../key_management_endpoints.py | 241 ++++++++++---- .../tool_management_endpoints.py | 213 ++++++++++-- litellm/proxy/prompts/prompt_endpoints.py | 161 ++++++---- .../response_polling/background_streaming.py | 117 +++++-- .../responses/file_search/emulated_handler.py | 286 +++++++++-------- litellm/responses/streaming_iterator.py | 227 +++++++++---- litellm/router_strategy/tag_based_routing.py | 174 ++++++---- ruff-strict-budget.json | 16 +- type-discipline-budget.json | 10 +- 19 files changed, 2261 insertions(+), 943 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 6e3cbdff9d0..c71ef7a0020 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 23919 + "limit": 21974 }, "reportArgumentType": { - "limit": 2580 + "limit": 2575 }, "reportAssignmentType": { "limit": 323 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 7573 + "limit": 7068 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5719 + "limit": 5697 }, "reportMissingTypeArgument": { - "limit": 15657 + "limit": 15627 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44832 + "limit": 44549 }, "reportUnknownLambdaType": { - "limit": 113 + "limit": 112 }, "reportUnknownMemberType": { - "limit": 39269 + "limit": 39156 }, "reportUnknownParameterType": { - "limit": 19988 + "limit": 19951 }, "reportUnknownVariableType": { - "limit": 30923 + "limit": 30798 }, "reportUnnecessaryCast": { "limit": 118 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 853 + "limit": 852 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 4f127f476c3..4c8e77d9feb 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,6 +1,9 @@ import json -from collections.abc import AsyncIterator, Iterator -from typing import Any, Final, 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 TypedDict from litellm import verbose_logger from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema @@ -9,7 +12,6 @@ from litellm.types.llms.openai import ( ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, ChatCompletionImageObject, - ChatCompletionRequest, ChatCompletionSystemMessage, ChatCompletionTextObject, ChatCompletionToolCallFunctionChunk, @@ -21,12 +23,79 @@ 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, ) +_JsonDict: TypeAlias = dict[str, object] +_JsonDictList: TypeAlias = list[_JsonDict] + + +class _ToolCallAccumulator(TypedDict): + name: str + arguments: str + + +class _GenAIFunctionCall(TypedDict): + name: str + args: Mapping[str, object] + + +class _GenAIPart(TypedDict, total=False): + text: str + functionCall: _GenAIFunctionCall + + +class _GenAIFunctionResponse(TypedDict, total=False): + name: str + response: object + + +class _GenAIRequestFunctionCall(TypedDict, total=False): + name: str + args: Mapping[str, object] + + +class _GenAIContentPart(TypedDict, total=False): + text: str + inline_data: Mapping[str, str] + functionResponse: _GenAIFunctionResponse + functionCall: _GenAIRequestFunctionCall + + +class _GenAIFunctionDeclaration(TypedDict, total=False): + name: str + description: str + parametersJsonSchema: object + + +class _GenAITool(TypedDict, total=False): + functionDeclarations: Sequence[_GenAIFunctionDeclaration] + + +class _GenAIFunctionCallingConfig(TypedDict, total=False): + mode: str + + +class _GenAIToolConfig(TypedDict, total=False): + functionCallingConfig: _GenAIFunctionCallingConfig + + +class _GenAISystemInstruction(TypedDict, total=False): + parts: Sequence[Mapping[str, str]] + + +_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({}) + class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ @@ -35,12 +104,12 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ sent_first_chunk: bool = False - # State tracking for accumulating partial tool calls - accumulated_tool_calls: dict[str, dict[str, Any]] + _parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads) - def __init__(self, completion_stream: Any): + 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) @@ -85,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 = [] + parts: Final = list[_GenAIPart]() for ( tool_call_index, tool_call_data, @@ -93,8 +162,10 @@ 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 = json.loads(tool_call_data["arguments"] or "{}") - function_call_part = { + 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", "args": parsed_args, @@ -172,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]: @@ -211,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, } @@ -273,9 +346,9 @@ class GoogleGenAIAdapter: def _add_generic_litellm_params_to_request( self, - completion_request_dict: dict[str, Any], + completion_request_dict: _JsonDict, litellm_params: GenericLiteLLMParams | None = None, - ) -> dict: + ) -> _JsonDict: """Add generic litellm params to request. e.g add api_base, api_key, api_version, etc. Args: @@ -287,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 @@ -295,7 +368,7 @@ class GoogleGenAIAdapter: def translate_completion_output_params_streaming( self, - completion_stream: Any, + completion_stream: object, ) -> AsyncIterator[bytes] | None: """Transform streaming completion output to Google GenAI format""" google_genai_wrapper: Final = GoogleGenAIStreamWrapper(completion_stream=completion_stream) @@ -304,15 +377,15 @@ class GoogleGenAIAdapter: def _transform_google_genai_tools_to_openai( self, - tools: list[dict[str, Any]], + tools: Sequence[_GenAITool], ) -> list[ChatCompletionToolParam]: """Transform Google GenAI tools to OpenAI tools format""" - openai_tools: Final[list[dict[str, Any]]] = [] + openai_tools: Final = list[_JsonDict]() for tool in tools: if "functionDeclarations" in tool: for func_decl in tool["functionDeclarations"]: - function_chunk: dict[str, Any] = { + function_chunk: _JsonDict = { "name": func_decl.get("name", ""), } @@ -321,7 +394,7 @@ class GoogleGenAIAdapter: if "parametersJsonSchema" in func_decl: function_chunk["parameters"] = func_decl["parametersJsonSchema"] - openai_tool = {"type": "function", "function": function_chunk} + openai_tool: _JsonDict = {"type": "function", "function": function_chunk} openai_tools.append(openai_tool) # normalize the tool schemas @@ -331,7 +404,7 @@ class GoogleGenAIAdapter: def _transform_google_genai_tool_config_to_openai( self, - tool_config: dict[str, Any], + tool_config: _GenAIToolConfig, ) -> ChatCompletionToolChoiceValues | None: """Transform Google GenAI tool_config to OpenAI tool_choice""" function_calling_config: Final = tool_config.get("functionCallingConfig", {}) @@ -345,20 +418,20 @@ class GoogleGenAIAdapter: def _transform_contents_to_messages( self, contents: list[dict[str, Any]], - system_instruction: dict[str, Any] | None = None, + system_instruction: _GenAISystemInstruction | None = None, ) -> list[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" messages: Final[list[AllMessageValues]] = [] # 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 @@ -461,7 +534,7 @@ class GoogleGenAIAdapter: def translate_completion_to_generate_content( self, response: ModelResponse, - ) -> dict[str, Any]: + ) -> _JsonDict: """ Transform litellm completion response to Google GenAI generate_content format @@ -484,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, Any]] = { + generate_content_response: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -524,7 +597,7 @@ class GoogleGenAIAdapter: self, response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper, - ) -> dict[str, Any] | None: + ) -> Mapping[str, object] | None: """ Transform streaming litellm completion chunk to Google GenAI generate_content format @@ -548,10 +621,10 @@ class GoogleGenAIAdapter: parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation(choice.delta, wrapper) else: parts = [] - finish_reason = getattr(choice, "finish_reason", None) + 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) @@ -560,7 +633,7 @@ class GoogleGenAIAdapter: return None # Create Google GenAI streaming format response - streaming_chunk: Final[dict[str, Any]] = { + streaming_chunk: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -596,10 +669,10 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, - message: Any, - ) -> list[dict[str, Any]]: + message: Message, + ) -> Sequence[_GenAIPart]: """Transform OpenAI message to Google GenAI parts format""" - parts: Final[list[dict[str, Any]]] = [] + parts: Final = list[_GenAIPart]() # Add text content if present if hasattr(message, "content") and message.content: @@ -607,16 +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 = json.loads(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 = { + function_call_part: _GenAIPart = { "functionCall": { - "name": tool_call.function.name or "undefined_tool_name", + "name": function.name or "undefined_tool_name", "args": args, } } @@ -625,28 +704,30 @@ class GoogleGenAIAdapter: return parts if parts else [{"text": ""}] def _transform_openai_delta_to_google_genai_parts_with_accumulation( - self, delta: Any, wrapper: GoogleGenAIStreamWrapper - ) -> list[dict[str, Any]]: + 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[dict[str, Any]]] = [] + 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"): continue # 3. Use `index` as the primary key for accumulation - tool_call_index = getattr(tool_call, "index", None) + tool_call_index: int | None = getattr(tool_call, "index", None) if tool_call_index is None: continue # Index is essential for tracking streaming tool calls @@ -658,8 +739,9 @@ 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: @@ -680,13 +762,13 @@ class GoogleGenAIAdapter: # 5. Attempt to parse arguments even if name hasn't arrived. try: # Attempt to parse the accumulated arguments string - parsed_args = json.loads(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. if accumulated_name: # If successful, create the part and clean up - function_call_part = {"functionCall": {"name": accumulated_name, "args": parsed_args}} + function_call_part: _GenAIPart = {"functionCall": {"name": accumulated_name, "args": parsed_args}} parts.append(function_call_part) # Remove the completed tool call from the accumulator @@ -714,7 +796,7 @@ class GoogleGenAIAdapter: return mapping.get(finish_reason, "STOP") - def _map_usage(self, usage: Any) -> 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 97e831f5822..c206849c86f 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -6,12 +6,14 @@ import random import time import uuid from collections import Counter -from collections.abc import Mapping, Sequence +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 +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload import httpx +from typing_extensions import Never, Required from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -29,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, @@ -48,7 +51,105 @@ _WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1" _WEBHOOK_PATH_LOGGING_BATCH: Final = "/v1/litellm/batch" _MAX_QUEUE_SIZE: Final = 10_000 _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0 -_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({}) + + +class _ModerationToolCall(TypedDict, total=False): + id: Required[str] + + +class _ModerationMessage(TypedDict, total=False): + content: str | None + tool_calls: Sequence[_ModerationToolCall] | None + + +class _ModerationChoice(TypedDict, total=False): + message: _ModerationMessage | None + + +class _ModerationResponse(TypedDict, total=False): + choices: Sequence[_ModerationChoice] + + +class _LogEventKwargs(TypedDict, total=False): + standard_logging_object: Required[StandardLoggingPayload] + litellm_call_id: 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: Mapping[str, object] | None + + +class _ModerationSourceMessage(TypedDict, total=False): + role: str + function_call: Mapping[str, object] | None + tool_calls: Sequence[_ModerationSourceToolCall | None] | None + + +class _FlattenedModerationMessage(TypedDict): + role: str | None + content: str + + +class _CorrelatablePayload(TypedDict): + id: str + + +class _SystemPromptCarrier(TypedDict, total=False): + messages: object + + +class _BlockFailurePayload(TypedDict, total=False): + id: object + model: object + model_group: object + model_id: str + model_parameters: object + startTime: float | None + endTime: float | None + completionStartTime: float | None + messages: object + metadata: StandardLoggingUserAPIKeyMetadata + response: str + status: str class _MalformedToolBlockingResponseError(Exception): @@ -143,7 +244,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): else {"Content-Type": "application/json"} ) - self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + self._periodic_flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task() @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -191,7 +292,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): params={"timeout": httpx.Timeout(5.0, connect=2.0)}, ) - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: """Start the periodic flush task only when an event loop is already running.""" try: loop: Final = asyncio.get_running_loop() @@ -212,7 +313,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Closing them here would close the shared connection pool for every other logger instance; let LiteLLM manage their lifecycle instead. """ - task: Final = getattr(self, "_periodic_flush_task", None) + task: Final[asyncio.Task[None] | None] = getattr(self, "_periodic_flush_task", None) if task is not None: task.cancel() @@ -253,7 +354,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod async def _guarded( - coro: Any, + coro: Awaitable[GenericGuardrailAPIInputs], inputs: GenericGuardrailAPIInputs, label: str, ) -> GenericGuardrailAPIInputs: @@ -371,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. @@ -400,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): request_data["_rubrik_logging_obj"] = logging_obj @staticmethod - def _normalize_tool_calls(tool_calls: Any) -> 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): @@ -427,7 +532,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}") @staticmethod - def _join_texts(texts: Any) -> str: + def _join_texts(texts: Sequence[str] | None) -> str: """Join response text segments into the single content string the webhook evaluates. Empty when there is no assistant text.""" if not texts: @@ -439,19 +544,22 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): tool_calls: Sequence[ChatCompletionMessageToolCall], content: str, request_id: str | None, - ) -> Mapping[str, Any]: + ) -> Mapping[str, object]: """Build an OpenAI ChatCompletion-format dict (assistant text + tool calls) for the after_completion webhook. ``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, Any]] = { + 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", @@ -467,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _flatten_messages_for_moderation(messages: Any) -> 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, @@ -488,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") @@ -506,8 +616,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _build_prompt_moderation_payload( inputs: GenericGuardrailAPIInputs, - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + request_data: Mapping[str, object], + ) -> Mapping[str, object]: """Build the bare OpenAI request the before_prompt webhook consumes. Unlike the after_completion envelope, this endpoint takes a raw OpenAI @@ -516,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, Any]] = { - "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. @@ -533,15 +635,19 @@ 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, Any] | None, - ) -> Mapping[str, Any]: + call_details: _RequestContextSource, + request_data: _RequestContextSource | None, + ) -> Mapping[str, object]: """Extract original request data from model_call_details for the response moderation service envelope. @@ -576,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: + 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 @@ -586,7 +692,7 @@ 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, Any], call_details: Mapping[str, Any]) -> 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"): @@ -596,7 +702,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Logging hooks --------------------------------------------------------- @staticmethod - def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | 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. @@ -610,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, Any], source: Mapping[str, Any]) -> 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. @@ -630,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): payload["id"] = correlated @staticmethod - def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> 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"]`` @@ -658,7 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): exc_info=True, ) - async def _prepare_log_payload(self, kwargs: Mapping[str, Any], 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) @@ -667,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() @@ -697,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, Any], 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: @@ -818,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**: @@ -860,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 = 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 " @@ -884,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. @@ -906,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, Any]: + ) -> _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") @@ -942,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: @@ -996,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Webhook services ------------------------------------------------------ - async def _post_json(self, endpoint: str, payload: Mapping[str, Any], 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: @@ -1006,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 = http_response.json() + result: Final[_ModerationResponse | None] = http_response.json() if not isinstance(result, dict): raise TypeError( f"{service_name} returned non-dict JSON " @@ -1021,9 +1133,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): async def _post_to_response_moderation_endpoint( self, - response_data: Mapping[str, Any], - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + response_data: Mapping[str, object], + request_data: Mapping[str, object], + ) -> _ModerationResponse: """Post the ``{request, response}`` envelope to the after_completion webhook and return its (possibly rewritten) response. @@ -1039,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "Response moderation service", ) - async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> 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). @@ -1047,7 +1159,7 @@ 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 @@ -1063,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: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 972ae1d9856..edd3fdb8c61 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, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast import litellm from litellm._logging import verbose_logger @@ -41,7 +41,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, @@ -51,6 +57,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, @@ -72,6 +80,8 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b # ``web_search_tool_result`` blocks to inject into the final response. WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" +_ResponseT = TypeVar("_ResponseT") + class _PlanMetadataView(TypedDict): websearch_native_blocks: Sequence[Mapping[str, object]] | None @@ -85,9 +95,96 @@ class _WebSearchSettingsView(TypedDict): websearch_interception_params: WebSearchInterceptionConfig +class _SearchToolLitellmParams(TypedDict, total=False): + search_provider: str | None + + class _SearchToolConfig(TypedDict, total=False): search_tool_name: str - litellm_params: Mapping[str, object] | None + litellm_params: _SearchToolLitellmParams | None + + +class _LitellmParamsProviderView(TypedDict, total=False): + custom_llm_provider: str + + +class _DeploymentCallKwargsView(TypedDict): + custom_llm_provider: str + litellm_params: _LitellmParamsProviderView + model: str + + +class _AcreateNamedParams(TypedDict, total=False): + metadata: Never + stop_sequences: Never + stream: bool | None + system: str | None + temperature: float | None + thinking: Never + tool_choice: Never + tools: Never + top_k: int | None + top_p: float | None + container: Never + + +class _AsearchNamedParams(TypedDict, total=False): + max_results: int | None + search_domain_filter: Never + max_tokens_per_page: int | None + country: str | None + api_key: str | None + api_base: str | None + timeout: float | None + extra_headers: Never + + +class _AcompletionNamedParams(TypedDict, total=False): + functions: Never + function_call: str | None + timeout: float | None + temperature: float | None + top_p: float | None + n: int | None + stream: bool | None + stream_options: Never + stop: Never + max_tokens: int | None + max_completion_tokens: int | None + modalities: Never + prediction: ChatCompletionPredictionContentParam | None + audio: ChatCompletionAudioParam | None + presence_penalty: float | None + frequency_penalty: float | None + logit_bias: Never + user: str | None + response_format: Never + seed: int | None + tools: Never + tool_choice: Never + parallel_tool_calls: bool | None + logprobs: bool | None + top_logprobs: int | None + deployment_id: str | None + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None + verbosity: Literal["low", "medium", "high"] | None + safety_identifier: str | None + service_tier: str | None + base_url: str | None + api_version: str | None + api_key: str | None + model_list: Never + extra_headers: Never + thinking: AnthropicThinkingParam | None + web_search_options: OpenAIWebSearchOptions | None + include_server_side_tool_invocations: bool | None + shared_session: "ClientSession | None" + enable_json_schema_validation: bool | None + + +_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {} +_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {} +_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {} class WebSearchInterceptionLogger(CustomLogger): @@ -275,12 +372,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 - custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get( + 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 = 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.get("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: @@ -903,17 +1005,18 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: + def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response if isinstance(response, dict): - existing = response.get("content") or [] + existing: Sequence[object] = response.get("content") or [] response["content"] = list(native_blocks) + list(existing) return response existing = getattr(response, "content", None) or [] + content_attribute: Final = "content" try: - response.content = list(native_blocks) + list(existing) + setattr(response, content_attribute, list(native_blocks) + list(existing)) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. @@ -1169,10 +1272,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( @@ -1180,9 +1283,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") @@ -1197,12 +1300,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 @@ -1344,12 +1449,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: @@ -1377,12 +1483,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, ) ) @@ -1422,7 +1531,7 @@ class WebSearchInterceptionLogger(CustomLogger): valid_token=user_api_key_auth, ) - team_id: Final = getattr(user_api_key_auth, "team_id", None) + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) if team_id: from litellm.proxy.proxy_server import ( prisma_client, @@ -1537,10 +1646,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.""" @@ -1548,8 +1657,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: @@ -1557,11 +1666,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 e4a4d23b438..47dfe8de292 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,12 +13,12 @@ Pattern Overview: """ import json -from collections.abc import Mapping +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 TypedDict, assert_never from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -61,6 +61,7 @@ if TYPE_CHECKING: ModifyResponseException, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -95,6 +96,48 @@ InputWriteBackTarget = ( ) +class _SSEDelta(TypedDict, total=False): + type: str + text: str + stop_reason: str | None + + +class _SSEEventData(TypedDict, total=False): + delta: _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 @@ -123,7 +166,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _build_streaming_usage_response( - responses_so_far: list[Any], + 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))) @@ -141,7 +184,7 @@ class AnthropicMessagesHandler(BaseTranslation): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: list[Any] | None = None, + responses_so_far: Sequence[object] | None = None, ) -> list[bytes]: """ Build an Anthropic SSE sequence delivering the guardrail block message @@ -159,7 +202,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]: @@ -184,7 +227,9 @@ class AnthropicMessagesHandler(BaseTranslation): ) return list(FakeAnthropicMessagesStreamIterator(response=block_response)) - def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[Any]) -> 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.""" @@ -234,7 +279,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _content_block_state( - responses_so_far: list[Any], + 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). @@ -260,7 +305,20 @@ class AnthropicMessagesHandler(BaseTranslation): return open_index, max_index @staticmethod - def _iter_sse_events(item: Any) -> list[dict]: + 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 @@ -268,22 +326,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]] = [] - 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 = 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.""" @@ -315,8 +366,8 @@ class AnthropicMessagesHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - ) -> Any: + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> Mapping[str, object]: """ Process input messages by applying guardrails to text content. """ @@ -467,8 +518,8 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _openai_system_message_to_anthropic( - message: dict[str, Any], - ) -> dict[str, Any] | None: # mutable-ok: API message payload + 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") if isinstance(content, str): @@ -477,14 +528,14 @@ class AnthropicMessagesHandler(BaseTranslation): ) # mutable-ok: API message payload if not isinstance(content, list): return None - blocks: Final[list[dict[str, Any]]] = [] # mutable-ok: API message payload + blocks: Final[list[dict[str, object]]] = [] # mutable-ok: API message payload for block in content: if not isinstance(block, dict) or block.get("type") != "text": continue text = block.get("text") if not isinstance(text, str) or not text: continue - anthropic_block: dict[str, Any] = { # mutable-ok: API message payload + anthropic_block: dict[str, object] = { # mutable-ok: API message payload "type": "text", "text": text, } # mutable-ok: API message payload @@ -514,7 +565,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 @@ -602,7 +653,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _extract_midturn_system_text( - message: dict[str, Any], # mutable-ok: API message payload + message: Mapping[str, object], msg_idx: int, ) -> ExtractedInput: """Match the adapter's filtering so positional guardrail write-back stays aligned.""" @@ -636,7 +687,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _extract_input_text_and_images( cls, - message: dict[str, Any], + message: Mapping[str, object], msg_idx: int, skip_system_message: bool = False, skip_tool_message: bool = False, @@ -696,7 +747,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)),) @@ -707,7 +758,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _extract_tool_result( cls, - content_item: Mapping[str, Any], + content_item: Mapping[str, object], msg_idx: int, content_idx: int, ) -> ExtractedInput: @@ -736,7 +787,7 @@ class AnthropicMessagesHandler(BaseTranslation): ) @staticmethod - def _image_sources(block: Mapping[str, Any]) -> tuple[str, ...]: + def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]: source: Final = block.get("source") if not isinstance(source, Mapping): return () @@ -746,7 +797,7 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: list[dict[str, Any]], + messages: Sequence[_WritableMessage], responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: @@ -788,10 +839,10 @@ class AnthropicMessagesHandler(BaseTranslation): self, response: "AnthropicMessagesResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - user_api_key_dict: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> Any: + ) -> "AnthropicMessagesResponse": """ Process output response by applying guardrails to text content and tool calls. @@ -869,10 +920,10 @@ class AnthropicMessagesHandler(BaseTranslation): self, responses_so_far: list[Any], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - user_api_key_dict: Any | None = None, + 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. @@ -950,8 +1001,8 @@ class AnthropicMessagesHandler(BaseTranslation): def _prepare_request_data( self, request_data: dict | None, - response: Any, - user_api_key_dict: Any | None, + response: object, + user_api_key_dict: "UserAPIKeyAuth | None", key: str, ) -> dict: """Ensure request_data has the response/responses_so_far key and metadata.""" @@ -968,7 +1019,7 @@ class AnthropicMessagesHandler(BaseTranslation): return request_data @staticmethod - def _get_response_content(response: Any) -> 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 [] @@ -978,7 +1029,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]], @@ -986,21 +1037,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, Any] = {} - if isinstance(content_block, dict): - block_type = content_block.get("type") - block_dict = cast(dict[str, Any], 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( @@ -1012,12 +1052,27 @@ 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], images_to_check: list[str], tool_calls_to_check: list["ChatCompletionToolCallChunk"], - response: Any, + response: object, ) -> "GenericGuardrailAPIInputs": """Build GenericGuardrailAPIInputs with optional images, tool calls, model.""" inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) @@ -1034,7 +1089,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. @@ -1105,7 +1160,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", "") @@ -1117,7 +1172,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. @@ -1168,7 +1223,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: @@ -1212,7 +1267,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _extract_output_text_and_images( self, - content_block: dict[str, Any], + content_block: Mapping[str, object], content_idx: int, texts_to_check: list[str], images_to_check: list[str], @@ -1235,7 +1290,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, @@ -1260,7 +1315,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"): @@ -1276,14 +1331,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, Any], 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 dbeac453791..d230b438086 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,8 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers: """ import re -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, 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 NotRequired, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -27,11 +29,11 @@ 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, AllAnthropicToolsValues, - AnthopicMessagesAssistantMessageParam, - AnthropicMessagesUserMessageParam, ) from litellm.types.llms.openai import ChatCompletionToolParam from litellm.types.utils import ModelResponse @@ -82,6 +84,69 @@ _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: str + max_tokens: int + timeout: float + litellm_metadata: Mapping[str, object] + user: NotRequired[str] + allowed_model_region: NotRequired[str] + + +class _SummaryAcompletion(Protocol): + def __call__( + self, *, messages: Sequence[Mapping[str, object]], **kwargs: Unpack[_SummaryCallKwargs] + ) -> "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.""" @@ -157,11 +222,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), @@ -347,7 +412,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( @@ -399,40 +464,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: @@ -446,7 +528,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. @@ -465,8 +547,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)`` @@ -481,19 +563,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. @@ -600,7 +684,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, @@ -623,7 +707,7 @@ def _count_effective_tokens( try: openai_shape = adapter.translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", messages_without_compaction, ) ) @@ -679,17 +763,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. @@ -704,16 +789,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 [ @@ -736,7 +823,7 @@ def _extract_summary_text(raw: str | None) -> str | None: def _system_to_openai_message( system: str | list[dict[str, Any]] | None, -) -> dict[str, Any] | 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 @@ -747,17 +834,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 @@ -773,7 +862,7 @@ def _build_summary_messages( try: openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", stripped, ) ) @@ -785,7 +874,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) @@ -809,7 +898,7 @@ def _is_user_message(msg: object) -> bool: return isinstance(msg, dict) and msg.get("role") == "user" -def _append_text_to_content(content: Any, extra_text: str) -> Any: +def _append_text_to_content(content: object, extra_text: str) -> object: """Append ``extra_text`` to an OpenAI-shape message ``content`` field. Handles the two common shapes: ``str`` and ``list`` of content parts. @@ -820,16 +909,17 @@ def _append_text_to_content(content: Any, extra_text: str) -> Any: if isinstance(content, str): return f"{content}\n\n{extra_text}" if isinstance(content, list): - return [*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}] 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"]: @@ -860,9 +950,8 @@ async def _call_summary_model( # the parent ``/v1/messages`` request. On timeout the caller catches the # exception and surfaces ``applied_edits[0].error = "summary_call_failed"``, # forwarding the request without compaction rather than hanging. - call_kwargs: Final[dict[str, Any]] = { + call_kwargs: Final[_SummaryCallKwargs] = { "model": summary_model, - "messages": summary_messages, "max_tokens": max_tokens, "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, "litellm_metadata": metadata, @@ -872,19 +961,23 @@ 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") - if end_user_id: + if isinstance(end_user_id, str) and end_user_id: call_kwargs["user"] = end_user_id if allowed_model_region is not None: call_kwargs["allowed_model_region"] = allowed_model_region - 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 @@ -900,7 +993,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 2e0ae30a192..6e720c058fb 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -1,8 +1,10 @@ +from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict import httpx -from httpx._types import RequestFiles +from httpx._types import FileTypes, RequestFiles +from typing_extensions import NotRequired import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION @@ -31,6 +33,31 @@ else: LiteLLMLoggingObj = Any +class _RunwayTaskResponse(TypedDict, total=False): + id: str + status: str + createdAt: str + completedAt: str + output: Sequence[str] | str + progress: int + failureCode: str + failure: str + + +class _RunwayVideoData(TypedDict): + id: str + object: Literal["video"] + status: str + created_at: int + output_url: NotRequired[str] + completed_at: NotRequired[int] + progress: NotRequired[int] + error: NotRequired[Mapping[str, str]] + model: NotRequired[str] + size: NotRequired[str] + seconds: NotRequired[str] + + class RunwayMLVideoConfig(BaseVideoConfig): """ Configuration class for RunwayML video generation. @@ -44,6 +71,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. @@ -68,7 +99,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: """ Map OpenAI parameters to RunwayML format. @@ -78,37 +109,44 @@ class RunwayMLVideoConfig(BaseVideoConfig): - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ - mapped_params: Final[dict[str, Any]] = {} + 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, @@ -163,7 +201,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]: @@ -179,17 +217,15 @@ class RunwayMLVideoConfig(BaseVideoConfig): "duration": 5 } """ - # Build the request data - request_data: Final[dict[str, Any]] = { + # Build the request data with the mapped parameters merged in + request_data: Final = { "model": model, "promptText": prompt, + **video_create_optional_request_params, } - # Add mapped parameters - request_data.update(video_create_optional_request_params) - # RunwayML uses JSON body, no files multipart - files_list: Final[list[tuple[str, Any]]] = [] + files_list: Final[Sequence[tuple[str, FileTypes]]] = [] # Append the specific endpoint for video generation full_api_base: Final = f"{api_base}/image_to_video" @@ -216,10 +252,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): We map this to OpenAI VideoObject format. """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_RunwayVideoData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -229,9 +265,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Add optional fields if present if "output" in response_data and response_data["output"]: # RunwayML returns output as array of URLs when task succeeds - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) + output: Final = response_data["output"] + video_data["output_url"] = output if isinstance(output, str) else output[0] if "completedAt" in response_data: video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) @@ -254,7 +289,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): if "duration" in request_data: video_data["seconds"] = str(request_data["duration"]) - video_obj: Final = VideoObject(**video_data) + video_obj: Final = VideoObject.model_validate(video_data) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) @@ -326,20 +361,18 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Get task status to retrieve video URL url: Final = f"{api_base}/tasks/{encoded_video_id}" - params: Final[dict[str, Any]] = {} + 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 @@ -373,7 +406,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 @@ -402,7 +435,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 @@ -421,7 +454,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video remix request for RunwayML API. @@ -448,7 +481,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video list request for RunwayML API. @@ -484,9 +517,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Construct the URL for task cancellation url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel" - data: Final[dict[str, Any]] = {} - - return url, data + return url, dict[str, str]() def transform_video_delete_response( self, @@ -494,7 +525,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): logging_obj: LiteLLMLoggingObj, ) -> VideoObject: """Transform the RunwayML video delete/cancel response.""" - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_obj: Final = VideoObject( id=response_data.get("id", ""), @@ -524,9 +555,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): url: Final = f"{api_base}/tasks/{encoded_video_id}" # Empty dict for GET request (no body) - data: Final[dict[str, Any]] = {} - - return url, data + return url, dict[str, str]() def transform_video_status_retrieve_response( self, @@ -537,10 +566,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ Transform the RunwayML video status retrieve response. """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_RunwayVideoData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -549,9 +578,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Add optional fields if present if "output" in response_data and response_data["output"]: - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) + output: Final = response_data["output"] + video_data["output_url"] = output if isinstance(output, str) else output[0] if "completedAt" in response_data: video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) @@ -565,14 +593,14 @@ class RunwayMLVideoConfig(BaseVideoConfig): "message": response_data.get("failure", "Video generation failed"), } - video_obj: Final = VideoObject(**video_data) + video_obj: Final = VideoObject.model_validate(video_data) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj - def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers): raise NotImplementedError("video create character is not supported for RunwayML") def transform_video_create_character_response(self, raw_response, logging_obj): diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3538fc5b1a7..131ee41ea8b 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,12 +5,13 @@ import json import os import re import time -from collections.abc import Callable, Iterable, Iterator -from typing import Any, Final +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from typing import Final, TypedDict import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from typing_extensions import Required import litellm from litellm._uuid import uuid @@ -50,9 +51,10 @@ from litellm.types.llms.openai import ( HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, + OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GenerateContentResponseBody from litellm.types.utils import LlmProviders, ModelResponse from ..common_utils import VertexAIError @@ -62,6 +64,47 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +class _OpenAIBatchRequestBody(TypedDict, total=False): + model: str + messages: Sequence[AllMessageValues] + + +class _OpenAIBatchJsonlEntry(TypedDict, total=False): + custom_id: Required[object] + body: _OpenAIBatchRequestBody + + +class _VertexBatchOutputRequest(TypedDict, total=False): + labels: Mapping[str, str] + + +class _VertexBatchResponse(GenerateContentResponseBody, total=False): + modelVersion: str + + +class _VertexBatchOutputRow(TypedDict, total=False): + request: _VertexBatchOutputRequest + status: str + processed_time: str + response: _VertexBatchResponse + + +class _GcsObjectMetadata(TypedDict, total=False): + purpose: OpenAIFilesPurpose + + +class _GcsObjectResponse(GcsBucketResponse, total=False): + metadata: _GcsObjectMetadata + + +def _parse_gcs_object_response(raw_response: Response) -> _GcsObjectResponse: + return raw_response.json() + + +def _parse_vertex_batch_output_row(line: str) -> _VertexBatchOutputRow: + return json.loads(line) + + def _sanitize_gcp_label_value(value: str) -> str: """ Sanitize a string to meet GCP label value constraints. @@ -106,7 +149,7 @@ def _decode_gcp_label_value_chunks(values: list[str]) -> str | None: return None -def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) -> None: +def _litellm_batch_custom_id_labels(custom_id: object) -> Mapping[str, str]: """ Store OpenAI batch custom_id for Vertex batch correlation. @@ -115,15 +158,19 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) round-trip correlation in batch output transforms. """ custom_id_str: Final = str(custom_id) - labels["litellm_custom_id"] = _sanitize_gcp_label_value(custom_id_str) raw_label_chunks: Final = _encode_gcp_label_value_chunks(custom_id_str) - labels["litellm_custom_id_raw"] = raw_label_chunks[0] - for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1): - labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk + return { + "litellm_custom_id": _sanitize_gcp_label_value(custom_id_str), + "litellm_custom_id_raw": raw_label_chunks[0], + **{ + f"litellm_custom_id_raw_{index}": raw_label_chunk + for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1) + }, + } -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: - """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, str]) -> str: + """Prefer encoded custom_id when present (see _litellm_batch_custom_id_labels).""" raw: Final = labels.get("litellm_custom_id_raw") if raw: raw_chunks: Final = [str(raw)] @@ -141,9 +188,9 @@ def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: def _openai_batch_jsonl_entry_to_vertex_wrapped_request( - openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> dict[str, Any]: + openai_entry: _OpenAIBatchJsonlEntry, + map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]], +) -> Mapping[str, object]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -151,11 +198,11 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ - openai_request_body: Final = openai_entry.get("body") or {} + openai_request_body: Final[_OpenAIBatchRequestBody] = openai_entry.get("body") or {} vertex_request_body: Final = _transform_request_body( - messages=openai_request_body.get("messages", []), + messages=[*openai_request_body.get("messages", [])], model=openai_request_body.get("model", ""), - optional_params=map_openai_to_vertex_params(openai_request_body), + optional_params=dict(map_openai_to_vertex_params(openai_request_body)), custom_llm_provider="vertex_ai", litellm_params={}, cached_content=None, @@ -163,9 +210,10 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( custom_id: Final = openai_entry.get("custom_id") if custom_id is not None: - if "labels" not in vertex_request_body: - vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) + vertex_request_body["labels"] = { + **vertex_request_body.get("labels", {}), + **_litellm_batch_custom_id_labels(custom_id), + } return {"request": vertex_request_body} @@ -186,7 +234,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: ``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited JSONL. """ - content: Any = openai_file_content + content: FileTypes | str = openai_file_content if isinstance(content, tuple): content = content[1] @@ -241,7 +289,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: def _iter_openai_jsonl_entries( openai_file_content: FileTypes, -) -> Iterator[dict[str, Any]]: +) -> Iterator[_OpenAIBatchJsonlEntry]: for line in _iter_openai_jsonl_lines(openai_file_content): yield json.loads(line) @@ -257,7 +305,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -308,7 +356,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: Sequence[_OpenAIBatchJsonlEntry], ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job @@ -396,8 +444,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, - openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + openai_request_body: _OpenAIBatchRequestBody, + ) -> Mapping[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ @@ -409,7 +457,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): _model: Final = openai_request_body.get("model", "") vertex_params: Final = config.map_openai_params( model=_model, - non_default_params=openai_request_body, + non_default_params=dict(openai_request_body), optional_params={}, drop_params=False, ) @@ -463,10 +511,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Transform VertexAI File upload response into OpenAI-style FileObject """ - response_json: Final = raw_response.json() + response_json: Final = _parse_gcs_object_response(raw_response) try: - response_object: Final = GcsBucketResponse(**response_json) + response_object: Final = _GcsObjectResponse(**response_json) except Exception as e: raise VertexAIError( status_code=raw_response.status_code, @@ -523,7 +571,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - response_json: Final = raw_response.json() + response_json: Final = _parse_gcs_object_response(raw_response) gcs_id = response_json.get("id", "") gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" return OpenAIFileObject( @@ -682,7 +730,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # discriminating fields. Anything else (e.g. a binary file whose # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. - first_row: Final = json.loads(first_line) + first_row: Final = _parse_vertex_batch_output_row(first_line) is_vertex_batch_output: Final = ( "request" in first_row and "response" in first_row @@ -723,7 +771,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): for line in itertools.chain([first_line], lines): try: openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), + vertex_output=_parse_vertex_batch_output_row(line), vertex_gemini_config=vertex_gemini_config, logging_obj=batch_transform_logging_obj, mock_httpx_response=mock_httpx_response, @@ -742,11 +790,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _transform_single_vertex_batch_output_to_openai( self, - vertex_output: dict[str, Any], + vertex_output: _VertexBatchOutputRow, vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index b2fa538b1cc..e6fcff548d6 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -6,8 +6,11 @@ Admins use the management endpoints to read and update input_policy / output_pol """ import uuid +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 @@ -23,33 +26,109 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient -def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow: +class _ToolTableRecord(Protocol): + tool_name: str + input_policy: str | None + output_policy: str | None + + +class _TokenRelationRecord(Protocol): + token: str | None + key_alias: str | None + + +class _TeamRelationRecord(Protocol): + team_id: str | None + team_alias: str | None + + +class _ObjectPermissionRecord(Protocol): + object_permission_id: str + blocked_tools: Sequence[str] | None + verification_tokens: Sequence[_TokenRelationRecord] | None + teams: Sequence[_TeamRelationRecord] | None + + +class _ToolTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | None = None, + ) -> Sequence[_ToolTableRecord]: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _ToolTableRecord | None: ... + + async def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _ObjectPermissionTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, + ) -> Sequence[_ObjectPermissionRecord]: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _ObjectPermissionRecord | None: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _ModelDumpMethod(Protocol): + def __call__(self) -> Mapping: ... + + +_ROW_DICT: Final = TypeAdapter(dict) + + +class _ToolTableHolder(Protocol): + @property + def table(self) -> _ToolTable: ... + + +class _ObjectPermissionTableHolder(Protocol): + @property + def table(self) -> _ObjectPermissionTable: ... + + +def _tool_table(repo: _ToolTableHolder) -> _ToolTable: + return repo.table + + +def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable: + return repo.table + + +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", ""), @@ -87,7 +166,7 @@ async def batch_upsert_tools( if not data: return now: Final = datetime.now(timezone.utc) - table: Final = ToolRepository(prisma_client).table + table: Final = _tool_table(ToolRepository(prisma_client)) for item in data: tool_name = item.get("tool_name", "") origin = item.get("origin") or "user_defined" @@ -132,8 +211,8 @@ async def list_tools( ) -> list[LiteLLM_ToolTableRow]: """Return all tools, optionally filtered by input_policy.""" try: - where: Final = {"input_policy": input_policy} if input_policy is not None else {} - rows: Final = await ToolRepository(prisma_client).table.find_many( + where: Final[Mapping[str, str]] = {"input_policy": input_policy} if input_policy is not None else {} + rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many( where=where, order={"created_at": "desc"}, ) @@ -149,7 +228,7 @@ async def get_tool( ) -> LiteLLM_ToolTableRow | None: """Return a single tool row by tool_name.""" try: - row: Final = await ToolRepository(prisma_client).table.find_unique( + row: Final = await _tool_table(ToolRepository(prisma_client)).find_unique( where={"tool_name": tool_name}, ) if row is None: @@ -172,7 +251,7 @@ async def update_tool_policy( _updated_by: Final = updated_by or "system" now: Final = datetime.now(timezone.utc) - create_data: Final[dict] = { + create_data: Final[Mapping[str, str | datetime]] = { "tool_id": str(uuid.uuid4()), "tool_name": tool_name, "input_policy": input_policy or "untrusted", @@ -182,16 +261,18 @@ async def update_tool_policy( "created_at": now, "updated_at": now, } - update_data: Final[dict] = { - "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 ToolRepository(prisma_client).table.upsert( + await _tool_table(ToolRepository(prisma_client)).upsert( where={"tool_name": tool_name}, data={ "create": create_data, @@ -214,7 +295,7 @@ async def get_tools_by_names( if not tool_names: return {} try: - rows: Final = await ToolRepository(prisma_client).table.find_many( + rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many( where={"tool_name": {"in": tool_names}}, ) return { @@ -239,7 +320,7 @@ async def list_overrides_for_tool( """ out: Final[list[ToolPolicyOverrideRow]] = [] try: - perms: Final = await ObjectPermissionRepository(prisma_client).table.find_many( + perms: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many( where={"blocked_tools": {"has": tool_name}}, include={ "verification_tokens": True, @@ -248,8 +329,8 @@ async def list_overrides_for_tool( ) for perm in perms: op_id = getattr(perm, "object_permission_id", None) or "" - tokens = getattr(perm, "verification_tokens", []) or [] - teams = getattr(perm, "teams", []) or [] + tokens: Sequence[_TokenRelationRecord] = getattr(perm, "verification_tokens", []) or [] + teams: Sequence[_TeamRelationRecord] = getattr(perm, "teams", []) or [] for t in tokens: out.append( ToolPolicyOverrideRow( @@ -302,7 +383,7 @@ class ToolPolicyRegistry: try: tools: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: ToolRepository(prisma_client).table.find_many(), + lambda: _tool_table(ToolRepository(prisma_client)).find_many(), reason="sync_tool_policy_from_db_tools_lookup_failure", ) self._tool_input_policies = { @@ -314,13 +395,13 @@ class ToolPolicyRegistry: perms: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: ObjectPermissionRepository(prisma_client).table.find_many(), + lambda: _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many(), reason="sync_tool_policy_from_db_perms_lookup_failure", ) 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) @@ -352,10 +433,12 @@ class ToolPolicyRegistry: """ if not tool_names: return {} - blocked: Final[set] = 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: @@ -385,18 +468,17 @@ async def add_tool_to_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique( where={"object_permission_id": object_permission_id}, ) 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 ObjectPermissionRepository(prisma_client).table.update( + await _object_permission_table(ObjectPermissionRepository(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: @@ -413,18 +495,17 @@ async def remove_tool_from_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique( where={"object_permission_id": object_permission_id}, ) 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 ObjectPermissionRepository(prisma_client).table.update( + await _object_permission_table(ObjectPermissionRepository(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 a416a197ab8..7dfc80d2b2a 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 Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -85,14 +85,6 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( if TYPE_CHECKING: from prisma import models as prisma_models from prisma import types as prisma_types - from prisma.actions import ( - LiteLLM_InvitationLinkActions, - LiteLLM_OrganizationMembershipActions, - LiteLLM_TeamMembershipActions, - LiteLLM_TeamTableActions, - LiteLLM_UserTableActions, - LiteLLM_VerificationTokenActions, - ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.proxy_server import PrismaClient @@ -100,55 +92,151 @@ if TYPE_CHECKING: router: Final = APIRouter() +_PrismaTableT = TypeVar("_PrismaTableT", covariant=True) + + +class _TableActions(Protocol[_PrismaTableT]): + async def find_unique( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> "_PrismaTableT | None": ... + + async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + skip: int | None = None, + take: int | None = None, + ) -> "Sequence[_PrismaTableT]": ... + + async def create(self, *, data: Mapping[str, object]) -> "_PrismaTableT": ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + + +class _PrismaTableHolder(Protocol[_PrismaTableT]): + @property + def table(self) -> "_TableActions[_PrismaTableT]": ... + + +def _typed_table(holder: "_PrismaTableHolder[_PrismaTableT]") -> "_TableActions[_PrismaTableT]": + return holder.table + + +class _LenientTableActions(Protocol[_PrismaTableT]): + async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + skip: int | None = None, + take: int | None = None, + ) -> "Sequence[_PrismaTableT] | None": ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _LenientTableHolder(Protocol[_PrismaTableT]): + @property + def table(self) -> "_LenientTableActions[_PrismaTableT]": ... + + +def _lenient_table(holder: "_LenientTableHolder[_PrismaTableT]") -> "_LenientTableActions[_PrismaTableT]": + return holder.table + + +class _UserDeleteRow(Protocol): + user_id: str + user_email: str | None + + @property + def teams(self) -> Sequence[str]: ... + + def json(self, *, exclude_none: bool) -> str: ... + + +class _TeamCleanupRow(Protocol): + team_id: str + members_with_roles: str + + def model_dump(self) -> Mapping[str, object]: ... + def _user_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]": - user_table: Final[LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table - return user_table +) -> "_TableActions[prisma_models.LiteLLM_UserTable]": + return _typed_table(UserRepository(prisma_client)) + + +def _user_table_lenient( + prisma_client: "PrismaClient | None", +) -> "_LenientTableActions[prisma_models.LiteLLM_UserTable]": + return _lenient_table(UserRepository(prisma_client)) + + +def _user_delete_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[_UserDeleteRow]": + return _typed_table(UserRepository(prisma_client)) def _team_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]": - team_table: Final[LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table - return team_table +) -> "_TableActions[prisma_models.LiteLLM_TeamTable]": + return _typed_table(TeamRepository(prisma_client)) + + +def _team_cleanup_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[_TeamCleanupRow]": + return _typed_table(TeamRepository(prisma_client)) def _verification_token_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]": - token_table: Final[LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]] = ( - VerificationTokenRepository(prisma_client).table - ) - return token_table +) -> "_TableActions[prisma_models.LiteLLM_VerificationToken]": + return _typed_table(VerificationTokenRepository(prisma_client)) + + +def _verification_token_table_lenient( + prisma_client: "PrismaClient | None", +) -> "_LenientTableActions[prisma_models.LiteLLM_VerificationToken]": + return _lenient_table(VerificationTokenRepository(prisma_client)) + + +def _organization_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[prisma_models.LiteLLM_OrganizationTable]": + return _typed_table(OrganizationRepository(prisma_client)) def _organization_membership_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]": - membership_table: Final[LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]] = ( - OrganizationMembershipRepository(prisma_client).table - ) - return membership_table +) -> "_TableActions[prisma_models.LiteLLM_OrganizationMembership]": + return _typed_table(OrganizationMembershipRepository(prisma_client)) def _invitation_link_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink]": - invitation_table: LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink] = InvitationLinkRepository( - prisma_client - ).table - return invitation_table +) -> "_TableActions[prisma_models.LiteLLM_InvitationLink]": + return _typed_table(InvitationLinkRepository(prisma_client)) def _team_membership_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]": - team_membership_table: Final[LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]] = ( - TeamMembershipRepository(prisma_client).table - ) - return team_membership_table +) -> "_TableActions[prisma_models.LiteLLM_TeamMembership]": + return _typed_table(TeamMembershipRepository(prisma_client)) def _hash_password_in_dict(data: dict) -> None: @@ -234,7 +322,7 @@ async def _check_duplicate_user_field( if case_insensitive: where_clause[field_name]["mode"] = "insensitive" - existing_user: Final = await UserRepository(prisma_client).table.find_first(where=where_clause) + existing_user: Final = await _user_table_lenient(prisma_client).find_first(where=where_clause) if existing_user is not None: existing_value: Final = getattr(existing_user, field_name, value) @@ -650,7 +738,7 @@ async def ui_get_available_role( def get_team_from_list( - team_list: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, + team_list: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None, team_id: str, ) -> LiteLLM_TeamTable | LiteLLM_TeamMembership | None: if team_list is None: @@ -732,18 +820,59 @@ def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKey ) -async def _get_user_info_teams( - prisma_client: Any, +_TeamIdList: TypeAlias = list[str] + + +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: _TeamIdList, + table_name: Literal["team"], + query_type: Literal["find_all"], + ) -> "Sequence[TeamListResponseObject] | None": ... + + +async def _get_user_info_row( + prisma_client: "_UserInfoDataClient", + user_id: str, +) -> "prisma_models.LiteLLM_UserTable | None": + return await prisma_client.get_data(user_id=user_id) + + +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[Any], list[Any] | None]: +) -> tuple[Sequence[TeamListResponseObject], Sequence[TeamListResponseObject] | None]: """Fetch and merge teams from membership + user.teams field.""" from litellm.proxy.management_endpoints.team_endpoints import list_team - team_list: list[Any] = [] - team_id_list: list[str] = [] - teams_1: Final = await list_team( http_request=Request( scope={"type": "http", "path": "/user/info"}, @@ -752,11 +881,10 @@ async def _get_user_info_teams( user_api_key_dict=user_api_key_dict, ) - if teams_1 is not None and isinstance(teams_1, list): - team_list = teams_1 - team_id_list = [team.team_id for team in teams_1] + team_list: Final = teams_1 if teams_1 is not None and isinstance(teams_1, list) else list[TeamListResponseObject]() + team_id_list: Final = [team.team_id for team in team_list] - teams_2: list[Any] | 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): @@ -767,7 +895,7 @@ async def _get_user_info_teams( ) elif user_api_key_dict.user_id is not None and user_id is None: caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id) - caller_team_ids: Final = getattr(caller_user_info, "teams", None) + 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, @@ -804,9 +932,9 @@ def _redact_scim_enterprise_metadata( def _build_user_info_response( user_id: str | None, user_info: Any | None, - keys: list[LiteLLM_VerificationToken] | None, - team_list: list[Any], - teams_1: list[Any] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, + team_list: Sequence[TeamListResponseObject], + teams_1: Sequence[TeamListResponseObject] | None, ) -> UserInfoResponse: """Create UserInfoResponse while filtering sensitive fields.""" if user_info is None and keys is not None: @@ -814,7 +942,7 @@ def _build_user_info_response( user_info = {"spend": spend} returned_keys: Final = _process_keys_for_user_info(keys=keys, all_teams=teams_1) - team_list.sort(key=lambda x: getattr(x, "team_alias", "") or "") + sorted_team_list: Final = sorted(team_list, key=lambda x: getattr(x, "team_alias", "") or "") _user_info: Final = user_info.model_dump() if isinstance(user_info, BaseModel) else user_info if isinstance(_user_info, dict): @@ -825,7 +953,7 @@ def _build_user_info_response( user_id=user_id, user_info=_user_info, keys=returned_keys, - teams=team_list, + teams=sorted_team_list, ) @@ -870,9 +998,9 @@ async def user_info( user_id = user_api_key_dict.user_id ## GET USER ROW ## - user_info = None + user_info: prisma_models.LiteLLM_UserTable | None = None if user_id is not None: - user_info = await prisma_client.get_data(user_id=user_id) + user_info = await _get_user_info_row(prisma_client, user_id) if user_info is None: raise HTTPException( @@ -888,11 +1016,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, @@ -1058,6 +1182,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 @@ -1081,22 +1211,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[list] = 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 = 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 @@ -1121,8 +1254,8 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): def _process_keys_for_user_info( - keys: list[LiteLLM_VerificationToken] | None, - all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, + all_teams: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None, ): from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.proxy_server import general_settings, litellm_master_key_hash @@ -1212,7 +1345,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, @@ -1768,7 +1901,10 @@ async def bulk_user_update( # Apply update transformations (reuse existing logic) data_json: Final[dict] = data.user_updates.model_dump(exclude_unset=True) - non_default_values: Final = _update_internal_user_params(data_json=data_json, data=data.user_updates) + _raw_update_values: Final[Mapping[str, object]] = _update_internal_user_params( + data_json=data_json, data=data.user_updates + ) + non_default_values: Final = dict[str, object](_raw_update_values) # Remove user identification fields since we're updating by user_id non_default_values.pop("user_id", None) @@ -1780,7 +1916,7 @@ async def bulk_user_update( try: # Perform bulk database update - await UserRepository(prisma_client).table.update_many( + await _user_table_lenient(prisma_client).update_many( where={}, data=non_default_values, # Update all users ) @@ -1885,7 +2021,7 @@ async def get_user_key_counts( # Get count for each user_id individually for user_id in user_ids: - count = await VerificationTokenRepository(prisma_client).table.count( + count = await _verification_token_table_lenient(prisma_client).count( where={ "user_id": user_id, "OR": [ @@ -2122,7 +2258,7 @@ async def get_users( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) - users: Sequence[prisma_models.LiteLLM_UserTable] | None = await UserRepository(prisma_client).table.find_many( + users: Sequence[prisma_models.LiteLLM_UserTable] | None = await _user_table_lenient(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -2130,7 +2266,7 @@ async def get_users( ) # Get total count of user rows - total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count: Final[int] = await _user_table_lenient(prisma_client).count(where=where_conditions) # Get key count for each user if users is not None: @@ -2256,7 +2392,7 @@ async def delete_user( # check that all teams passed exist for user_id in data.user_ids: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user_row = await _user_delete_table(prisma_client).find_unique(where={"user_id": user_id}) if user_row is None: raise HTTPException( @@ -2308,8 +2444,8 @@ async def delete_user( ) ## CLEANUP MEMBERS_WITH_ROLES - fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}}) - teams_to_update = [] + fetch_all_teams = await _team_cleanup_table(prisma_client).find_many(where={"team_id": {"in": user_row.teams}}) + teams_to_update = list[_TeamCleanupRow]() for team in fetch_all_teams: is_member_in_team, new_team_members = _cleanup_members_with_roles( existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()), @@ -2327,7 +2463,7 @@ async def delete_user( ## update teams for team in teams_to_update: - await TeamRepository(prisma_client).table.update( + await _team_cleanup_table(prisma_client).update( where={"team_id": team.team_id}, data={"members_with_roles": team.members_with_roles}, ) @@ -2382,14 +2518,14 @@ async def add_internal_user_to_organization( try: # Check if organization_id exists - organization_row: Final = await OrganizationRepository(prisma_client).table.find_unique( + organization_row: Final = await _organization_table(prisma_client).find_unique( where={"organization_id": organization_id} ) if organization_row is None: raise Exception(f"Organization not found, passed organization_id={organization_id}") # Create a new organization membership entry - new_membership: Final = await OrganizationMembershipRepository(prisma_client).table.create( + new_membership: Final = await _organization_membership_table(prisma_client).create( data={ "user_id": user_id, "organization_id": organization_id, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2a385c4c42a..8fe388d2fea 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,9 +18,10 @@ 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 Any, Final, Literal, Optional, Protocol, TypeVar, cast +from typing import Any, Final, Literal, Optional, Protocol, TypeAlias, TypeVar, cast import fastapi import yaml @@ -88,6 +89,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) +from litellm.proxy.management_helpers import object_permission_utils from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, @@ -144,6 +146,7 @@ from litellm.types.utils import ( ) _PrismaRowT = TypeVar("_PrismaRowT") +_PrismaRowCoT: Final = TypeVar("_PrismaRowCoT", covariant=True) _RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel) @@ -169,7 +172,7 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): *, where: Mapping[str, object] | None = None, include: Mapping[str, object] | None = None, - order: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, skip: int | None = None, take: int | None = None, ) -> list[_PrismaRowT]: ... @@ -189,17 +192,73 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): data: Mapping[str, object], ) -> _PrismaRowT | None: ... + async def upsert( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _PrismaRowT: ... -class _UserRowLike(Protocol): - user_id: str | None - user_email: str | None - user_alias: str | None - def model_dump(self) -> Mapping[str, object]: ... +class _PrismaTableHolder(Protocol[_PrismaRowT]): + @property + def table(self) -> _PrismaTableActions[_PrismaRowT]: ... + +def _typed_table(holder: _PrismaTableHolder[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]: + return holder.table + + +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() + + +class _PrismaTableLenient(Protocol[_PrismaRowCoT]): + async def find_unique(self, *, where: Mapping[str, object]) -> _PrismaRowCoT: ... + + async def find_many(self, *, where: Mapping[str, object] | None = None) -> Sequence[_PrismaRowCoT] | None: ... + + +class _PrismaTableLenientHolder(Protocol[_PrismaRowCoT]): + @property + def table(self) -> _PrismaTableLenient[_PrismaRowCoT]: ... + + +def _lenient_table(holder: _PrismaTableLenientHolder[_PrismaRowCoT]) -> _PrismaTableLenient[_PrismaRowCoT]: + return holder.table + + +def _prisma_table_lenient( + repository: BaseRepository[_RepositoryModelT], +) -> _PrismaTableLenient[_RepositoryModelT]: + return _lenient_table(repository) + + +def _jsonify_for_db(client: PrismaClient, data: Mapping[str, object]) -> Mapping[str, object]: + return client.jsonify_object(dict[str, object](data)) + + class _TxTables(Protocol): litellm_proxymodeltable: _PrismaTableActions[object] @@ -207,21 +266,82 @@ class _TxTables(Protocol): def _prisma_table( repository: BaseRepository[_RepositoryModelT], ) -> _PrismaTableActions[_RepositoryModelT]: - return repository.table + return _typed_table(repository) def _deleted_verification_token_table( prisma_client: PrismaClient, ) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]: - return DeletedVerificationTokenRepository(prisma_client).table + return _typed_table(DeletedVerificationTokenRepository(prisma_client)) def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]: - return CredentialsRepository(prisma_client).table + return _typed_table(CredentialsRepository(prisma_client)) def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]: - return ConfigRepository(prisma_client).table + return _typed_table(ConfigRepository(prisma_client)) + + +def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]: + return _typed_table(DeprecatedVerificationTokenRepository(prisma_client)) + + +_StringList: TypeAlias = list[str] + + +class _CreatedUserRow(Protocol): + models: _StringList + + +def _created_user_row(user_row: "_CreatedUserRow | None") -> "_CreatedUserRow | None": + return user_row + + +async def _query_raw_text_rows(prisma_client: PrismaClient, sql: str, *params: object) -> Sequence[Mapping[str, str]]: + return await prisma_client.db.query_raw(sql, *params) + + +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 _SpendCache(Protocol): + async def async_get_cache(self, key: str) -> float | None: ... + + +def _spend_cache(cache: _SpendCache) -> _SpendCache: + return cache + + +class _ObjectPermissionUtils(Protocol): + @property + def attach_object_permission_to_dict( + self, + ) -> Callable[..., Awaitable[Mapping[str, object]]]: ... + + +def _object_permission_utils(module: _ObjectPermissionUtils) -> _ObjectPermissionUtils: + return module + + +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: @@ -886,7 +1006,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 and key in [ "max_budget", @@ -984,9 +1104,9 @@ async def _common_key_generation_helper( soft_budget=data.soft_budget, model_max_budget=data.model_max_budget or {}, ) - new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + new_budget: Final = _jsonify_for_db(prisma_client, budget_row.json(exclude_none=True)) - _budget: Final = await BudgetRepository(prisma_client).table.create( + _budget: Final = await _prisma_table(BudgetRepository(prisma_client)).create( data={ **new_budget, "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -1655,11 +1775,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: @@ -1686,7 +1806,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): @@ -1855,11 +1975,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: @@ -1887,7 +2007,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) @@ -1959,7 +2081,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(): @@ -2737,13 +2859,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: @@ -2774,7 +2896,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) @@ -2927,14 +3051,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, @@ -2980,7 +3106,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( @@ -3058,7 +3184,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) @@ -3089,14 +3215,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, @@ -3223,7 +3351,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], ) @@ -3435,7 +3563,7 @@ async def _get_model_max_budget_current_spend( virtual_key_model_spend_cache_key = ( f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}" ) - current_spend: float | None = await user_api_key_cache.async_get_cache( + current_spend: float | None = await _spend_cache(user_api_key_cache).async_get_cache( key=virtual_key_model_spend_cache_key, ) if current_spend is None: @@ -3444,7 +3572,7 @@ async def _get_model_max_budget_current_spend( f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:" f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}" ) - current_spend = await user_api_key_cache.async_get_cache( + current_spend = await _spend_cache(user_api_key_cache).async_get_cache( key=virtual_key_model_spend_cache_key, ) try: @@ -3635,7 +3763,7 @@ async def info_key_fn( hashed_key: str | None = key if key is not None: hashed_key = _hash_token_if_needed(token=key) - key_info = await VerificationTokenRepository(prisma_client).table.find_unique( + key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_key}, include={"litellm_budget_table": True}, ) @@ -3945,7 +4073,7 @@ async def generate_key_helper_fn( if table_name is None or table_name == "user": # do not auto-create users for `/key/generate` ## CREATE USER (If necessary) if query_type == "insert_data": - user_row = await prisma_client.insert_data(data=user_data, table_name="user") + user_row = _created_user_row(await prisma_client.insert_data(data=user_data, table_name="user")) if user_row is None: raise Exception("Failed to create user") @@ -4232,7 +4360,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) @@ -4355,13 +4483,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( @@ -4376,14 +4503,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, @@ -4449,7 +4576,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: @@ -4605,7 +4732,7 @@ async def _insert_deprecated_key( try: revoke_at: Final = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds) - await DeprecatedVerificationTokenRepository(prisma_client).table.upsert( + await _deprecated_verification_token_table(prisma_client).upsert( where={"token": old_token_hash}, data={ "create": { @@ -4704,11 +4831,13 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) - updated_token: Final[Mapping[str, object] | None] = await VerificationTokenRepository(prisma_client).table.update( + updated_token: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_api_key}, data=with_settings_updated_at(jsonified_update_data), ) - updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {} + updated_token_dict: Final = ( + dict[str, object](_as_object_dict(dict(updated_token))) if updated_token is not None else dict[str, object]() + ) updated_token_dict["key"] = new_token updated_token_dict["token_id"] = updated_token_dict.pop("token") @@ -5247,7 +5376,7 @@ async def validate_key_list_check( if key_hash: try: - key_info: Final = await VerificationTokenRepository(prisma_client).table.find_unique( + key_info: Final = await _prisma_table_lenient(VerificationTokenRepository(prisma_client)).find_unique( where={"token": key_hash}, ) except Exception: @@ -5278,7 +5407,7 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Final[list[BaseModel] | None] = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _prisma_table_lenient(TeamRepository(prisma_client)).find_many( where={"team_id": {"in": complete_user_info.teams}} ) if teams is None: @@ -5653,7 +5782,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final = await _query_raw_text_rows(prisma_client, count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -5666,7 +5795,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final = await _query_raw_text_rows(prisma_client, aliases_sql, *aliases_params) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -5952,7 +6081,7 @@ async def _list_key_helper( # Fetch keys with pagination if use_deleted_table: - keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many( + keys = await _deleted_verification_token_table(prisma_client).find_many( where=where, skip=skip, take=size, @@ -5966,7 +6095,7 @@ async def _list_key_helper( ), ) else: - keys = await VerificationTokenRepository(prisma_client).table.find_many( + keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where=where, skip=skip, take=size, @@ -5995,13 +6124,13 @@ async def _list_key_helper( total_pages: Final = -(-total_count // size) # Ceiling division # Fetch user information if expand includes "user" - user_map = {} + user_map: Mapping[str, LiteLLM_UserTable] = {} if expand and "user" in expand: user_ids: Final = [key.user_id for key in keys if key.user_id] created_by_ids: Final = [key.created_by for key in keys if key.created_by] all_ids: Final = list(set(user_ids + created_by_ids)) # Remove duplicates if all_ids: - users: Final[Sequence[_UserRowLike]] = await UserRepository(prisma_client).table.find_many( + users: Final = await _prisma_table(UserRepository(prisma_client)).find_many( where={"user_id": {"in": all_ids}} ) user_map = {user.user_id: user for user in users} @@ -6014,10 +6143,14 @@ async def _list_key_helper( key_dict = key.model_dump() except Exception: # Fallback for Pydantic v1 compatibility - key_dict = key.dict() + key_dict = dict[str, object](_legacy_model_dict(key)) # Attach object_permission if object_permission_id is set (only for non-deleted keys) if not use_deleted_table: - key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) + key_dict = dict[str, object]( + await _object_permission_utils(object_permission_utils).attach_object_permission_to_dict( + key_dict, prisma_client + ) + ) # Include user information if expand includes "user" if expand and "user" in expand: @@ -6025,7 +6158,7 @@ async def _list_key_helper( try: key_dict["user"] = user_map[key.user_id].model_dump() except Exception: - key_dict["user"] = user_map[key.user_id].dict() + key_dict["user"] = _legacy_model_dict(user_map[key.user_id]) if key.created_by and key.created_by in user_map: created_by_user = user_map[key.created_by] key_dict["created_by_user"] = { @@ -6039,7 +6172,7 @@ async def _list_key_helper( # Use deleted key type to preserve deleted_at, deleted_by, etc. key_list.append(LiteLLM_DeletedVerificationToken.model_validate(key_dict)) else: - key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object + key_list.append(UserAPIKeyAuth.model_validate(key_dict)) # Return full key object else: _token = key_dict.get("token") key_list.append(cast(str, _token)) # Return only the token diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 0fdedafb2bf..dfb1422308f 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -10,8 +10,9 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a """ import uuid +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Annotated, Any, Final +from typing import TYPE_CHECKING, Annotated, Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, Field, TypeAdapter @@ -49,6 +50,144 @@ from litellm.types.tool_management import ( ToolUsageLogsResponse, ) + +class _DailyToolSpendRecord(Protocol): + date: str + tool_name: str + spend: float + request_count: int + + +class _SpendLogToolIndexRecord(Protocol): + request_id: str + + +class _SpendLogRecord(Protocol): + request_id: str + startTime: datetime + model: str | None + spend: float | None + total_tokens: int | None + messages: object + proxy_server_request: object + + +class _VerificationTokenRecord(Protocol): + object_permission_id: str | None + + +class _TeamRecord(Protocol): + object_permission_id: str | None + + +class _DailyToolSpendTable(Protocol): + async def group_by( + self, + *, + by: Sequence[str], + sum: Mapping[str, bool], + where: Mapping[str, object], + order: Mapping[str, object], + take: int, + ) -> Sequence[object] | None: ... + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Sequence[Mapping[str, str]], + ) -> Sequence[_DailyToolSpendRecord]: ... + + +class _SpendLogToolIndexTable(Protocol): + async def count(self, *, where: Mapping[str, object]) -> int: ... + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str], + skip: int, + take: int, + ) -> Sequence[_SpendLogToolIndexRecord]: ... + + +class _SpendLogsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_SpendLogRecord]: ... + + +class _VerificationTokenTable(Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> _VerificationTokenRecord | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _TeamTable(Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _ObjectPermissionTable(Protocol): + async def create(self, *, data: Mapping[str, str | Sequence[str]]) -> object: ... + + async def delete(self, *, where: Mapping[str, object]) -> object: ... + + +class _DailyToolSpendTableHolder(Protocol): + @property + def table(self) -> _DailyToolSpendTable: ... + + +class _SpendLogToolIndexTableHolder(Protocol): + @property + def table(self) -> _SpendLogToolIndexTable: ... + + +class _SpendLogsTableHolder(Protocol): + @property + def table(self) -> _SpendLogsTable: ... + + +class _VerificationTokenTableHolder(Protocol): + @property + def table(self) -> _VerificationTokenTable: ... + + +class _TeamTableHolder(Protocol): + @property + def table(self) -> _TeamTable: ... + + +class _ObjectPermissionTableHolder(Protocol): + @property + def table(self) -> _ObjectPermissionTable: ... + + +def _daily_tool_spend_table(repo: _DailyToolSpendTableHolder) -> _DailyToolSpendTable: + return repo.table + + +def _spend_log_tool_index_table(repo: _SpendLogToolIndexTableHolder) -> _SpendLogToolIndexTable: + return repo.table + + +def _spend_logs_table(repo: _SpendLogsTableHolder) -> _SpendLogsTable: + return repo.table + + +def _verification_token_table(repo: _VerificationTokenTableHolder) -> _VerificationTokenTable: + return repo.table + + +def _team_table(repo: _TeamTableHolder) -> _TeamTable: + return repo.table + + +def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable: + return repo.table + + router: Final = APIRouter() TOOL_POLICY_OPTIONS: Final = ToolPolicyOptionsResponse( @@ -154,6 +293,7 @@ class _TopToolRow(BaseModel): _TOP_TOOL_ROWS: Final = TypeAdapter(list[_TopToolRow]) +_PARSED_JSON: Final = TypeAdapter(object) @router.get( @@ -201,7 +341,7 @@ async def get_tool_spend( end_str: Final = end_day.strftime("%Y-%m-%d") date_window: Final = {"date": {"gte": start_str, "lte": end_str}} - table: Final = DailyToolSpendRepository(prisma_client).table + table: Final = _daily_tool_spend_table(DailyToolSpendRepository(prisma_client)) top_tools: Final = _TOP_TOOL_ROWS.validate_python( await table.group_by( by=["tool_name"], @@ -222,7 +362,7 @@ async def get_tool_spend( for row in top_tools ] - daily_rows: Final = ( + daily_rows: Final[Sequence[_DailyToolSpendRecord]] = ( await table.find_many( where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}}, order=[{"date": "asc"}, {"spend": "desc"}], @@ -270,23 +410,23 @@ async def get_tool_detail( raise HTTPException(status_code=500, detail=str(e)) -def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None: +def _input_snippet_for_tool_log(sl: _SpendLogRecord | None, max_len: int = 200) -> str | None: """Short snippet from messages or proxy_server_request for tool usage log row.""" if sl is None: return None - messages: Final = getattr(sl, "messages", None) + messages: Final[object] = getattr(sl, "messages", None) if messages is not None: s = _snippet_str(messages, max_len) if s: return s - psr = getattr(sl, "proxy_server_request", None) + psr: object = getattr(sl, "proxy_server_request", None) if not psr: return None if isinstance(psr, str): import json try: - psr = json.loads(psr) + psr = _PARSED_JSON.validate_python(json.loads(psr)) except Exception: return _snippet_str(psr, max_len) if isinstance(psr, dict): @@ -299,7 +439,7 @@ def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None: return _snippet_str(psr, max_len) -def _snippet_str(text: Any, max_len: int = 200) -> str | None: +def _snippet_str(text: object, max_len: int = 200) -> str | None: if text is None: return None if isinstance(text, str): @@ -344,10 +484,9 @@ async def get_tool_usage_logs( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - where: Final[dict] = {"tool_name": tool_name} + start_time_filter: datetime | None = None + end_time_filter: datetime | None = None if start_date or end_date: - start_time_filter: datetime | None = None - end_time_filter: datetime | None = None if start_date: try: start_time_filter = datetime.strptime(start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S").replace( @@ -362,15 +501,15 @@ async def get_tool_usage_logs( ) except ValueError: pass - if start_time_filter is not None or end_time_filter is not None: - where["start_time"] = {} - if start_time_filter is not None: - where["start_time"]["gte"] = start_time_filter - if end_time_filter is not None: - where["start_time"]["lte"] = end_time_filter + start_time_range: Final[Mapping[str, datetime]] = { + key: value for key, value in (("gte", start_time_filter), ("lte", end_time_filter)) if value is not None + } + where: Final[Mapping[str, str | Mapping[str, datetime]]] = ( + {"tool_name": tool_name, "start_time": start_time_range} if start_time_range else {"tool_name": tool_name} + ) - total: Final = await SpendLogToolIndexRepository(prisma_client).table.count(where=where) - index_rows: Final = await SpendLogToolIndexRepository(prisma_client).table.find_many( + total: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).count(where=where) + index_rows: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, @@ -380,7 +519,9 @@ async def get_tool_usage_logs( if not request_ids: return ToolUsageLogsResponse(logs=[], total=total, page=page, page_size=page_size) - spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}}) + spend_logs = await _spend_logs_table(SpendLogsRepository(prisma_client)).find_many( + where={"request_id": {"in": request_ids}} + ) log_by_id: Final = {s.request_id: s for s in spend_logs} logs_out: Final[list[ToolUsageLogEntry]] = [] @@ -449,23 +590,29 @@ async def _resolve_key_hash_to_object_permission_id( hashed: Final = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) if not hashed: return None - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) + row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed} + ) if row is None: return None - op_id: Final = getattr(row, "object_permission_id", None) + op_id: Final[str | None] = getattr(row, "object_permission_id", None) if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count: Final = await VerificationTokenRepository(prisma_client).table.update_many( + updated_count: Final = await _verification_token_table(VerificationTokenRepository(prisma_client)).update_many( where={"token": hashed, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) + await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete( + where={"object_permission_id": new_id} + ) + row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed} + ) return getattr(row, "object_permission_id", None) if row else None return new_id @@ -478,23 +625,25 @@ async def _resolve_team_id_to_object_permission_id( if not team_id or not team_id.strip(): return None team_id_clean: Final = team_id.strip() - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) + row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean}) if row is None: return None - op_id: Final = getattr(row, "object_permission_id", None) + op_id: Final[str | None] = getattr(row, "object_permission_id", None) if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count: Final = await TeamRepository(prisma_client).table.update_many( + updated_count: Final = await _team_table(TeamRepository(prisma_client)).update_many( where={"team_id": team_id_clean, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) + await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete( + where={"object_permission_id": new_id} + ) + row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean}) return getattr(row, "object_permission_id", None) if row else None return new_id diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index d8e9f8dfaee..d1dc1482401 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -3,8 +3,9 @@ CRUD ENDPOINTS FOR PROMPTS """ import tempfile +from collections.abc import Mapping, Sequence from pathlib import Path -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, Protocol, cast from fastapi import ( APIRouter, @@ -15,7 +16,7 @@ from fastapi import ( Response, UploadFile, ) -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -38,8 +39,47 @@ from litellm.types.prompts.init_prompts import ( ) from litellm.types.proxy.prompt_endpoints import TestPromptRequest +if TYPE_CHECKING: + from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry + from litellm.proxy.utils import PrismaClient + + +class _PromptRecord(Protocol): + id: str + version: int + environment: str | None + + +class _PromptTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str] | None = None, + take: int | None = None, + distinct: Sequence[str] | None = None, + ) -> Sequence[_PromptRecord]: ... + + async def create(self, *, data: Mapping[str, str | int | None]) -> _PromptRecord: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> _PromptRecord: ... + + async def delete_many(self, *, where: Mapping[str, str]) -> object: ... + + +class _PromptTableHolder(Protocol): + @property + def table(self) -> _PromptTable: ... + + +def _prompt_table(repo: _PromptTableHolder) -> _PromptTable: + return repo.table + + router: Final = APIRouter() +_PARSED_VALUE: Final = TypeAdapter(object) + def get_base_prompt_id(prompt_id: str) -> str: """ @@ -132,7 +172,7 @@ def construct_versioned_prompt_id(prompt_id: str, version: int | None = None) -> return f"{base_id}.v{version}" -def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: dict[str, Any]) -> str: +def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Mapping[str, object]) -> str: """ Find the latest version of a prompt from available prompt IDs. @@ -198,7 +238,9 @@ def get_latest_prompt_versions(prompts: list[PromptSpec]) -> list[PromptSpec]: return list(latest_prompts.values()) -async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment: str = "development") -> int: +async def get_next_version_for_prompt( + prisma_client: "PrismaClient", prompt_id: str, environment: str = "development" +) -> int: """ Get the next version number for a prompt in a specific environment. @@ -210,7 +252,7 @@ async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment Returns: Next version number (1 if no versions exist, max_version + 1 otherwise) """ - existing_prompts: Final = await PromptRepository(prisma_client).table.find_many( + existing_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where={"prompt_id": prompt_id, "environment": environment} ) @@ -431,10 +473,10 @@ async def get_prompt_versions( # Query DB for versions versioned_prompts: Final = [] if prisma_client is not None: - where_clause: Final[dict[str, Any]] = {"prompt_id": base_prompt_id} - if environment: - where_clause["environment"] = environment - db_prompts: Final = await PromptRepository(prisma_client).table.find_many( + where_clause: Final[Mapping[str, str]] = ( + {"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id} + ) + db_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=where_clause, order={"version": "desc"}, ) @@ -590,7 +632,7 @@ async def get_prompt_info( # Query all environments this prompt exists in (lightweight: distinct on environment) all_environments: list[str] = [] if prisma_client is not None: - all_prompt_rows: Final = await PromptRepository(prisma_client).table.find_many( + all_prompt_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where={"prompt_id": base_prompt_id}, distinct=["environment"], ) @@ -602,13 +644,16 @@ async def get_prompt_info( prompt_spec = None requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None if environment and prisma_client is not None: - where_clause: Final[dict[str, Any]] = { - "prompt_id": base_prompt_id, - "environment": environment, + where_clause: Final[Mapping[str, str | int]] = { + key: value + for key, value in ( + ("prompt_id", base_prompt_id), + ("environment", environment), + ("version", requested_version), + ) + if value is not None } - if requested_version is not None: - where_clause["version"] = requested_version - env_prompts: Final = await PromptRepository(prisma_client).table.find_many( + env_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=where_clause, order={"version": "desc"}, take=1, @@ -721,7 +766,7 @@ async def create_prompt( ) # Store prompt in db with version - prompt_db_entry: Final = await PromptRepository(prisma_client).table.create( + prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create( data={ "prompt_id": request.prompt_id, "version": new_version, @@ -811,7 +856,9 @@ async def update_prompt( ) # Check if any version of this prompt exists (in any environment) - existing_prompts = await PromptRepository(prisma_client).table.find_many(where={"prompt_id": base_prompt_id}) + existing_prompts = await _prompt_table(PromptRepository(prisma_client)).find_many( + where={"prompt_id": base_prompt_id} + ) if not existing_prompts: raise HTTPException( @@ -835,7 +882,7 @@ async def update_prompt( ) # Store new version in db - prompt_db_entry: Final = await PromptRepository(prisma_client).table.create( + prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create( data={ "prompt_id": base_prompt_id, "version": new_version, @@ -936,12 +983,12 @@ async def delete_prompt( base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) # Build delete filter; scope to environment if provided - delete_where: Final[dict[str, Any]] = {"prompt_id": base_prompt_id} - if environment: - delete_where["environment"] = environment + delete_where: Final[Mapping[str, str]] = ( + {"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id} + ) # Delete versions from the database (scoped to environment if provided) - await PromptRepository(prisma_client).table.delete_many(where=delete_where) + await _prompt_table(PromptRepository(prisma_client)).delete_many(where=delete_where) # Remove matching prompts from memory — scope to environment if provided if environment: @@ -967,7 +1014,9 @@ async def delete_prompt( raise HTTPException(status_code=500, detail=str(e)) -def _reload_prompt_in_registry(registry: Any, versioned_id: str, updated_prompt_spec: PromptSpec) -> PromptSpec: +def _reload_prompt_in_registry( + registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec +) -> PromptSpec: """Remove stale entry and re-initialize the prompt in the in-memory registry.""" if versioned_id in registry.IN_MEMORY_PROMPTS: del registry.IN_MEMORY_PROMPTS[versioned_id] @@ -1033,14 +1082,13 @@ async def patch_prompt( requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None # Build query to find the exact row by composite unique key - find_where: Final[dict[str, Any]] = { - "prompt_id": base_prompt_id, - "environment": env, + find_where: Final[Mapping[str, str | int]] = { + key: value + for key, value in (("prompt_id", base_prompt_id), ("environment", env), ("version", requested_version)) + if value is not None } - if requested_version is not None: - find_where["version"] = requested_version - db_rows: Final = await PromptRepository(prisma_client).table.find_many( + db_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=find_where, order={"version": "desc"}, take=1, @@ -1084,15 +1132,18 @@ async def patch_prompt( raise HTTPException(status_code=400, detail="litellm_params cannot be None") # Build update data dict - update_data: Final[dict[str, Any]] = { - "litellm_params": updated_litellm_params.model_dump_json(), - "prompt_info": updated_prompt_info.model_dump_json(), + update_data: Final[Mapping[str, str]] = { + key: value + for key, value in ( + ("litellm_params", updated_litellm_params.model_dump_json()), + ("prompt_info", updated_prompt_info.model_dump_json()), + ("created_by", user_api_key_dict.user_id), + ) + if value } - if user_api_key_dict.user_id: - update_data["created_by"] = user_api_key_dict.user_id # Update by primary key (id) to target exactly one row - updated_prompt_db_entry: Final = await PromptRepository(prisma_client).table.update( + updated_prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).update( where={"id": target_row.id}, data=update_data, ) @@ -1216,23 +1267,25 @@ async def test_prompt( # Use ProxyBaseLLMRequestProcessing to go through all proxy logic base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) - result: Final = await base_llm_response_processor.base_process_llm_request( - request=fastapi_request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="acompletion", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, + result: Final = _PARSED_VALUE.validate_python( + await base_llm_response_processor.base_process_llm_request( + request=fastapi_request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) ) if isinstance(result, BaseModel): @@ -1257,7 +1310,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 383ada5a1bc..a8ec8884d9b 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,9 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from typing import Any, Final, cast +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 TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -20,6 +23,56 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus +if TYPE_CHECKING: + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + +_JsonDict: TypeAlias = dict[str, object] +_JsonList: TypeAlias = list[object] + + +class _OutputItem(TypedDict, total=False): + id: str + content: Sequence[object] + + +class _TerminalResponse(TypedDict, total=False): + status: ResponsesAPIStatus + error: _JsonDict + usage: _JsonDict + reasoning: _JsonDict + tool_choice: object + tools: _JsonList + model: str + instructions: str + temperature: float + top_p: float + max_output_tokens: int + previous_response_id: str + text: _JsonDict + truncation: str + parallel_tool_calls: bool + user: str + store: bool + incomplete_details: _JsonDict + output: Sequence[_OutputItem] + + +class _StreamEvent(TypedDict, total=False): + type: str + item: _OutputItem + item_id: str + content_index: int + delta: str + part: object + response: _TerminalResponse + + +class _StreamEventParser: + parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) + async def background_streaming_task( polling_id: str, @@ -29,16 +82,16 @@ async def background_streaming_task( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, general_settings: dict, - llm_router, - proxy_config, - proxy_logging_obj, + llm_router: "Router | None", + proxy_config: "ProxyConfig", + proxy_logging_obj: "ProxyLogging", select_data_generator, user_model, - user_temperature, - user_request_timeout, - user_max_tokens, - user_api_base, - version, + user_temperature: float | None, + user_request_timeout: float | None, + user_max_tokens: int | None, + user_api_base: str | None, + version: str | None, ): """ Background task to stream response and update cache @@ -69,7 +122,7 @@ async def background_streaming_task( # Make streaming request. # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. - response: Final = await processor.base_process_llm_request( + response: Final[StreamingResponse] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -91,8 +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, dict[str, Any]]] = {} # Track output items by ID - accumulated_text: Final = {} # Track accumulated text deltas by (item_id, content_index) + 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 @@ -121,7 +176,7 @@ async def background_streaming_task( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None - _event_to_status: Final = { + _event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = { "response.completed": "completed", "response.failed": "failed", "response.incomplete": "incomplete", @@ -162,7 +217,7 @@ async def background_streaming_task( break try: - event = json.loads(chunk_data) + event: _StreamEvent = _StreamEventParser.parse(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -181,9 +236,8 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update the output item with new content - if "content" not in output_items[item_id]: - output_items[item_id]["content"] = [] - output_items[item_id]["content"].append(content_part) + added_item = output_items[item_id] + added_item["content"] = (*added_item.get("content", ()), content_part) state_dirty = True elif event_type == "response.output_text.delta": @@ -201,12 +255,14 @@ async def background_streaming_task( accumulated_text[key] += delta # Update the content in output_items - if "content" in output_items[item_id]: - content_list = output_items[item_id]["content"] + 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 - if isinstance(content_list[content_index], dict): - content_list[content_index]["text"] = accumulated_text[key] + 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": @@ -217,10 +273,14 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update with final content from event - if "content" in output_items[item_id]: - content_list = output_items[item_id]["content"] + done_item = output_items[item_id] + if "content" in done_item: + content_list = done_item["content"] if content_index < len(content_list): - content_list[content_index] = content_part + 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": @@ -248,12 +308,9 @@ async def background_streaming_task( # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming response_data = event.get("response", {}) - terminal_status = cast( - ResponsesAPIStatus, - response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ), + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), ) # Extract error for failed and incomplete responses diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7854b17a06f..205189a6043 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -14,16 +14,19 @@ Flow: import json import time import uuid -from collections.abc import Iterable -from typing import Any, Final, cast +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, TypedDict, cast from litellm._internal_context import is_internal_call from litellm._logging import verbose_logger -from litellm.types.llms.openai import ResponseOutputItem, ResponsesAPIResponse +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.vector_stores import VectorStoreSearchResult +if TYPE_CHECKING: + from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig + # Keep ToolParam broad so we stay compatible with both dict and Pydantic forms -ToolParam = Any +ToolParam = object FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" @@ -35,7 +38,7 @@ FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" def should_use_emulated_file_search( tools: Iterable[ToolParam] | None, - provider_config: Any, # BaseResponsesAPIConfig + provider_config: "BaseResponsesAPIConfig | None", ) -> bool: """Return True when there is a file_search tool and the provider can't handle it natively.""" if not tools: @@ -51,7 +54,7 @@ def should_use_emulated_file_search( # --------------------------------------------------------------------------- -def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: +def _build_function_tool(vector_store_ids: Sequence[str]) -> Mapping[str, object]: """ Create a Responses API function-tool definition that describes file search. The function accepts one or more natural-language queries (like OpenAI's native @@ -94,27 +97,26 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: } +def _file_search_tool_vector_store_ids(tool: object) -> Sequence[str] | None: + if not (isinstance(tool, dict) and tool.get("type") == "file_search"): + return None + return tool.get("vector_store_ids") or [] + + def _replace_file_search_tools( tools: Iterable[ToolParam] | None, -) -> tuple[list[dict[str, Any]], list[str]]: +) -> tuple[Sequence[object], Sequence[str]]: """ Replace all file_search tools with a single function tool. Returns: (new_tools_list, all_vector_store_ids) """ - non_file_search: Final[list[dict[str, Any]]] = [] - vector_store_ids: Final[list[str]] = [] - - for tool in tools or []: - if isinstance(tool, dict) and tool.get("type") == "file_search": - ids = tool.get("vector_store_ids") or [] - vector_store_ids.extend(ids) - else: - non_file_search.append(tool) + ids_and_tools: Final = tuple((_file_search_tool_vector_store_ids(tool), tool) for tool in tools or ()) # Deduplicate while preserving order - unique_ids: Final[list[str]] = list(dict.fromkeys(vector_store_ids)) + unique_ids: Final = list(dict.fromkeys(vs_id for ids, _ in ids_and_tools if ids is not None for vs_id in ids)) + non_file_search: Final = [tool for ids, tool in ids_and_tools if ids is None] if unique_ids: non_file_search.append(_build_function_tool(unique_ids)) @@ -127,9 +129,9 @@ def _replace_file_search_tools( async def _run_vector_searches( - queries: list[str], - vector_store_ids: list[str], -) -> tuple[list[str], list[VectorStoreSearchResult]]: + queries: Sequence[str], + vector_store_ids: Sequence[str], +) -> tuple[Sequence[str], Sequence[VectorStoreSearchResult]]: """ Run `asearch` against all vector stores for all queries and collect results. @@ -172,7 +174,7 @@ async def _run_vector_searches( # --------------------------------------------------------------------------- -def _get_field(result: Any, key: str, default: Any = None) -> Any: +def _get_field(result: object, key: str, default: object = None) -> object: """Read a field from either a dict/TypedDict or an attribute-based object.""" if isinstance(result, dict): return result.get(key, default) @@ -180,7 +182,7 @@ def _get_field(result: Any, key: str, default: Any = None) -> Any: def _format_search_results_as_tool_output( - results: list[VectorStoreSearchResult], + results: Sequence[VectorStoreSearchResult], ) -> str: """Serialize search results into a string to pass back as the tool's output.""" if not results: @@ -191,7 +193,8 @@ def _format_search_results_as_tool_output( score = _get_field(result, "score") file_id = _get_field(result, "file_id") filename = _get_field(result, "filename") - content_items = _get_field(result, "content") or [] + raw_content = _get_field(result, "content") + content_items = raw_content if isinstance(raw_content, list) else [] text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] text = " ".join(t for t in text_chunks if t) @@ -209,9 +212,24 @@ def _format_search_results_as_tool_output( return "\n\n".join(parts) +def _format_result_for_include(result: VectorStoreSearchResult) -> Mapping[str, object]: + file_id: Final = _get_field(result, "file_id") or "" + raw_content: Final = _get_field(result, "content") + content_items: Final = raw_content if isinstance(raw_content, list) else [] + text_chunks: Final = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] + text: Final = " ".join(t for t in text_chunks if t) + return { + "file_id": file_id, + "filename": _get_field(result, "filename") or "", + "score": _get_field(result, "score"), + "text": text, + "attributes": _get_field(result, "attributes") or {}, + } + + def _build_search_results_for_include( - results: list[VectorStoreSearchResult], -) -> list[dict[str, Any]]: + results: Sequence[VectorStoreSearchResult], +) -> Sequence[Mapping[str, object]]: """ Convert VectorStoreSearchResult objects to the format expected in file_search_call.search_results (mirrors OpenAI's include= format). @@ -220,30 +238,15 @@ def _build_search_results_for_include( behaviour of OpenAI's native file_search which surfaces every relevant chunk even when multiple chunks originate from the same document. """ - formatted: Final[list[dict[str, Any]]] = [] - for result in results: - file_id = _get_field(result, "file_id") or "" - content_items = _get_field(result, "content") or [] - text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] - text = " ".join(t for t in text_chunks if t) - formatted.append( - { - "file_id": file_id, - "filename": _get_field(result, "filename") or "", - "score": _get_field(result, "score"), - "text": text, - "attributes": _get_field(result, "attributes") or {}, - } - ) - return formatted + return [_format_result_for_include(result) for result in results] def _build_file_search_call_output( call_id: str, - queries: list[str], - results: list[VectorStoreSearchResult] | None = None, + queries: Sequence[str], + results: Sequence[VectorStoreSearchResult] | None = None, include_search_results: bool = False, -) -> dict[str, Any]: +) -> Mapping[str, object]: """Build the file_search_call output item (mirrors OpenAI's format). Args: @@ -266,39 +269,34 @@ def _build_file_search_call_output( def _build_file_citation_annotations( - results: list[VectorStoreSearchResult], + results: Sequence[VectorStoreSearchResult], text: str, -) -> list[dict[str, Any]]: +) -> Sequence[Mapping[str, object]]: """ Build file_citation annotations for the text. Each result with a file_id gets a citation at the end of the text. """ - annotations: Final[list[dict[str, Any]]] = [] index: Final = len(text) # cite at end of text block - seen_file_ids: Final[set] = set() + id_filename_pairs: Final = tuple( + (_get_field(result, "file_id"), _get_field(result, "filename")) for result in results + ) + first_filename_by_id: Final = {file_id: filename for file_id, filename in reversed(id_filename_pairs) if file_id} - for result in results: - file_id = _get_field(result, "file_id") - filename = _get_field(result, "filename") - if not file_id or file_id in seen_file_ids: - continue - seen_file_ids.add(file_id) - annotations.append( - { - "type": "file_citation", - "index": index, - "file_id": file_id, - "filename": filename or "", - } - ) - - return annotations + return [ + { + "type": "file_citation", + "index": index, + "file_id": file_id, + "filename": first_filename_by_id[file_id] or "", + } + for file_id in dict.fromkeys(file_id for file_id, _ in id_filename_pairs if file_id) + ] def _build_message_output( response_text: str, - results: list[VectorStoreSearchResult], -) -> dict[str, Any]: + results: Sequence[VectorStoreSearchResult], +) -> Mapping[str, object]: """Build the message output item with optional file_citation annotations.""" annotations: Final = _build_file_citation_annotations(results, response_text) return { @@ -330,8 +328,8 @@ def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str: def _synthesize_responses_api_response( original_response: ResponsesAPIResponse, - file_search_call_output: dict[str, Any], - message_output: dict[str, Any], + file_search_call_output: Mapping[str, object], + message_output: Mapping[str, object], first_response: ResponsesAPIResponse | None = None, ) -> ResponsesAPIResponse: """ @@ -343,21 +341,20 @@ def _synthesize_responses_api_response( synthesized _hidden_params so that billing callbacks see the total cost of both provider calls that the emulated flow makes. """ - synthesized_output: Final[list[dict[str, Any]]] = [file_search_call_output, message_output] synthesized: Final = ResponsesAPIResponse( id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"), object="response", created_at=getattr(original_response, "created_at", int(time.time())), status="completed", model=getattr(original_response, "model", ""), - output=cast(list[ResponseOutputItem | dict[str, Any]], synthesized_output), + output=[dict(file_search_call_output), dict(message_output)], usage=getattr(original_response, "usage", None), error=None, ) if hasattr(original_response, "_hidden_params"): hidden: Final = dict(getattr(original_response, "_hidden_params") or {}) if first_response is not None and hasattr(first_response, "_hidden_params"): - first_hidden: Final = getattr(first_response, "_hidden_params") or {} + first_hidden: Final[object] = getattr(first_response, "_hidden_params", None) or {} first_cost: Final = ( first_hidden.get("response_cost") if isinstance(first_hidden, dict) @@ -382,9 +379,10 @@ async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover def _prepare_emulated_file_search_call( - kwargs: dict[str, Any], -) -> tuple[bool, dict[str, Any]]: - include_items: Final[list[str]] = list(kwargs.get("include") or []) + kwargs: Mapping[str, object], +) -> tuple[bool, Mapping[str, object]]: + raw_include: Final = kwargs.get("include") + include_items: Final[Sequence[str]] = raw_include if isinstance(raw_include, list) else [] include_search_results: Final = "file_search_call.results" in include_items original_stream: Final = kwargs.get("stream") @@ -398,7 +396,7 @@ def _prepare_emulated_file_search_call( return include_search_results, updated_kwargs -def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[str, str]: +def _extract_tool_call_fields(tool_call: object, fallback_call_id: str) -> tuple[str, str]: """Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item.""" if isinstance(tool_call, dict): call_id = str(tool_call.get("call_id") or tool_call.get("id") or fallback_call_id) @@ -410,7 +408,13 @@ def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[st return call_id, raw_args -def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: +class _FileSearchArguments(TypedDict, total=False): + queries: Sequence[str] + query: str + vector_store_id: str + + +def _resolve_queries_from_args(args: _FileSearchArguments, input: object) -> Sequence[str]: """Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks.""" queries_from_call: Final = args.get("queries") if not queries_from_call: @@ -422,76 +426,96 @@ def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: return queries_from_call -async def _execute_file_search_tool_calls( - file_search_calls: list[Any], - all_vs_ids: list[str], - input: Any, +def _parse_file_search_arguments(raw_args: str) -> _FileSearchArguments: + if not isinstance(raw_args, str): + return raw_args + try: + return json.loads(raw_args) + except json.JSONDecodeError: + return {} + + +async def _execute_single_file_search_call( + tool_call: object, + all_vs_ids: Sequence[str], + input: object, file_search_call_id: str, -) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult]]: +) -> tuple[Mapping[str, object], Sequence[str], Sequence[VectorStoreSearchResult]]: + call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) + args: Final = _parse_file_search_arguments(raw_args) + queries_from_call: Final = _resolve_queries_from_args(args, input) + + vs_id_arg: Final = args.get("vector_store_id") + vs_ids_for_call: Final = [vs_id_arg] if vs_id_arg else all_vs_ids + + queries, results = await _run_vector_searches( + queries=queries_from_call, + vector_store_ids=vs_ids_for_call, + ) + + return ( + { + "type": "function_call_output", + "call_id": call_id, + "output": _format_search_results_as_tool_output(results), + }, + queries, + results, + ) + + +async def _execute_file_search_tool_calls( + file_search_calls: Sequence[object], + all_vs_ids: Sequence[str], + input: object, + file_search_call_id: str, +) -> tuple[Sequence[Mapping[str, object]], Sequence[str], Sequence[VectorStoreSearchResult]]: """Run the vector search for each file_search tool_call and collect results.""" - tool_results: Final[list[dict[str, Any]]] = [] - all_queries: Final[list[str]] = [] - all_results: Final[list[VectorStoreSearchResult]] = [] + per_call: Final = tuple( + [ + await _execute_single_file_search_call( + tool_call=tool_call, + all_vs_ids=all_vs_ids, + input=input, + file_search_call_id=file_search_call_id, + ) + for tool_call in file_search_calls + ] + ) - for tool_call in file_search_calls: - call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) - - try: - args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args - except json.JSONDecodeError: - args = {} - - queries_from_call = _resolve_queries_from_args(args, input) - - vs_id_arg = args.get("vector_store_id") - vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids - - queries, results = await _run_vector_searches( - queries=queries_from_call, - vector_store_ids=vs_ids_for_call, - ) - all_queries.extend(queries) - all_results.extend(results) - - tool_results.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": _format_search_results_as_tool_output(results), - } - ) - - return tool_results, all_queries, all_results + return ( + [tool_result for tool_result, _, _ in per_call], + [query for _, queries, _ in per_call for query in queries], + [result for _, _, results in per_call for result in results], + ) def _build_follow_up_input( - input: Any, + input: object, first_response: ResponsesAPIResponse, - tool_results: list[dict[str, Any]], -) -> list[Any]: + tool_results: Sequence[Mapping[str, object]], +) -> Sequence[object]: """Assemble the follow-up call input: original messages + first-response output + tool results. Including all output items (text blocks, reasoning, non-file-search calls) ensures providers like Anthropic that emit text before the tool call have complete conversation context. Serializes Pydantic model instances to plain dicts so the transformation layer can call .get(). """ - original_input_items: Final = ( - list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}] + original_input_items: Final[tuple[object, ...]] = ( + tuple(input) if isinstance(input, (list, tuple)) else ({"role": "user", "content": str(input)},) + ) + first_response_output_items: Final[tuple[object, ...]] = tuple( + _item + if isinstance(_item, dict) + else (_item.model_dump(exclude_none=True) if hasattr(_item, "model_dump") else _item) + for _item in first_response.output ) - first_response_output_items: Final[list[Any]] = [] - for _item in first_response.output: - if isinstance(_item, dict): - first_response_output_items.append(_item) - elif hasattr(_item, "model_dump"): - first_response_output_items.append(_item.model_dump(exclude_none=True)) - else: - first_response_output_items.append(_item) - return original_input_items + first_response_output_items + tool_results + return [*original_input_items, *first_response_output_items, *tool_results] async def aresponses_with_emulated_file_search( - input: Any, + input: object, model: str, tools: Iterable[ToolParam] | None = None, # Pass-through params — forwarded as-is to the underlying aresponses call @@ -504,7 +528,7 @@ async def aresponses_with_emulated_file_search( runs vector search, and synthesizes an OpenAI-format response. """ # Determine whether caller wants search_results populated in the output. - _include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) + _include_search_results, call_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) # 1. Replace file_search tools with function tool transformed_tools, all_vs_ids = _replace_file_search_tools(tools) @@ -521,7 +545,7 @@ async def aresponses_with_emulated_file_search( input=input, model=model, tools=transformed_tools or None, - **kwargs, + **call_kwargs, ), ) finally: @@ -585,7 +609,7 @@ async def aresponses_with_emulated_file_search( input=follow_up_input, model=model, tools=None, # no tools needed for the answer step - **kwargs, + **call_kwargs, ), ) finally: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d7f6ece5cd1..1385e329e93 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 +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,12 +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, ) + from litellm.types.router import LiteLLM_Params @lru_cache(maxsize=1) @@ -69,6 +71,79 @@ 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 _LoadsJsonValue(Protocol): + def __call__(self, s: str | bytes, /) -> object: ... + + +class _LoadsJsonDict(Protocol): + def __call__(self, s: str | bytes, /) -> _MutableJsonObject: ... + + +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_loads_json_value(fn: _LoadsJsonValue) -> _LoadsJsonValue: + return fn + + +def _typed_loads_json_dict(fn: _LoadsJsonDict) -> _LoadsJsonDict: + return fn + + +def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams: + return fn + + +def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr: + return fn + + +_LOADS_JSON_VALUE: Final = _typed_loads_json_value(json.loads) +_LOADS_JSON_DICT: Final = _typed_loads_json_dict(json.loads) + +_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache" +_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text" + + def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None: model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None model_id: Final = model_info.get("id") if _is_json_object(model_info) else None @@ -185,7 +260,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), @@ -228,7 +303,7 @@ class BaseResponsesAPIStreamingIterator: try: # Parse the JSON chunk - parsed_chunk: Final = json.loads(chunk) + parsed_chunk: Final = _LOADS_JSON_VALUE(chunk) # Format as ResponsesAPIStreamingResponse if isinstance(parsed_chunk, dict): @@ -514,7 +589,7 @@ class BaseResponsesAPIStreamingIterator: if response_obj is None: return - caching_handler: Final = 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 @@ -532,8 +607,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( - 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 @@ -586,12 +664,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, @@ -1043,8 +1124,8 @@ class _HasModelDumpJson(Protocol): def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... -def _dump_response_object(obj: Any) -> dict[str, Any]: - if hasattr(obj, "model_dump"): +def _dump_response_object(obj: object) -> Mapping[str, object]: + if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): return obj @@ -1073,21 +1154,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( @@ -1117,7 +1197,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() @@ -1134,15 +1214,19 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []): + 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( @@ -1200,7 +1284,8 @@ def _build_synthetic_response_events( ] sequence_number = 0 - for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): + output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or [] + for output_index, output_item in enumerate(output_items): output_item_payload = _dump_response_object(output_item) item_id = str(output_item_payload.get("id") or transformed.id) item_type = output_item_payload.get("type") @@ -1214,7 +1299,9 @@ def _build_synthetic_response_events( ) if item_type == "message": - for content_index, part in enumerate(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( openai_types.ContentPartAddedEvent( @@ -1261,7 +1348,9 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): + 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): @@ -1354,7 +1443,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[Any] | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, authorized_model: str | None = None, ): @@ -1363,16 +1452,16 @@ 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[Any] = guardrail_callbacks or [] + self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] # 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: @@ -1380,7 +1469,7 @@ class ResponsesWebSocketStreaming: event = event.decode("utf-8") if isinstance(event, str): try: - event_obj = json.loads(event) + event_obj = _LOADS_JSON_DICT(event) except (json.JSONDecodeError, TypeError): return else: @@ -1393,7 +1482,7 @@ class ResponsesWebSocketStreaming: """Extract user input content from response.create for logging.""" try: if isinstance(message, str): - msg_obj = json.loads(message) + msg_obj = _LOADS_JSON_DICT(message) elif _is_json_object(message): msg_obj = message else: @@ -1463,7 +1552,7 @@ class ResponsesWebSocketStreaming: # masked response.completed. if self.output_guardrail_callbacks: try: - _evt_type = json.loads(response_str).get("type") + _evt_type = _LOADS_JSON_DICT(response_str).get("type") except (json.JSONDecodeError, TypeError): _evt_type = None if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES: @@ -1485,7 +1574,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. @@ -1527,7 +1616,7 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = json.loads(message) + msg_obj: Final = _LOADS_JSON_DICT(message) except (json.JSONDecodeError, TypeError): return message @@ -1553,7 +1642,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 @@ -1655,11 +1744,12 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final = _LOADS_JSON_DICT(response_str) except (json.JSONDecodeError, TypeError): 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": @@ -1679,7 +1769,7 @@ class ResponsesWebSocketStreaming: continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text(text, pii_tokens) + unmasked = unmask_pii_text(text, pii_tokens) if unmasked != text: content_block["text"] = unmasked modified = True @@ -1688,7 +1778,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(delta, pii_tokens) + unmasked = unmask_pii_text(delta, pii_tokens) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) @@ -1711,7 +1801,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final[Mapping[str, object]] = json.loads(response_str) + evt_obj: Final = _LOADS_JSON_DICT(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1859,7 +1949,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, @@ -1871,10 +1961,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 @@ -1894,7 +1985,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): @@ -1937,7 +2028,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. @@ -1952,7 +2043,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 @@ -2009,10 +2100,10 @@ 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 = json.loads(raw_message) + msg_obj: Final = _LOADS_JSON_DICT(raw_message) except json.JSONDecodeError: await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None @@ -2022,7 +2113,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 @@ -2038,13 +2129,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 @@ -2062,7 +2153,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. @@ -2085,7 +2176,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": [...], ...}} @@ -2194,7 +2285,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. @@ -2202,7 +2293,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) @@ -2216,7 +2307,7 @@ class ManagedResponsesWebSocketHandler: continue if chunk_type == "response.completed" and completed_event is None: try: - completed_event = json.loads(serialized) + completed_event = _LOADS_JSON_DICT(serialized) except Exception: pass try: @@ -2228,7 +2319,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: @@ -2293,13 +2384,15 @@ 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 = call_kwargs.pop("model", None) + requested_model: Final = _typed_pops_optional_str(call_kwargs.pop)("model", None) if requested_model is None or requested_model == self.model_group: model = self.model else: model = 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 bbe97613c57..209dea87cc5 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -10,10 +10,10 @@ Use this to route requests between Teams import re from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload from litellm._logging import verbose_logger -from litellm.types.router import RouterErrors +from litellm.types.router import DeploymentTypedDict, RouterErrors if TYPE_CHECKING: from litellm.router import Router as _Router @@ -23,9 +23,68 @@ else: LitellmRouter = Any +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 _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 _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 _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: ... + + +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( - tag_regexes: list[str], - header_strings: list[str], + tag_regexes: Sequence[str], + header_strings: Sequence[str], ) -> str | None: """ Test compiled regex patterns against "Header-Name: value" strings. @@ -46,7 +105,9 @@ def _is_valid_deployment_tag_regex( return None -def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], match_any: bool = True) -> bool: +def is_valid_deployment_tag( + deployment_tags: Sequence[str], request_tags: Sequence[str], match_any: bool = True +) -> bool: """ Check if a tag is valid, the matching can be either any or all based on `match_any` flag """ @@ -73,7 +134,7 @@ def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], def _match_deployment( - deployment: Any, + deployment: _DeploymentLike, request_tags: list[str] | None, header_strings: list[str], match_any: bool, @@ -90,8 +151,8 @@ def _match_deployment( ran and failed, so the regex cannot override strict-tag policy. """ litellm_params: Final = deployment.get("litellm_params", {}) - deployment_tags: Final[list[str] | None] = litellm_params.get("tags") - deployment_tag_regex: Final[list[str] | None] = litellm_params.get("tag_regex") + deployment_tags: Final[Sequence[str] | None] = litellm_params.get("tags") + deployment_tag_regex: Final[Sequence[str] | None] = litellm_params.get("tag_regex") # 1. Exact tag match (existing behaviour). if deployment_tags and request_tags: @@ -162,38 +223,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[ def _exclude_deployments( - deployments: Sequence[Any] | Mapping[Any, Any], + deployments: _DeploymentPool, excluded_set: frozenset[str], -) -> list[Any]: +) -> 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: Sequence[Any] | Mapping[Any, Any], + deployments: _DeploymentPool, required_set: frozenset[str], -) -> tuple[Any, ...]: +) -> 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: Sequence[Any] | Mapping[Any, Any], -) -> tuple[Any, ...]: + 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: Sequence[Any] | Mapping[Any, Any]) -> frozenset[str]: +def _known_tag_values(deployments: _DeploymentPool) -> frozenset[str]: return frozenset( 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: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -217,7 +278,7 @@ def _unknown_required_tag_hides_an_answer( def _chain_allows_fail_open( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -228,12 +289,12 @@ def _chain_allows_fail_open( def _trusted_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, inherited_required_set: frozenset[str] | None, -) -> tuple[Any, ...]: +) -> 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 @@ -260,8 +321,8 @@ def _trusted_only_pool( def _resolve_or_fail_open( - pool: Sequence[Any], - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + pool: Sequence[_DeploymentLike], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -269,7 +330,7 @@ def _resolve_or_fail_open( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: if pool: return tuple(pool) if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed): @@ -289,7 +350,7 @@ def _resolve_or_fail_open( def _resolve_constraint_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -297,7 +358,7 @@ def _resolve_constraint_only_pool( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: pool: Final = ( _require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set) if required_set @@ -319,8 +380,8 @@ def _resolve_constraint_only_pool( def _all_deployments_or_fallback( llm_router_instance: LitellmRouter, model: str, - fallback: Sequence[Any] | Mapping[Any, Any], -) -> Sequence[Any] | Mapping[Any, Any]: + 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 @@ -330,7 +391,7 @@ def _all_deployments_or_fallback( def _chain_tag_filtering_override( llm_router_instance: LitellmRouter, model: str, - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + 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 @@ -392,10 +453,10 @@ def _tag_known_to_group( 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. @@ -473,25 +534,25 @@ async def get_deployments_for_tag( request_tags, ) - new_healthy_deployments: Final[list[Any]] = [] - default_deployments: Final[list[Any]] = [] - 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", @@ -507,10 +568,10 @@ async def get_deployments_for_tag( "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( @@ -545,10 +606,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 = [] - 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 @@ -562,7 +624,7 @@ async def get_deployments_for_tag( def _get_tags_from_request_kwargs( - request_kwargs: dict[Any, Any] | None = None, + request_kwargs: _RequestKwargsLike | None = None, metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", ) -> list[str]: """ @@ -577,12 +639,12 @@ def _get_tags_from_request_kwargs( if request_kwargs is None: return [] if metadata_variable_name in request_kwargs: - metadata: Final = request_kwargs[metadata_variable_name] or {} + metadata: Final[_MetadataLike] = request_kwargs[metadata_variable_name] or {} tags = metadata.get("tags", []) - return tags if tags is not None else [] + return list(tags) if tags is not None else [] elif "litellm_params" in request_kwargs: - litellm_params: Final = request_kwargs["litellm_params"] or {} - _metadata: Final = litellm_params.get(metadata_variable_name, {}) or {} + litellm_params: Final[_NestedLitellmParamsLike] = request_kwargs["litellm_params"] or {} + _metadata: Final[_MetadataLike] = litellm_params.get(metadata_variable_name, {}) or {} tags = _metadata.get("tags", []) - return tags if tags is not None else [] + return list(tags) if tags is not None else [] return [] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index dff010bfd30..90033af024b 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3058 + "limit": 3036 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2022 + "limit": 2020 }, "ANN202": { "limit": 855 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1384 + "limit": 1286 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 505 }, "B009": { - "limit": 64 + "limit": 63 }, "B010": { "limit": 190 @@ -123,7 +123,7 @@ "limit": 12 }, "PERF403": { - "limit": 34 + "limit": 33 }, "PIE804": { "limit": 18 @@ -180,7 +180,7 @@ "limit": 8 }, "RUF019": { - "limit": 38 + "limit": 36 }, "RUF046": { "limit": 4 @@ -201,7 +201,7 @@ "limit": 58 }, "SIM102": { - "limit": 321 + "limit": 318 }, "SIM103": { "limit": 119 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1224 + "limit": 1214 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fdacf375844..83def5fe2e3 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23003 + "limit": 22780 }, "LIT002": { - "limit": 27146 + "limit": 27144 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1077 + "limit": 1069 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16731 + "limit": 16725 }, "LIT011": { - "limit": 5596 + "limit": 5590 } } From 64e993773da57f5bd39dc60b1d45b2eaafc5acd6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:25:22 +0000 Subject: [PATCH 02/16] chore(typing): clear 1.5k basedpyright Any errors across 54 files Retypes the 54 highest-density reportAny/reportExplicitAny sources with real types instead of shuffling the ceilings around: typed prisma table Protocols so the untyped client surface stops at the query, local TypedDicts for JSON and dict payloads, concrete chunk and logging types on the streaming and callback surfaces, and 3-argument getattr with a Callable annotation where an SDK object is genuinely duck-typed No cast(), no type: ignore, no noqa, no suppression comments, and no new Any annotations. Whole-tree basedpyright drops 1,941 errors with no rule rising anywhere, and all three budget files are ratcheted so the cleared headroom cannot silently grow back Adds a GDC regression test pinning the named AttributeError that the typed credential accessor now raises when with_gdch_audience is missing --- basedpyright-code-budget.json | 26 +-- .../proxy/audit_logging_endpoints.py | 108 +++++++---- litellm/_lazy_imports.py | 79 ++++---- .../transformation.py | 46 ++--- litellm/caching/valkey_semantic_cache.py | 71 ++++--- .../bitbucket/bitbucket_client.py | 49 +++-- .../compression_interception/handler.py | 84 +++++---- litellm/integrations/custom_logger.py | 28 +-- .../opik_payload_builder/payload_builders.py | 12 +- .../prometheus_helpers/prometheus_api.py | 54 +++++- .../websearch_interception/handler.py | 16 +- .../llm_response_utils/response_metadata.py | 31 ++- .../model_response_utils.py | 66 +++++-- .../litellm_core_utils/streaming_handler.py | 8 +- litellm/litellm_core_utils/url_utils.py | 25 +-- .../llms/anthropic/batches/transformation.py | 59 ++++-- .../context_management/dispatcher.py | 81 +++++--- .../responses_adapters/handler.py | 40 ++-- .../llms/anthropic/skills/transformation.py | 30 ++- litellm/llms/azure/files/handler.py | 13 +- litellm/llms/bedrock/realtime/handler.py | 104 ++++++++-- litellm/llms/chatgpt/chat/streaming_utils.py | 41 +++- .../llms/compactifai/chat/transformation.py | 33 +++- .../llms/custom_httpx/container_handler.py | 178 ++++++++++++------ litellm/llms/gdc/chat/transformation.py | 45 ++++- .../llms/infinity/rerank/transformation.py | 29 ++- .../litellm_proxy/skills/code_execution.py | 76 ++++++-- litellm/llms/oci/chat/cohere.py | 38 +++- litellm/llms/ollama/completion/handler.py | 36 +++- .../llms/openai/containers/transformation.py | 72 +++++-- .../vector_stores/rag_api/transformation.py | 127 ++++++++++--- litellm/models/base.py | 2 +- .../mcp_server/oauth2_flow_backfill.py | 63 ++++++- .../_experimental/mcp_server/toolset_db.py | 93 +++++++-- .../proxy/agent_endpoints/agent_registry.py | 81 +++++--- .../proxy/client/cli/commands/credentials.py | 51 ++++- litellm/proxy/client/cli/commands/teams.py | 52 +++-- litellm/proxy/common_utils/callback_utils.py | 14 +- litellm/proxy/common_utils/get_routes.py | 65 ++++--- .../proxy/common_utils/user_api_key_cache.py | 48 ++--- litellm/proxy/db/routing_prisma_wrapper.py | 36 ++-- .../guardrail_hooks/custom_code/sandbox.py | 14 +- .../hiddenlayer/hiddenlayer.py | 33 ++-- .../llm_as_a_judge/__init__.py | 130 ++++++++++--- .../guardrail_hooks/noma/noma_v2.py | 26 +-- .../promptguard/promptguard.py | 20 +- .../jwt_key_mapping_endpoints.py | 76 ++++++-- .../response_polling/background_streaming.py | 47 ++++- .../search_endpoints/search_tool_registry.py | 57 ++++-- litellm/rag/rag_query.py | 65 +++++-- litellm/repositories/base_repository.py | 51 +++-- .../repositories/credentials_repository.py | 58 ++++-- litellm/repositories/team_repository.py | 76 ++++++-- litellm/rust_bridge/responses_websocket.py | 23 ++- .../secret_managers/secret_manager_handler.py | 60 +++++- ruff-strict-budget.json | 18 +- .../gdc/chat/test_gdc_chat_transformation.py | 19 ++ type-discipline-budget.json | 12 +- 58 files changed, 2172 insertions(+), 823 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index b4c324a2c4c..58b8ae01cb0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19955 + "limit": 18773 }, "reportArgumentType": { "limit": 2566 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 6049 + "limit": 5774 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5663 + "limit": 5640 }, "reportMissingTypeArgument": { - "limit": 15555 + "limit": 15498 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44655 + "limit": 44589 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 39017 + "limit": 38882 }, "reportUnknownParameterType": { - "limit": 19885 + "limit": 19806 }, "reportUnknownVariableType": { - "limit": 30572 + "limit": 30456 }, "reportUnnecessaryCast": { - "limit": 117 + "limit": 116 }, "reportUnnecessaryComparison": { - "limit": 699 + "limit": 698 }, "reportUnnecessaryContains": { "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 836 + "limit": 835 }, "reportUntypedBaseClass": { "limit": 0 @@ -138,9 +138,9 @@ "limit": 139 }, "reportUnusedImport": { - "limit": 545 + "limit": 544 }, "reportUnusedVariable": { - "limit": 146 + "limit": 142 } } diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index 18ac29b9781..21bde06e86a 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -7,7 +7,9 @@ GET - /audit/{id} - Get audit log by id GET - /audit - Get all audit logs """ -from typing import Any, Dict, List, Optional +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol #### AUDIT LOGGING #### from fastapi import APIRouter, Depends, HTTPException, Query @@ -15,6 +17,7 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import ( AuditLogResponse, PaginatedAuditLogResponse, ) +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -22,7 +25,44 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() -def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]: +class _AuditLogFields(TypedDict): + """Columns of the `LiteLLM_AuditLog` table, as returned by `model_dump()`.""" + + id: ReadOnly[str] + updated_at: ReadOnly[datetime] + changed_by: ReadOnly[str] + changed_by_api_key: ReadOnly[str] + action: ReadOnly[str] + table_name: ReadOnly[str] + object_id: ReadOnly[str] + before_value: ReadOnly[dict[str, object] | None] + updated_values: ReadOnly[dict[str, object] | None] + + +class _AuditLogRecord(Protocol): + """Row of the `LiteLLM_AuditLog` table as materialised by the Prisma client.""" + + def model_dump(self) -> _AuditLogFields: ... + + +class _AuditLogTable(Protocol): + """The `litellm_auditlog` accessor of the Prisma client.""" + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str], + skip: int, + take: int, + ) -> Sequence[_AuditLogRecord]: ... + + async def count(self, *, where: Mapping[str, object]) -> int: ... + + async def find_unique(self, *, where: Mapping[str, str]) -> _AuditLogRecord | None: ... + + +def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, object]: """ Build an OR condition that matches a value inside a JSON column at the given key, checking both before_value and updated_values. @@ -53,33 +93,33 @@ async def get_audit_logs( page: int = Query(1, ge=1), page_size: int = Query(10, ge=1, le=100), # Filter parameters - changed_by: Optional[str] = Query( + changed_by: str | None = Query( None, description="Filter by user or system that performed the action" ), - changed_by_api_key: Optional[str] = Query( + changed_by_api_key: str | None = Query( None, description="Filter by API key hash that performed the action" ), - action: Optional[str] = Query( + action: str | None = Query( None, description="Filter by action type (create, update, delete)" ), - table_name: Optional[str] = Query( + table_name: str | None = Query( None, description="Filter by table name that was modified" ), - object_id: Optional[str] = Query( + object_id: str | None = Query( None, description="Filter by ID of the object that was modified" ), - start_date: Optional[str] = Query(None, description="Filter logs after this date"), - end_date: Optional[str] = Query(None, description="Filter logs before this date"), - object_team_id: Optional[str] = Query( + start_date: str | None = Query(None, description="Filter logs after this date"), + end_date: str | None = Query(None, description="Filter logs before this date"), + object_team_id: str | None = Query( None, description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)", ), - object_key_hash: Optional[str] = Query( + object_key_hash: str | None = Query( None, description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)", ), # Sorting parameters - sort_by: Optional[str] = Query( + sort_by: str | None = Query( None, description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')", ), @@ -102,7 +142,7 @@ async def get_audit_logs( ) # Build filter conditions - where_conditions: Dict[str, Any] = {} + where_conditions: Final[dict[str, object]] = {} if changed_by: where_conditions["changed_by"] = changed_by if changed_by_api_key: @@ -114,33 +154,31 @@ async def get_audit_logs( if object_id: where_conditions["object_id"] = object_id if start_date or end_date: - date_filter: Dict[str, Any] = {} - if start_date: - date_filter["gte"] = start_date - if end_date: - date_filter["lte"] = end_date + date_filter: Final[Mapping[str, str]] = { + bound: bound_value for bound, bound_value in (("gte", start_date), ("lte", end_date)) if bound_value + } where_conditions["updated_at"] = date_filter # JSON field filters (PostgreSQL only) — each filter is AND'd with the # others, but checks both before_value and updated_values internally (OR). - if object_team_id: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("team_id", object_team_id) - ] - if object_key_hash: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("token", object_key_hash) + if object_team_id or object_key_hash: + where_conditions["AND"] = [ + _build_json_field_or_condition(json_key, json_value) + for json_key, json_value in ( + ("team_id", object_team_id), + ("token", object_key_hash), + ) + if json_value ] # Build sort conditions - order_by: Dict[str, Any] = {} - if sort_by and isinstance(sort_by, str): - order_by[sort_by] = sort_order - else: - order_by["updated_at"] = sort_order # Default sort by updated_at + sort_column: Final[str] = sort_by if sort_by and isinstance(sort_by, str) else "updated_at" + order_by: Final[Mapping[str, str]] = {sort_column: sort_order} + + audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog # Get paginated results - audit_logs = await prisma_client.db.litellm_auditlog.find_many( + audit_logs: Final = await audit_log_table.find_many( where=where_conditions, order=order_by, skip=(page - 1) * page_size, @@ -148,8 +186,8 @@ async def get_audit_logs( ) # Get total count for pagination - total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions) - total_pages = -(-total_count // page_size) # Ceiling division + total_count: Final = await audit_log_table.count(where=where_conditions) + total_pages: Final = -(-total_count // page_size) # Ceiling division # Return paginated response return PaginatedAuditLogResponse( @@ -198,8 +236,10 @@ async def get_audit_log_by_id( detail={"message": CommonProxyErrors.db_not_connected_error.value}, ) + audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog + # Get the audit log by ID - audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id}) + audit_log: Final = await audit_log_table.find_unique(where={"id": id}) if audit_log is None: raise HTTPException( diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 933464d3f23..dff3b11e353 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -17,8 +17,14 @@ until they're actually needed. import importlib import sys -from collections.abc import Callable -from typing import Any, Final, cast +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Any, Final, cast + +if TYPE_CHECKING: + import httpx + import tiktoken + + from .caching.llm_caching_handler import LLMClientCache as LLMClientCacheType # Import all the data structures that define what can be lazy-loaded # These are just lists of names and maps of where to find them @@ -54,7 +60,7 @@ from ._lazy_imports_registry import ( ) -def get_litellm_globals() -> dict: +def get_litellm_globals() -> dict[str, object]: """ Get the globals dictionary of the litellm module. @@ -64,7 +70,7 @@ def get_litellm_globals() -> dict: return sys.modules["litellm"].__dict__ -def _get_utils_globals() -> dict: +def _get_utils_globals() -> dict[str, object]: """ Get the globals dictionary of the utils module. @@ -74,14 +80,19 @@ def _get_utils_globals() -> dict: return sys.modules["litellm.utils"].__dict__ +def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None": + """Read the configured `litellm.request_timeout` used for the module level http clients.""" + return litellm_globals.get("request_timeout") + + # These are special lazy loaders for things that are used internally # They're separate from the main lazy import system because they have specific use cases # Lazy loader for default encoding - avoids importing heavy tiktoken library at startup -_default_encoding: Any | None = None +_default_encoding: "tiktoken.Encoding | None" = None -def _get_default_encoding() -> Any: +def _get_default_encoding() -> "tiktoken.Encoding": """ Lazily load and cache the default OpenAI encoding. @@ -100,10 +111,10 @@ def _get_default_encoding() -> Any: # Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time -_get_modified_max_tokens_func: Any | None = None +_get_modified_max_tokens_func: Callable[..., int | None] | None = None -def _get_modified_max_tokens() -> Any: +def _get_modified_max_tokens() -> Callable[..., int | None]: """ Lazily load and cache the get_modified_max_tokens function. @@ -124,10 +135,10 @@ def _get_modified_max_tokens() -> Any: # Lazy loader for token_counter to avoid importing token_counter module at module import time -_token_counter_new_func: Any | None = None +_token_counter_new_func: Callable[..., int] | None = None -def _get_token_counter_new() -> Any: +def _get_token_counter_new() -> Callable[..., int]: """ Lazily load and cache the token_counter function (aliased as token_counter_new). @@ -154,10 +165,10 @@ def _get_token_counter_new() -> Any: # This registry maps attribute names (like "ModelResponse") to handler functions # It's built once the first time someone accesses a lazy-loaded attribute # Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...} -_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None +_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], object]] | None = None -def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: +def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]: """ Build the registry that maps attribute names to their handler functions. @@ -206,7 +217,7 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: return _LAZY_IMPORT_REGISTRY -def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any: +def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object: """ Generic function that handles lazy importing for most attributes. @@ -255,7 +266,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # Step 6: Get the actual attribute from the module # Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class - value: Final = getattr(module, attr_name) + value: Final[object] = getattr(module, attr_name) # Step 7: Cache it so we don't have to import again next time _globals[name] = value @@ -272,62 +283,62 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # The registry (above) maps attribute names to these handler functions. -def _lazy_import_utils(name: str) -> Any: +def _lazy_import_utils(name: str) -> object: """Handler for utils module attributes (ModelResponse, token_counter, etc.)""" return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils") -def _lazy_import_cost_calculator(name: str) -> Any: +def _lazy_import_cost_calculator(name: str) -> object: """Handler for cost calculator functions (completion_cost, cost_per_token, etc.)""" return _generic_lazy_import(name, _COST_CALCULATOR_IMPORT_MAP, "Cost calculator") -def _lazy_import_token_counter(name: str) -> Any: +def _lazy_import_token_counter(name: str) -> object: """Handler for token counter utilities""" return _generic_lazy_import(name, _TOKEN_COUNTER_IMPORT_MAP, "Token counter") -def _lazy_import_bedrock_types(name: str) -> Any: +def _lazy_import_bedrock_types(name: str) -> object: """Handler for Bedrock type aliases""" return _generic_lazy_import(name, _BEDROCK_TYPES_IMPORT_MAP, "Bedrock types") -def _lazy_import_types_utils(name: str) -> Any: +def _lazy_import_types_utils(name: str) -> object: """Handler for types from litellm.types.utils (BudgetConfig, ImageObject, etc.)""" return _generic_lazy_import(name, _TYPES_UTILS_IMPORT_MAP, "Types utils") -def _lazy_import_caching(name: str) -> Any: +def _lazy_import_caching(name: str) -> object: """Handler for caching classes (Cache, DualCache, RedisCache, etc.)""" return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching") -def _lazy_import_dotprompt(name: str) -> Any: +def _lazy_import_dotprompt(name: str) -> object: """Handler for dotprompt integration globals""" return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt") -def _lazy_import_types(name: str) -> Any: +def _lazy_import_types(name: str) -> object: """Handler for type classes (GuardrailItem, etc.)""" return _generic_lazy_import(name, _TYPES_IMPORT_MAP, "Types") -def _lazy_import_llm_configs(name: str) -> Any: +def _lazy_import_llm_configs(name: str) -> object: """Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)""" return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config") -def _lazy_import_litellm_logging(name: str) -> Any: +def _lazy_import_litellm_logging(name: str) -> object: """Handler for litellm_logging module (Logging, modify_integration)""" return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging") -def _lazy_import_llm_provider_logic(name: str) -> Any: +def _lazy_import_llm_provider_logic(name: str) -> object: """Handler for LLM provider logic functions (get_llm_provider, etc.)""" return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") -def _lazy_import_utils_module(name: str) -> Any: +def _lazy_import_utils_module(name: str) -> object: """ Handler for utils module lazy imports. @@ -355,7 +366,7 @@ def _lazy_import_utils_module(name: str) -> Any: module = importlib.import_module(module_path) # Get the actual attribute from the module - value: Final = getattr(module, attr_name) + value: Final[object] = getattr(module, attr_name) # Cache it so we don't have to import again next time _globals[name] = value @@ -370,7 +381,7 @@ def _lazy_import_utils_module(name: str) -> Any: # These handlers have custom logic that doesn't fit the generic pattern -def _lazy_import_llm_client_cache(name: str) -> Any: +def _lazy_import_llm_client_cache(name: str) -> object: """ Handler for LLM client cache - has special logic for singleton instance. @@ -387,7 +398,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: # Import the class module: Final = importlib.import_module("litellm.caching.llm_caching_handler") - LLMClientCache: Final = getattr(module, "LLMClientCache") + LLMClientCache: Final[type[LLMClientCacheType]] = getattr(module, "LLMClientCache") # If they want the class itself, return it if name == "LLMClientCache": @@ -403,7 +414,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}") -def _lazy_import_http_handlers(name: str) -> Any: +def _lazy_import_http_handlers(name: str) -> object: """ Handler for HTTP clients - has special logic for creating client instances. @@ -419,8 +430,8 @@ def _lazy_import_http_handlers(name: str) -> Any: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client # Get timeout from module config (if set) - timeout = _globals.get("request_timeout") - params: Final = {"timeout": timeout, "client_alias": "module level aclient"} + async_timeout: Final = _get_module_level_client_timeout(_globals) + params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"} # Create the client instance provider_id: Final = cast(Any, "litellm_module_level_client") @@ -437,8 +448,8 @@ def _lazy_import_http_handlers(name: str) -> Any: # Create a sync HTTP client from litellm.llms.custom_httpx.http_handler import HTTPHandler - timeout = _globals.get("request_timeout") - sync_client: Final = HTTPHandler(timeout=timeout) + sync_timeout: Final = _get_module_level_client_timeout(_globals) + sync_client: Final = HTTPHandler(timeout=sync_timeout) # Cache it _globals["module_level_client"] = sync_client diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 15cf77708f9..90afa5adf9e 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -17,12 +17,16 @@ A2A Streaming Events: - Artifact update (kind: "artifact-update") - Content/artifact delivery """ +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from uuid import uuid4 from litellm._logging import verbose_logger +if TYPE_CHECKING: + from litellm.types.utils import Choices + class A2AStreamingContext: """ @@ -30,7 +34,7 @@ class A2AStreamingContext: Tracks task_id, context_id, and message accumulation. """ - def __init__(self, request_id: str, input_message: dict[str, Any]): + def __init__(self, request_id: str, input_message: Mapping[str, object]): self.request_id = request_id self.task_id = str(uuid4()) self.context_id = str(uuid4()) @@ -46,7 +50,7 @@ class A2ACompletionBridgeTransformation: """ @staticmethod - def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str: + def _extract_text_from_a2a_parts(parts: Sequence[Mapping[str, object]]) -> str: """Extract text from A2A parts (with or without explicit ``kind``).""" content_parts: Final[list[str]] = [] for part in parts: @@ -62,16 +66,16 @@ class A2ACompletionBridgeTransformation: @staticmethod def get_forward_metadata( - a2a_message: dict[str, Any], + a2a_message: Mapping[str, object], params: dict[str, Any] | None = None, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """ Merge A2A metadata from MessageSendParams and the message for downstream providers. Forwarded once on the LangGraph run payload (``metadata``), not duplicated on each input message — see ``apply_forward_metadata_to_completion_params``. """ - merged: Final[dict[str, Any]] = {} + merged: Final[dict[str, object]] = {} if params and isinstance(params.get("metadata"), dict): merged.update(params["metadata"]) message_metadata: Final = a2a_message.get("metadata") @@ -81,8 +85,8 @@ class A2ACompletionBridgeTransformation: @staticmethod def apply_forward_metadata_to_completion_params( - completion_params: dict[str, Any], - a2a_message: dict[str, Any], + completion_params: dict[str, object], + a2a_message: Mapping[str, object], params: dict[str, Any] | None = None, ) -> None: """ @@ -104,8 +108,8 @@ class A2ACompletionBridgeTransformation: # ``extra_body.metadata`` so the configured keys remain authoritative # and an A2A caller cannot overwrite server-set run metadata. existing_metadata: Final = extra_body.get("metadata") - existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {} - merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict} + existing_dict: Final[dict[str, object]] = existing_metadata if isinstance(existing_metadata, dict) else {} + merged_metadata: Final[dict[str, object]] = {**forward_metadata, **existing_dict} extra_body = {**extra_body, "metadata": merged_metadata} completion_params["extra_body"] = extra_body @@ -114,7 +118,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def a2a_message_to_openai_messages( a2a_message: dict[str, Any], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform an A2A message to OpenAI message format. @@ -124,8 +128,8 @@ class A2ACompletionBridgeTransformation: Returns: List of OpenAI-format messages """ - role: Final = a2a_message.get("role", "user") - parts = a2a_message.get("parts", []) + role: Final[object] = a2a_message.get("role", "user") + parts: Sequence[Mapping[str, object]] = a2a_message.get("parts", []) # Map A2A roles to OpenAI roles openai_role = role @@ -143,7 +147,7 @@ class A2ACompletionBridgeTransformation: # Do not attach A2A message.metadata here — the completion bridge forwards it # once at run level via extra_body.metadata (LangGraph POST /runs/wait shape). - openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content} + openai_message: Final[dict[str, object]] = {"role": openai_role, "content": content} verbose_logger.debug( "A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content) @@ -155,7 +159,7 @@ class A2ACompletionBridgeTransformation: def openai_response_to_a2a_response( response: Any, request_id: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform a LiteLLM ModelResponse to A2A SendMessageResponse format. @@ -169,7 +173,7 @@ class A2ACompletionBridgeTransformation: # Extract content from response content = "" if hasattr(response, "choices") and response.choices: - choice: Final = response.choices[0] + choice: Final[Choices] = response.choices[0] if hasattr(choice, "message") and choice.message: content = choice.message.content or "" @@ -182,7 +186,7 @@ class A2ACompletionBridgeTransformation: } # Build A2A response - a2a_response: Final = { + a2a_response: Final[dict[str, object]] = { "jsonrpc": "2.0", "id": request_id, "result": a2a_message, @@ -200,7 +204,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def create_task_event( ctx: A2AStreamingContext, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create the initial task event with status 'submitted'. @@ -235,7 +239,7 @@ class A2ACompletionBridgeTransformation: state: str, final: bool = False, message_text: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create a status update event. @@ -245,7 +249,7 @@ class A2ACompletionBridgeTransformation: final: Whether this is the final event message_text: Optional message text for 'working' status """ - status: Final[dict[str, Any]] = { + status: Final[dict[str, object]] = { "state": state, "timestamp": A2ACompletionBridgeTransformation._get_timestamp(), } @@ -277,7 +281,7 @@ class A2ACompletionBridgeTransformation: def create_artifact_update_event( ctx: A2AStreamingContext, text: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create an artifact update event with content. diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index 737d212a89d..ac4d4033546 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -17,6 +17,7 @@ RedisSemanticCache since those are backend agnostic. import asyncio import hashlib import os +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any, Final @@ -62,7 +63,7 @@ class ValkeySemanticCache(RedisSemanticCache): sync_client: Redis | None = None, async_client: AsyncRedis | None = None, embedding_max_input_tokens: int | None = None, - **kwargs: Any, + **kwargs: object, ): if similarity_threshold is None: raise ValueError("similarity_threshold must be provided, passed None") @@ -115,7 +116,7 @@ class ValkeySemanticCache(RedisSemanticCache): return hashlib.sha256(str(key).encode("utf-8")).hexdigest() @staticmethod - def _embedding_to_bytes(embedding: list[float]) -> bytes: + def _embedding_to_bytes(embedding: Sequence[float]) -> bytes: return pack_vector(embedding) def _index_schema(self, dim: int) -> tuple[TagField, VectorField]: @@ -189,7 +190,9 @@ class ValkeySemanticCache(RedisSemanticCache): def _doc_key(self, key: str) -> str: return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}" - def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict: + def _doc_mapping( + self, key: str, prompt: str, value_str: str, embedding: Sequence[float] + ) -> Mapping[str | bytes, str | bytes]: return { self.CACHE_KEY_FIELD_NAME: self._scope_tag(key), self.PROMPT_FIELD_NAME: prompt, @@ -205,30 +208,49 @@ class ValkeySemanticCache(RedisSemanticCache): ) return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2) + async def _async_search(self, key: str, embedding: Sequence[float]) -> object: + """Run the KNN query on the async client, stopping the untyped search surface here.""" + return await self.async_client.ft(self.index_name).search( + self._knn_query(key), + query_params={"vec": self._embedding_to_bytes(embedding)}, + ) + @classmethod - def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None: - docs: Final = getattr(search_result, "docs", []) + def _first_hit(cls, search_result: object) -> _ValkeyCacheHit | None: + docs: Final[Sequence[object]] = getattr(search_result, "docs", []) if not docs: return None doc: Final = docs[0] + response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME) + distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME) return _ValkeyCacheHit( - response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)), - distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)), + response=str(response_field), + distance=float(distance_field), ) - def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any: + @staticmethod + def _record_similarity(kwargs: dict[str, Any], similarity: float) -> None: + """Stamp the semantic-similarity score onto the request metadata carried in ``kwargs``.""" + kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity + + @staticmethod + def _embedding_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: + """The request metadata forwarded to the embedding call.""" + return kwargs.get("metadata") + + def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: object) -> object: if hit is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None similarity: Final = 1 - hit.distance - kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity + self._record_similarity(kwargs, similarity) if similarity < self.similarity_threshold: return None return self._get_cache_logic(cached_response=hit.response) - def set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + def set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -247,12 +269,12 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in Valkey semantic-cache set_cache: {e}") - def get_cache(self, key: str, **kwargs: Any) -> Any: + def get_cache(self, key: str, **kwargs: object) -> object: print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None embedding: Final = self._get_embedding(prompt) @@ -265,9 +287,9 @@ class ValkeySemanticCache(RedisSemanticCache): return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: print_verbose(f"Error in Valkey semantic-cache get_cache: {e}") - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) - async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -275,7 +297,7 @@ class ValkeySemanticCache(RedisSemanticCache): print_verbose("No prompt provided for semantic caching") return - embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs)) await self._ensure_index_async(len(embedding)) doc_key: Final = self._doc_key(key) @@ -286,31 +308,28 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}") - async def async_get_cache(self, key: str, **kwargs: Any) -> Any: + async def async_get_cache(self, key: str, **kwargs: object) -> object: print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None - embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs)) await self._ensure_index_async(len(embedding)) - search_result: Final = await self.async_client.ft(self.index_name).search( - self._knn_query(key), - query_params={"vec": self._embedding_to_bytes(embedding)}, - ) + search_result: Final[object] = await self._async_search(key, embedding) return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}") - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) - async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None: try: await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list]) except Exception as e: print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}") - async def _index_info(self) -> dict: + async def _index_info(self) -> Mapping[str, object]: return await self.async_client.ft(self.index_name).info() diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index e06e5ab358f..c6967ef3340 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -4,11 +4,33 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. import base64 import urllib.parse -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final + +from typing_extensions import ReadOnly, TypedDict from litellm.llms.custom_httpx.http_handler import HTTPHandler +class BitBucketSrcEntry(TypedDict, total=False): + """One entry of a BitBucket ``src`` directory listing.""" + + type: ReadOnly[str] + path: ReadOnly[str] + + +class BitBucketSrcListing(TypedDict, total=False): + """A page of a BitBucket ``src`` directory listing.""" + + values: ReadOnly[Sequence[BitBucketSrcEntry]] + + +class BitBucketBranchListing(TypedDict, total=False): + """A page of a BitBucket ``refs/branches`` listing.""" + + values: ReadOnly[Sequence[Mapping[str, object]]] + + def _sanitize_file_path(file_path: str) -> str: """Reject path traversal and URL-encode each path segment.""" if "#" in file_path or "?" in file_path: @@ -31,7 +53,7 @@ class BitBucketClient: - Branch-specific file fetching """ - def __init__(self, config: dict[str, Any]): + def __init__(self, config: Mapping[str, object]): """ Initialize the BitBucket client. @@ -135,16 +157,13 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() - files: Final = [] + data: Final[BitBucketSrcListing] = response.json() - for item in data.get("values", []): - if item.get("type") == "commit_file": - file_path = item.get("path", "") - if file_path.endswith(file_extension): - files.append(file_path) - - return files + return [ + file_path + for item in data.get("values", []) + if item.get("type") == "commit_file" and (file_path := item.get("path", "")).endswith(file_extension) + ] except Exception as e: # Check if it's an HTTP error @@ -162,7 +181,7 @@ class BitBucketClient: else: raise Exception(f"Error listing files in '{directory_path}': {e}") - def get_repository_info(self) -> dict[str, Any]: + def get_repository_info(self) -> Mapping[str, object]: """ Get information about the repository. @@ -191,7 +210,7 @@ class BitBucketClient: except Exception: return False - def get_branches(self) -> list[dict[str, Any]]: + def get_branches(self) -> Sequence[Mapping[str, object]]: """ Get list of branches in the repository. @@ -204,12 +223,12 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() + data: Final[BitBucketBranchListing] = response.json() return data.get("values", []) except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str) -> dict[str, Any] | None: + def get_file_metadata(self, file_path: str) -> Mapping[str, object] | None: """ Get metadata about a file (size, last modified, etc.). diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 7ea60053e6f..219718f6771 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,7 +7,10 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid -from typing import Any, Final, cast +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, cast + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.compression import compress @@ -26,6 +29,19 @@ LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve" _CACHE_TTL_SECONDS: Final = 15 * 60 +class _AgenticLoopParams(TypedDict, total=False): + """The ``agentic_loop_params`` entry the agentic loop driver records on the logging object.""" + + model: ReadOnly[str] + + +class _AgenticLoopLoggingObj(Protocol): + """Logging object view exposing the untyped call details this handler reads.""" + + @property + def model_call_details(self) -> Mapping[str, _AgenticLoopParams]: ... + + def _compression_savings_from_counts( original_tokens: object, compressed_tokens: object ) -> CompressionSavingsMetadata | None: @@ -78,7 +94,7 @@ class CompressionInterceptionLogger(CustomLogger): compression_trigger: int = 200_000, compression_target: int | None = None, embedding_model: str | None = None, - embedding_model_params: dict[str, Any] | None = None, + embedding_model_params: dict[str, object] | None = None, ): super().__init__() self.enabled = enabled @@ -101,7 +117,7 @@ class CompressionInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( litellm_settings: dict[str, Any], - callback_specific_params: dict[str, Any], + callback_specific_params: Mapping[str, object], ) -> "CompressionInterceptionLogger": compression_params: CompressionInterceptionConfig = {} if "compression_interception_params" in litellm_settings: @@ -115,7 +131,9 @@ class CompressionInterceptionLogger(CustomLogger): ) return CompressionInterceptionLogger.from_config_yaml(compression_params) - async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, Any], call_type: CallTypes | None + ) -> dict[str, object] | None: if not self.enabled: return None if call_type is not None and call_type != CallTypes.anthropic_messages: @@ -145,7 +163,7 @@ class CompressionInterceptionLogger(CustomLogger): cache: Final = cast(dict[str, str], compressed.get("cache", {})) skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason")) - compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", [])) + compressed_tools: Final = cast(list[dict[str, object]], compressed.get("tools", [])) # Only mutate kwargs when compression actually produced a result. # If compression was a no-op (below trigger, invalid tool sequence, etc.), @@ -156,7 +174,7 @@ class CompressionInterceptionLogger(CustomLogger): kwargs["messages"] = compressed["messages"] if compressed_tools: kwargs["tools"] = self._merge_tools( - existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")), + existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")), compressed_tools=compressed_tools, ) call_id = cast(str | None, kwargs.get("litellm_call_id")) @@ -189,14 +207,14 @@ class CompressionInterceptionLogger(CustomLogger): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, - messages: list[dict], - tools: list[dict] | None, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, stream: bool, custom_llm_provider: str, - kwargs: dict, - ) -> tuple[bool, dict]: + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: if not self.enabled: return False, {} if not self._has_retrieval_tool(tools): @@ -214,19 +232,19 @@ class CompressionInterceptionLogger(CustomLogger): async def async_build_agentic_loop_plan( self, - tools: dict, + tools: Mapping[str, object], model: str, - messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, - anthropic_messages_optional_request_params: dict, - logging_obj: Any, + messages: list[dict[str, object]], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: Mapping[str, object], + logging_obj: _AgenticLoopLoggingObj | None, stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], ) -> AgenticLoopPlan: self._prune_expired_cache() - tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", [])) - thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", [])) + tool_calls: Final = cast(list[dict[str, object]], tools.get("tool_calls", [])) + thinking_blocks: Final = cast(list[dict[str, object]], tools.get("thinking_blocks", [])) call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) cache: Final = self._get_cache(call_id=call_id) @@ -269,7 +287,7 @@ class CompressionInterceptionLogger(CustomLogger): full_model_name = model if logging_obj is not None: agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) - full_model_name = cast(str, agentic_params.get("model", model)) + full_model_name = agentic_params.get("model", model) request_patch: Final = AgenticLoopRequestPatch( model=full_model_name, @@ -304,15 +322,15 @@ class CompressionInterceptionLogger(CustomLogger): return {} return cache_entry[0] - def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None: + def _resolve_call_id(self, logging_obj: _AgenticLoopLoggingObj | None, kwargs: Mapping[str, object]) -> str | None: if logging_obj is not None: logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) if isinstance(logging_call_id, str) and logging_call_id: return logging_call_id kwargs_call_id: Final = kwargs.get("litellm_call_id") - return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None) + return kwargs_call_id if isinstance(kwargs_call_id, str) else None - def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str: + def _resolve_retrieval_content(self, tool_call: Mapping[str, object], cache: Mapping[str, str]) -> str: raw_input: Final = tool_call.get("input", {}) key = "" if isinstance(raw_input, dict): @@ -323,7 +341,9 @@ class CompressionInterceptionLogger(CustomLogger): return cache[key] return f"[compressed content key '{key}' not found]" - def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + def _extract_retrieval_tool_calls( + self, response: object + ) -> tuple[list[dict[str, object]], list[dict[str, object]]]: if isinstance(response, dict): content = response.get("content", []) else: @@ -332,8 +352,8 @@ class CompressionInterceptionLogger(CustomLogger): if not isinstance(content, list): return [], [] - tool_calls: Final[list[dict[str, Any]]] = [] - thinking_blocks: Final[list[dict[str, Any]]] = [] + tool_calls: Final[list[dict[str, object]]] = [] + thinking_blocks: Final[list[dict[str, object]]] = [] for block in content: if isinstance(block, dict): @@ -380,13 +400,13 @@ class CompressionInterceptionLogger(CustomLogger): return tool_calls, thinking_blocks - def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + def _prepare_followup_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]: internal_keys: Final = {"litellm_logging_obj"} return { k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys } - def _has_retrieval_tool(self, tools: Any) -> bool: + def _has_retrieval_tool(self, tools: object) -> bool: if not isinstance(tools, list): return False for tool in tools: @@ -402,9 +422,9 @@ class CompressionInterceptionLogger(CustomLogger): def _merge_tools( self, - existing_tools: list[dict[str, Any]] | None, - compressed_tools: list[dict[str, Any]], - ) -> list[dict[str, Any]]: + existing_tools: Sequence[Mapping[str, object]] | None, + compressed_tools: Sequence[Mapping[str, object]], + ) -> list[Mapping[str, object]]: merged: Final = list(existing_tools or []) if self._has_retrieval_tool(merged): return merged diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a0c78674ac8..417edc77e9c 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -2,7 +2,7 @@ # On success, logs events to Promptlayer import re import traceback -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import BaseModel @@ -103,11 +103,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return [] callbacks: Final = AllCallbacks() - callback_info: Final = getattr(callbacks, lookup_name, None) + callback_info: Final[object] = getattr(callbacks, lookup_name, None) if callback_info is None: return [] - params: Final = getattr(callback_info, "litellm_callback_params", None) + params: Final[Sequence[str] | None] = getattr(callback_info, "litellm_callback_params", None) if not params: return [] @@ -783,7 +783,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac - Converting to string and then truncating the logged content catches this 2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user """ - field_value: Final = standard_logging_object.get(field_name) + field_value: Final[object] = standard_logging_object.get(field_name) if field_value: str_value: Final = str(field_value) if len(str_value) > max_length: @@ -937,8 +937,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -969,8 +969,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -991,7 +991,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> Any: + ) -> object: """Recursively redact inline base64 from any nested structure with a max recursion depth limit.""" if depth > max_depth: verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth) @@ -1022,16 +1022,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def _process_messages( self, - messages: list[Any], + messages: list[object], max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> list[dict[str, Any]]: - filtered_messages: Final[list[dict[str, Any]]] = [] + ) -> list[dict[str, object]]: + filtered_messages: Final[list[dict[str, object]]] = [] for msg in messages: if not isinstance(msg, dict): continue - contents: Any = msg.get("content") + contents: object = msg.get("content") if isinstance(contents, list): - cleaned: list[Any] = [] + cleaned: list[object] = [] for c in contents: if self._should_keep_content(content=c): cleaned.append(self._redact_base64(value=c, max_depth=max_depth)) diff --git a/litellm/integrations/opik/opik_payload_builder/payload_builders.py b/litellm/integrations/opik/opik_payload_builder/payload_builders.py index e40d72ea542..855b84ba4c8 100644 --- a/litellm/integrations/opik/opik_payload_builder/payload_builders.py +++ b/litellm/integrations/opik/opik_payload_builder/payload_builders.py @@ -17,12 +17,12 @@ def build_trace_payload( end_time: datetime, input_data: Any, output_data: Any, - metadata: dict[str, Any], + metadata: dict[str, object], tags: list[str], thread_id: str | None, ) -> types.TracePayload: """Build a complete trace payload.""" - trace_name: Final = response_obj.get("object", "unknown type") + trace_name: Final[str] = response_obj.get("object", "unknown type") return types.TracePayload( project_name=project_name, @@ -47,7 +47,7 @@ def build_span_payload( end_time: datetime, input_data: Any, output_data: Any, - metadata: dict[str, Any], + metadata: dict[str, object], tags: list[str], usage: dict[str, int], provider: str | None = None, @@ -56,9 +56,9 @@ def build_span_payload( """Build a complete span payload.""" span_id: Final = utils.create_uuid7() - model: Final = response_obj.get("model", "unknown-model") - obj_type: Final = response_obj.get("object", "unknown-object") - created: Final = response_obj.get("created", 0) + model: Final[str] = response_obj.get("model", "unknown-model") + obj_type: Final[str] = response_obj.get("object", "unknown-object") + created: Final[int] = response_obj.get("created", 0) span_name: Final = f"{model}_{obj_type}_{created}" _logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id) diff --git a/litellm/integrations/prometheus_helpers/prometheus_api.py b/litellm/integrations/prometheus_helpers/prometheus_api.py index 9f77f87a670..f677db17648 100644 --- a/litellm/integrations/prometheus_helpers/prometheus_api.py +++ b/litellm/integrations/prometheus_helpers/prometheus_api.py @@ -4,9 +4,13 @@ Helper functions to query prometheus API import json import time +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta from typing import Final +from httpx import Response +from typing_extensions import ReadOnly, TypedDict + from litellm import get_secret from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -19,9 +23,45 @@ PROMETHEUS_SELECTED_INSTANCE: Final[str | None] = get_secret("PROMETHEUS_SELECTE async_http_handler: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) +class _PrometheusSeries(TypedDict): + """One time series in a Prometheus query result, with its labels and its ``[timestamp, value]`` samples.""" + + metric: ReadOnly[Mapping[str, str]] + values: ReadOnly[Sequence[Sequence[float | str]]] + + +class _PrometheusResultData(TypedDict): + """The ``data`` envelope of a Prometheus query response.""" + + result: ReadOnly[Sequence[_PrometheusSeries]] + + +class _PrometheusQueryResponse(TypedDict): + """The JSON body returned by the Prometheus ``/api/v1/query`` and ``/api/v1/query_range`` endpoints.""" + + data: ReadOnly[_PrometheusResultData] + + +class _DailySpend(TypedDict): + """One day of spend, in the shape ``get_daily_spend_from_prometheus`` returns.""" + + date: ReadOnly[str] + spend: ReadOnly[float] + + +def _query_body(response: Response) -> _PrometheusQueryResponse: + """Read the untyped JSON body of a Prometheus query response.""" + return response.json() + + +def _query_series(response: Response) -> Sequence[_PrometheusSeries]: + """Read the time series list out of the untyped JSON body of a Prometheus query response.""" + return response.json()["data"]["result"] + + async def get_metric_from_prometheus( metric_name: str, -): +) -> Sequence[_PrometheusSeries]: # Get the start of the current day in Unix timestamp if PROMETHEUS_URL is None: raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env") @@ -31,13 +71,13 @@ async def get_metric_from_prometheus( response: Final = await async_http_handler.get( f"{PROMETHEUS_URL}/api/v1/query", params={"query": query, "time": now} ) # End of the day - _json_response: Final = response.json() + _json_response: Final = _query_body(response) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] + results: Final = _query_series(response) return results -async def get_fallback_metric_from_prometheus(): +async def get_fallback_metric_from_prometheus() -> str: """ Gets fallback metrics from prometheus for the last 24 hours """ @@ -96,7 +136,7 @@ def _quote_promql_string_literal(value: str) -> str: return json.dumps(value, ensure_ascii=False) -async def get_daily_spend_from_prometheus(api_key: str | None): +async def get_daily_spend_from_prometheus(api_key: str | None) -> Sequence[_DailySpend]: """ Expected Response Format: [ @@ -133,9 +173,9 @@ async def get_daily_spend_from_prometheus(api_key: str | None): } response: Final = await async_http_handler.get(url, params=params) - _json_response: Final = response.json() + _json_response: Final = _query_body(response) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] + results: Final = _query_series(response) formatted_results: Final = [] for result in results: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index e59ef0449d0..e1ec3360351 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, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, TypeVar, cast from typing_extensions import ReadOnly @@ -74,6 +74,10 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b # ``web_search_tool_result`` blocks to inject into the final response. WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" +_RESPONSE_CONTENT_FIELD: Final = "content" + +_ResponseT: Final = TypeVar("_ResponseT") + class _PlanMetadataView(TypedDict): websearch_native_blocks: Sequence[Mapping[str, object]] | None @@ -926,17 +930,17 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: + def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response if isinstance(response, dict): - existing = response.get("content") or [] - response["content"] = list(native_blocks) + list(existing) + existing = response.get(_RESPONSE_CONTENT_FIELD) or [] + response[_RESPONSE_CONTENT_FIELD] = list(native_blocks) + list(existing) return response - existing = getattr(response, "content", None) or [] + existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or [] try: - response.content = list(native_blocks) + list(existing) + setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing)) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 44fed944d2a..1737e2b8cb0 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,6 +1,9 @@ import datetime +from collections.abc import Mapping from typing import Any, Final +import httpx + from litellm.constants import LITELLM_DETAILED_TIMING from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base @@ -25,11 +28,7 @@ class ResponseMetadata: @property def supports_response_time(self) -> bool: """Check if response type supports timing metrics""" - return ( - isinstance(self.result, ModelResponse) - or isinstance(self.result, EmbeddingResponse) - or isinstance(self.result, TranscriptionResponse) - ) + return isinstance(self.result, (ModelResponse, EmbeddingResponse, TranscriptionResponse)) def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict) -> None: """Set hidden parameters on the response""" @@ -45,14 +44,14 @@ class ResponseMetadata: result=self.result, litellm_model_name=model, router_model_id=model_id ), "additional_headers": process_response_headers( - self._get_value_from_hidden_params("additional_headers") or {}, + self._get_additional_headers_from_hidden_params() or {}, preserve_litellm_internal_headers=True, ), "litellm_model_name": model, } self._update_hidden_params(new_params) - def _update_hidden_params(self, new_params: dict) -> None: + def _update_hidden_params(self, new_params: Mapping[str, object]) -> None: """ Update hidden params - handles when self._hidden_params is a dict or HiddenParams object """ @@ -64,12 +63,12 @@ class ResponseMetadata: for key, value in new_params.items(): setattr(self._hidden_params, key, value) - def _get_value_from_hidden_params(self, key: str) -> Any | None: - """Get value from hidden params - handles when self._hidden_params is a dict or HiddenParams object""" + def _get_additional_headers_from_hidden_params(self) -> httpx.Headers | dict[str, str] | None: + """Get `additional_headers` from hidden params - handles when self._hidden_params is a dict or HiddenParams object""" if isinstance(self._hidden_params, dict): - return self._hidden_params.get(key, None) + return self._hidden_params.get("additional_headers", None) elif isinstance(self._hidden_params, HiddenParams): - return getattr(self._hidden_params, key, None) + return getattr(self._hidden_params, "additional_headers", None) def set_timing_metrics( self, @@ -96,7 +95,7 @@ class ResponseMetadata: ######################################################### # 2. Add LiteLLM overhead duration ######################################################### - llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") + llm_api_duration_ms: Final[float | None] = logging_obj.model_call_details.get("llm_api_duration_ms") if llm_api_duration_ms is not None: overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) self._update_hidden_params( @@ -108,7 +107,7 @@ class ResponseMetadata: ######################################################### # 3. Add callback processing duration ######################################################### - callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None) + callback_duration_ms: Final[float | None] = getattr(logging_obj, "callback_duration_ms", None) if callback_duration_ms is not None: self._update_hidden_params( { @@ -136,17 +135,17 @@ class ResponseMetadata: # 5. Detailed per-phase timing (opt-in via env var) ######################################################### if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None: - detailed: Final[dict] = { + detailed: Final[dict[str, float]] = { "timing_llm_api_ms": round(llm_api_duration_ms, 4), } # message copy time from Logging.__init__() - msg_copy_ms: Final = getattr(logging_obj, "message_copy_duration_ms", None) + msg_copy_ms: Final[float | None] = getattr(logging_obj, "message_copy_duration_ms", None) if msg_copy_ms is not None: detailed["timing_message_copy_ms"] = round(msg_copy_ms, 4) # pre-processing = time from request start to LLM API call start - api_call_start: Final = logging_obj.model_call_details.get("api_call_start_time") + api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time") if api_call_start is not None and start_time is not None: pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000 detailed["timing_pre_processing_ms"] = round(pre_ms, 4) diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index ea4be1c856f..7d412229a96 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -2,10 +2,54 @@ Utility functions for ModelResponse and ModelResponseStream objects. """ -from typing import Any, Final +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, Protocol from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream +_NO_EXTRA_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _HasModelExtra(Protocol): + """A Pydantic model, seen through the extra fields it collected.""" + + @property + def model_extra(self) -> Mapping[str, object] | None: ... + + +class _StreamingChoice(_HasModelExtra, Protocol): + """The streaming-choice fields this emptiness check reads.""" + + @property + def finish_reason(self) -> object: ... + + @property + def logprobs(self) -> object: ... + + @property + def enhancements(self) -> object: ... + + @property + def delta(self) -> Delta | None: ... + + +def _extra_fields(model: _HasModelExtra) -> Mapping[str, object]: + """The dynamically added fields Pydantic stored on ``model``.""" + return model.model_extra or _NO_EXTRA_FIELDS + + +def _attribute(obj: object, name: str) -> object: + """The named attribute of ``obj``, or ``None`` when it is absent.""" + attribute: Final[object] = getattr(obj, name, None) + return attribute + + +def _has_callable_attribute(obj: object, name: str) -> bool: + """Whether the named attribute of ``obj`` is callable.""" + attribute: Final[object] = getattr(obj, name) + return callable(attribute) + def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: """ @@ -41,7 +85,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: # Check model_extra for dynamically added fields (this is where Pydantic stores them) if hasattr(model_response, "model_extra") and model_response.model_extra: - for extra_field_name, extra_field_value in model_response.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(model_response).items(): if _has_meaningful_content(extra_field_value): return False @@ -57,7 +101,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: continue # Check if any other field has meaningful content - model_response_value = getattr(model_response, model_response_field, None) + model_response_value = _attribute(model_response, model_response_field) if _has_meaningful_content(model_response_value): return False @@ -71,7 +115,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: return True -def _has_meaningful_content(value: Any) -> bool: +def _has_meaningful_content(value: object) -> bool: """ Check if a value contains meaningful content. @@ -102,7 +146,7 @@ def _has_meaningful_content(value: Any) -> bool: return True -def _is_choice_non_empty(choice: Any) -> bool: +def _is_choice_non_empty(choice: _StreamingChoice) -> bool: """ Deep check if a choice contains any meaningful content. @@ -131,7 +175,7 @@ def _is_choice_non_empty(choice: Any) -> bool: # Check model_extra for dynamically added fields on the choice if hasattr(choice, "model_extra") and choice.model_extra: - for extra_field_name, extra_field_value in choice.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(choice).items(): # Skip certain structural fields that are just default/None placeholders if extra_field_name == "index" and extra_field_value == 0: continue @@ -147,7 +191,7 @@ def _is_choice_non_empty(choice: Any) -> bool: # Skip private attributes, methods, and known empty fields if ( attr_name.startswith("_") - or callable(getattr(choice, attr_name)) + or _has_callable_attribute(choice, attr_name) or attr_name.startswith("model_") or attr_name in { @@ -160,7 +204,7 @@ def _is_choice_non_empty(choice: Any) -> bool: ): continue - attr_value = getattr(choice, attr_name, None) + attr_value = _attribute(choice, attr_name) if _has_meaningful_content(attr_value): return True @@ -179,7 +223,7 @@ def _is_delta_non_empty(delta: Delta) -> bool: """ # Check model_extra for dynamically added fields (this is where Pydantic stores them) if hasattr(delta, "model_extra") and delta.model_extra: - for extra_field_name, extra_field_value in delta.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(delta).items(): # Even structural fields are meaningful if they have actual content if _has_meaningful_content(extra_field_value): return True @@ -187,10 +231,10 @@ def _is_delta_non_empty(delta: Delta) -> bool: # Check all regular attributes of the delta object for attr_name in dir(delta): # Skip private attributes, methods, and Pydantic-specific fields - if attr_name.startswith("_") or callable(getattr(delta, attr_name)) or attr_name.startswith("model_"): + if attr_name.startswith("_") or _has_callable_attribute(delta, attr_name) or attr_name.startswith("model_"): continue - attr_value = getattr(delta, attr_name, None) + attr_value = _attribute(delta, attr_name) if _has_meaningful_content(attr_value): return True diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 485091bccd0..300ba427cda 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -92,7 +92,7 @@ def print_verbose(print_statement: object): @dataclass(frozen=True, slots=True) class _ProviderChunkParsed: - response_obj: dict[str, Any] + response_obj: dict[str, object] @dataclass(frozen=True, slots=True) @@ -1242,7 +1242,7 @@ class CustomStreamWrapper: for key, value in anthropic_response_obj["provider_specific_fields"].items(): setattr(model_response, key, value) - response_obj = cast(dict[str, Any], anthropic_response_obj) + response_obj = cast(dict[str, object], anthropic_response_obj) elif self.model == "replicate" or self.custom_llm_provider == "replicate": response_obj = self.handle_replicate_chunk(chunk) completion_obj["content"] = response_obj["text"] @@ -1398,7 +1398,7 @@ class CustomStreamWrapper: if not isinstance(chunk, str): raise ValueError(f"chunk is not a string: {chunk}") response_obj = cast( - dict[str, Any], + dict[str, object], litellm.CodestralTextCompletionConfig()._chunk_parser(chunk), ) completion_obj["content"] = response_obj["text"] @@ -2462,7 +2462,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: prompt_tokens: int = 0 completion_tokens: int = 0 - latest_usage_chunk = None + latest_usage_chunk: Usage | Mapping[str, int] | None = None prompt_tokens_details: PromptTokensDetailsWrapper | None = None completion_tokens_details: CompletionTokensDetailsWrapper | None = None cache_creation_token_details: CacheCreationTokenDetails | None = None diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0a59eaa75d3..53ae25b928c 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -20,6 +20,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): """ import socket +from collections.abc import Sequence from ipaddress import ip_address, ip_network from typing import Any, Final from urllib.parse import quote, urlparse, urlunparse @@ -44,7 +45,7 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" -def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str: +def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. ``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986 @@ -64,7 +65,7 @@ def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") - return quote(value_str, safe="") -def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: +def encode_url_path_segments(value: object, *, field_name: str = "path") -> str: """Percent-encode a user-controlled URL path made of multiple segments. Empty segments are rejected, so leading, trailing, or consecutive slashes @@ -77,9 +78,9 @@ def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: if value_str == "": raise ValueError(f"{field_name} is required") - encoded_segments: Final = [] - for segment in value_str.split("/"): - encoded_segments.append(encode_url_path_segment(segment, field_name=field_name)) + encoded_segments: Final = tuple( + encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/") + ) return "/".join(encoded_segments) @@ -202,7 +203,7 @@ def _format_host_header(hostname: str, port: int, default_port: int) -> str: return f"{bracketed}:{port}" -def _sockaddr_host(sockaddr: Any) -> str: +def _sockaddr_host(sockaddr: Sequence[object]) -> str: """Return the host element of a ``getaddrinfo`` sockaddr as ``str``. ``getaddrinfo`` with ``IPPROTO_TCP`` returns AF_INET / AF_INET6 sockaddrs @@ -285,7 +286,7 @@ def validate_url(url: str) -> tuple[str, str]: raise SSRFError(f"No addresses found for '{hostname}'") if not is_allowlisted: - for family, type_, proto, canonname, sockaddr in addrinfo: + for _family, _type, _proto, _canonname, sockaddr in addrinfo: resolved_ip = _sockaddr_host(sockaddr) if _is_blocked_ip(resolved_ip): raise SSRFError( @@ -363,7 +364,7 @@ def assert_same_origin(candidate_url: str, expected_url: str) -> None: _MAX_REDIRECTS: Final = 10 -def _extract_redirect_url(response: Any, request_url: str) -> str: +def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: """Extract and resolve the redirect target from a response's Location header.""" location: Final = response.headers.get("location") if not isinstance(location, str) or not location: @@ -372,7 +373,7 @@ def _extract_redirect_url(response: Any, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: Any, url: str, **kwargs: Any) -> Any: +def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -398,7 +399,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = client.get( + response: httpx.Response = client.get( validated_url, headers={**caller_headers, "Host": original_host}, follow_redirects=False, @@ -412,7 +413,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: raise SSRFError("Too many redirects") -async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: +async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) @@ -421,7 +422,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = await client.get( + response: httpx.Response = await client.get( validated_url, headers={**caller_headers, "Host": original_host}, follow_redirects=False, diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 3f8fd2c27f4..3071c3ddd58 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -1,9 +1,11 @@ import json import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast import httpx from httpx import Headers, Response +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig @@ -19,6 +21,29 @@ else: LoggingClass = Any +class AnthropicBatchRequestCounts(TypedDict, total=False): + """The ``request_counts`` object of an Anthropic Message Batch.""" + + processing: ReadOnly[int] + succeeded: ReadOnly[int] + errored: ReadOnly[int] + canceled: ReadOnly[int] + expired: ReadOnly[int] + + +class AnthropicMessageBatch(TypedDict, total=False): + """The fields of an Anthropic Message Batch that map onto an OpenAI Batch.""" + + id: ReadOnly[str] + processing_status: ReadOnly[str] + created_at: ReadOnly[str | None] + ended_at: ReadOnly[str | None] + expires_at: ReadOnly[str | None] + cancel_initiated_at: ReadOnly[str | None] + archived_at: ReadOnly[str | None] + request_counts: ReadOnly[AnthropicBatchRequestCounts] + + class AnthropicBatchesConfig(BaseBatchesConfig): def __init__(self): from ..chat.transformation import AnthropicConfig @@ -83,7 +108,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): create_batch_data: CreateBatchRequest, optional_params: dict, litellm_params: dict, - ) -> bytes | str | dict[str, Any]: + ) -> bytes | str | dict[str, object]: """ Transform the batch creation request to Anthropic format. @@ -133,7 +158,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): batch_id: str, optional_params: dict, litellm_params: dict, - ) -> bytes | str | dict[str, Any]: + ) -> bytes | str | dict[str, object]: """ Transform batch retrieval request for Anthropic. @@ -152,7 +177,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): ) -> LiteLLMBatch: """Transform Anthropic MessageBatch retrieval response to LiteLLM format.""" try: - response_data: Final = raw_response.json() + response_data: Final[AnthropicMessageBatch] = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Anthropic batch response: {e}") @@ -161,18 +186,20 @@ class AnthropicBatchesConfig(BaseBatchesConfig): processing_status: Final = response_data.get("processing_status", "in_progress") # Map Anthropic processing_status to OpenAI status - status_mapping: dict[ - str, - Literal[ - "validating", - "failed", - "in_progress", - "finalizing", - "completed", - "expired", - "cancelling", - "cancelled", - ], + status_mapping: Final[ + Mapping[ + str, + Literal[ + "validating", + "failed", + "in_progress", + "finalizing", + "completed", + "expired", + "cancelling", + "cancelled", + ], + ] ] = { "in_progress": "in_progress", "canceling": "cancelling", @@ -279,7 +306,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): if not line: continue try: - response_json = json.loads(line) + response_json: Mapping[str, Mapping[str, dict[str, object]]] = json.loads(line) # Update model_response with the parsed JSON completion_response = response_json["result"]["message"] transformed_response = self.anthropic_chat_config.transform_parsed_response( diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py index 41795fa0f32..8e06b73a92c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -1,8 +1,10 @@ """Dispatch ``context_management`` edits to registered polyfill editors.""" import inspect -from collections.abc import Awaitable, Callable -from typing import Any, Final, cast +from collections.abc import Awaitable, Callable, Mapping +from typing import TYPE_CHECKING, Final, TypeAlias, TypedDict, cast + +from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm.types.llms.anthropic import AppliedEdit @@ -11,31 +13,54 @@ from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112 from .result import PolyfillResult -EditorFn = Callable[..., Any] +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router -_EDITOR_REGISTRY: Final[dict[str, EditorFn]] = { +AnthropicMessages: TypeAlias = list[dict[str, object]] +AnthropicSystem: TypeAlias = str | list[dict[str, object]] | None +AnthropicTools: TypeAlias = list[dict[str, object]] | None +EditSpec: TypeAlias = dict[str, object] +ContextManagementSpec: TypeAlias = EditSpec | list[EditSpec] | None +SyncEditorReturn: TypeAlias = tuple[AnthropicMessages, AppliedEdit | None] +EditorFn: TypeAlias = Callable[..., object] + + +class EditorKwargs(TypedDict): + """The keyword payload every registered editor accepts.""" + + model: ReadOnly[str] + messages: ReadOnly[AnthropicMessages] + tools: ReadOnly[AnthropicTools] + system: ReadOnly[AnthropicSystem] + edit_spec: ReadOnly[EditSpec] + + +_EDITOR_REGISTRY: Final[Mapping[str, EditorFn]] = { CLEAR_TOOL_USES_EDIT_TYPE: apply_clear_tool_uses_20250919, COMPACT_EDIT_TYPE: apply_compact_20260112, } -def _normalize_spec( - spec: dict[str, Any] | list[dict[str, Any]] | None, -) -> list[dict[str, Any]] | None: +def _map_openai_spec(spec: list[EditSpec]) -> EditSpec | None: + """Translate the OpenAI list form into the Anthropic-native dict form.""" + # Local import to avoid an import cycle at module load. + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + return AnthropicConfig.map_openai_context_management_to_anthropic(spec) + + +def _normalize_spec(spec: ContextManagementSpec) -> list[EditSpec] | None: """Accept Anthropic-native dict form or OpenAI list form; return edits list.""" - if isinstance(spec, list): - # Local import to avoid an import cycle at module load. - from litellm.llms.anthropic.chat.transformation import AnthropicConfig + normalized: Final = _map_openai_spec(spec) if isinstance(spec, list) else spec - spec = AnthropicConfig.map_openai_context_management_to_anthropic(spec) - - edits: Final = spec.get("edits") if isinstance(spec, dict) else None + edits: Final = normalized.get("edits") if isinstance(normalized, dict) else None if not edits or not isinstance(edits, list): return None return [edit for edit in edits if isinstance(edit, dict)] -def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: +def _wrap_editor_return(raw: object, *, fallback_system: AnthropicSystem) -> PolyfillResult: """Coerce an editor's native return shape into a ``PolyfillResult``. v0 sync editors (e.g. ``clear_tool_uses_20250919``) return a 2-tuple @@ -46,7 +71,7 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: return raw # Legacy 2-tuple return — sync editors don't mutate ``system``, so # carry the caller's value forward. - messages, applied = cast(tuple[list[dict[str, Any]], Any], raw) + messages, applied = cast(SyncEditorReturn, raw) return PolyfillResult( messages=messages, system=fallback_system, @@ -57,13 +82,13 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: async def apply_context_management( *, model: str, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, - system: Any, - context_management_spec: dict[str, Any] | list[dict[str, Any]] | None, - litellm_metadata: dict[str, Any] | None = None, - llm_router: Any = None, - user_api_key_auth: Any = None, + messages: AnthropicMessages, + tools: AnthropicTools, + system: AnthropicSystem, + context_management_spec: ContextManagementSpec, + litellm_metadata: Mapping[str, object] | None = None, + llm_router: "Router | None" = None, + user_api_key_auth: "UserAPIKeyAuth | None" = None, ) -> PolyfillResult: """Run edits in order; return a single ``PolyfillResult``. @@ -92,7 +117,7 @@ async def apply_context_management( ) continue - kwargs: dict[str, Any] = { + kwargs: EditorKwargs = { "model": model, "messages": current_messages, "tools": tools, @@ -102,10 +127,12 @@ async def apply_context_management( # Only async editors accept these — passing them to sync v0 editors # would break their signature. if inspect.iscoroutinefunction(editor): - kwargs["litellm_metadata"] = litellm_metadata - kwargs["llm_router"] = llm_router - kwargs["user_api_key_auth"] = user_api_key_auth - raw_result = await cast(Callable[..., Awaitable[Any]], editor)(**kwargs) + raw_result = await cast(Callable[..., Awaitable[PolyfillResult]], editor)( + **kwargs, + litellm_metadata=litellm_metadata, + llm_router=llm_router, + user_api_key_auth=user_api_key_auth, + ) else: raw_result = editor(**kwargs) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index e6d8686b466..d99c7f556fb 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -33,18 +33,18 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, def _build_responses_kwargs( *, max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, @@ -134,22 +134,22 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, - **kwargs, + **kwargs: object, ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: responses_kwargs: Final = _build_responses_kwargs( max_tokens=max_tokens, @@ -185,23 +185,23 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, _is_async: bool = False, - **kwargs, + **kwargs: object, ) -> ( AnthropicMessagesResponse | AsyncIterator[bytes] diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 566322bbdd6..20f3a3d74f0 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -2,9 +2,10 @@ Anthropic Skills API configuration and transformations """ -from typing import Any, Final +from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -23,6 +24,25 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +class _SkillPayload(TypedDict): + """The JSON body Anthropic returns for a single skill, before ``Skill`` validates it.""" + + id: ReadOnly[str] + created_at: ReadOnly[str] + source: ReadOnly[str] + updated_at: ReadOnly[str] + display_title: NotRequired[ReadOnly[str | None]] + latest_version: NotRequired[ReadOnly[str | None]] + type: NotRequired[ReadOnly[str]] + + +class _DeleteSkillPayload(TypedDict): + """The JSON body Anthropic returns for a skill deletion, before ``DeleteSkillResponse`` validates it.""" + + id: ReadOnly[str] + type: NotRequired[ReadOnly[str]] + + class AnthropicSkillsConfig(BaseSkillsAPIConfig): """Anthropic-specific Skills API configuration""" @@ -104,7 +124,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final[_SkillPayload] = raw_response.json() verbose_logger.debug("Transforming create skill response: %s", response_json) return Skill(**response_json) @@ -122,7 +142,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): url: Final = self.get_complete_url(api_base=api_base, endpoint="skills") # Build query parameters - query_params: Final[dict[str, Any]] = {} + query_params: Final[dict[str, int | str]] = {} if "limit" in list_params and list_params["limit"]: query_params["limit"] = list_params["limit"] if "page" in list_params and list_params["page"]: @@ -168,7 +188,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final[_SkillPayload] = raw_response.json() verbose_logger.debug("Transforming get skill response: %s", response_json) return Skill(**response_json) @@ -193,7 +213,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteSkillResponse: """Transform Anthropic response to DeleteSkillResponse""" - response_json: Final = raw_response.json() + response_json: Final[_DeleteSkillPayload] = raw_response.json() verbose_logger.debug("Transforming delete skill response: %s", response_json) return DeleteSkillResponse(**response_json) diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index 4f93896699f..f3ed2e6d6ed 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -40,6 +40,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): data.pop("expires_after", None) return data + @staticmethod + def _to_openai_file_object(response: FileObject) -> OpenAIFileObject: + """Re-wrap the SDK's file object as litellm's, carrying every field across.""" + return OpenAIFileObject(**response.model_dump()) + async def acreate_file( self, create_file_data: CreateFileRequest, @@ -48,7 +53,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): verbose_logger.debug("create_file_data=%s", create_file_data) response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data)) verbose_logger.debug("create_file_response=%s", response) - return OpenAIFileObject(**response.model_dump()) + return self._to_openai_file_object(response) def create_file( self, @@ -61,7 +66,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: int | None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, litellm_params: dict | None = None, - ) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]: + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -84,7 +89,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): response: Final = cast(AzureOpenAI | OpenAI, openai_client).files.create( **self._prepare_create_file_data(create_file_data) ) - return OpenAIFileObject(**response.model_dump()) + return self._to_openai_file_object(response) async def afile_content( self, @@ -105,7 +110,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): api_version: str | None = None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, litellm_params: dict | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3eeb3cb9fc6..e795e907eb4 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,9 +7,10 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json -from typing import Any, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -18,9 +19,74 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig +if TYPE_CHECKING: + from litellm.types.realtime import RealtimeResponseTransformInput + _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) +class _ClientWebSocket(Protocol): + """The client-facing websocket surface used by the Bedrock realtime bridge.""" + + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + + +class _BedrockInputStream(Protocol): + """The write half of a Bedrock bidirectional stream.""" + + async def send(self, chunk: object) -> None: ... + + async def close(self) -> None: ... + + +class _BedrockPayloadPart(Protocol): + """A single Bedrock bidirectional output payload.""" + + bytes_: bytes | None + + +class _BedrockOutputChunk(Protocol): + """A chunk read off the read half of a Bedrock bidirectional stream.""" + + value: _BedrockPayloadPart | None + + +class _BedrockOutputStream(Protocol): + """The read half of a Bedrock bidirectional stream.""" + + async def receive(self) -> _BedrockOutputChunk | None: ... + + +class _BedrockBidirectionalStream(Protocol): + """The bidirectional stream returned by ``invoke_model_with_bidirectional_stream``.""" + + input_stream: _BedrockInputStream + + async def await_output(self) -> tuple[object, _BedrockOutputStream]: ... + + +class _ClientSessionPayload(TypedDict, total=False): + """The ``session`` body of a client ``session.update`` frame.""" + + modalities: ReadOnly[object] + + +class _ClientRealtimeFrame(TypedDict, total=False): + """The fields read off a client realtime frame.""" + + type: ReadOnly[str] + session: ReadOnly[_ClientSessionPayload] + + +def _decode_client_frame(payload: str) -> _ClientRealtimeFrame: + """Decode a client realtime frame into the fields this bridge reads.""" + return json.loads(payload) + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -30,7 +96,7 @@ class BedrockRealtime(BaseAWSLLM): async def async_realtime( self, model: str, - websocket: Any, + websocket: _ClientWebSocket, logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, @@ -46,7 +112,7 @@ class BedrockRealtime(BaseAWSLLM): aws_sts_endpoint: str | None = None, aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, - **kwargs, + **kwargs: object, ): """ Establish bidirectional streaming connection with Bedrock Nova Sonic. @@ -118,13 +184,16 @@ class BedrockRealtime(BaseAWSLLM): ) bedrock_client: Final = BedrockRuntimeClient(config=config) + async def open_bidirectional_stream() -> _BedrockBidirectionalStream: + return await bedrock_client.invoke_model_with_bidirectional_stream( + InvokeModelWithBidirectionalStreamOperationInput(model_id=model) + ) + transformation_config: Final = BedrockRealtimeConfig() try: # Initialize the bidirectional stream - bedrock_stream: Final = await bedrock_client.invoke_model_with_bidirectional_stream( - InvokeModelWithBidirectionalStreamOperationInput(model_id=model) - ) + bedrock_stream: Final = await open_bidirectional_stream() verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") @@ -132,7 +201,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") # Track state for transformation - session_state: Final = { + session_state: Final[RealtimeResponseTransformInput] = { "current_output_item_id": None, "current_response_id": None, "current_conversation_id": None, @@ -182,11 +251,11 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_client_to_bedrock( self, - client_ws: Any, - bedrock_stream: Any, + client_ws: _ClientWebSocket, + bedrock_stream: _BedrockBidirectionalStream, transformation_config: BedrockRealtimeConfig, model: str, - session_state: dict, + session_state: "RealtimeResponseTransformInput", logging_obj: LiteLLMLogging | None = None, ): """Forward messages from client WebSocket to Bedrock stream.""" @@ -195,10 +264,11 @@ class BedrockRealtime(BaseAWSLLM): InvokeModelWithBidirectionalStreamInputChunk, ) + def build_input_chunk(payload: bytes) -> object: + return InvokeModelWithBidirectionalStreamInputChunk(value=BidirectionalInputPayloadPart(bytes_=payload)) + async def send_to_bedrock(bedrock_message: str) -> None: - event: Final = InvokeModelWithBidirectionalStreamInputChunk( - value=BidirectionalInputPayloadPart(bytes_=bedrock_message.encode("utf-8")) - ) + event: Final = build_input_chunk(bedrock_message.encode("utf-8")) await bedrock_stream.input_stream.send(event) verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) @@ -223,7 +293,7 @@ class BedrockRealtime(BaseAWSLLM): client_message_type: str | None = None requested_modalities: list[str] | None = None with contextlib.suppress(Exception): - parsed_client_message = json.loads(message) + parsed_client_message = _decode_client_frame(message) client_message_type = parsed_client_message.get("type") if client_message_type == "session.update": requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( @@ -246,12 +316,12 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_bedrock_to_client( self, - bedrock_stream: Any, - client_ws: Any, + bedrock_stream: _BedrockBidirectionalStream, + client_ws: _ClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging, - session_state: dict, + session_state: "RealtimeResponseTransformInput", ): """Forward messages from Bedrock stream to client WebSocket.""" try: diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index 57f679947f6..f8c54c0a1d8 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -4,7 +4,32 @@ Streaming utilities for ChatGPT provider. Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API. """ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final, Protocol + +from litellm.types.utils import Delta + + +class ChatGPTStreamChoice(Protocol): + """Streaming choice as read by :class:`ChatGPTToolCallNormalizer`.""" + + @property + def delta(self) -> Delta | None: ... + + +class ChatGPTStreamChunk(Protocol): + """Streaming chunk as read by :class:`ChatGPTToolCallNormalizer`.""" + + @property + def choices(self) -> Sequence[ChatGPTStreamChoice]: ... + + +class ChatGPTChunkStream(Protocol): + """Sync/async chunk source wrapped by :class:`ChatGPTToolCallNormalizer`.""" + + def __next__(self) -> ChatGPTStreamChunk: ... + + async def __anext__(self) -> ChatGPTStreamChunk: ... class ChatGPTToolCallNormalizer: @@ -20,36 +45,36 @@ class ChatGPTToolCallNormalizer: chunks to the consumer. """ - def __init__(self, stream: Any): + def __init__(self, stream: ChatGPTChunkStream): self._stream = stream self._seen_ids: dict[str, int] = {} # tool_call_id -> assigned_index self._next_index: int = 0 self._last_id: str | None = None # tracks which tool call the next delta belongs to - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: return getattr(self._stream, name) - def __iter__(self): + def __iter__(self) -> "ChatGPTToolCallNormalizer": return self - def __aiter__(self): + def __aiter__(self) -> "ChatGPTToolCallNormalizer": return self - def __next__(self): + def __next__(self) -> ChatGPTStreamChunk: while True: chunk = next(self._stream) result = self._normalize(chunk) if result is not None: return result - async def __anext__(self): + async def __anext__(self) -> ChatGPTStreamChunk: while True: chunk = await self._stream.__anext__() result = self._normalize(chunk) if result is not None: return result - def _normalize(self, chunk: Any) -> Any: + def _normalize(self, chunk: ChatGPTStreamChunk) -> ChatGPTStreamChunk | None: """Fix tool_calls in the chunk. Returns None to skip duplicate chunks.""" if not chunk.choices: return chunk diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index 44e1ab15801..3d189911e9e 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -2,13 +2,16 @@ CompactifAI chat completion transformation """ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import OpenAIError from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig @@ -21,6 +24,18 @@ else: LiteLLMLoggingObj = Any +class CompactifAIResponseFields(TypedDict, total=False): + """The chat completion fields of a CompactifAI response body.""" + + id: ReadOnly[str] + choices: ReadOnly[Sequence[Mapping[str, object]]] + created: ReadOnly[int] + model: ReadOnly[str] + system_fingerprint: ReadOnly[str | None] + usage: ReadOnly[Mapping[str, object]] + object: ReadOnly[str] + + class CompactifAIChatConfig(OpenAIGPTConfig): """ Configuration class for CompactifAI chat completions. @@ -45,11 +60,11 @@ class CompactifAIChatConfig(OpenAIGPTConfig): raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, - request_data: dict, - messages: list, - optional_params: dict, - litellm_params: dict, - encoding: Any, + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -79,14 +94,18 @@ class CompactifAIChatConfig(OpenAIGPTConfig): message["content"] = tool_calls[0]["function"].get("arguments", "") message["tool_calls"] = None - returned_response: Final = ModelResponse(**response_json) + response_fields: Final[CompactifAIResponseFields] = response_json + + returned_response: Final = ModelResponse(**response_fields) # Set model name with provider prefix returned_response.model = f"compactifai/{model}" return returned_response - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers + ) -> BaseLLMException: """ Get the appropriate error class for CompactifAI errors. Since CompactifAI is OpenAI-compatible, we use OpenAI error handling. diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index 7690351e3b2..343cf043187 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -6,11 +6,13 @@ endpoint defined in endpoints.json, eliminating the need for individual handler """ import json -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping, Sequence from pathlib import Path from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import BaseModel +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -32,22 +34,54 @@ if TYPE_CHECKING: from litellm.llms.base_llm.containers.transformation import BaseContainerConfig +class EndpointConfig(TypedDict): + """One endpoint entry of ``litellm/containers/endpoints.json``.""" + + name: ReadOnly[str] + async_name: ReadOnly[str] + path: ReadOnly[str] + method: ReadOnly[str] + path_params: ReadOnly[Sequence[str]] + query_params: ReadOnly[Sequence[str]] + response_type: ReadOnly[str] + is_multipart: NotRequired[ReadOnly[bool]] + returns_binary: NotRequired[ReadOnly[bool]] + + +class EndpointsConfig(TypedDict): + """The parsed ``litellm/containers/endpoints.json`` document.""" + + endpoints: ReadOnly[Sequence[EndpointConfig]] + + +class ContainerErrorDetail(TypedDict, total=False): + """The ``error`` object of a container API error body.""" + + message: ReadOnly[str] + + +class ContainerResponseBody(TypedDict, total=False): + """The fields this handler reads off a container API JSON body.""" + + error: ReadOnly[ContainerErrorDetail] + + # Response type mapping -RESPONSE_TYPES: Final[dict[str, type]] = { +RESPONSE_TYPES: Final[Mapping[str, type[BaseModel]]] = { "ContainerFileListResponse": ContainerFileListResponse, "ContainerFileObject": ContainerFileObject, "DeleteContainerFileResponse": DeleteContainerFileResponse, } -def _load_endpoints_config() -> dict: +def _load_endpoints_config() -> EndpointsConfig: """Load the endpoints configuration from JSON file.""" config_path: Final = Path(__file__).parent.parent.parent / "containers" / "endpoints.json" with open(config_path) as f: return json.load(f) -def _get_endpoint_config(endpoint_name: str) -> dict | None: +def _get_endpoint_config(endpoint_name: str) -> EndpointConfig | None: """Get config for a specific endpoint by name.""" config: Final = _load_endpoints_config() for endpoint in config["endpoints"]: @@ -56,10 +90,15 @@ def _get_endpoint_config(endpoint_name: str) -> dict | None: return None +def _response_model(response_type_name: str) -> type[BaseModel] | None: + """The pydantic model a container endpoint's ``response_type`` names.""" + return RESPONSE_TYPES.get(response_type_name) + + def _build_url( api_base: str, path_template: str, - path_params: dict[str, str], + path_params: Mapping[str, object], ) -> str: """Build the full URL by substituting path parameters. @@ -89,22 +128,18 @@ def _build_url( def _build_query_params( - query_param_names: list, - kwargs: dict[str, Any], -) -> dict[str, str]: + query_param_names: Sequence[str], + kwargs: Mapping[str, object], +) -> dict[str, object]: """Build query parameters from kwargs.""" - params: Final = {} - for param_name in query_param_names: - value = kwargs.get(param_name) - if value is not None: - params[param_name] = str(value) if not isinstance(value, str) else value - return params + supplied: Final = ((param_name, kwargs.get(param_name)) for param_name in query_param_names) + return {name: value if isinstance(value, str) else str(value) for name, value in supplied if value is not None} def _prepare_multipart_file_upload( file: Any, - headers: dict[str, Any], -) -> tuple: + headers: dict[str, object], +) -> tuple[dict[str, tuple[str, bytes, str]], dict[str, object]]: """ Prepare file and headers for multipart upload. @@ -129,6 +164,52 @@ def _prepare_multipart_file_upload( return files, headers_copy +def _request_headers( + container_provider_config: "BaseContainerConfig", + extra_headers: dict[str, object] | None, + litellm_params: GenericLiteLLMParams, +) -> dict[str, object]: + """The provider auth headers for a container request.""" + return container_provider_config.validate_environment( + headers=extra_headers or {}, + api_key=litellm_params.get("api_key", None), + ) + + +def _request_api_base( + container_provider_config: "BaseContainerConfig", + litellm_params: GenericLiteLLMParams, +) -> str: + """The provider base URL for a container request.""" + return container_provider_config.get_complete_url( + api_base=litellm_params.get("api_base", None), + litellm_params=dict(litellm_params), + ) + + +def _sync_http_client( + client: HTTPHandler | AsyncHTTPHandler | None, + litellm_params: GenericLiteLLMParams, +) -> HTTPHandler: + """The sync HTTP client for a container request, reusing the caller's when usable.""" + if client is None or not isinstance(client, HTTPHandler): + return _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) + return client + + +def _async_http_client( + client: HTTPHandler | AsyncHTTPHandler | None, + litellm_params: GenericLiteLLMParams, +) -> AsyncHTTPHandler: + """The async HTTP client for a container request, reusing the caller's when usable.""" + if client is None or not isinstance(client, AsyncHTTPHandler): + return get_async_httpx_client( + llm_provider=litellm.LlmProviders.OPENAI, + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + return client + + class GenericContainerHandler: """ Generic handler for container file API endpoints. @@ -143,13 +224,13 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, _is_async: bool = False, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, - ) -> Any | Coroutine[Any, Any, Any]: + **kwargs: object, + ) -> Any | Coroutine[object, object, Any]: """ Generic handler for any container file endpoint. @@ -196,11 +277,11 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, + **kwargs: object, ) -> Any: """Synchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) @@ -208,23 +289,14 @@ class GenericContainerHandler: raise ValueError(f"Unknown endpoint: {endpoint_name}") # Get HTTP client - if client is None or not isinstance(client, HTTPHandler): - http_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) - else: - http_client = client + http_client: Final = _sync_http_client(client, litellm_params) # Build request - headers = container_provider_config.validate_environment( - headers=extra_headers or {}, - api_key=litellm_params.get("api_key", None), - ) + headers = _request_headers(container_provider_config, extra_headers, litellm_params) if extra_headers: headers.update(extra_headers) - api_base: Final = container_provider_config.get_complete_url( - api_base=litellm_params.get("api_base", None), - litellm_params=dict(litellm_params), - ) + api_base: Final = _request_api_base(container_provider_config, litellm_params) # Build URL with path params path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} @@ -275,11 +347,11 @@ class GenericContainerHandler: return response.content # Check for error response - response_json: Final = response.json() + response_json: Final[ContainerResponseBody] = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg: Final = response_json.get("error", {}).get("message", str(response_json)) + error_msg: Final = response_json["error"].get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, @@ -287,7 +359,7 @@ class GenericContainerHandler: ) # Parse response - response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"]) + response_type: Final = _response_model(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json @@ -301,11 +373,11 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, + **kwargs: object, ) -> Any: """Asynchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) @@ -313,26 +385,14 @@ class GenericContainerHandler: raise ValueError(f"Unknown endpoint: {endpoint_name}") # Get HTTP client - if client is None or not isinstance(client, AsyncHTTPHandler): - http_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.OPENAI, - params={"ssl_verify": litellm_params.get("ssl_verify", None)}, - ) - else: - http_client = client + http_client: Final = _async_http_client(client, litellm_params) # Build request - headers = container_provider_config.validate_environment( - headers=extra_headers or {}, - api_key=litellm_params.get("api_key", None), - ) + headers = _request_headers(container_provider_config, extra_headers, litellm_params) if extra_headers: headers.update(extra_headers) - api_base: Final = container_provider_config.get_complete_url( - api_base=litellm_params.get("api_base", None), - litellm_params=dict(litellm_params), - ) + api_base: Final = _request_api_base(container_provider_config, litellm_params) # Build URL with path params path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} @@ -383,11 +443,11 @@ class GenericContainerHandler: return response.content # Check for error response - response_json: Final = response.json() + response_json: Final[ContainerResponseBody] = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg: Final = response_json.get("error", {}).get("message", str(response_json)) + error_msg: Final = response_json["error"].get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, @@ -395,7 +455,7 @@ class GenericContainerHandler: ) # Parse response - response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"]) + response_type: Final = _response_model(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 2d0322bf10f..03037512551 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -6,11 +6,25 @@ import json import os import re import threading -from typing import Any, Final +from collections.abc import Callable +from typing import Any, Final, Protocol from urllib.parse import urlsplit import litellm from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig +from litellm.types.llms.openai import AllMessageValues + + +class _GDCHAudienceCredentials(Protocol): + """A GDCH service account credential already bound to an audience, ready to mint a bearer token.""" + + @property + def valid(self) -> bool: ... + + @property + def token(self) -> str: ... + + def refresh(self, request: object) -> None: ... class GDCGeminiConfig(OpenAILikeChatConfig): @@ -21,7 +35,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._creds_lock = threading.Lock() - self._gdch_creds_cache: dict = {} + self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {} def get_supported_openai_params(self, model: str) -> list: return [ @@ -110,7 +124,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return f"{api_base}/v1/projects/{project}/locations/{location}/chat/completions" - def _read_env_bool(self, val: Any, env_var: str, default: bool = True) -> bool | str: + def _read_env_bool(self, val: bool | str | None, env_var: str, default: bool = True) -> bool | str: def _parse(s: str) -> bool | str: cleaned: Final = s.strip().lower() if cleaned in ("false", "0", "no", "off"): @@ -129,7 +143,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return default return _parse(_env_val) - def _fetch_auth(self, gdch_creds: Any, ssl_verify: bool | str) -> None: + def _fetch_auth(self, gdch_creds: _GDCHAudienceCredentials, ssl_verify: bool | str) -> None: import requests from google.auth.transport import requests as auth_requests @@ -138,13 +152,24 @@ class GDCGeminiConfig(OpenAILikeChatConfig): auth_request: Final = auth_requests.Request(session=auth_session) gdch_creds.refresh(auth_request) - def _cached_fetch_token(self, creds: Any, audience: str, ssl_verify: bool | str, api_key: str | None = None) -> str: + def _with_gdch_audience(self, creds: object, audience: str) -> _GDCHAudienceCredentials: + """The credential rebound to ``audience``, which GDCH requires before a token refresh.""" + bind_audience: Final[Callable[[str], _GDCHAudienceCredentials] | None] = getattr( + creds, "with_gdch_audience", None + ) + if bind_audience is None: + raise AttributeError("GDC credentials must expose with_gdch_audience to be bound to a request audience") + return bind_audience(audience) + + def _cached_fetch_token( + self, creds: object, audience: str, ssl_verify: bool | str, api_key: str | None = None + ) -> str: # Key cache by both audience and credential identity to prevent cross-caller contamination cache_key: Final = (audience.rstrip("/"), api_key or str(id(creds))) with self._creds_lock: if cache_key not in self._gdch_creds_cache: - self._gdch_creds_cache[cache_key] = creds.with_gdch_audience(audience.rstrip("/")) + self._gdch_creds_cache[cache_key] = self._with_gdch_audience(creds, audience.rstrip("/")) gdch_creds: Final = self._gdch_creds_cache[cache_key] @@ -155,7 +180,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return token - def _load_creds_from_key(self, api_key: str) -> tuple[Any, bool]: + def _load_creds_from_key(self, api_key: str) -> tuple[object | None, bool]: import google.auth try: @@ -175,7 +200,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): self, headers: dict, model: str, - messages: list[Any], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, @@ -230,7 +255,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): if self._read_env_bool(litellm_params.get("gdc_token_caching"), "GDC_TOKEN_CACHING", default=False): token = self._cached_fetch_token(creds, audience, ssl_verify, api_key) else: - gdch_creds: Final = creds.with_gdch_audience(audience) + gdch_creds: Final = self._with_gdch_audience(creds, audience) self._fetch_auth(gdch_creds, ssl_verify) token = gdch_creds.token headers["Authorization"] = f"Bearer {token}" @@ -252,7 +277,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): def transform_request( self, model: str, - messages: list[Any], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, headers: dict, diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index ebcbf1b5a07..df4adf0c9a2 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -4,9 +4,11 @@ Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` Why separate file? Make it easy to see how transformation works """ +from collections.abc import Sequence from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._uuid import uuid @@ -25,6 +27,31 @@ from litellm.types.rerank import ( from ..common_utils import InfinityError +class _InfinityRerankUsage(TypedDict, extra_items=ReadOnly[int]): + """The token counters Infinity reports in the ``usage`` block of a rerank response.""" + + +class _InfinityRerankResult(TypedDict): + """One scored document in an Infinity ``/v1/rerank`` response.""" + + index: ReadOnly[int] + relevance_score: ReadOnly[float] + document: ReadOnly[str] + + +class _InfinityRerankResponse(TypedDict): + """The JSON body returned by Infinity's ``/v1/rerank`` endpoint.""" + + id: ReadOnly[NotRequired[str]] + usage: ReadOnly[NotRequired[_InfinityRerankUsage]] + results: ReadOnly[Sequence[_InfinityRerankResult]] + + +def _parse_rerank_response(raw_response: httpx.Response) -> _InfinityRerankResponse: + """Read the untyped JSON body of an Infinity rerank response.""" + return raw_response.json() + + class InfinityRerankConfig(CohereRerankConfig): def get_complete_url( self, @@ -80,7 +107,7 @@ class InfinityRerankConfig(CohereRerankConfig): No transformation required, Infinity follows Cohere API response format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _parse_rerank_response(raw_response) except Exception: raise InfinityError(message=raw_response.text, status_code=raw_response.status_code) diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index d435994ce20..f51213ca1c3 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -13,12 +13,57 @@ Generated files are returned directly in the response - no separate storage need import base64 import json +from collections.abc import Sequence from enum import Enum -from typing import Any, Final +from typing import Any, Final, Protocol + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger +class _ToolCallFunction(Protocol): + """Function payload of an assistant tool call.""" + + name: str | None + arguments: str + + +class _ToolCall(Protocol): + """Tool call requested by the assistant on a chat completion choice.""" + + id: str + function: _ToolCallFunction + + +class _AssistantMessage(Protocol): + """Assistant message carried by a chat completion choice.""" + + content: str | None + tool_calls: Sequence[_ToolCall] | None + + +class _CompletionChoice(Protocol): + """Single choice of a chat completion response.""" + + finish_reason: str + message: _AssistantMessage + + +class _SandboxFile(TypedDict): + """File generated inside the sandbox during a code execution run.""" + + name: ReadOnly[str] + mime_type: ReadOnly[str] + content_base64: ReadOnly[str] + + +class _CodeExecutionArguments(TypedDict): + """Arguments the model passes to the `litellm_code_execution` tool.""" + + code: NotRequired[ReadOnly[str]] + + class LiteLLMInternalTools(str, Enum): """ Enum for internal LiteLLM tools that are injected into requests. @@ -30,7 +75,7 @@ class LiteLLMInternalTools(str, Enum): CODE_EXECUTION = "litellm_code_execution" -def get_litellm_code_execution_tool() -> dict[str, Any]: +def get_litellm_code_execution_tool() -> dict[str, object]: """ Returns the litellm_code_execution tool definition in OpenAI format. @@ -51,7 +96,7 @@ def get_litellm_code_execution_tool() -> dict[str, Any]: } -def get_litellm_code_execution_tool_anthropic() -> dict[str, Any]: +def get_litellm_code_execution_tool_anthropic() -> dict[str, object]: """ Returns the litellm_code_execution tool definition in Anthropic/messages API format. @@ -103,7 +148,7 @@ class CodeExecutionHandler: skill_files: dict[str, bytes], skill_id: str | None = None, **kwargs, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Execute an LLM call with automatic code execution handling. @@ -134,8 +179,8 @@ class CodeExecutionHandler: ) current_messages: Final = list(messages) - generated_files: Final[list[dict[str, Any]]] = [] # Files returned directly - execution_results: Final[list[dict]] = [] + generated_files: Final[list[dict[str, object]]] = [] # Files returned directly + execution_results: Final[list[dict[str, object]]] = [] executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) response: Any = None # Initialize to avoid possibly unbound error @@ -151,11 +196,12 @@ class CodeExecutionHandler: **kwargs, ) - assistant_message = response.choices[0].message - stop_reason = response.choices[0].finish_reason + choice: _CompletionChoice = response.choices[0] + assistant_message = choice.message + stop_reason: str = choice.finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -190,8 +236,8 @@ class CodeExecutionHandler: if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: # Execute code in sandbox try: - args = json.loads(tool_call.function.arguments) - code = args.get("code", "") + args: _CodeExecutionArguments = json.loads(tool_call.function.arguments) + code: str = args.get("code", "") verbose_logger.debug("CodeExecutionHandler: Executing code (%s chars)", len(code)) @@ -202,13 +248,15 @@ class CodeExecutionHandler: verbose_logger.debug("CodeExecutionHandler: Execution result: %s", exec_result) + sandbox_files: Sequence[_SandboxFile] = exec_result["files"] + execution_results.append( { "iteration": iteration, "success": exec_result["success"], "output": exec_result["output"], "error": exec_result["error"], - "files": [f["name"] for f in exec_result["files"]], + "files": [f["name"] for f in sandbox_files], } ) @@ -216,9 +264,9 @@ class CodeExecutionHandler: tool_result = exec_result["output"] or "" # Collect generated files (returned directly, no storage) - if exec_result["files"]: + if sandbox_files: tool_result += "\n\nGenerated files:" - for f in exec_result["files"]: + for f in sandbox_files: file_content = base64.b64decode(f["content_base64"]) # Add to generated files list (returned in response) generated_files.append( diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 7ae438fd4cd..18f75f0e45c 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -8,10 +8,12 @@ response parsing, and streaming chunk parsing for models served with import datetime import json -from typing import Any, Final +from collections.abc import Sequence +from typing import Any, Final, TypedDict import httpx from pydantic import ValidationError +from typing_extensions import ReadOnly from litellm.llms.oci.chat.generic import ( _normalize_oci_finish_reason, @@ -46,6 +48,20 @@ from litellm.types.utils import ( ) +class _OpenAIToolCallFunction(TypedDict, total=False): + """The ``function`` block of an OpenAI-format assistant tool call.""" + + name: ReadOnly[str | None] + arguments: ReadOnly[str | dict[str, object]] + + +class _OpenAIToolCall(TypedDict, total=False): + """An entry of an OpenAI-format assistant message's ``tool_calls``.""" + + id: ReadOnly[str | None] + function: ReadOnly[_OpenAIToolCallFunction] + + def _extract_text_content(content: Any) -> str: """Return the plain-text representation of a message content value.""" if content is None: @@ -78,10 +94,10 @@ def adapt_messages_to_cohere_standard( """ # First pass: build tool_call_id → CohereToolCall so tool-result messages can # reference the originating call by name and parameters. - tool_call_lookup: Final[dict[str, CohereToolCall]] = {} + tool_call_lookup: Final[dict[str | None, CohereToolCall]] = {} for msg in messages: if msg.get("role") == "assistant": - tool_calls_raw: Any = msg.get("tool_calls") or [] + tool_calls_raw: Sequence[_OpenAIToolCall] = msg.get("tool_calls") or [] for tc in tool_calls_raw: tc_id = tc.get("id", "") raw_args = tc.get("function", {}).get("arguments", "{}") @@ -150,8 +166,22 @@ def adapt_messages_to_cohere_standard( return chat_history +class _OpenAIToolDefinitionFunction(TypedDict, total=False): + """The ``function`` block of an OpenAI-format tool definition.""" + + name: ReadOnly[str] + description: ReadOnly[str] + parameters: ReadOnly[dict[str, object]] + + +class _OpenAIToolDefinition(TypedDict, total=False): + """An entry of an OpenAI-format ``tools`` array.""" + + function: ReadOnly[_OpenAIToolDefinitionFunction] + + def adapt_tool_definitions_to_cohere_standard( - tools: list[dict[str, Any]], + tools: Sequence[_OpenAIToolDefinition], ) -> list[CohereTool]: """Adapt OpenAI-format tool definitions to the OCI Cohere format. diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 6e490f3ff15..449952217b5 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -4,16 +4,32 @@ Ollama /chat/completion calls handled in llm_http_handler.py [TODO]: migrate embeddings to a base handler as well. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, TypedDict + +from typing_extensions import NotRequired, ReadOnly import litellm from litellm.types.utils import EmbeddingResponse +class TokenEncoder(Protocol): + """The tokenizer surface used to estimate prompt tokens.""" + + def encode(self, text: str, /) -> Sequence[int]: ... + + +class OllamaEmbeddingResponse(TypedDict): + """Body of an Ollama ``/api/embed`` response.""" + + embeddings: ReadOnly[list[list[float]]] + prompt_eval_count: ReadOnly[NotRequired[int]] + + def _prepare_ollama_embedding_payload( - model: str, prompts: list[str], optional_params: dict[str, Any] -) -> dict[str, Any]: - data: Final[dict[str, Any]] = {"model": model, "input": prompts} + model: str, prompts: list[str], optional_params: Mapping[str, object] +) -> dict[str, object]: + data: Final[dict[str, object]] = {"model": model, "input": prompts} special_optional_params: Final = ["truncate", "options", "keep_alive", "dimensions"] for k, v in optional_params.items(): @@ -27,12 +43,12 @@ def _prepare_ollama_embedding_payload( def _process_ollama_embedding_response( - response_json: dict, + response_json: OllamaEmbeddingResponse, prompts: list[str], model: str, model_response: EmbeddingResponse, logging_obj: Any, - encoding: Any, + encoding: TokenEncoder | None, ) -> EmbeddingResponse: output_data: Final = [] embeddings: Final[list[list[float]]] = response_json["embeddings"] @@ -72,7 +88,7 @@ async def ollama_aembeddings( model_response: EmbeddingResponse, optional_params: dict, logging_obj: Any, - encoding: Any, + encoding: TokenEncoder | None, ): if not api_base.endswith("/api/embed"): api_base += "/api/embed" @@ -80,7 +96,7 @@ async def ollama_aembeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = await litellm.module_level_aclient.post(url=api_base, json=data) - response_json: Final = response.json() + response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( response_json=response_json, @@ -99,7 +115,7 @@ def ollama_embeddings( optional_params: dict, model_response: EmbeddingResponse, logging_obj: Any, - encoding: Any = None, + encoding: TokenEncoder | None = None, ): if not api_base.endswith("/api/embed"): api_base += "/api/embed" @@ -107,7 +123,7 @@ def ollama_embeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = litellm.module_level_client.post(url=api_base, json=data) - response_json: Final = response.json() + response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( response_json=response_json, diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 6fc50458aa3..44bd401c115 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -1,6 +1,8 @@ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( @@ -11,9 +13,11 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.containers.main import ( ContainerCreateOptionalRequestParams, ContainerFileListResponse, + ContainerFileObject, ContainerListResponse, ContainerObject, DeleteContainerResult, + ExpiresAfter, ) from litellm.types.router import GenericLiteLLMParams @@ -32,6 +36,46 @@ else: BaseLLMException = Any +class OpenAIContainerPayload(TypedDict): + """The JSON body OpenAI returns for a single container.""" + + id: ReadOnly[str] + object: ReadOnly[Literal["container"]] + created_at: ReadOnly[int] + status: ReadOnly[str] + expires_after: ReadOnly[ExpiresAfter | None] + last_active_at: ReadOnly[int | None] + name: ReadOnly[str | None] + + +class OpenAIContainerListPayload(TypedDict): + """The JSON body OpenAI returns for a page of containers.""" + + object: ReadOnly[Literal["list"]] + data: ReadOnly[list[ContainerObject]] + first_id: ReadOnly[str | None] + last_id: ReadOnly[str | None] + has_more: ReadOnly[bool] + + +class OpenAIContainerDeletedPayload(TypedDict): + """The JSON body OpenAI returns for a deleted container.""" + + id: ReadOnly[str] + object: ReadOnly[Literal["container.deleted"]] + deleted: ReadOnly[bool] + + +class OpenAIContainerFileListPayload(TypedDict): + """The JSON body OpenAI returns for a page of container files.""" + + object: ReadOnly[Literal["list"]] + data: ReadOnly[list[ContainerFileObject]] + first_id: ReadOnly[str | None] + last_id: ReadOnly[str | None] + has_more: ReadOnly[bool] + + class OpenAIContainerConfig(BaseContainerConfig): """Configuration class for OpenAI container API.""" @@ -87,7 +131,7 @@ class OpenAIContainerConfig(BaseContainerConfig): def transform_container_create_request( self, name: str, - container_create_optional_request_params: dict, + container_create_optional_request_params: Mapping[str, object], litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: @@ -111,7 +155,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container creation response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerPayload] = raw_response.json() # Transform the response data container_obj: Final = ContainerObject(**response_data) @@ -140,7 +184,7 @@ class OpenAIContainerConfig(BaseContainerConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """Transform the container list request for OpenAI API. @@ -151,7 +195,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = api_base # Prepare query parameters - params: Final = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if limit is not None: @@ -171,7 +215,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerListResponse: """Transform the OpenAI container list response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerListPayload] = raw_response.json() # Transform the response data container_list: Final = ContainerListResponse(**response_data) @@ -191,7 +235,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} return url, data @@ -201,7 +245,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container retrieve response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerPayload] = raw_response.json() # Transform the response data container_obj: Final = ContainerObject(**response_data) @@ -224,7 +268,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} return url, data @@ -234,7 +278,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteContainerResult: """Transform the OpenAI container delete response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerDeletedPayload] = raw_response.json() # Transform the response data delete_result: Final = DeleteContainerResult(**response_data) @@ -250,7 +294,7 @@ class OpenAIContainerConfig(BaseContainerConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """Transform the container file list request for OpenAI API. @@ -262,7 +306,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files") # Prepare query parameters - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if limit is not None: @@ -282,7 +326,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerFileListResponse: """Transform the OpenAI container file list response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerFileListPayload] = raw_response.json() # Transform the response data file_list: Final = ContainerFileListResponse(**response_data) @@ -308,7 +352,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content") # No query parameters needed - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} return url, params diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 962dfe52c0a..deedf237a50 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url @@ -25,6 +27,68 @@ else: LiteLLMLoggingObj = Any +class _AuthHeadersView(TypedDict): + """Auth headers read out of an untyped ``BaseVectorStoreAuthCredentials``.""" + + headers: ReadOnly[Mapping[str, str]] + + +class _RagContext(TypedDict, total=False): + """One context entry of a Vertex RAG ``:retrieveContexts`` response.""" + + text: ReadOnly[str] + sourceUri: ReadOnly[str] + sourceDisplayName: ReadOnly[str] + pageSpan: ReadOnly[Mapping[str, object]] + score: ReadOnly[float] + + +class _RagContexts(TypedDict, total=False): + """The ``contexts`` envelope wrapping the context list.""" + + contexts: ReadOnly[Sequence[_RagContext]] + + +class _RetrieveContextsBody(TypedDict, total=False): + """Body of a Vertex RAG ``:retrieveContexts`` response.""" + + contexts: ReadOnly[_RagContexts] + + +class _RetrieveContextsView(TypedDict): + """Typed view over the untyped ``:retrieveContexts`` JSON payload.""" + + body: ReadOnly[_RetrieveContextsBody] + + +class _RagCorpusBody(TypedDict, total=False): + """Body of a Vertex RAG ``ragCorpora`` create response.""" + + name: ReadOnly[str] + display_name: ReadOnly[str] + createTime: ReadOnly[str | int | float] + labels: ReadOnly[Mapping[str, str]] + + +class _RagCorpusView(TypedDict): + """Typed view over the untyped ``ragCorpora`` create JSON payload.""" + + body: ReadOnly[_RagCorpusBody] + + +class _RagSearchQuery(TypedDict, total=False): + """The ``query`` block of a Vertex RAG ``:retrieveContexts`` request.""" + + text: ReadOnly[str] + rag_retrieval_config: ReadOnly[Mapping[str, object]] + + +class _LoggedQueryView(TypedDict): + """The search query recovered from the logging object's call details.""" + + search_query: ReadOnly[str] + + class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Configuration for Vertex AI Vector Store RAG API @@ -35,7 +99,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() - def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: # Get credentials and project info vertex_credentials: Final = self.get_vertex_ai_credentials(dict(litellm_params)) vertex_project: Final = self.get_vertex_ai_project(dict(litellm_params)) @@ -60,20 +124,23 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): "write": [("POST", "/ragCorpora")], } - def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: + def validate_environment( + self, headers: dict[str, str], litellm_params: GenericLiteLLMParams | None + ) -> dict[str, str]: """ Validate and set up authentication for Vertex AI RAG API """ litellm_params = litellm_params or GenericLiteLLMParams() auth_headers: Final = self.get_auth_credentials(litellm_params.model_dump()) - headers.update(auth_headers.get("headers", {})) + auth_view: Final[_AuthHeadersView] = {"headers": auth_headers.get("headers", {})} + headers.update(auth_view["headers"]) return headers def get_complete_url( self, api_base: str | None, - litellm_params: dict, + litellm_params: dict[str, object], ) -> str: """ Get the Base endpoint for Vertex AI RAG API @@ -95,9 +162,9 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: + litellm_params: dict[str, object], + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: """ Transform search request for Vertex AI RAG API """ @@ -120,35 +187,34 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Just the corpus ID, construct full path full_rag_corpus = f"projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{vector_store_id}" - # Build the request body for Vertex AI RAG API - request_body: Final[dict[str, Any]] = { - "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, - "query": {"text": query}, - } - ######################################################### # Update logging object with details of the request ######################################################### litellm_logging_obj.model_call_details["query"] = query # Add optional parameters + rag_retrieval_config: Final[dict[str, object]] = {} max_num_results: Final = vector_store_search_optional_params.get("max_num_results") if max_num_results is not None: - request_body["query"]["rag_retrieval_config"] = {"top_k": max_num_results} + rag_retrieval_config["top_k"] = max_num_results # Add filters if provided - filters: Final = vector_store_search_optional_params.get("filters") + filters: Final[object] = vector_store_search_optional_params.get("filters") if filters is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["filter"] = filters + rag_retrieval_config["filter"] = filters # Add ranking options if provided - ranking_options: Final = vector_store_search_optional_params.get("ranking_options") + ranking_options: Final[object] = vector_store_search_optional_params.get("ranking_options") if ranking_options is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["ranking"] = ranking_options + rag_retrieval_config["ranking"] = ranking_options + + query_body: Final[_RagSearchQuery] = ( + {"text": query, "rag_retrieval_config": rag_retrieval_config} if rag_retrieval_config else {"text": query} + ) + request_body: Final[dict[str, object]] = { + "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, + "query": query_body, + } return url, request_body @@ -159,12 +225,13 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG API response to standard vector store search response """ try: - response_json: Final = response.json() + response_view: Final[_RetrieveContextsView] = {"body": response.json()} + response_json: Final = response_view["body"] # Extract contexts from Vertex AI response - handle nested structure contexts: Final = response_json.get("contexts", {}).get("contexts", []) # Transform contexts to standard format - search_results: Final = [] + search_results: Final[list[VectorStoreSearchResult]] = [] for context in contexts: content = [ VectorStoreResultContent( @@ -202,9 +269,12 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): ) search_results.append(result) + query_view: Final[_LoggedQueryView] = { + "search_query": litellm_logging_obj.model_call_details.get("query", "") + } return VectorStoreSearchResponse( object="vector_store.search_results.page", - search_query=litellm_logging_obj.model_call_details.get("query", ""), + search_query=query_view["search_query"], data=search_results, ) @@ -219,14 +289,14 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: """ Transform create request for Vertex AI RAG Corpus """ url: Final = f"{api_base}/ragCorpora" # Base URL for creating RAG corpus # Build the request body for Vertex AI RAG Corpus creation - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "display_name": vector_store_create_optional_params.get("name", "litellm-vector-store"), "description": "Vector store created via LiteLLM", } @@ -243,7 +313,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG Corpus creation response to standard vector store response """ try: - response_json: Final = response.json() + response_view: Final[_RagCorpusView] = {"body": response.json()} + response_json: Final = response_view["body"] # Extract the corpus ID from the response name corpus_name: Final = response_json.get("name", "") diff --git a/litellm/models/base.py b/litellm/models/base.py index 7eedf10212e..8125bfd0205 100644 --- a/litellm/models/base.py +++ b/litellm/models/base.py @@ -33,6 +33,6 @@ class DomainModel(BaseModel): return cls(**record.dict()) return cls(**dict(record)) - def to_db_dict(self, exclude_unset: bool = False) -> dict[str, Any]: + def to_db_dict(self, exclude_unset: bool = False) -> dict[str, object]: """Convert domain model to a dictionary for database operations.""" return self.model_dump(exclude_none=True, exclude_unset=exclude_unset) diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py index c09106273e1..150900e7ff2 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py @@ -36,7 +36,10 @@ a healed fleet has no null rows and the backfill exits after one query. import json from collections import Counter -from typing import Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import Final, Literal, Protocol + +from pydantic import JsonValue from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials @@ -55,9 +58,59 @@ BackfillRule = Literal[ _BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill" -def _decrypted_credentials(raw_credentials: Any) -> MCPCredentials | None: +class _MCPServerRow(Protocol): + """The ``LiteLLM_MCPServerTable`` columns this backfill reads.""" + + @property + def server_id(self) -> str: ... + + @property + def authorization_url(self) -> str | None: ... + + @property + def registration_url(self) -> str | None: ... + + @property + def token_url(self) -> str | None: ... + + @property + def credentials(self) -> str | Mapping[str, JsonValue] | None: ... + + +class _MCPUserCredentialRow(Protocol): + """The ``LiteLLM_MCPUserCredentials`` columns this backfill reads.""" + + @property + def server_id(self) -> str: ... + + @property + def credential_b64(self) -> str: ... + + +class _MCPServerTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPServerRow]: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> object: ... + + +class _MCPUserCredentialsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPUserCredentialRow]: ... + + +def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: + """The MCP server table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpservertable + + +def _mcp_user_credentials_table(prisma_client: PrismaClient) -> _MCPUserCredentialsTable: + """The per-user MCP credential table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpusercredentials + + +def _decrypted_credentials(raw_credentials: str | Mapping[str, JsonValue] | None) -> MCPCredentials | None: if raw_credentials is None: return None + parsed: JsonValue | Mapping[str, JsonValue] if isinstance(raw_credentials, str): try: parsed = json.loads(raw_credentials) @@ -92,14 +145,14 @@ def classify_null_flow_row( async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]: """Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable ones, warn on the ambiguous ones, and return counts per rule.""" - null_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpservertable.find_many( + null_rows: Final[Sequence[_MCPServerRow]] = await _mcp_server_table(prisma_client).find_many( where={"auth_type": "oauth2", "oauth2_flow": None}, ) if not null_rows: return {} server_ids: Final = [row.server_id for row in null_rows] - token_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpusercredentials.find_many( + token_rows: Final[Sequence[_MCPUserCredentialRow]] = await _mcp_user_credentials_table(prisma_client).find_many( where={"server_id": {"in": server_ids}}, ) server_ids_with_oauth_tokens: Final[set[str]] = { @@ -141,7 +194,7 @@ async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[Backfi stamped_flows: Final = {flow for _, (flow, _) in classified if flow is not None} for stamped_flow in stamped_flows: server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow] - await prisma_client.db.litellm_mcpservertable.update_many( + await _mcp_server_table(prisma_client).update_many( where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None}, data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR}, ) diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 9672383a572..ecaaf35e817 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -1,5 +1,9 @@ import json -from typing import Final +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -7,18 +11,73 @@ from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import MCPToolsetRepository from litellm.types.mcp_server.mcp_toolset import ( MCPToolset, + MCPToolsetTool, NewMCPToolsetRequest, UpdateMCPToolsetRequest, ) -def _toolset_from_row(row) -> MCPToolset: +class MCPToolsetFields(TypedDict): + """The ``MCPToolset`` constructor keywords a toolset row expands into.""" + + toolset_id: ReadOnly[str] + toolset_name: ReadOnly[str] + description: NotRequired[ReadOnly[str | None]] + tools: NotRequired[ReadOnly[list[MCPToolsetTool]]] + created_at: NotRequired[ReadOnly[datetime | None]] + created_by: NotRequired[ReadOnly[str | None]] + updated_at: NotRequired[ReadOnly[datetime | None]] + updated_by: NotRequired[ReadOnly[str | None]] + + +class MCPToolsetRowData(TypedDict): + """A toolset table row, whose ``tools`` column is stored as JSON.""" + + toolset_id: ReadOnly[str] + toolset_name: ReadOnly[str] + description: NotRequired[ReadOnly[str | None]] + tools: NotRequired[ReadOnly[str | list[MCPToolsetTool]]] + created_at: NotRequired[ReadOnly[datetime | None]] + created_by: NotRequired[ReadOnly[str | None]] + updated_at: NotRequired[ReadOnly[datetime | None]] + updated_by: NotRequired[ReadOnly[str | None]] + + +class MCPToolsetRow(Protocol): + """A row of the toolset table, as the prisma client returns it.""" + + def model_dump(self) -> MCPToolsetRowData: ... + + +class MCPToolsetTable(Protocol): + """The prisma table actions this module runs against the toolset table.""" + + async def create(self, data: Mapping[str, object]) -> MCPToolsetRow: ... + + async def find_unique(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... + + async def find_first(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... + + async def find_many(self, where: Mapping[str, object]) -> Sequence[MCPToolsetRow]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> MCPToolsetRow: ... + + async def delete(self, where: Mapping[str, object]) -> MCPToolsetRow: ... + + +def _toolset_table(prisma_client: PrismaClient) -> MCPToolsetTable: + """The toolset table actions of the prisma client.""" + return MCPToolsetRepository(prisma_client).table + + +def _toolset_from_row(row: MCPToolsetRow) -> MCPToolset: data: Final = row.model_dump() - tools = data.get("tools") or [] - if isinstance(tools, str): - tools = json.loads(tools) - data["tools"] = tools - return MCPToolset(**data) + tools: Final = data.get("tools") or [] + resolved: Final[MCPToolsetFields] = { + **data, + "tools": json.loads(tools) if isinstance(tools, str) else tools, + } + return MCPToolset(**resolved) async def create_mcp_toolset( @@ -31,7 +90,7 @@ async def create_mcp_toolset( data_dict["tools"] = json.dumps(data_dict.get("tools", [])) data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - row: Final = await MCPToolsetRepository(prisma_client).table.create(data=data_dict) + row: Final = await _toolset_table(prisma_client).create(data=data_dict) return _toolset_from_row(row) @@ -39,7 +98,7 @@ async def get_mcp_toolset( prisma_client: PrismaClient, toolset_id: str, ) -> MCPToolset | None: - row: Final = await MCPToolsetRepository(prisma_client).table.find_unique(where={"toolset_id": toolset_id}) + row: Final = await _toolset_table(prisma_client).find_unique(where={"toolset_id": toolset_id}) if row is None: return None return _toolset_from_row(row) @@ -47,13 +106,11 @@ async def get_mcp_toolset( async def list_mcp_toolsets( prisma_client: PrismaClient, - toolset_ids: list[str] | None = None, -) -> list[MCPToolset]: + toolset_ids: Sequence[str] | None = None, +) -> Sequence[MCPToolset]: try: - where = {} - if toolset_ids is not None: - where = {"toolset_id": {"in": toolset_ids}} - rows: Final = await MCPToolsetRepository(prisma_client).table.find_many(where=where) + where: Final[Mapping[str, object]] = {} if toolset_ids is None else {"toolset_id": {"in": toolset_ids}} + rows: Final = await _toolset_table(prisma_client).find_many(where=where) return [_toolset_from_row(r) for r in rows] except Exception as e: verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e) @@ -64,7 +121,7 @@ async def get_mcp_toolset_by_name( prisma_client: PrismaClient, toolset_name: str, ) -> MCPToolset | None: - row: Final = await MCPToolsetRepository(prisma_client).table.find_first(where={"toolset_name": toolset_name}) + row: Final = await _toolset_table(prisma_client).find_first(where={"toolset_name": toolset_name}) if row is None: return None return _toolset_from_row(row) @@ -80,7 +137,7 @@ async def update_mcp_toolset( data_dict["tools"] = json.dumps(data_dict["tools"]) data_dict["updated_by"] = touched_by try: - row: Final = await MCPToolsetRepository(prisma_client).table.update( + row: Final = await _toolset_table(prisma_client).update( where={"toolset_id": data.toolset_id}, data=data_dict, ) @@ -98,7 +155,7 @@ async def delete_mcp_toolset( toolset_id: str, ) -> MCPToolset | None: try: - row: Final = await MCPToolsetRepository(prisma_client).table.delete(where={"toolset_id": toolset_id}) + row: Final = await _toolset_table(prisma_client).delete(where={"toolset_id": toolset_id}) except Exception as e: from prisma.errors import RecordNotFoundError diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 64de6827679..edb945fc674 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -12,6 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -83,9 +84,17 @@ class AgentTableClient(Protocol): async def delete(self, where: Mapping[str, object]) -> AgentRecord: ... +class _AgentsRepositoryView(Protocol): + @property + def table(self) -> AgentTableClient: ... + + +def _agents_table_of(repository: _AgentsRepositoryView) -> AgentTableClient: + return repository.table + + def agents_table(prisma_client: PrismaClient) -> AgentTableClient: - table: Final[AgentTableClient] = AgentsRepository(prisma_client).table - return table + return _agents_table_of(AgentsRepository(prisma_client)) class ObjectPermissionGrantRecord(Protocol): @@ -99,9 +108,17 @@ class ObjectPermissionTableClient(Protocol): async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... +class _ObjectPermissionRepositoryView(Protocol): + @property + def table(self) -> ObjectPermissionTableClient: ... + + +def _object_permission_table_of(repository: _ObjectPermissionRepositoryView) -> ObjectPermissionTableClient: + return repository.table + + def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient: - table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table - return table + return _object_permission_table_of(ObjectPermissionRepository(prisma_client)) class GrantMigrationResult(NamedTuple): @@ -283,19 +300,21 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {}) + litellm_params_dict: Final[Mapping[str, object]] = ( + litellm_params_obj.model_dump() + if isinstance(litellm_params_obj, SupportsModelDump) + else (dict(litellm_params_obj) if litellm_params_obj else {}) + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[Mapping[str, object]] = ( + agent_card_params_obj.model_dump() + if isinstance(agent_card_params_obj, SupportsModelDump) + else (dict(agent_card_params_obj) if agent_card_params_obj else {}) + ) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Handle object_permission (MCP tool access for agent) @@ -386,15 +405,13 @@ class AgentRegistry: The patched agent """ try: - existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) - if existing_agent is not None: - existing_agent = dict(existing_agent) - - if existing_agent is None: + existing_record: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) + if existing_record is None: raise Exception(f"Agent with ID {agent_id} not found") + existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, Any]] = {} + update_data: Final[dict[str, object]] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): @@ -418,7 +435,7 @@ class AgentRegistry: update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else [] if agent.get("object_permission") is not None: agent_copy: Final = dict(augment_agent) - existing_object_permission_id: Final = existing_agent.get("object_permission_id") + existing_object_permission_id: Final = existing_record.object_permission_id object_permission_id: Final = await handle_update_object_permission_common( agent_copy, existing_object_permission_id, @@ -460,19 +477,21 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {}) + litellm_params_dict: Final[Mapping[str, object]] = ( + litellm_params_obj.model_dump() + if isinstance(litellm_params_obj, SupportsModelDump) + else (dict(litellm_params_obj) if litellm_params_obj else {}) + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[Mapping[str, object]] = ( + agent_card_params_obj.model_dump() + if isinstance(agent_card_params_obj, SupportsModelDump) + else (dict(agent_card_params_obj) if agent_card_params_obj else {}) + ) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Serialize static_headers for update diff --git a/litellm/proxy/client/cli/commands/credentials.py b/litellm/proxy/client/cli/commands/credentials.py index c550b39d33f..2beff05b375 100644 --- a/litellm/proxy/client/cli/commands/credentials.py +++ b/litellm/proxy/client/cli/commands/credentials.py @@ -1,14 +1,42 @@ import json +from collections.abc import Mapping, Sequence from typing import Final, Literal import click import requests import rich from rich.table import Table +from typing_extensions import NotRequired, ReadOnly, TypedDict from ...credentials import CredentialsManagementClient +class _CliContext(TypedDict): + """Values the top-level CLI group stores on the click context.""" + + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +class _CliContextView(TypedDict): + obj: ReadOnly[_CliContext] + + +class _CredentialRow(TypedDict): + """Single credential entry as returned by ``GET /credentials``.""" + + credential_name: ReadOnly[NotRequired[str]] + credential_info: ReadOnly[NotRequired[Mapping[str, object]]] + + +class _CredentialRowsView(TypedDict): + rows: ReadOnly[Sequence[_CredentialRow]] + + +class _JsonBodyView(TypedDict): + body: ReadOnly[object] + + @click.group() def credentials(): """Manage credentials for the LiteLLM proxy server""" @@ -25,7 +53,8 @@ def credentials(): @click.pass_context def list(ctx: click.Context, output_format: Literal["table", "json"]): """List all credentials""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) response: Final = client.list() assert isinstance(response, dict) @@ -39,7 +68,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): table.add_column("Custom LLM Provider", style="green") # Add rows - for cred in response.get("credentials", []): + credential_rows: Final[_CredentialRowsView] = {"rows": response.get("credentials", [])} + for cred in credential_rows["rows"]: info = cred.get("credential_info", {}) table.add_row( str(cred.get("credential_name", "")), @@ -66,7 +96,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): @click.pass_context def create(ctx: click.Context, credential_name: str, info: str, values: str): """Create a new credential""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) try: credential_info: Final = json.loads(info) credential_values: Final = json.loads(values) @@ -79,8 +110,8 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -91,15 +122,16 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): @click.pass_context def delete(ctx: click.Context, credential_name: str): """Delete a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) try: response: Final = client.delete(credential_name) rich.print_json(data=response) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -110,6 +142,7 @@ def delete(ctx: click.Context, credential_name: str): @click.pass_context def get(ctx: click.Context, credential_name: str): """Get a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) response: Final = client.get(credential_name) rich.print_json(data=response) diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index e814ac84ebb..1212e194bde 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -1,21 +1,48 @@ """Team management commands for LiteLLM CLI.""" +from collections.abc import Mapping, Sequence from typing import Any, Final import click import requests from rich.console import Console from rich.table import Table +from typing_extensions import ReadOnly, TypedDict from litellm.proxy.client import Client +class _CliContext(TypedDict): + """The proxy connection settings the CLI group stores on the click context.""" + + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +def _cli_context(ctx: click.Context) -> _CliContext: + """The proxy connection settings the CLI group stored on the click context.""" + ctx_obj: Final[_CliContext] = ctx.obj + return ctx_obj + + +def _proxy_client(ctx: click.Context) -> Client: + """A proxy client for the base URL and API key on the click context.""" + ctx_obj: Final = _cli_context(ctx) + return Client(ctx_obj["base_url"], ctx_obj["api_key"]) + + +def _http_error_detail(error: requests.exceptions.HTTPError) -> object: + """The ``detail`` the proxy reported for a failed request.""" + error_body: Final[Mapping[str, object]] = error.response.json() + return error_body.get("detail", "Unknown error") + + @click.group() def teams(): """Manage teams and team assignments""" -def display_teams_table(teams: list[dict[str, Any]]) -> None: +def display_teams_table(teams: Sequence[dict[str, Any]]) -> None: """Display teams in a formatted table""" console: Final = Console() @@ -33,8 +60,8 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: for i, team in enumerate(teams): team_alias = team.get("team_alias") or "N/A" - team_id = team.get("team_id", "N/A") - models = team.get("models", []) + team_id: str = team.get("team_id", "N/A") + models: Sequence[str] = team.get("models", []) max_budget = team.get("max_budget") # Format models list @@ -64,7 +91,7 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: @click.pass_context def list(ctx: click.Context): """List teams that you belong to""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + client: Final = _proxy_client(ctx) try: # Use list() for simpler response structure (returns array directly) @@ -72,8 +99,7 @@ def list(ctx: click.Context): display_teams_table(teams) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) @@ -84,7 +110,7 @@ def list(ctx: click.Context): @click.pass_context def available(ctx: click.Context): """List teams that are available to join""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + client: Final = _proxy_client(ctx) try: teams: Final = client.teams.get_available() @@ -96,8 +122,7 @@ def available(ctx: click.Context): click.echo("No available teams to join.") except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) except Exception as e: click.echo(f"Error: {e}", err=True) raise click.Abort() @@ -108,8 +133,8 @@ def available(ctx: click.Context): @click.pass_context def assign_key(ctx: click.Context, team_id: str | None): """Assign your current CLI key to a team""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) - api_key: Final = ctx.obj["api_key"] + client: Final = _proxy_client(ctx) + api_key: Final = _cli_context(ctx)["api_key"] if not api_key: click.echo("No API key found. Please login first using 'litellm login'") @@ -145,7 +170,7 @@ def assign_key(ctx: click.Context, team_id: str | None): teams = client.teams.list() for team in teams: if team.get("team_id") == team_id: - models = team.get("models", []) + models: Sequence[str] = team.get("models", []) if models: click.echo(f"You can now access models: {', '.join(models)}") else: @@ -154,8 +179,7 @@ def assign_key(ctx: click.Context, team_id: str | None): except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 9379a8577a3..a28d03eebf0 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,6 +1,6 @@ import copy import os -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias @@ -525,8 +525,8 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( def sanitize_openai_provider_metadata( - metadata: dict[str, Any] | None, -) -> dict[str, str] | None: + metadata: Mapping[str, object] | None, +) -> Mapping[str, object] | None: """ Keep only provider-safe OpenAI metadata entries (string keys -> string values). @@ -644,7 +644,7 @@ def process_callback(_callback: str, callback_type: str, environment_variables: return {"name": _callback, "variables": env_vars_dict, "type": callback_type} -def normalize_callback_names(callbacks: Iterable[Any]) -> list[Any]: +def normalize_callback_names(callbacks: Iterable[object] | None) -> list[object]: if callbacks is None: return [] return [c.lower() if isinstance(c, str) else c for c in callbacks] @@ -674,7 +674,7 @@ def decrypt_callback_vars(metadata: Any) -> Any: return _transform_callback_vars(metadata, _decrypt_or_passthrough) -def _transform_callback_vars(metadata: Any, transform: Callable[[str, Any], Any]) -> Any: +def _transform_callback_vars(metadata: object, transform: Callable[[str, Any], Any]) -> object: if not isinstance(metadata, dict): return metadata out: Final = copy.deepcopy(metadata) @@ -704,7 +704,7 @@ def is_sensitive_callback_key( return _CALLBACK_VAR_MASKER.is_sensitive_key(key) -def _encrypt_if_plaintext(key: str, value: Any) -> Any: +def _encrypt_if_plaintext(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not is_sensitive_callback_key(key): @@ -725,7 +725,7 @@ def _encrypt_if_plaintext(key: str, value: Any) -> Any: return value -def _decrypt_or_passthrough(key: str, value: Any) -> Any: +def _decrypt_or_passthrough(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not value.startswith(_CALLBACK_VAR_ENCRYPTED_PREFIX): diff --git a/litellm/proxy/common_utils/get_routes.py b/litellm/proxy/common_utils/get_routes.py index 2118a6610b4..5aeb755071c 100644 --- a/litellm/proxy/common_utils/get_routes.py +++ b/litellm/proxy/common_utils/get_routes.py @@ -2,63 +2,78 @@ Utility class for getting routes from a FastAPI app. """ +from collections.abc import Sequence from typing import Any, Final from starlette.routing import BaseRoute +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger +class RouteInfo(TypedDict, total=False): + """One entry of the app's route listing.""" + + path: ReadOnly[object] + methods: ReadOnly[object] + name: ReadOnly[object] + endpoint: ReadOnly[str | None] + mounted_app: ReadOnly[bool] + + class GetRoutes: @staticmethod def get_app_routes( route: BaseRoute, endpoint_route: Any, - ) -> list[dict[str, Any]]: + ) -> Sequence[RouteInfo]: """ Get routes for a regular route. """ - routes: Final[list[dict[str, Any]]] = [] - route_info: Final = { + route_info: Final[RouteInfo] = { "path": getattr(route, "path", None), "methods": getattr(route, "methods", None), "name": getattr(route, "name", None), "endpoint": (endpoint_route.__name__ if getattr(route, "endpoint", None) else None), } - routes.append(route_info) - return routes + return [route_info] @staticmethod def get_routes_for_mounted_app( route: BaseRoute, - ) -> list[dict[str, Any]]: + ) -> Sequence[RouteInfo]: """ Get routes for a mounted sub-application. """ - routes: Final[list[dict[str, Any]]] = [] + routes: Final[list[RouteInfo]] = [] mount_path: Final = getattr(route, "path", "") - sub_app: Final = getattr(route, "app", None) - if sub_app and hasattr(sub_app, "routes"): - for sub_route in sub_app.routes: - # Get endpoint - either from endpoint attribute or app attribute - endpoint_func = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) + for sub_route in GetRoutes._mounted_app_routes(route): + endpoint_func: object = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) - if endpoint_func is not None: - sub_route_path = getattr(sub_route, "path", "") - full_path = mount_path.rstrip("/") + sub_route_path + if endpoint_func is not None: + sub_route_path = getattr(sub_route, "path", "") + full_path = mount_path.rstrip("/") + sub_route_path - route_info = { - "path": full_path, - "methods": getattr(sub_route, "methods", ["GET", "POST"]), - "name": getattr(sub_route, "name", None), - "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), - "mounted_app": True, - } - routes.append(route_info) + route_info: RouteInfo = { + "path": full_path, + "methods": getattr(sub_route, "methods", ["GET", "POST"]), + "name": getattr(sub_route, "name", None), + "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), + "mounted_app": True, + } + routes.append(route_info) return routes @staticmethod - def _safe_get_endpoint_name(endpoint_function: Any) -> str | None: + def _mounted_app_routes(route: BaseRoute) -> Sequence[BaseRoute]: + """The routes of the sub-application mounted at ``route``, if it mounts one.""" + sub_app: Final[object] = getattr(route, "app", None) + if sub_app and hasattr(sub_app, "routes"): + return getattr(sub_app, "routes") + return () + + @staticmethod + def _safe_get_endpoint_name(endpoint_function: object) -> str | None: """ Safely get the name of the endpoint function. """ @@ -66,7 +81,7 @@ class GetRoutes: if hasattr(endpoint_function, "__name__"): return getattr(endpoint_function, "__name__") elif hasattr(endpoint_function, "__class__") and hasattr(endpoint_function.__class__, "__name__"): - return getattr(endpoint_function.__class__, "__name__") + return endpoint_function.__class__.__name__ else: return None except Exception: diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 93d51bdd461..a8a71134034 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -40,30 +40,31 @@ class UserApiKeyCache(DualCache): @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... def get_cache( self, - key, - parent_otel_span=None, + key: object, + parent_otel_span: object = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, + **kwargs: object, ) -> Any | BaseModel | None: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) @@ -85,30 +86,31 @@ class UserApiKeyCache(DualCache): @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... async def async_get_cache( self, - key, - parent_otel_span=None, + key: object, + parent_otel_span: object = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, + **kwargs: object, ) -> Any | BaseModel | None: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) @@ -129,17 +131,17 @@ class UserApiKeyCache(DualCache): return None return decoded - def set_cache(self, key, value, local_only: bool = False, **kwargs): + def set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): + async def async_set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs) -> None: + async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None: """ Batch writes with the same Codec boundary as ``async_set_cache`` without ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 5aeb52be535..9cc4369fea6 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -6,11 +6,14 @@ otherwise PrismaClient uses the writer-only PrismaWrapper directly. import os from collections.abc import Callable -from typing import Any, Final +from datetime import timedelta +from typing import Any, Final, TypeAlias from litellm._logging import verbose_proxy_logger from litellm.proxy.db.prisma_client import PrismaWrapper +ConnectTimeout: TypeAlias = "int | timedelta | None" + # Per-model action methods that read from the database. These are routed to # the read replica when one is configured. _MODEL_READ_METHODS: Final = frozenset( @@ -31,6 +34,11 @@ _MODEL_READ_METHODS: Final = frozenset( _TOP_LEVEL_READ_METHODS: Final = frozenset({"query_first", "query_raw"}) +def _dynamic_attr(target: object, name: str) -> object: + """Fetch `name` off an unstubbed Prisma object as an opaque value.""" + return getattr(target, name) + + class _RoutedActions: """Per-model accessor that sends reads to the reader and writes to the writer. @@ -43,18 +51,18 @@ class _RoutedActions: def __init__( self, - writer_actions: Any, - reader_actions: Any, + writer_actions: object, + reader_actions: object, should_use_reader: Callable[[], bool], ): self._writer_actions = writer_actions self._reader_actions = reader_actions self._should_use_reader = should_use_reader - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: if name in _MODEL_READ_METHODS and self._should_use_reader(): - return getattr(self._reader_actions, name) - return getattr(self._writer_actions, name) + return _dynamic_attr(self._reader_actions, name) + return _dynamic_attr(self._writer_actions, name) class RoutingPrismaWrapper: @@ -135,7 +143,7 @@ class RoutingPrismaWrapper: return not self._reader_unavailable @staticmethod - async def _try_connect(client: PrismaWrapper, *args: Any, **kwargs: Any) -> Exception | None: + async def _try_connect(client: PrismaWrapper, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> Exception | None: if client.is_connected() is True: return None try: @@ -144,7 +152,7 @@ class RoutingPrismaWrapper: except Exception as e: return e - async def connect(self, *args: Any, **kwargs: Any) -> None: + async def connect(self, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> None: writer_error: Final = await self._try_connect(self._writer, *args, **kwargs) if writer_error is None: self._writer_unavailable = False @@ -176,7 +184,7 @@ class RoutingPrismaWrapper: writer_error, ) - async def disconnect(self, *args: Any, **kwargs: Any) -> None: + async def disconnect(self, *args: object, **kwargs: object) -> None: first_error: BaseException | None = None for client in (self._writer, self._reader): try: @@ -206,7 +214,7 @@ class RoutingPrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, - http_client: Any | None = None, + http_client: object | None = None, *, expected_generation: int | None = None, ) -> bool: @@ -245,7 +253,7 @@ class RoutingPrismaWrapper: ) return True - async def _recreate_reader(self, http_client: Any | None = None) -> None: + async def _recreate_reader(self, http_client: object | None = None) -> None: """Resolve the reader URL and recreate its Prisma client. IAM-enabled readers regenerate their token (host/port/user came from @@ -265,14 +273,14 @@ class RoutingPrismaWrapper: def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: - return getattr(self.read_target, name) - writer_attr: Final = getattr(self._writer, name) + return _dynamic_attr(self.read_target, name) + writer_attr: Final = _dynamic_attr(self._writer, name) # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / # tx are callables and stay on the writer untouched. if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"): try: - reader_attr: Final = getattr(self._reader, name) + reader_attr: Final = _dynamic_attr(self._reader, name) except AttributeError: return writer_attr return _RoutedActions(writer_attr, reader_attr, self._should_use_reader) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index da12222f233..707118bf016 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -15,6 +15,8 @@ restriction intact. """ import operator +from collections.abc import Callable +from types import CodeType from typing import Any, Final from RestrictedPython import ( @@ -58,7 +60,7 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): return self.node_contents_visit(node) -_INPLACE_OPS: Final[dict[str, Any]] = { +_INPLACE_OPS: Final[dict[str, Callable[[object, object], object]]] = { "+=": operator.iadd, "-=": operator.isub, "*=": operator.imul, @@ -75,7 +77,7 @@ _INPLACE_OPS: Final[dict[str, Any]] = { } -def _inplacevar_(op: str, x: Any, y: Any) -> Any: +def _inplacevar_(op: str, x: object, y: object) -> object: # RestrictedPython rewrites ``x += 1`` on a simple name into # ``x = _inplacevar_("+=", x, 1)``. The package deliberately ships no # default, so we dispatch through ``operator``'s in-place helpers, which @@ -86,7 +88,7 @@ def _inplacevar_(op: str, x: Any, y: Any) -> Any: return fn(x, y) -def _build_sandbox_builtins() -> dict[str, Any]: +def _build_sandbox_builtins() -> dict[str, object]: # ``limited_builtins`` overrides ``list``/``tuple``/``range`` from # ``safe_builtins`` with bounds-checking variants (e.g. ``limited_range`` # rejects ``range(10**18)``). ``utility_builtins`` adds ``set``, @@ -98,14 +100,14 @@ def _build_sandbox_builtins() -> dict[str, Any]: } -def build_sandbox_globals() -> dict[str, Any]: +def build_sandbox_globals() -> dict[str, object]: """Assemble the globals dict for executing guardrail code. Includes the LiteLLM-provided primitives (``regex_match``, ``http_get``, ``allow``/``block``/``modify``, etc.) plus the RestrictedPython guards that the compiled bytecode expects to find by name. """ - sandbox: Final[dict[str, Any]] = get_custom_code_primitives().copy() + sandbox: Final[dict[str, object]] = get_custom_code_primitives().copy() sandbox["__builtins__"] = _build_sandbox_builtins() sandbox["_getattr_"] = safer_getattr sandbox["_getitem_"] = default_guarded_getitem @@ -116,7 +118,7 @@ def build_sandbox_globals() -> dict[str, Any]: return sandbox -def compile_sandboxed(source: str, filename: str = "") -> Any: +def compile_sandboxed(source: str, filename: str = "") -> CodeType: """Compile guardrail source with RestrictedPython's AST transformer. Raises ``SyntaxError`` on either a Python syntax error or a restricted diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 507dd645953..0ae048c040d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -2,7 +2,7 @@ from __future__ import annotations import os from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse from uuid import uuid4 @@ -11,7 +11,7 @@ import requests from fastapi import HTTPException from httpx import HTTPStatusError from requests.auth import HTTPBasicAuth -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -36,24 +36,31 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options carried by this guardrail's forwarded keyword arguments.""" + + guardrail_name: ReadOnly[str | None] + supported_event_hooks: list[GuardrailEventHooks] | None + + class _HiddenlayerEvaluation(TypedDict, total=False): - action: str - threat_level: str + action: ReadOnly[str] + threat_level: ReadOnly[str] class _HiddenlayerAnalysisEntry(TypedDict, total=False): - name: str - detected: bool + name: ReadOnly[str] + detected: ReadOnly[bool] class _HiddenlayerModifiedSide(TypedDict): - messages: Any + messages: ReadOnly[Any] class _HiddenlayerResponse(TypedDict, total=False): - evaluation: _HiddenlayerEvaluation - analysis: Sequence[_HiddenlayerAnalysisEntry] - modified_data: Mapping[str, _HiddenlayerModifiedSide] + evaluation: ReadOnly[_HiddenlayerEvaluation] + analysis: ReadOnly[Sequence[_HiddenlayerAnalysisEntry]] + modified_data: ReadOnly[Mapping[str, _HiddenlayerModifiedSide]] class _LoggedCallMetadata(TypedDict, total=False): @@ -151,7 +158,7 @@ class HiddenlayerGuardrail(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, auth_url: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") @@ -356,7 +363,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, auth_url: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET") @@ -486,7 +493,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): self, payload: Any, input_type: Literal["request", "response"], - hl_headers: dict[str, str], + hl_headers: Mapping[str, str], ) -> httpx.Response: if input_type == "request": path = "detection/v2/request-evaluations" diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index e3f67f0024b..c5c36ae51c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -1,10 +1,11 @@ """LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria.""" -from collections.abc import Callable +from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -41,28 +42,87 @@ _parse_judge_verdict: Final = parse_json_verdict _extract_text_from_content: Final = extract_text_from_content +class _JudgeMessage(TypedDict): + """Chat message, as far as the judge prompt builder reads it.""" + + role: ReadOnly[NotRequired[str]] + content: ReadOnly[NotRequired[object]] + + +class _GuardrailOptions(TypedDict, total=False): + """Base :class:`CustomGuardrail` options forwarded untouched.""" + + mask_request_content: ReadOnly[bool] + mask_response_content: ReadOnly[bool] + violation_message_template: ReadOnly[str | None] + end_session_after_n_fails: ReadOnly[int | None] + on_violation: ReadOnly[str | None] + realtime_violation_message: ReadOnly[str | None] + on_sensitive_data: ReadOnly[str | None] + sensitive_data_route_to_model: ReadOnly[str | None] + sticky_session_routing: ReadOnly[bool] + run_in_parallel: ReadOnly[bool] + only_scan_new_messages: ReadOnly[bool] + + +class _RequestMessagesView(TypedDict): + messages: ReadOnly[Sequence[_JudgeMessage]] + + +class _RequestMetadataView(TypedDict): + metadata: ReadOnly[MutableMapping[str, object]] + + +class _OverallScoreView(TypedDict): + overall_score: ReadOnly[str | float] + + +class _JudgeModelView(TypedDict): + judge_model: ReadOnly[str] + + +class _CriteriaView(TypedDict): + criteria: ReadOnly[Sequence[Mapping[str, str | float]]] + + +class _OnFailureView(TypedDict): + on_failure: ReadOnly[Literal["block", "log"]] + + +class _ThresholdView(TypedDict): + overall_threshold: ReadOnly[str | float] + + +class _ModeView(TypedDict): + mode: ReadOnly[object] + + +class _DefaultOnView(TypedDict): + default_on: ReadOnly[object] + + def _get_litellm_param( litellm_params: "LitellmParams", guardrail: "Guardrail", key: str, - default: Any = None, + default: str | float | bool | None = None, ) -> Any: - val: Final = getattr(litellm_params, key, None) + val: Final[object] = getattr(litellm_params, key, None) if val is not None: return val raw: Final = guardrail.get("litellm_params") if isinstance(raw, dict) and key in raw: return raw[key] if raw is not None and not isinstance(raw, dict): - attr: Final = getattr(raw, key, None) + attr: Final[object] = getattr(raw, key, None) if attr is not None: return attr return default def _build_judge_prompt( - criteria: list[dict[str, Any]], - messages: list[dict[str, Any]], + criteria: Sequence[Mapping[str, object]], + messages: Sequence[_JudgeMessage], response_text: str, ) -> str: criteria_block: Final = "\n".join( @@ -87,13 +147,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): self, guardrail_name: str, judge_model: str, - criteria: list[dict[str, Any]], + criteria: Sequence[Mapping[str, object]], overall_threshold: float = 80.0, on_failure: Literal["block", "log"] = "block", event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, default_on: bool = False, router_provider: "Callable[[], Router | None] | None" = None, - **kwargs: Any, + **kwargs: Unpack[_GuardrailOptions], ) -> None: _event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None if event_hook is not None: @@ -121,9 +181,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): async def _run_judge( self, - messages: list[dict[str, Any]], + messages: Sequence[_JudgeMessage], response_text: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: judge_messages: Final = [ {"role": "system", "content": JUDGE_SYSTEM_PROMPT}, { @@ -162,10 +222,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): judge_result: dict[str, Any] = {} try: - messages: Final[list[dict[str, Any]]] = request_data.get("messages") or [] + request_messages: Final[_RequestMessagesView] = {"messages": request_data.get("messages") or []} try: - judge_result = await self._run_judge(messages, response_text) + judge_result = await self._run_judge(request_messages["messages"], response_text) except Exception as judge_err: verbose_logger.warning( "llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err @@ -174,7 +234,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): return inputs try: - overall_score: Final = max(0.0, min(100.0, float(judge_result.get("overall_score", 100)))) + raw_score: Final[_OverallScoreView] = {"overall_score": judge_result.get("overall_score", 100)} + overall_score: Final = max(0.0, min(100.0, float(raw_score["overall_score"]))) except (TypeError, ValueError): verbose_logger.warning("llm_as_a_judge: invalid overall_score from judge, failing open") return inputs @@ -189,7 +250,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): "threshold": self.overall_threshold, "verdicts": judge_result.get("verdicts", []), } - _metadata: Final = request_data.setdefault("metadata", {}) + request_metadata: Final[_RequestMetadataView] = {"metadata": request_data.setdefault("metadata", {})} + _metadata: Final = request_metadata["metadata"] existing: Final = _metadata.get("eval_information") if isinstance(existing, list): existing.append(eval_info) @@ -238,37 +300,45 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("llm_as_a_judge guardrail requires a guardrail_name") - judge_model: Final = _get_litellm_param(litellm_params, guardrail, "judge_model") - if not judge_model: + judge_model: Final[_JudgeModelView] = {"judge_model": _get_litellm_param(litellm_params, guardrail, "judge_model")} + if not judge_model["judge_model"]: raise ValueError("llm_as_a_judge guardrail requires judge_model in litellm_params") - criteria: Final = _get_litellm_param(litellm_params, guardrail, "criteria") or [] - if not criteria: + criteria: Final[_CriteriaView] = {"criteria": _get_litellm_param(litellm_params, guardrail, "criteria") or []} + if not criteria["criteria"]: raise ValueError("llm_as_a_judge guardrail requires at least one criterion") - weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria) + weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria["criteria"]) if abs(weight_total - 100) > 0.5: raise ValueError(f"llm_as_a_judge criterion weights must sum to 100 (got {weight_total})") - on_failure: Final = _get_litellm_param(litellm_params, guardrail, "on_failure", "block") - if on_failure not in _VALID_ON_FAILURE: - raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure}'") + on_failure: Final[_OnFailureView] = { + "on_failure": _get_litellm_param(litellm_params, guardrail, "on_failure", "block") + } + if on_failure["on_failure"] not in _VALID_ON_FAILURE: + raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure['on_failure']}'") - overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0)) + threshold: Final[_ThresholdView] = { + "overall_threshold": _get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0) + } + overall_threshold: Final = float(threshold["overall_threshold"]) - mode: Final = _get_litellm_param(litellm_params, guardrail, "mode") + mode: Final[_ModeView] = {"mode": _get_litellm_param(litellm_params, guardrail, "mode")} event_hook: GuardrailEventHooks | None = None - if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}: - event_hook = GuardrailEventHooks(mode) + if isinstance(mode["mode"], str) and mode["mode"] in {e.value for e in GuardrailEventHooks}: + event_hook = GuardrailEventHooks(mode["mode"]) + default_on: Final[_DefaultOnView] = { + "default_on": _get_litellm_param(litellm_params, guardrail, "default_on", False) + } instance: Final = LLMAsAJudgeGuardrail( guardrail_name=guardrail_name, - judge_model=judge_model, - criteria=criteria, + judge_model=judge_model["judge_model"], + criteria=criteria["criteria"], overall_threshold=overall_threshold, - on_failure=on_failure, + on_failure=on_failure["on_failure"], event_hook=event_hook, - default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)), + default_on=bool(default_on["default_on"]), ) litellm.logging_callback_manager.add_litellm_callback(instance) return instance diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index e9cd6addef8..704e2564ef5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -7,8 +7,9 @@ import enum import json import os +from collections.abc import Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, cast from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger @@ -36,6 +37,8 @@ _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail" _INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls") _DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname +_GuardrailJsonResponse: TypeAlias = Exception | str | dict[str, object] + _KEYS_DUPLICATING_SCAN_INPUTS: Final = ("messages", "input") _LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + ( "additional_args", @@ -119,7 +122,7 @@ class NomaV2Guardrail(CustomGuardrail): def _resolve_action_from_response( self, - response_json: dict, + response_json: Mapping[str, object], ) -> _Action: action: Final = response_json.get("action") if isinstance(action, str): @@ -165,10 +168,11 @@ class NomaV2Guardrail(CustomGuardrail): @staticmethod def _sanitize_payload_for_transport(payload: dict) -> dict: - def _default(obj: Any) -> Any: - if hasattr(obj, "model_dump"): + def _default(obj: object) -> object: + model_dump: Final[Callable[[], Mapping[str, object]] | None] = getattr(obj, "model_dump", None) + if model_dump is not None: try: - return obj.model_dump() + return model_dump() except Exception: pass return str(obj) @@ -178,7 +182,7 @@ class NomaV2Guardrail(CustomGuardrail): except (ValueError, TypeError): json_str = safe_dumps(payload) - safe_payload: Final = safe_json_loads(json_str, default={}) + safe_payload: Final[object] = safe_json_loads(json_str, default={}) if safe_payload == {} and payload: verbose_proxy_logger.warning( "Noma v2 guardrail: payload serialization failed, falling back to empty payload" @@ -196,7 +200,7 @@ class NomaV2Guardrail(CustomGuardrail): async def _call_noma_scan( self, payload: dict, - ) -> dict: + ) -> dict[str, object]: headers: Final[dict[str, str]] = {"Content-Type": "application/json"} authorization_header: Final = self._get_authorization_header() if authorization_header: @@ -215,7 +219,7 @@ class NomaV2Guardrail(CustomGuardrail): response.text, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[dict[str, object]] = response.json() verbose_proxy_logger.debug( "Noma v2 AIDR response parsed: %s", json.dumps(response_json, default=str), @@ -227,7 +231,7 @@ class NomaV2Guardrail(CustomGuardrail): request_data: dict, start_time: datetime, guardrail_status: GuardrailStatus, - guardrail_json_response: Any, + guardrail_json_response: _GuardrailJsonResponse, ) -> None: end_time: Final = datetime.now() duration: Final = (end_time - start_time).total_seconds() @@ -270,11 +274,11 @@ class NomaV2Guardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: start_time: Final = datetime.now() guardrail_status: GuardrailStatus = "success" - guardrail_json_response: Any = {} + guardrail_json_response: _GuardrailJsonResponse = {} dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data) if not isinstance(dynamic_params, dict): dynamic_params = {} - response_json: dict | None = None + response_json: dict[str, object] | None = None # Per-request dynamic params can override configured application context. application_id = self._get_non_empty_str(dynamic_params.get("application_id")) diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index 4775a8b3caa..155db816e94 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -7,8 +7,11 @@ before and after LLM calls. """ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing_extensions import ReadOnly, TypedDict + from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import ( @@ -20,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -34,6 +38,16 @@ _DEFAULT_API_BASE: Final = "https://api.promptguard.co" _GUARD_ENDPOINT: Final = "/api/v1/guard" +class PromptGuardResult(TypedDict, total=False): + """The fields this guardrail reads off a PromptGuard Guard API response.""" + + decision: ReadOnly[str] + threat_type: ReadOnly[str] + event_id: ReadOnly[str] + confidence: ReadOnly[float] + redacted_messages: ReadOnly[list[AllMessageValues]] + + class PromptGuardMissingCredentials(Exception): pass @@ -96,7 +110,7 @@ class PromptGuardGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, object], input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: @@ -114,7 +128,7 @@ class PromptGuardGuardrail(CustomGuardrail): direction: Final = "input" if input_type == "request" else "output" - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "messages": messages, "direction": direction, } @@ -144,7 +158,7 @@ class PromptGuardGuardrail(CustomGuardrail): timeout=10.0, ) response.raise_for_status() - result: Final = response.json() + result: Final[PromptGuardResult] = response.json() except Exception as exc: verbose_proxy_logger.error("PromptGuard API error: %s", str(exc)) if self.block_on_error: diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 41e52f05c01..7b45afa0ad2 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,4 +1,6 @@ -from typing import Final +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query @@ -18,7 +20,59 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository router: Final = APIRouter() -def _to_response(mapping) -> JWTKeyMappingResponse: +class _JWTKeyMappingRecord(Protocol): + """A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read.""" + + @property + def id(self) -> str: ... + + @property + def jwt_claim_name(self) -> str: ... + + @property + def jwt_claim_value(self) -> str: ... + + @property + def description(self) -> str | None: ... + + @property + def is_active(self) -> bool: ... + + @property + def created_at(self) -> datetime: ... + + @property + def updated_at(self) -> datetime: ... + + @property + def created_by(self) -> str | None: ... + + @property + def updated_by(self) -> str | None: ... + + +class _JWTKeyMappingTable(Protocol): + """The Prisma table actions these endpoints issue against the JWT key mapping table.""" + + async def create(self, *, data: Mapping[str, object]) -> _JWTKeyMappingRecord: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _JWTKeyMappingRecord: ... + + async def delete(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ... + + async def find_many(self, *, skip: int, take: int, order: Mapping[str, str]) -> Sequence[_JWTKeyMappingRecord]: ... + + async def count(self) -> int: ... + + +def _mapping_table(prisma_client: object) -> _JWTKeyMappingTable: + """View the JWT key mapping repository's untyped Prisma table through the actions used here.""" + return JWTKeyMappingRepository(prisma_client).table + + +def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse: """Convert a Prisma mapping object to a safe response (no hashed token).""" return JWTKeyMappingResponse( id=mapping.id, @@ -62,7 +116,7 @@ async def create_jwt_key_mapping( if data.description is not None: create_data["description"] = data.description - new_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.create(data=create_data) + new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data) # Invalidate cache cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}" @@ -110,7 +164,7 @@ async def update_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) + old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -118,9 +172,7 @@ async def update_jwt_key_mapping( cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - updated_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.update( - where={"id": data.id}, data=update_data - ) + updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data) # Invalidate new cache key if claim fields changed cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}" @@ -159,7 +211,7 @@ async def delete_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) + old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -167,7 +219,7 @@ async def delete_jwt_key_mapping( cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - await JWTKeyMappingRepository(prisma_client).table.delete(where={"id": data.id}) + await _mapping_table(prisma_client).delete(where={"id": data.id}) return {"status": "success"} except HTTPException: raise @@ -195,12 +247,12 @@ async def list_jwt_key_mappings( try: skip: Final = (page - 1) * size - mappings: Final = await JWTKeyMappingRepository(prisma_client).table.find_many( + mappings: Final = await _mapping_table(prisma_client).find_many( skip=skip, take=size, order={"created_at": "desc"}, ) - total_count: Final = await JWTKeyMappingRepository(prisma_client).table.count() + total_count: Final = await _mapping_table(prisma_client).count() return { "mappings": [_to_response(m) for m in mappings], "total_count": total_count, @@ -232,7 +284,7 @@ async def info_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": id}) + mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": id}) if mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") return _to_response(mapping) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 020698dabd9..3c5d889d23f 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,8 +10,8 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import Sequence -from typing import TYPE_CHECKING, Final, TypedDict, cast +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, TypedDict from fastapi import Request, Response from fastapi.responses import StreamingResponse @@ -38,6 +38,36 @@ class _StreamOutputItem(TypedDict, total=False): content: ReadOnly[Sequence[_StreamContentPart | None]] +class _StreamTerminalResponse(TypedDict, total=False): + """Fields of the ``response`` payload carried by a terminal streaming event.""" + + status: ReadOnly[ResponsesAPIStatus] + tool_choice: ReadOnly[object] + model: ReadOnly[str] + instructions: ReadOnly[str] + temperature: ReadOnly[float] + top_p: ReadOnly[float] + max_output_tokens: ReadOnly[int] + previous_response_id: ReadOnly[str] + truncation: ReadOnly[str] + parallel_tool_calls: ReadOnly[bool] + user: ReadOnly[str] + store: ReadOnly[bool] + output: ReadOnly[Sequence[_StreamOutputItem]] + + +class _StreamEvent(TypedDict, total=False): + """One decoded ``data:`` frame of an OpenAI Responses streaming body.""" + + type: ReadOnly[str] + item: ReadOnly[_StreamOutputItem] + item_id: ReadOnly[str] + part: ReadOnly[_StreamContentPart] + content_index: ReadOnly[int] + delta: ReadOnly[str] + response: ReadOnly[_StreamTerminalResponse] + + async def background_streaming_task( polling_id: str, data, @@ -139,7 +169,7 @@ async def background_streaming_task( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None - _event_to_status: Final = { + _event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = { "response.completed": "completed", "response.failed": "failed", "response.incomplete": "incomplete", @@ -180,7 +210,7 @@ async def background_streaming_task( break try: - event = json.loads(chunk_data) + event: _StreamEvent = json.loads(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -288,12 +318,9 @@ async def background_streaming_task( # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming response_data = event.get("response", {}) - terminal_status = cast( - ResponsesAPIStatus, - response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ), + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), ) # Extract error for failed and incomplete responses diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index fe4794f3ba1..b25263e4c64 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -2,8 +2,9 @@ Search Tool Registry for managing search tool configurations. """ +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Final +from typing import Final, Protocol from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -13,6 +14,40 @@ from litellm.repositories.table_repositories import SearchToolsRepository from litellm.types.search import SearchTool +class SearchToolRecord(Protocol): + search_tool_id: str + search_tool_name: str + created_at: datetime + updated_at: datetime + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + +class SearchToolTableClient(Protocol): + async def create(self, data: Mapping[str, object]) -> SearchToolRecord: ... + + async def find_unique(self, where: Mapping[str, object]) -> SearchToolRecord | None: ... + + async def find_many(self, order: Mapping[str, str] | None = None) -> Sequence[SearchToolRecord]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> SearchToolRecord: ... + + async def delete(self, where: Mapping[str, object]) -> SearchToolRecord: ... + + +class _SearchToolsRepositoryView(Protocol): + @property + def table(self) -> SearchToolTableClient: ... + + +def _search_tools_table_of(repository: _SearchToolsRepositoryView) -> SearchToolTableClient: + return repository.table + + +def _search_tools_table(prisma_client: PrismaClient) -> SearchToolTableClient: + return _search_tools_table_of(SearchToolsRepository(prisma_client)) + + class SearchToolRegistry: """ Handles adding, removing, and getting search tools in DB + in memory. @@ -22,7 +57,7 @@ class SearchToolRegistry: pass @staticmethod - def _convert_prisma_to_dict(prisma_obj) -> dict: + def _convert_prisma_to_dict(prisma_obj: SearchToolRecord) -> dict: """ Convert Prisma result to dict with datetime objects as ISO format strings. @@ -35,9 +70,9 @@ class SearchToolRegistry: result: Final = dict(prisma_obj) # Convert datetime objects to ISO format strings if "created_at" in result and result["created_at"]: - result["created_at"] = result["created_at"].isoformat() + result["created_at"] = prisma_obj.created_at.isoformat() if "updated_at" in result and result["updated_at"]: - result["updated_at"] = result["updated_at"].isoformat() + result["updated_at"] = prisma_obj.updated_at.isoformat() return result ########################################################### @@ -61,7 +96,7 @@ class SearchToolRegistry: search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {})) # Create search tool in DB - created_search_tool: Final = await SearchToolsRepository(prisma_client).table.create( + created_search_tool: Final = await _search_tools_table(prisma_client).create( data={ "search_tool_name": search_tool_name, "litellm_params": litellm_params, @@ -95,7 +130,7 @@ class SearchToolRegistry: """ try: # Get search tool before deletion for response - existing_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + existing_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_id": search_tool_id} ) @@ -103,7 +138,7 @@ class SearchToolRegistry: raise Exception(f"Search tool with ID {search_tool_id} not found") # Delete from DB - await SearchToolsRepository(prisma_client).table.delete(where={"search_tool_id": search_tool_id}) + await _search_tools_table(prisma_client).delete(where={"search_tool_id": search_tool_id}) return { "message": f"Search tool {search_tool_id} deleted successfully", @@ -131,7 +166,7 @@ class SearchToolRegistry: search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {})) # Update in DB - updated_search_tool: Final = await SearchToolsRepository(prisma_client).table.update( + updated_search_tool: Final = await _search_tools_table(prisma_client).update( where={"search_tool_id": search_tool_id}, data={ "search_tool_name": search_tool_name, @@ -163,7 +198,7 @@ class SearchToolRegistry: try: search_tools_from_db: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: SearchToolsRepository(prisma_client).table.find_many( + lambda: _search_tools_table(prisma_client).find_many( order={"created_at": "desc"}, ), reason="get_all_search_tools_from_db_lookup_failure", @@ -194,7 +229,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + search_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_id": search_tool_id} ) @@ -222,7 +257,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + search_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_name": search_tool_name} ) diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 16b8f82815c..255faf94402 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,11 +1,45 @@ +from collections.abc import Sequence from typing import Any, Final +from typing_extensions import NotRequired, ReadOnly, TypedDict + from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.utils import ModelResponse -from litellm.types.vector_stores import ( - VectorStoreResultContent, - VectorStoreSearchResponse, -) +from litellm.types.vector_stores import VectorStoreSearchResponse + + +class _ResultContentView(TypedDict): + """Content entry carried by a vector store search result.""" + + type: ReadOnly[NotRequired[str]] + text: ReadOnly[str] + + +class _SearchResultView(TypedDict): + """Vector store search result, as far as :class:`RAGQuery` reads it.""" + + content: ReadOnly[NotRequired[Sequence[_ResultContentView]]] + text: ReadOnly[NotRequired[str]] + + +class _SearchDataView(TypedDict): + results: ReadOnly[Sequence[_SearchResultView]] + + +class _ContextChunksView(TypedDict): + chunks: ReadOnly[Sequence[_SearchResultView | str | None]] + + +class _RerankResultView(TypedDict): + index: ReadOnly[NotRequired[int]] + + +class _RerankResultsView(TypedDict): + results: ReadOnly[Sequence[_RerankResultView]] + + +class _MessageView(TypedDict): + message: ReadOnly[object] class RAGQuery: @@ -42,9 +76,10 @@ class RAGQuery: """ context_content = RAGQuery.CONTENT_PREFIX_STRING - for chunk in context_chunks: + chunks: Final[_ContextChunksView] = {"chunks": context_chunks} + for chunk in chunks["chunks"]: if isinstance(chunk, dict): - result_content: list[VectorStoreResultContent] | None = chunk.get("content") + result_content: Sequence[_ResultContentView] | None = chunk.get("content") if result_content: for content_item in result_content: content_text: str | None = content_item.get("text") @@ -64,14 +99,15 @@ class RAGQuery: def add_search_results_to_response( response: ModelResponse, search_results: VectorStoreSearchResponse, - rerank_results: Any | None = None, + rerank_results: object = None, ) -> ModelResponse: """ Add search results to the response choices. """ if hasattr(response, "choices") and response.choices: for choice in response.choices: - message = getattr(choice, "message", None) + message_view: _MessageView = {"message": getattr(choice, "message", None)} + message = message_view["message"] if message is not None: # Get existing provider_specific_fields or create new dict provider_fields = getattr(message, "provider_specific_fields", None) or {} @@ -91,7 +127,8 @@ class RAGQuery: ) -> list[str | dict[str, Any]]: """Extract text documents from vector store search response.""" documents: Final[list[str | dict[str, Any]]] = [] - for result in search_response.get("data", []): + search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} + for result in search_data["results"]: content_list = result.get("content", []) for content in content_list: if content.get("type") == "text" and content.get("text"): @@ -99,11 +136,13 @@ class RAGQuery: return documents @staticmethod - def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[Any]: + def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]: """Get the original search results corresponding to the top reranked results.""" - top_chunks: Final = [] - original_results: Final = search_response.get("data", []) - for result in rerank_response.get("results", []): + top_chunks: Final[list[_SearchResultView]] = [] + search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} + original_results: Final = search_data["results"] + reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])} + for result in reranked["results"]: index = result.get("index") if index is not None and index < len(original_results): top_chunks.append(original_results[index]) diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index 7008099fe8c..8dfe3ed962e 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -3,7 +3,7 @@ Base repository class with common functionality. """ from abc import ABC, abstractmethod -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Any, Final, Generic, Protocol, TypeVar, runtime_checkable from pydantic import BaseModel @@ -24,6 +24,22 @@ class SupportsDict(Protocol): DbRecord = Mapping[str, object] | SupportsModelDump | SupportsDict | Sequence[tuple[str, object]] +class PrismaCrudActions(Protocol): + """The Prisma table actions reached by the generic repository CRUD helpers.""" + + async def find_unique(self, *, where: Mapping[str, object]) -> DbRecord | None: ... + + find_many: Callable[..., Awaitable[Sequence[DbRecord]]] + + async def create(self, *, data: Mapping[str, object]) -> DbRecord: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> DbRecord | None: ... + + async def delete(self, *, where: Mapping[str, object]) -> DbRecord | None: ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + def record_to_dict(record: DbRecord) -> Mapping[str, object]: """Project a database record into a mapping of column name to value.""" if isinstance(record, SupportsModelDump): @@ -38,7 +54,7 @@ def record_to_dict(record: DbRecord) -> Mapping[str, object]: class BaseRepository(ABC, Generic[T]): """Abstract base class for all repositories.""" - def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper + def __init__(self, prisma_client: object): self._prisma_client = prisma_client @property @@ -53,6 +69,11 @@ class BaseRepository(ABC, Generic[T]): """Return the Prisma table for this repository.""" ... + @property + def _crud_actions(self) -> PrismaCrudActions: + """View ``table`` through the action surface the CRUD helpers below use.""" + return self.table + @property @abstractmethod def model_class(self) -> type[T]: @@ -71,18 +92,18 @@ class BaseRepository(ABC, Generic[T]): async def find_by_id(self, id_value: str, id_field: str = "id") -> T | None: """Find a record by its primary key.""" - record: Final = await self.table.find_unique(where={id_field: id_value}) + record: Final = await self._crud_actions.find_unique(where={id_field: id_value}) return self._to_model(record) async def find_many( self, - where: dict[str, Any] | None = None, + where: Mapping[str, object] | None = None, skip: int | None = None, take: int | None = None, - order: dict[str, str] | None = None, + order: Mapping[str, str] | None = None, ) -> list[T]: """Find multiple records matching the criteria.""" - kwargs: Final[dict[str, Any]] = {} + kwargs: Final[dict[str, object]] = {} if where: kwargs["where"] = where if skip is not None: @@ -92,31 +113,31 @@ class BaseRepository(ABC, Generic[T]): if order: kwargs["order"] = order - records: Final = await self.table.find_many(**kwargs) + records: Final = await self._crud_actions.find_many(**kwargs) return self._to_model_list(records) - async def create(self, data: dict[str, Any]) -> T: + async def create(self, data: Mapping[str, object]) -> T: """Create a new record.""" - record: Final = await self.table.create(data=data) + record: Final = await self._crud_actions.create(data=data) model: Final = self._to_model(record) assert model is not None return model - async def update(self, id_value: str, data: dict[str, Any], id_field: str = "id") -> T | None: + async def update(self, id_value: str, data: Mapping[str, object], id_field: str = "id") -> T | None: """Update an existing record.""" - record: Final = await self.table.update(where={id_field: id_value}, data=data) + record: Final = await self._crud_actions.update(where={id_field: id_value}, data=data) return self._to_model(record) async def delete(self, id_value: str, id_field: str = "id") -> T | None: """Delete a record by its primary key.""" - record: Final = await self.table.delete(where={id_field: id_value}) + record: Final = await self._crud_actions.delete(where={id_field: id_value}) return self._to_model(record) - async def count(self, where: dict[str, Any] | None = None) -> int: + async def count(self, where: Mapping[str, object] | None = None) -> int: """Count records matching the criteria.""" - return await self.table.count(where=where) + return await self._crud_actions.count(where=where) async def exists(self, id_value: str, id_field: str = "id") -> bool: """Check if a record exists.""" - record: Final = await self.table.find_unique(where={id_field: id_value}) + record: Final = await self._crud_actions.find_unique(where={id_field: id_value}) return record is not None diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index 9fdb6e4aca7..b22f24e0de5 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -6,11 +6,41 @@ credential values is the caller's responsibility (see ``CredentialHelperUtils``) so reads return the stored values verbatim. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.credentials import CredentialItem from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync +if TYPE_CHECKING: + from prisma.models import LiteLLM_CredentialsTable + + +class _CredentialsDb(Protocol): + @property + def litellm_credentialstable(self) -> object: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _CredentialsDb: ... + + +class _CredentialsActions(Protocol): + """Prisma table actions used by :class:`CredentialsRepository`.""" + + async def find_many(self) -> "Sequence[LiteLLM_CredentialsTable]": ... + + async def create(self, *, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable": ... + + async def find_unique(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ... + + async def update( + self, *, where: Mapping[str, object], data: Mapping[str, object] + ) -> "LiteLLM_CredentialsTable | None": ... + + async def delete(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ... + class CredentialsRepository: """Repository for credentials database operations, keyed by credential name.""" @@ -19,7 +49,7 @@ class CredentialsRepository: self._prisma_client = prisma_client @property - def prisma_client(self) -> Any: + def prisma_client(self) -> _PrismaClientView: if self._prisma_client is None: raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client @@ -31,6 +61,10 @@ class CredentialsRepository: table_name="litellm_credentialstable", ) + @property + def _credentials_table(self) -> _CredentialsActions: + return self.table + @staticmethod def _to_model(record: Any) -> CredentialItem | None: if record is None: @@ -42,18 +76,20 @@ class CredentialsRepository: credential_info=data.get("credential_info") or {}, ) - async def find_all(self) -> Any: - return await self.table.find_many() + async def find_all(self) -> "Sequence[LiteLLM_CredentialsTable]": + return await self._credentials_table.find_many() - async def create(self, data: dict[str, Any]) -> Any: - return await self.table.create(data=data) + async def create(self, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable": + return await self._credentials_table.create(data=data) async def find_by_name(self, credential_name: str) -> CredentialItem | None: - record: Final = await self.table.find_unique(where={"credential_name": credential_name}) + record: Final = await self._credentials_table.find_unique(where={"credential_name": credential_name}) return self._to_model(record) - async def update_by_name(self, credential_name: str, data: dict[str, Any]) -> Any: - return await self.table.update(where={"credential_name": credential_name}, data=data) + async def update_by_name( + self, credential_name: str, data: Mapping[str, object] + ) -> "LiteLLM_CredentialsTable | None": + return await self._credentials_table.update(where={"credential_name": credential_name}, data=data) - async def delete_by_name(self, credential_name: str) -> Any: - return await self.table.delete(where={"credential_name": credential_name}) + async def delete_by_name(self, credential_name: str) -> "LiteLLM_CredentialsTable | None": + return await self._credentials_table.delete(where={"credential_name": credential_name}) diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 7efd32288e4..68490797348 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -3,9 +3,10 @@ Team repository for database operations on LiteLLM_TeamTable. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from contextlib import AbstractAsyncContextManager from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from pydantic import TypeAdapter @@ -13,12 +14,45 @@ from litellm.models.team import LiteLLM_TeamTable, Member from litellm.repositories.base_repository import ( BaseRepository, DbRecord, + PrismaCrudActions, record_to_dict, ) if TYPE_CHECKING: from prisma import Prisma + +class _TeamTables(Protocol): + """The Prisma tables this repository reaches, on the client or inside a transaction.""" + + litellm_teamtable: PrismaCrudActions + litellm_deletedteamtable: PrismaCrudActions + + +class _TeamArrays(Protocol): + """The string array columns of a team row, which the domain model leaves untyped.""" + + @property + def members(self) -> Sequence[str]: ... + + @property + def admins(self) -> Sequence[str]: ... + + @property + def models(self) -> Sequence[str]: ... + + +def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays: + """View a team's untyped list columns as sequences of ids.""" + return team + + +class _TeamDatabase(_TeamTables, Protocol): + """The Prisma client surface used for team reads, writes, and archival transactions.""" + + def tx(self) -> AbstractAsyncContextManager[_TeamTables]: ... + + _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) _JSON_ENCODED_TEAM_FIELDS: Final = ( "metadata", @@ -34,12 +68,16 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" @property - def table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper - return self.prisma_client.db.litellm_teamtable + def _db(self) -> _TeamDatabase: + return self.prisma_client.db @property - def deleted_table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper - return self.prisma_client.db.litellm_deletedteamtable + def table(self) -> Any: # any-ok: callers reach model-specific actions this repository does not use + return self._db.litellm_teamtable + + @property + def deleted_table(self) -> PrismaCrudActions: + return self._db.litellm_deletedteamtable @property def model_class(self) -> type[LiteLLM_TeamTable]: @@ -75,8 +113,8 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) if not rows: return None - raw_value: Final = rows[0]["members_with_roles"] - parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value + raw_value: Final[object] = rows[0]["members_with_roles"] + parsed: Final[object] = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] return _MEMBERS_WITH_ROLES_ADAPTER.validate_python(parsed) @@ -86,24 +124,24 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): async def find_by_alias(self, team_alias: str) -> LiteLLM_TeamTable | None: """Find a team by alias.""" - records: Final = await self.table.find_many(where={"team_alias": team_alias}) + records: Final = await self._crud_actions.find_many(where={"team_alias": team_alias}) if records: return self._to_model(records[0]) return None async def find_by_organization_id(self, organization_id: str) -> list[LiteLLM_TeamTable]: """Find all teams belonging to an organization.""" - records: Final = await self.table.find_many(where={"organization_id": organization_id}) + records: Final = await self._crud_actions.find_many(where={"organization_id": organization_id}) return self._to_model_list(records) async def find_by_member(self, user_id: str) -> list[LiteLLM_TeamTable]: """Find all teams where user is a member.""" - records: Final = await self.table.find_many(where={"members": {"has": user_id}}) + records: Final = await self._crud_actions.find_many(where={"members": {"has": user_id}}) return self._to_model_list(records) async def find_by_admin(self, user_id: str) -> list[LiteLLM_TeamTable]: """Find all teams where user is an admin.""" - records: Final = await self.table.find_many(where={"admins": {"has": user_id}}) + records: Final = await self._crud_actions.find_many(where={"admins": {"has": user_id}}) return self._to_model_list(records) async def create_team( @@ -232,7 +270,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): archive_data["litellm_changed_by"] = litellm_changed_by archive_data["deleted_at"] = datetime.utcnow() - async with self.prisma_client.db.tx() as tx: + async with self._db.tx() as tx: await tx.litellm_deletedteamtable.create(data=archive_data) await tx.litellm_teamtable.delete(where={"team_id": team_id}) @@ -293,7 +331,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"members": {"push": user_id}}, ) @@ -310,7 +348,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - members: Final = [m for m in team.members if m != user_id] + members: Final = [m for m in _team_arrays(team).members if m != user_id] return await self.update(team_id, {"members": members}, id_field="team_id") async def add_admin(self, team_id: str, user_id: str) -> LiteLLM_TeamTable | None: @@ -318,7 +356,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"admins": {"push": user_id}}, ) @@ -335,7 +373,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - admins: Final = [a for a in team.admins if a != user_id] + admins: Final = [a for a in _team_arrays(team).admins if a != user_id] return await self.update(team_id, {"admins": admins}, id_field="team_id") async def add_models(self, team_id: str, models: list[str]) -> LiteLLM_TeamTable | None: @@ -343,7 +381,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"models": {"push": models}}, ) @@ -360,5 +398,5 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - current_models: Final = [m for m in team.models if m not in models] + current_models: Final = [m for m in _team_arrays(team).models if m not in models] return await self.update(team_id, {"models": current_models}, id_field="team_id") diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 20fc2634a8a..c57a2b14dcb 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Awaitable from dataclasses import dataclass -from typing import Any, Final, Protocol +from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK @@ -12,6 +13,16 @@ from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds +class RustResponsesWebSocketSocket(Protocol): + """Open socket handle handed back by the native bridge.""" + + def send_text(self, text: str) -> Awaitable[None]: ... + + def recv_text(self) -> Awaitable[str | None]: ... + + def close(self) -> Awaitable[None]: ... + + class RustResponsesWebSocketConnection(Protocol): @classmethod def connect( @@ -19,7 +30,7 @@ class RustResponsesWebSocketConnection(Protocol): url: str, headers: dict[str, str], timeout_seconds: float | None, - ) -> Any: + ) -> Awaitable[RustResponsesWebSocketSocket]: raise NotImplementedError @@ -32,7 +43,7 @@ _UNSET: Final[_Unset] = _Unset() @dataclass(slots=True) class _RustResponsesWebSocketState: - connection: Any = None + connection: type[RustResponsesWebSocketConnection] | None = None _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() @@ -40,13 +51,13 @@ _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() def set_rust_responses_websocket( *, - connection: Any = _UNSET, + connection: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET, ) -> None: if not isinstance(connection, _Unset): _STATE.connection = connection -def load_rust_responses_websocket() -> Any: +def load_rust_responses_websocket() -> type[RustResponsesWebSocketConnection] | None: if _STATE.connection is not None: return _STATE.connection native_bridge: Final = get_native_bridge() @@ -59,7 +70,7 @@ def load_rust_responses_websocket() -> Any: class _ConnectionAdapter: - def __init__(self, connection: Any): + def __init__(self, connection: RustResponsesWebSocketSocket): self._connection = connection async def send(self, text: str) -> None: diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index c77d2505d0e..a85382f863e 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -6,14 +6,62 @@ Handles retrieving secrets from different secret management systems. import base64 import os -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Protocol, overload import litellm from litellm._logging import print_verbose -from litellm.types.secret_managers.main import KeyManagementSystem +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem -def _is_base64(s): +class _VaultSecret(Protocol): + """The secret object returned by the Azure Key Vault and Infisical clients.""" + + @property + def value(self) -> str | None: ... + + @property + def secret_value(self) -> str | None: ... + + +class _KmsPlaintext(Protocol): + """The decrypted payload the AWS KMS client exposes under the ``Plaintext`` response key.""" + + def decode(self, encoding: str) -> str | None: ... + + +class _KmsDecryptResponse(Protocol): + """The decrypt response returned by the Google KMS and AWS KMS clients.""" + + @property + def plaintext(self) -> bytes: ... + + def __getitem__(self, key: str) -> _KmsPlaintext: ... + + +class _SecretManagerClient(Protocol): + """The untyped secret manager client surface reached by ``get_secret_from_manager``.""" + + def get_secret(self, secret_name: str, /) -> _VaultSecret: ... + + @overload + def decrypt(self, *, request: Mapping[str, object]) -> _KmsDecryptResponse: ... + + @overload + def decrypt(self, *, CiphertextBlob: bytes) -> _KmsDecryptResponse: ... + + def sync_read_secret( + self, + *, + secret_name: str, + primary_secret_name: str | None = None, + optional_params: Mapping[str, object] | None = None, + ) -> str | None: ... + + def get_secret_from_google_secret_manager(self, secret_name: str, /) -> str | None: ... + + +def _is_base64(s: str | bytes) -> bool: """Check if a string is valid base64.""" import binascii @@ -24,10 +72,10 @@ def _is_base64(s): def get_secret_from_manager( - client: Any, + client: _SecretManagerClient, key_manager: str, secret_name: str, - key_management_settings: Any | None = None, + key_management_settings: KeyManagementSettings | None = None, ) -> str | None: """ Get a secret from the configured secret manager. @@ -56,7 +104,7 @@ def get_secret_from_manager( elif ( key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient" ): - encrypted_secret: Any = os.getenv(secret_name) + encrypted_secret: str | bytes | None = os.getenv(secret_name) if encrypted_secret is None: raise ValueError("Google KMS requires the encrypted secret to be in the environment!") b64_flag: Final = _is_base64(encrypted_secret) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 5c312dcf1c8..ceaf1ec504f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 3020 + "limit": 3012 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 827 + "limit": 818 }, "ANN201": { - "limit": 2017 + "limit": 2014 }, "ANN202": { - "limit": 852 + "limit": 851 }, "ANN204": { - "limit": 711 + "limit": 707 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1188 + "limit": 1061 }, "ASYNC230": { "limit": 11 @@ -198,7 +198,7 @@ "limit": 22 }, "SIM101": { - "limit": 58 + "limit": 57 }, "SIM102": { "limit": 317 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1212 + "limit": 1202 }, "TRY002": { "limit": 524 @@ -249,7 +249,7 @@ "limit": 113 }, "TRY300": { - "limit": 859 + "limit": 858 }, "UP028": { "limit": 2 diff --git a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py b/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py index d106cf7ea21..a768f8ccd8c 100644 --- a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py +++ b/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py @@ -313,6 +313,25 @@ class TestGDCGeminiConfig: api_base=TEST_API_BASE, ) + def test_validate_environment_credentials_missing_audience_binding_are_named(self): + config = GDCGeminiConfig() + creds_without_audience_binding = MagicMock(spec=[]) + + with patch( + "google.auth.load_credentials_from_dict", + return_value=(creds_without_audience_binding, None), + ): + with pytest.raises(AttributeError, match="must expose with_gdch_audience"): + config.validate_environment( + headers={}, + model=TEST_MODEL, + messages=[], + optional_params={}, + litellm_params={"vertex_project": TEST_PROJECT}, + api_key=TEST_API_KEY, + api_base=TEST_API_BASE, + ) + def test_validate_environment_string_false_disables_token_caching(self): config = GDCGeminiConfig() mock_creds = MagicMock() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 20cd1165577..890184fad78 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22805 + "limit": 22728 }, "LIT002": { - "limit": 26878 + "limit": 26860 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1069 + "limit": 1066 }, "LIT007": { "limit": 0 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16695 + "limit": 16673 }, "LIT011": { - "limit": 5588 + "limit": 5586 }, "LIT012": { - "limit": 4519 + "limit": 4512 } } From a90ad5fe5c6dc3d837a0f28f32b1cc9e2ce93c33 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:35:15 -0700 Subject: [PATCH 03/16] feat(bedrock): honor streaming buffer/sampling config for unbuffered post_call scans --- .../guardrail_hooks/bedrock_guardrails.py | 46 ++++- .../guardrails/guardrail_initializers.py | 4 + litellm/types/guardrails.py | 32 ++++ .../test_bedrock_guardrails.py | 172 ++++++++++++++++++ 4 files changed, 253 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index dd76a27c80f..d7f222db35d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -52,7 +52,12 @@ from litellm.proxy.guardrails.anthropic_sse import ( model_response_text, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks +from litellm.types.guardrails import ( + BedrockChecksConfigModel, + BedrockGuardrailStreamingParams, + GuardrailEventHooks, + LitellmParams, +) from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksMessage, @@ -221,9 +226,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + streaming_buffer_until_moderated: bool | None = None, + streaming_sampling_rate: int | None = None, + streaming_end_of_stream_only: bool | None = None, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self._set_streaming_params( + BedrockGuardrailStreamingParams.from_extras( + { + "streaming_buffer_until_moderated": streaming_buffer_until_moderated, + "streaming_sampling_rate": streaming_sampling_rate, + "streaming_end_of_stream_only": streaming_end_of_stream_only, + } + ) + ) self.guardrailIdentifier = guardrailIdentifier self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" @@ -278,6 +295,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): list(self.checks.keys()) if self.checks else None, ) + def _set_streaming_params(self, streaming_params: BedrockGuardrailStreamingParams) -> None: + self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated + self.streaming_sampling_rate = streaming_params.streaming_sampling_rate + self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only + + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + super().update_in_memory_litellm_params(litellm_params) + self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra)) + + def _streams_incrementally(self) -> bool: + return not self.streaming_buffer_until_moderated and not self.mask_response_content + @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: return [ @@ -2660,6 +2689,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Collect content from the stream and run the bedrock OUTPUT scan (post_call only validates the response). """ + if self._streams_incrementally(): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + async for streamed_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + guardrail_to_apply=self, + buffer_until_moderated_default=False, + ): + yield streamed_chunk + return + # Import here to avoid circular imports from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.main import stream_chunk_builder diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 47aea62f4c2..76dea1b7784 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -11,6 +11,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): BedrockGuardrail, ) + streaming_params: Final = BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra) _bedrock_callback: Final = BedrockGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, @@ -38,6 +39,9 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only, only_scan_new_messages=litellm_params.only_scan_new_messages or False, + streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated, + streaming_sampling_rate=streaming_params.streaming_sampling_rate, + streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 9be78757511..c5398160c69 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from enum import Enum from typing import Any, Final, Literal @@ -550,6 +551,37 @@ class BedrockGuardrailConfigModel(BaseModel): ) +class BedrockGuardrailStreamingParams(BaseModel): + streaming_buffer_until_moderated: bool = Field( + default=True, + description="If True (default), withhold every streamed chunk until the end-of-stream " + "ApplyGuardrail scan passes, so no flagged content reaches the client before a block. " + "If False, chunks stream through unbuffered, so flagged content can reach the client " + "before the scan finishes; a flagged scan still ends the stream, with a block message " + "when disable_exception_on_block is true and an in-stream error frame otherwise.", + ) + streaming_sampling_rate: int = Field( + default=5, + ge=1, + description="When not buffering and not end-of-stream-only, scan the accumulated response " + "every Nth streamed chunk. Each sampled scan is a full ApplyGuardrail call that delays " + "that chunk, so lower values add latency and AWS text-unit cost.", + ) + streaming_end_of_stream_only: bool = Field( + default=False, + description="When not buffering, skip per-chunk sampling and run one ApplyGuardrail scan " + "on the assembled response at end of stream. Combined with " + "streaming_buffer_until_moderated=false the full response streams live before the scan " + "and the scan result lands in guardrail_information; a flagged response still ends the " + "stream with a block message (disable_exception_on_block=true) or an error frame.", + ) + + @classmethod + def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams": + source: Final[Mapping[str, object]] = extras or {} + return cls.model_validate({name: source[name] for name in cls.model_fields if source.get(name) is not None}) + + class LakeraV2GuardrailConfigModel(BaseModel): """Configuration parameters for the Lakera AI v2 guardrail""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 36b356e34d0..ca4b0a65e0b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5345,3 +5345,175 @@ def test_initialize_bedrock_forwards_aws_external_id(): assert guardrail.optional_params["aws_external_id"] == "external-id-123" finally: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, guardrail) + + +def _chat_chunk(content: str, finish_reason: str | None) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content=content, role="assistant"), + finish_reason=finish_reason, + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + +def _streaming_litellm_params(**extras): + from litellm.types.guardrails import LitellmParams + + return LitellmParams( + guardrail="bedrock", + mode="post_call", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + **extras, + ) + + +def test_initialize_bedrock_wires_streaming_flags(): + from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock + + configured = initialize_bedrock( + _streaming_litellm_params( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=3, + streaming_end_of_stream_only=True, + ), + {"guardrail_name": "bedrock-streaming"}, + ) + defaulted = initialize_bedrock( + _streaming_litellm_params(), + {"guardrail_name": "bedrock-defaults"}, + ) + for registered in (configured, defaulted): + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, registered) + + assert configured.streaming_buffer_until_moderated is False + assert configured.streaming_sampling_rate == 3 + assert configured.streaming_end_of_stream_only is True + assert defaulted.streaming_buffer_until_moderated is True + assert defaulted.streaming_sampling_rate == 5 + assert defaulted.streaming_end_of_stream_only is False + + +def test_initialize_bedrock_rejects_non_positive_sampling_rate(): + from pydantic import ValidationError + + from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock + + with pytest.raises(ValidationError): + initialize_bedrock( + _streaming_litellm_params(streaming_sampling_rate=0), + {"guardrail_name": "bedrock-bad-rate"}, + ) + + +def test_update_in_memory_litellm_params_round_trips_streaming_flags(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-update", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + ) + + guardrail.update_in_memory_litellm_params( + _streaming_litellm_params( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=7, + streaming_end_of_stream_only=True, + ) + ) + assert guardrail.streaming_buffer_until_moderated is False + assert guardrail.streaming_sampling_rate == 7 + assert guardrail.streaming_end_of_stream_only is True + + guardrail.update_in_memory_litellm_params(_streaming_litellm_params()) + assert guardrail.streaming_buffer_until_moderated is True + assert guardrail.streaming_sampling_rate == 5 + assert guardrail.streaming_end_of_stream_only is False + + +async def _run_streaming_hook_recording_order(guardrail: BedrockGuardrail) -> list: + events = [] + minimal = {"action": "NONE", "assessments": [], "outputs": []} + + async def record_scan(*args, **kwargs): + events.append("scan") + return minimal + + async def mock_stream(): + yield _chat_chunk("Hello", None) + yield _chat_chunk(" world", None) + yield _chat_chunk("", "stop") + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + ): + content = chunk.choices[0].delta.content if chunk.choices else None + events.append(("chunk", content)) + return events + + +@pytest.mark.asyncio +async def test_unbuffered_end_of_stream_hook_yields_chunks_before_scan(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-audit-mode", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + streaming_buffer_until_moderated=False, + streaming_end_of_stream_only=True, + ) + + events = await _run_streaming_hook_recording_order(guardrail) + + scan_index = events.index("scan") + chunk_events = [e for e in events if e != "scan"] + assert events.count("scan") == 1 + assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[: scan_index] + assert ("chunk", "Hello") in events[:scan_index] + assert ("chunk", " world") in events[:scan_index] + assert len(chunk_events) == 3 + + +@pytest.mark.asyncio +async def test_buffered_default_hook_scans_before_any_chunk(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-buffered-default", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + + events = await _run_streaming_hook_recording_order(guardrail) + + assert events[0] == "scan" + assert all(e == "scan" or e[0] == "chunk" for e in events) + assert len([e for e in events if e != "scan"]) >= 1 + + +@pytest.mark.asyncio +async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-mask-buffered", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + mask_response_content=True, + streaming_buffer_until_moderated=False, + streaming_end_of_stream_only=True, + ) + + assert guardrail._streams_incrementally() is False + events = await _run_streaming_hook_recording_order(guardrail) + assert events[0] == "scan" From 64eec53fd844684df2f953558c405c66eb27ce2d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:25:08 -0700 Subject: [PATCH 04/16] fix(guardrails): surface post-flush stream blocks as in-stream error frames and keep guardrail_information in spend logs A guardrail block or failed scan that fires after SSE chunks have been flushed can no longer set an HTTP status, so raising HTTPException there silently truncated the stream. _emit_streaming_http_error now routes post-flush failures through the endpoint translation's build_stream_error_items, emitting the surface-correct error frame on chat completions (data: {error}), /v1/messages (event: error), and /v1/responses (ErrorEvent with the next sequence number). Pre-flush blocks still raise with a real HTTP status. Successful flags-on scans also logged metadata.guardrail_information as null: the chat handler planted litellm_metadata on a route whose bucket is metadata, flipping the bucket for every later write, and responses streams fired their spend log before the eos scan ran. The chat handler now merges user_api_key metadata through get_or_create_metadata_bucket, and deferred stream-complete logging is armed for aresponses like it already was for anthropic_messages. --- .../chat/guardrail_translation/handler.py | 15 ++ .../guardrail_translation/base_translation.py | 48 ++++ .../chat/guardrail_translation/handler.py | 32 +-- .../guardrail_translation/handler.py | 35 +++ litellm/proxy/common_request_processing.py | 14 +- .../guardrail_hooks/bedrock_guardrails.py | 4 +- .../unified_guardrail/unified_guardrail.py | 67 ++++- litellm/responses/streaming_iterator.py | 21 +- .../openai/test_moderations.py | 32 +-- .../test_openai_moderation_streaming.py | 32 ++- .../test_bedrock_guardrails.py | 73 ++++++ .../test_unified_guardrail.py | 235 +++++++++++++++++- .../test_deferred_guardrail_logging.py | 69 +++++ .../proxy/test_common_request_processing.py | 16 +- 14 files changed, 608 insertions(+), 85 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b9ca18c7843..6cb2e568d0f 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -58,6 +58,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -165,6 +167,19 @@ class AnthropicMessagesHandler(BaseTranslation): return self._block_continuation_chunks(exc, responses_so_far or []) return self._standalone_block_chunks(exc) + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames + + message, _ = serialize_http_exception_detail(exc.detail) + return list(anthropic_sse_error_frames(message)) + def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: import uuid diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..b07a6b986d7 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,8 +1,11 @@ from abc import ABC, abstractmethod +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Final, Optional if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -73,6 +76,31 @@ class BaseTranslation(ABC): return transformed + @staticmethod + def merge_user_api_key_metadata_into_request( + request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place + user_api_key_dict: Optional["UserAPIKeyAuth"], + ) -> None: + """ + Add the prefixed ``user_api_key_*`` metadata to the request's resolved + metadata bucket without overwriting existing keys. + + Writes must go through ``get_or_create_metadata_bucket``: creating a + ``litellm_metadata`` key on a route whose bucket is ``metadata`` (chat + completions) flips the bucket for every later metadata write, and spend + logging never sees those writes (e.g. guardrail_information). + """ + from litellm.litellm_core_utils.core_helpers import ( + get_or_create_metadata_bucket, + ) + + user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if not user_metadata: + return + _, metadata_bucket = get_or_create_metadata_bucket(request_data) + for key, value in user_metadata.items(): + metadata_bucket.setdefault(key, value) + @abstractmethod async def process_input_messages( self, @@ -147,6 +175,26 @@ class BaseTranslation(ABC): """ return None + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + """ + Build the stream items that surface a guardrail HTTPException (a block + with the default exception-on-block config, or a failed scan) after the + response has already started streaming, in this endpoint's wire format. + + Called only once chunks have been sent: the HTTP status is gone, so the + failure must travel as an in-stream error frame. ``responses_so_far`` + holds the chunks the client has already received, for formats whose + error frame continues the stream (e.g. sequence numbers). + + Returns None when the format has no in-stream error frame; the caller + then re-raises ``exc``. Override in endpoint subclasses. + """ + return None + def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None: """ Convert request data to OpenAI-spec structured messages. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index e411dc497fc..e61fb719c98 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,6 +14,7 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Union, cast import litellm @@ -46,6 +47,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail @@ -381,11 +384,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "response" not in request_data: request_data["response"] = response - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -554,11 +553,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "responses" not in request_data: request_data["responses"] = responses_so_far - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -590,6 +585,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + import json + + from litellm.proxy.common_request_processing import sse_error_payload + + _, error_obj = sse_error_payload(exc) + return [f"data: {json.dumps({'error': error_obj})}\n\n".encode()] + @staticmethod def _accumulate_string_content_by_choice_index( responses_so_far: list["ModelResponseStream"], @@ -652,10 +659,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): request_data = {"responses": responses_so_far} elif "responses" not in request_data: request_data["responses"] = responses_so_far - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if responses_so_far and getattr(responses_so_far[0], "model", None): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..50d7f7452a8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -48,6 +48,8 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam, + ErrorEvent, + ErrorEventError, OpenAIMcpServerTool, ResponsesAPIStreamEvents, ) @@ -59,6 +61,8 @@ from litellm.types.responses.main import ( from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth @@ -80,6 +84,14 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: + sequence_numbers: Final = ( + item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) + for item in reversed(responses_so_far or []) + ) + return next((n + 1 for n in sequence_numbers if isinstance(n, int)), 0) + + class OpenAIResponsesHandler(BaseTranslation): """ Handler for processing OpenAI Responses API with guardrails. @@ -620,6 +632,29 @@ class OpenAIResponsesHandler(BaseTranslation): } return responses_so_far[-1].get("type") in terminal_types + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + + message, _ = serialize_http_exception_detail(exc.detail) + return [ + ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=_next_stream_sequence_number(responses_so_far), + error=ErrorEventError( + type="guardrail_error", + code=str(exc.status_code), + message=message, + param=None, + ), + ) + ] + def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: """ Get the string so far from the responses so far. diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ff6c8d1b1f8..6c7c77610b3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -502,7 +502,7 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch return logging_obj -def _serialize_http_exception_detail( +def serialize_http_exception_detail( detail: object, ) -> tuple[str, dict | None]: """ @@ -803,7 +803,7 @@ async def _buffer_first_chunk_honoring_disconnect( raise _ClientDisconnectedBeforeFirstChunk() -def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: +def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: """Build the ProxyException-shaped ``{"error": ...}`` body used in SSE error frames. Matches ``ProxyException.to_dict()`` so streaming and non-streaming error frames @@ -812,7 +812,7 @@ def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: # Preserve status code from HTTPException (e.g. guardrail blocks) error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start") - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) @@ -927,7 +927,7 @@ async def create_response( # Unexpected error consuming first chunk. verbose_proxy_logger.exception("Error consuming first chunk from generator: %s", e) - error_status, error_obj = _sse_error_payload(e) + error_status, error_obj = sse_error_payload(e) async def error_gen_message() -> AsyncGenerator[str, None]: for frame in _sse_error_frames(error_obj): @@ -1104,7 +1104,7 @@ async def open_sse_before_first_byte( # would never fire and the failure would go unaudited. The hook # also gets to sanitize what reaches the client, by returning or # raising a replacement, so its answer decides the frame. - _, error_obj = _sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) + _, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) for frame in _sse_error_frames(error_obj): yield frame.encode() return @@ -2374,7 +2374,7 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete elif ( _post_call_guardrails_active - and route_type == "anthropic_messages" + and route_type in ("anthropic_messages", "aresponses") and self._is_streaming_response(response) ): from litellm.litellm_core_utils.logging_worker import ( @@ -3245,7 +3245,7 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raw_detail: Final = _getattr_object(e, "detail", str(e)) - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(e, "provider_specific_fields", None) or {} if structured_fields: merged_fields: dict | None = {**existing_fields, **structured_fields} diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d7f222db35d..951cbc13290 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -42,7 +42,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_request_processing import _serialize_http_exception_detail +from litellm.proxy.common_request_processing import serialize_http_exception_detail from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, @@ -2760,7 +2760,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if not raw_sse or (not is_block and not headers_flushed): raise - block_message, _ = _serialize_http_exception_detail(block_detail) + block_message, _ = serialize_http_exception_detail(block_detail) for error_frame in anthropic_sse_error_frames( block_message if is_block else f"{block_exc.status_code}: {block_message}" ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e95e97bfe74..6647ac4c293 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -57,6 +57,9 @@ class _EndpointTranslation(Protocol): @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + @property + def build_stream_error_items(self) -> "Callable[..., Sequence[object] | None]": ... + def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: return translation @@ -408,14 +411,32 @@ class UnifiedLLMGuardrails(CustomLogger): call_type: str | None, responses_so_far: Sequence[object], request_data: dict, + endpoint_translation: _EndpointTranslation | None = None, + stream_started: bool = False, + responses_yielded: Sequence[object] | None = None, ) -> AsyncGenerator[object, None]: - """Surface a mid-stream HTTPException. For A2A call types the response has - already started, so emit an in-stream JSON-RPC error chunk; otherwise - re-raise so the proxy can report it. + """Surface a mid-stream HTTPException (a guardrail block with the default + exception-on-block config, or a failed scan). + + A2A call types emit an in-stream JSON-RPC error chunk. For other call + types, once chunks have already reached the client the HTTP status is + gone, so the failure is delegated to the endpoint translation's + ``build_stream_error_items`` and travels as an in-stream error frame in + that endpoint's wire format. Before the first chunk (or when the format + has no in-stream error frame) the exception is re-raised so the proxy + can report it with a real HTTP status. """ if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data)) return + if stream_started and endpoint_translation is not None: + error_items: Final = endpoint_translation.build_stream_error_items( + exc, responses_so_far=list(responses_yielded) if responses_yielded is not None else None + ) + if error_items is not None: + for error_item in error_items: + yield error_item + return raise exc def _build_transform_chunk( @@ -586,7 +607,15 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk raise _StreamTerminated() except HTTPException as e: - async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data): + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): yield error_item raise _StreamTerminated() @@ -1070,11 +1099,17 @@ class UnifiedLLMGuardrails(CustomLogger): return except HTTPException as e: # Response already started (we already yielded chunks); cannot send 400. - # For A2A, yield an in-stream JSON-RPC error so the client sees it. - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - return - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=chunks_yielded, + responses_yielded=responses_yielded, + ): + yield error_item + return chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1133,7 +1168,13 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk return except HTTPException as e: - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - else: - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): + yield error_item diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 368fd481e63..94c93ae7c6b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -422,15 +422,20 @@ class BaseResponsesAPIStreamingIterator: end_time: Final = datetime.now() if is_async: - asyncio.create_task( - self.logging_obj.dispatch_success_handlers( - logging_response, - start_time=self.start_time, - end_time=end_time, - cache_hit=self._completed_response_cache_hit, - prefer_async_handlers=True, - ) + logging_coroutine: Final = self.logging_obj.dispatch_success_handlers( + logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + prefer_async_handlers=True, ) + deferred_dispatch_armed: Final = getattr(self.logging_obj, "_on_deferred_stream_complete", None) is not None + if deferred_dispatch_armed: + # End-of-stream guardrail scans write guardrail_information after + # the terminal event; dispatching now would snapshot metadata early. + self.logging_obj._deferred_stream_complete_args = (logging_coroutine,) + else: + asyncio.create_task(logging_coroutine) else: run_async_function( async_function=self.logging_obj.async_success_handler, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 112bc5e6e49..2b43720a126 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -482,23 +482,25 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException when processing streaming harmful content - from fastapi import HTTPException + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - async def _drain(): - result_chunks = [] - async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - result_chunks.append(chunk) + result_chunks = [] + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + result_chunks.append(chunk) - with pytest.raises(HTTPException) as exc_info: - await _drain() - - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + frame = result_chunks[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 914af0e2368..476d443d8d8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -161,19 +161,27 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException - with pytest.raises(HTTPException) as exc_info: - async for ( - _ - ) in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - pass + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + collected = [] + async for ( + chunk + ) in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + collected.append(chunk) + + frame = collected[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index ca4b0a65e0b..235a5c0c09b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5517,3 +5517,76 @@ async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): assert guardrail._streams_incrementally() is False events = await _run_streaming_hook_recording_order(guardrail) assert events[0] == "scan" + + +@pytest.mark.asyncio +async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_truncating(): + """Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream + scan used to raise after SSE headers were flushed, so the client saw a + silently truncated stream. The unified hook must emit the chat in-stream + error frame instead.""" + from litellm.llms import load_guardrail_translation_mappings + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, + ) + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + streaming_end_of_stream_only=True, + streaming_buffer_until_moderated=False, + guardrail_name="bedrock-eos", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "actionReason": "Guardrail blocked.", + "outputs": [{"text": "Sorry, the model cannot answer this question."}], + "assessments": [ + {"topicPolicy": {"topics": [{"name": "Forbidden topic", "type": "DENY", "action": "BLOCKED"}]}} + ], + } + + def _chunk(content, finish_reason=None): + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta={"content": content, "role": "assistant"}, + finish_reason=finish_reason, + ) + ], + ) + + async def _mock_stream(): + yield _chunk("the forbidden ") + yield _chunk("topic answer", finish_reason="stop") + + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + try: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response) + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"), + response=_mock_stream(), + request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"}, + ): + out.append(item) + finally: + unified_module.endpoint_guardrail_translation_mappings = None + + assert len(out) == 3 + assert isinstance(out[0], ModelResponseStream) + frame = out[-1] + assert isinstance(frame, bytes) + payload = json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8b9ecfbbeee..0b32558a00a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -948,19 +948,24 @@ class TestStreamingTransform: assert streamed == "ABCDEFGHIJ" @pytest.mark.asyncio - async def test_incremental_diff_underflow_raises(self): + async def test_incremental_diff_underflow_emits_error_frame(self): """A transform shorter than what was already streamed cannot retract - bytes: it raises HTTPException(stream_transform_underflow).""" + bytes. Chunks have already been flushed by then, so the underflow + surfaces as the in-stream error frame, not an unraisable HTTPException.""" + import json as _json + # First sample emits "ABCDEF" (6 chars); second sample shrinks to 3. guardrail = _StreamingTextGuardrail(shrink_to="ABC", shrink_after=1) chunks = [_stream_chunk("abcdef"), _stream_chunk("ghij")] - with pytest.raises(unified_module.HTTPException) as exc_info: - await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) - assert exc_info.value.status_code == 400 - assert exc_info.value.detail["error"] == "stream_transform_underflow" + frame = out[-1] + assert isinstance(frame, bytes) + payload = _json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "stream_transform_underflow" + assert payload["error"]["code"] == "400" @pytest.mark.asyncio async def test_incremental_diff_final_chunk_preserves_finish_reason(self): @@ -1747,3 +1752,221 @@ class TestAppliedGuardrailsReflectsExecution: async def test_ordinary_guardrail_is_auto_marked_applied(self): data = await self._run(_AutoLoggingGuardrail()) assert "auto-logging" in _applied_guardrails(data) + + +class _EosHttpBlockingGuardrail(CustomGuardrail): + """Raises the bedrock-shaped block HTTPException at end-of-stream scan time.""" + + def __init__(self): + super().__init__(guardrail_name="eos-http-block") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + raise unified_module.HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "bedrock_guardrail_response": "BLOCKED_TOPIC", + }, + ) + + +def _anthropic_sse_event(event_type, data): + import json as _json + + return f"event: {event_type}\ndata: {_json.dumps(data)}\n\n".encode() + + +def _anthropic_message_chunks(texts): + head = [ + _anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + _anthropic_sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ] + deltas = [ + _anthropic_sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + for text in texts + ] + tail = [ + _anthropic_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + _anthropic_sse_event("message_stop", {"type": "message_stop"}), + ] + return head + deltas + tail + + +class TestStreamingHttpErrorFrames: + """A post-flush end-of-stream guardrail block (HTTPException) must surface as + the endpoint's in-stream error frame instead of an unhandled raise that + silently truncates the SSE stream (PR #38722 defect 1).""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_block_emits_data_error_frame(self): + import json as _json + + guardrail = _EosHttpBlockingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out[:2] == chunks + frame = out[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + + @pytest.mark.asyncio + async def test_messages_eos_block_emits_anthropic_error_event(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages" + ) + + raw = b"".join(c for c in out if isinstance(c, bytes)).decode() + assert "hello " in raw + assert "event: error" in raw + assert "Violated guardrail policy" in raw + assert "guardrail_error" in raw + + @pytest.mark.asyncio + async def test_responses_eos_block_emits_error_event_with_next_sequence(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = [ + {"type": "response.created", "sequence_number": 0}, + {"type": "response.output_text.delta", "sequence_number": 1, "delta": "hello"}, + { + "type": "response.completed", + "sequence_number": 2, + "response": { + "model": "gpt-4", + "output": [{"type": "message", "content": [{"type": "output_text", "text": "hello"}]}], + }, + }, + ] + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses" + ) + + assert chunks[0] in out and chunks[1] in out + assert chunks[2] not in out + error_event = out[-1] + assert error_event.type == "error" + assert error_event.sequence_number == 2 + assert error_event.error.message == "Violated guardrail policy" + assert error_event.error.code == "400" + assert error_event.error.type == "guardrail_error" + + @pytest.mark.asyncio + async def test_pre_flush_block_still_raises_http_exception(self): + guardrail = _EosHttpBlockingGuardrail() + guardrail.streaming_buffer_until_moderated = True + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + with pytest.raises(unified_module.HTTPException) as exc_info: + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Violated guardrail policy" + + +class _AuditRecordingGuardrail(CustomGuardrail): + """Successful scan that records guardrail_information, like a flags-on audit.""" + + def __init__(self): + super().__init__(guardrail_name="audit-recorder") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"action": "NONE"}, + request_data=request_data, + guardrail_status="success", + ) + return inputs + + +class TestStreamingGuardrailInformationBucket: + """guardrail_information written during a chat streaming end-of-stream scan + must land in the request's ``metadata`` bucket that spend logging snapshots. + Regression for PR #38722 defect 2: the chat handler used to plant a + ``litellm_metadata`` key first, flipping the bucket so every later + guardrail_information write was diverted and /spend/logs showed null.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_scan_writes_guardrail_information_to_metadata(self): + guardrail = _AuditRecordingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + async def _mock_stream(): + for chunk in chunks: + yield chunk + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", user_id="user-1", request_route="/v1/chat/completions" + ) + request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}} + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_mock_stream(), + request_data=request_data, + ): + out.append(item) + + assert "litellm_metadata" not in request_data + recorded = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(recorded) == 1 + assert recorded[0]["guardrail_name"] == "audit-recorder" + assert recorded[0]["guardrail_status"] == "success" + assert request_data["metadata"]["user_api_key_user_id"] == "user-1" diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index e70fc61de30..b7317d33b84 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1228,3 +1228,72 @@ class TestFireDeferredStreamLogging: assert info is not None, "guardrail_information should be populated" assert len(info) == 1 assert info[0]["guardrail_name"] == "info-writer" + + +class TestResponsesIteratorDeferredLogging: + """Regression for PR #38722 defect 2 on /v1/responses streams: when the + proxy arms _on_deferred_stream_complete, the responses streaming iterator + must store the logging coroutine for ProxyLogging._fire_deferred_stream_logging + (which runs AFTER end-of-stream guardrail scans write guardrail_information) + instead of dispatching immediately with a premature metadata snapshot.""" + + def _iterator(self, logging_obj): + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + iterator = object.__new__(BaseResponsesAPIStreamingIterator) + iterator.logging_obj = logging_obj + iterator.start_time = None + iterator.completed_response = None + iterator._completed_response_logged = False + iterator._completed_response_cache_hit = None + iterator._persist_completed_response_before_logging = False + return iterator + + def _logging_obj(self): + recorded = {} + + async def dispatch_success_handlers(result=None, **kwargs): + recorded["dispatched"] = True + + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = dispatch_success_handlers + return logging_obj, recorded + + @pytest.mark.asyncio + async def test_armed_iterator_stores_deferred_coroutine(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = MagicMock() + iterator = self._iterator(logging_obj) + + with patch("asyncio.create_task") as mock_create_task: + iterator._log_completed_response(is_async=True) + + mock_create_task.assert_not_called() + args = logging_obj._deferred_stream_complete_args + assert isinstance(args, tuple) and len(args) == 1 + assert "dispatched" not in recorded + await args[0] + assert recorded["dispatched"] is True + + @pytest.mark.asyncio + async def test_unarmed_iterator_dispatches_immediately(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = None + iterator = self._iterator(logging_obj) + + created = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + iterator._log_completed_response(is_async=True) + + assert len(created) == 1 + await created[0] + assert recorded["dispatched"] is True diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 71d4666416d..99f03d23378 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1656,32 +1656,32 @@ class TestCommonRequestProcessingHelpers: assert payload["error"]["message"] == "MCP request blocked: no rewritable argument field present" assert payload["error"]["provider_specific_fields"]["error"]["code"] == "panw_prisma_airs_blocked" - async def test_serialize_http_exception_detail_helper(self): + async def testserialize_http_exception_detail_helper(self): """Direct unit coverage for the L1 helper across all branches.""" from litellm.proxy.common_request_processing import ( - _serialize_http_exception_detail, + serialize_http_exception_detail, ) import json as _json - assert _serialize_http_exception_detail("plain") == ("plain", None) + assert serialize_http_exception_detail("plain") == ("plain", None) - msg, fields = _serialize_http_exception_detail({"error": "Violated", "extra": "x"}) + msg, fields = serialize_http_exception_detail({"error": "Violated", "extra": "x"}) assert msg == "Violated" assert fields == {"error": "Violated", "extra": "x"} - msg, fields = _serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) + msg, fields = serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) assert msg == "blocked" assert fields == {"error": {"message": "blocked", "code": "x"}} - msg, fields = _serialize_http_exception_detail({"message": "top-level"}) + msg, fields = serialize_http_exception_detail({"message": "top-level"}) assert msg == "top-level" assert fields == {"message": "top-level"} - msg, fields = _serialize_http_exception_detail({"weird": ["a", "b"]}) + msg, fields = serialize_http_exception_detail({"weird": ["a", "b"]}) assert msg == _json.dumps({"weird": ["a", "b"]}) assert fields == {"weird": ["a", "b"]} - assert _serialize_http_exception_detail(42) == ("42", None) + assert serialize_http_exception_detail(42) == ("42", None) async def test_create_streaming_response_first_chunk_error_string_code(self): """ From f60ccf623460fcf029f3c22179d867eaf5fdd99a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 02:23:12 -0700 Subject: [PATCH 05/16] fix(guardrails): match deferred stream dispatch shape per stream owner and defer passthrough logging until guardrail eos --- litellm/proxy/common_request_processing.py | 132 +++++++++++------ .../streaming_handler.py | 41 ++++-- .../test_deferred_guardrail_logging.py | 123 ++++++++++++++++ .../test_streaming_handler_interrupt.py | 133 ++++++++++++++++-- .../proxy/test_common_request_processing.py | 2 +- .../test_proxy_logging_hook_detection.py | 29 ++-- 6 files changed, 381 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6c7c77610b3..3fca267e321 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2341,53 +2341,14 @@ class ProxyBaseLLMRequestProcessing: if requested_model_from_client: self.data["_litellm_client_requested_model"] = requested_model_from_client - # Streaming: attach a closure that fires after all guardrail - # end-of-stream blocks complete. CSW.__anext__ stores the - # assembled response on logging_obj; the outer consumer - # (ProxyLogging._fire_deferred_stream_logging) fires the - # closure after the full streaming pipeline finishes. - # The closure runs non-apply_guardrail hooks on the - # assembled response, then fires success logging. - # Only for CustomStreamWrapper — raw async generators from - # passthrough routes bypass CSW and would orphan the closure. - from litellm.litellm_core_utils.streaming_handler import ( - CustomStreamWrapper, - ) - - if _post_call_guardrails_active and isinstance(response, CustomStreamWrapper): - # Intentionally a live reference (not a copy) — mirrors - # ProxyLogging.post_call_success_hook which also mutates - # data["guardrail_to_apply"] during iteration. - _captured_data: Final = self.data - _captured_user_api_key_dict: Final = user_api_key_dict - _captured_logging_obj: Final = logging_obj - - async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: - await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( - captured_data=_captured_data, - captured_user_api_key_dict=_captured_user_api_key_dict, - captured_logging_obj=_captured_logging_obj, - assembled_response=assembled_response, - cache_hit=cache_hit, - ) - - logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete - elif ( - _post_call_guardrails_active - and route_type in ("anthropic_messages", "aresponses") - and self._is_streaming_response(response) - ): - from litellm.litellm_core_utils.logging_worker import ( - GLOBAL_LOGGING_WORKER, + if _post_call_guardrails_active: + self._arm_deferred_stream_dispatch( + response=response, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, ) - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) - - logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete - if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): @@ -3057,6 +3018,87 @@ class ProxyBaseLLMRequestProcessing: except Exception as e: verbose_proxy_logger.exception("Error firing deferred logging: %s", e) + def _arm_deferred_stream_dispatch( + self, + response: object, + route_type: str, + user_api_key_dict: "UserAPIKeyAuth", + logging_obj: LiteLLMLoggingObj, + ) -> None: + """ + Streaming with post-call guardrails active: attach a closure that + ProxyLogging._fire_deferred_stream_logging fires after all guardrail + end-of-stream blocks complete, so the spend log sees + guardrail_information. + + Three closure shapes, matching who owns logging for the stream: + - CustomStreamWrapper (chat completions) stores + (assembled_response, cache_hit); the closure also runs + non-apply_guardrail post-call hooks via + _run_deferred_stream_guardrails. + - Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares + its inner CustomStreamWrapper's logging_obj, so it stores the same + (assembled_response, cache_hit) shape; the closure only dispatches + success logging, matching the route's pre-existing hook surface. + - Native anthropic_messages/aresponses iterators store a single + ready-made logging coroutine to enqueue. + + Raw async generators from passthrough routes bypass all three and + would orphan the closure, so they are not armed here. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + if isinstance(response, CustomStreamWrapper): + # Intentionally a live reference (not a copy) — mirrors + # ProxyLogging.post_call_success_hook which also mutates + # data["guardrail_to_apply"] during iteration. + _captured_data: Final = self.data + _captured_user_api_key_dict: Final = user_api_key_dict + _captured_logging_obj: Final = logging_obj + + async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=_captured_data, + captured_user_api_key_dict=_captured_user_api_key_dict, + captured_logging_obj=_captured_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + return + + if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): + return + + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + if isinstance(response, LiteLLMCompletionStreamingIterator): + _captured_bridge_logging_obj: Final = logging_obj + + async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: + await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + + logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete + return + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + async def _on_deferred_native_stream_complete( + logging_coroutine: Coroutine[object, object, object], + ) -> None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + + logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete + @staticmethod async def _run_deferred_stream_guardrails( captured_data: dict, diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 5ad41b00890..022a1ecbac4 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -65,6 +65,19 @@ class PassThroughStreamingHandler: route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler ) raw_bytes: Final[list[bytes]] = [] + + def _build_logging_coroutine() -> Coroutine[None, None, None]: + return resolved_route_streaming_logging( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body or {}, + endpoint_type=endpoint_type, + start_time=start_time, + raw_bytes=raw_bytes, + end_time=datetime.now(), + ) + logging_scheduled = False model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection( request_body=request_body, @@ -114,6 +127,21 @@ class PassThroughStreamingHandler: ) if pending: yield pending + # Stream completed cleanly. When the proxy armed deferred + # dispatch (post-call guardrails active), park the logging + # coroutine on logging_obj instead of enqueueing now, so + # ProxyLogging._fire_deferred_stream_logging fires it after + # guardrail end-of-stream blocks populate guardrail_information. + # Disconnect/exception paths skip this and fall through to the + # immediate enqueue in ``finally`` to keep partial billing + # (LIT-2642). + if ( + getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None + and raw_bytes + and response.status_code < 400 + ): + logging_scheduled = True + litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),) except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) raise @@ -128,18 +156,7 @@ class PassThroughStreamingHandler: if not logging_scheduled and raw_bytes and response.status_code < 400: logging_scheduled = True try: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=resolved_route_streaming_logging( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body or {}, - endpoint_type=endpoint_type, - start_time=start_time, - raw_bytes=raw_bytes, - end_time=datetime.now(), - ) - ) + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index b7317d33b84..99fae277bdd 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1297,3 +1297,126 @@ class TestResponsesIteratorDeferredLogging: assert len(created) == 1 await created[0] assert recorded["dispatched"] is True + + +class TestArmDeferredStreamDispatch: + """Regression for PR #38722: the closure shape armed on logging_obj must + match the args the stream's logging owner stores. Bridged /v1/responses + (LiteLLMCompletionStreamingIterator) shares its inner CustomStreamWrapper's + logging_obj, which stores (assembled_response, cache_hit); arming the + single-coroutine native closure there made _fire_deferred_stream_logging + raise TypeError inside the streaming hook, leaking an in-stream 500 error + frame on every streamed /v1/responses request.""" + + def _processor(self): + return ProxyBaseLLMRequestProcessing(data={"model": "gpt-test"}) + + def _dispatch_recording_logging_obj(self): + recorded = {} + + async def dispatch_success_handlers( + result=None, start_time=None, end_time=None, cache_hit=None, prefer_async_handlers=False + ): + recorded["result"] = result + recorded["cache_hit"] = cache_hit + recorded["prefer_async_handlers"] = prefer_async_handlers + + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = dispatch_success_handlers + logging_obj._on_deferred_stream_complete = None + logging_obj._deferred_stream_complete_args = None + return logging_obj, recorded + + @pytest.mark.asyncio + async def test_bridged_responses_iterator_gets_csw_arg_shape(self): + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + logging_obj, recorded = self._dispatch_recording_logging_obj() + bridged = object.__new__(LiteLLMCompletionStreamingIterator) + + self._processor()._arm_deferred_stream_dispatch( + response=bridged, + route_type="aresponses", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + async def test_native_stream_closure_enqueues_single_coroutine(self): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, _ = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + closure = logging_obj._on_deferred_stream_complete + assert closure is not None + + async def _logging_coroutine(): + return None + + coro = _logging_coroutine() + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue: + await closure(coro) + mock_enqueue.assert_called_once_with(async_coroutine=coro) + coro.close() + + @pytest.mark.asyncio + async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + logging_obj, recorded = self._dispatch_recording_logging_obj() + csw = object.__new__(CustomStreamWrapper) + processor = self._processor() + + monkeypatch.setattr( # test-quality-ok: empty the process-global callback registry so no ambient guardrail runs + litellm, "callbacks", [] + ) + processor._arm_deferred_stream_dispatch( + response=csw, + route_type="acompletion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + assembled = object() + await logging_obj._on_deferred_stream_complete(assembled, False) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + + def test_non_native_route_generator_not_armed(self): + logging_obj, _ = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="acompletion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assert logging_obj._on_deferred_stream_complete is None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 1d82a5dfc6e..56c89fed79a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -28,12 +28,21 @@ def _make_streaming_response(chunks): return mock +def _unarmed_logging_obj(): + """Real Logging objects only carry _on_deferred_stream_complete when the + proxy arms deferred dispatch; a bare MagicMock's auto-attribute is truthy + and would spuriously trigger the deferral branch.""" + obj = MagicMock() + obj._on_deferred_stream_complete = None + return obj + + @pytest.mark.asyncio async def test_chunk_processor_logs_on_normal_completion(): chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] response = _make_streaming_response(chunks) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -66,7 +75,7 @@ async def test_chunk_processor_logs_on_client_disconnect(): chunks = [b"event-1", b"event-2", b"event-3"] response = _make_streaming_response(chunks) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -104,7 +113,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er response = _make_streaming_response(chunks) response.status_code = 403 - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -134,7 +143,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er async def test_chunk_processor_does_not_schedule_logging_when_no_chunks(): response = _make_streaming_response([]) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -189,7 +198,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker(): async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "claude-3-haiku"}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=EndpointType.GENERIC, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -230,7 +239,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne gen = PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "claude-3-haiku"}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=EndpointType.GENERIC, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -246,7 +255,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne def _logging_obj_with_write_once_cst(): """Build a MagicMock that mirrors the real Logging behavior: _update_completion_start_time latches self.completion_start_time so the write-once guard actually latches.""" - obj = MagicMock() + obj = _unarmed_logging_obj() obj.completion_start_time = None def _update(*, completion_start_time): @@ -301,7 +310,7 @@ async def test_chunk_processor_does_not_reset_completion_start_time_on_later_chu response = _make_streaming_response(chunks) real_first = datetime(2020, 1, 1, 0, 0, 0) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() # Simulate first-chunk stamp having already landed (e.g. under contention or a # prior wrapper that already set it): later chunks must be no-ops. mock_logging_obj.completion_start_time = real_first @@ -387,7 +396,7 @@ async def _collect_openai_passthrough_chunks(chunks, endpoint_type): async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "gpt-4o-mini", "stream": True}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -517,3 +526,109 @@ def test_convert_raw_bytes_survives_truncated_multibyte_sequence(): lines = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) assert any('"type": "message_delta"' in line for line in lines) + + +@pytest.mark.asyncio +async def test_chunk_processor_defers_logging_until_fire_when_armed(): + """Regression for PR #38722: native /v1/messages streams route through + chunk_processor, which enqueued the spend log the moment the stream ended, + racing the guardrail end-of-stream scan and logging + guardrail_information as null. With deferred dispatch armed, the completed + stream must park the logging coroutine on logging_obj and only enqueue it + when ProxyLogging._fire_deferred_stream_logging fires after the scan.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + chunks = [b"event-1", b"event-2"] + response = _make_streaming_response(chunks) + + logging_obj = _unarmed_logging_obj() + logging_obj._deferred_stream_complete_args = None + + enqueued = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + async_coroutine.close() + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, + "ensure_initialized_and_enqueue", + side_effect=_capture, + ) as mock_enqueue: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=gen, + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + received = [] + async for chunk in gen: + received.append(chunk) + await asyncio.sleep(0) + + assert received == chunks + mock_enqueue.assert_not_called() + parked = logging_obj._deferred_stream_complete_args + assert isinstance(parked, tuple) and len(parked) == 1 + assert asyncio.iscoroutine(parked[0]) + + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_called_once() + + +@pytest.mark.asyncio +async def test_chunk_processor_enqueues_immediately_on_disconnect_even_when_armed(): + """Client disconnects never reach _fire_deferred_stream_logging, so parking + the coroutine there would lose the partial-usage spend log (LIT-2642); the + disconnect path must keep enqueueing immediately.""" + chunks = [b"event-1", b"event-2", b"event-3"] + response = _make_streaming_response(chunks) + + logging_obj = _unarmed_logging_obj() + + async def _armed_closure(logging_coroutine): + raise AssertionError("deferred closure must not fire on disconnect") + + logging_obj._on_deferred_stream_complete = _armed_closure + logging_obj._deferred_stream_complete_args = None + + enqueued = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + async_coroutine.close() + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, + "ensure_initialized_and_enqueue", + side_effect=_capture, + ) as mock_enqueue: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ) + await gen.__anext__() + await gen.aclose() + + mock_enqueue.assert_called_once() + assert logging_obj._deferred_stream_complete_args is None diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 99f03d23378..0fea239625a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1656,7 +1656,7 @@ class TestCommonRequestProcessingHelpers: assert payload["error"]["message"] == "MCP request blocked: no rewritable argument field present" assert payload["error"]["provider_specific_fields"]["error"]["code"] == "panw_prisma_airs_blocked" - async def testserialize_http_exception_detail_helper(self): + async def test_serialize_http_exception_detail_helper(self): """Direct unit coverage for the L1 helper across all branches.""" from litellm.proxy.common_request_processing import ( serialize_http_exception_detail, diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 542572e1e56..9f1321aec2c 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -346,14 +346,14 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions @pytest.mark.asyncio -async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch): +async def test_unified_guardrail_iterator_accepts_explicit_guardrail(): """ The dispatch passes each guardrail explicitly instead of through a shared request_data key, so chaining two unified-routed guardrails cannot drop - all but the last one. + all but the last one. The block fires after the deltas were already + flushed to the client, so it surfaces as a trailing in-stream error frame + rather than a raised HTTPException. """ - from fastapi import HTTPException - from litellm.proxy.utils import unified_guardrail guardrail = _content_filter_guardrail("BLOCK") @@ -367,14 +367,19 @@ async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch for chunk in _anthropic_stream_chunks(["the", " zebra runs"]): yield chunk - with pytest.raises(HTTPException): - async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), - response=fake_stream(), - request_data=request_data, - guardrail_to_apply=guardrail, - ): - pass + delivered = [] + async for item in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + response=fake_stream(), + request_data=request_data, + guardrail_to_apply=guardrail, + ): + delivered.append(item) + + raw = b"".join(c for c in delivered if isinstance(c, bytes)).decode() + assert "event: error" in raw + assert "guardrail_error" in raw + assert raw.index("guardrail_error") > raw.index(" zebra runs") @pytest.mark.asyncio From 229970c5005fefd540d87c292750fe6c3c7ea6a4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 02:30:34 -0700 Subject: [PATCH 06/16] fix(guardrails): unwrap HiddenParamsAsyncIteratorWrapper before deferred dispatch class sniffing --- litellm/proxy/common_request_processing.py | 11 ++++-- .../test_deferred_guardrail_logging.py | 34 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3fca267e321..69c2cb3f0f0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3045,10 +3045,17 @@ class ProxyBaseLLMRequestProcessing: Raw async generators from passthrough routes bypass all three and would orphan the closure, so they are not armed here. + + The router wraps iterators that cannot carry _hidden_params in + HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the + unwrapped inner iterator. """ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper - if isinstance(response, CustomStreamWrapper): + unwrapped: Final = response._inner if isinstance(response, HiddenParamsAsyncIteratorWrapper) else response + + if isinstance(unwrapped, CustomStreamWrapper): # Intentionally a live reference (not a copy) — mirrors # ProxyLogging.post_call_success_hook which also mutates # data["guardrail_to_apply"] during iteration. @@ -3075,7 +3082,7 @@ class ProxyBaseLLMRequestProcessing: LiteLLMCompletionStreamingIterator, ) - if isinstance(response, LiteLLMCompletionStreamingIterator): + if isinstance(unwrapped, LiteLLMCompletionStreamingIterator): _captured_bridge_logging_obj: Final = logging_obj async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 99fae277bdd..8fde4cc9d5e 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1352,6 +1352,40 @@ class TestArmDeferredStreamDispatch: assert recorded["cache_hit"] is False assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio + async def test_router_wrapped_bridged_iterator_gets_csw_arg_shape(self): + """The router wraps iterators without _hidden_params in + HiddenParamsAsyncIteratorWrapper before the proxy arms deferral, so + every production streamed /v1/responses reaches arming wrapped; + sniffing the wrapper instead of the inner iterator armed the 1-arg + native closure against the CSW's 2-arg stored shape and leaked a + TypeError 500 frame into the stream.""" + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.router_utils.add_retry_fallback_headers import ( + HiddenParamsAsyncIteratorWrapper, + ) + + logging_obj, recorded = self._dispatch_recording_logging_obj() + wrapped = HiddenParamsAsyncIteratorWrapper(object.__new__(LiteLLMCompletionStreamingIterator)) + + self._processor()._arm_deferred_stream_dispatch( + response=wrapped, + route_type="aresponses", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio async def test_native_stream_closure_enqueues_single_coroutine(self): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER From 9c653be72cd87509385149e10b907df833de4840 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 06:42:59 -0700 Subject: [PATCH 07/16] chore(budgets): drop merge headroom above staging ceilings --- basedpyright-code-budget.json | 20 ++++++++++---------- ruff-strict-budget.json | 32 ++++++++++++++++---------------- type-discipline-budget.json | 6 +++--- 3 files changed, 29 insertions(+), 29 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0a71a81693d..7b114dcbb9a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -6,19 +6,19 @@ "limit": 2384 }, "reportAssignmentType": { - "limit": 320 + "limit": 319 }, "reportAttributeAccessIssue": { - "limit": 488 + "limit": 480 }, "reportCallIssue": { - "limit": 114 + "limit": 112 }, "reportConstantRedefinition": { "limit": 40 }, "reportDeprecated": { - "limit": 213 + "limit": 212 }, "reportDuplicateImport": { "limit": 19 @@ -30,7 +30,7 @@ "limit": 7 }, "reportGeneralTypeIssues": { - "limit": 154 + "limit": 101 }, "reportIncompatibleMethodOverride": { "limit": 56 @@ -45,7 +45,7 @@ "limit": 35 }, "reportInvalidTypeForm": { - "limit": 35 + "limit": 34 }, "reportInvalidTypeVarUse": { "limit": 2 @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1823 + "limit": 1808 }, "reportRedeclaration": { "limit": 8 @@ -93,7 +93,7 @@ "limit": 197 }, "reportTypedDictNotRequiredAccess": { - "limit": 26 + "limit": 25 }, "reportUndefinedVariable": { "limit": 0 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 834 + "limit": 829 }, "reportUntypedBaseClass": { "limit": 0 @@ -135,7 +135,7 @@ "limit": 21 }, "reportUnusedFunction": { - "limit": 139 + "limit": 138 }, "reportUnusedImport": { "limit": 543 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 61bef07fcb3..3e8a97b67a9 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,10 +9,10 @@ "limit": 818 }, "ANN201": { - "limit": 2005 + "limit": 2003 }, "ANN202": { - "limit": 848 + "limit": 845 }, "ANN204": { "limit": 700 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 733 + "limit": 655 }, "ASYNC230": { "limit": 11 @@ -33,13 +33,13 @@ "limit": 2 }, "B006": { - "limit": 177 + "limit": 176 }, "B008": { "limit": 503 }, "B009": { - "limit": 59 + "limit": 52 }, "B010": { "limit": 190 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2923 + "limit": 2917 }, "C401": { "limit": 8 @@ -78,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 312 + "limit": 311 }, "D419": { "limit": 6 @@ -108,7 +108,7 @@ "limit": 3 }, "F401": { - "limit": 17 + "limit": 13 }, "LOG015": { "limit": 5 @@ -117,7 +117,7 @@ "limit": 1 }, "PERF102": { - "limit": 27 + "limit": 23 }, "PERF401": { "limit": 12 @@ -171,28 +171,28 @@ "limit": 3 }, "RET504": { - "limit": 176 + "limit": 175 }, "RUF012": { - "limit": 241 + "limit": 239 }, "RUF015": { "limit": 8 }, "RUF019": { - "limit": 38 + "limit": 32 }, "RUF046": { "limit": 4 }, "RUF059": { - "limit": 67 + "limit": 66 }, "RUF100": { "limit": 0 }, "S110": { - "limit": 218 + "limit": 217 }, "S112": { "limit": 22 @@ -201,7 +201,7 @@ "limit": 57 }, "SIM102": { - "limit": 317 + "limit": 315 }, "SIM103": { "limit": 119 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1165 + "limit": 1117 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 70376a56b9a..5c29d527709 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -21,16 +21,16 @@ "limit": 0 }, "LIT008": { - "limit": 950 + "limit": 945 }, "LIT009": { "limit": 0 }, "LIT010": { - "limit": 16598 + "limit": 16564 }, "LIT011": { - "limit": 5579 + "limit": 5577 }, "LIT012": { "limit": 4502 From c4663a4ae96be8977fcb2f8d8d335418cd9252b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 06:51:54 -0700 Subject: [PATCH 08/16] chore(budgets): drop stale PLW0133 entry graduated to ruff.toml --- ruff-strict-budget.json | 3 --- 1 file changed, 3 deletions(-) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 3e8a97b67a9..92148dddd76 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -152,9 +152,6 @@ "PLW0127": { "limit": 57 }, - "PLW0133": { - "limit": 1 - }, "PLW0602": { "limit": 215 }, From c09f76a079835e11500c09ce9cbc47d24c41e329 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 07:15:00 -0700 Subject: [PATCH 09/16] refactor(mcp): drop always-true isinstance guard on typed token-endpoint response --- .../outbound_credentials/client_credentials.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index da00abfe604..ad18d1bb10f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -19,9 +19,9 @@ Implements the client-credentials behavior contract for the v2 resolver: identity. The token-endpoint POST is injected (``M2MTokenEndpointPost``) so the grant orchestration is -testable without a live IdP; ``post_client_credentials_grant`` is the httpx edge and the one -place the untyped response boundary is contained. Failures are values: the source returns -``Result[OAuthToken, CredError]``; only the httpx edge touches exceptions. +testable without a live IdP; ``post_client_credentials_grant`` is the httpx edge. Failures are +values: the source returns ``Result[OAuthToken, CredError]``; only the httpx edge touches +exceptions. """ from __future__ import annotations @@ -95,18 +95,17 @@ async def post_client_credentials_grant( ) -> TokenEndpointOutcome: """POST the grant to the token endpoint and classify the transport outcome. - The httpx edge: litellm's handler is partially typed (and raises ``HTTPStatusError`` itself on - a 4xx/5xx), so the untyped boundary is contained here and every field the caller reads comes - out of a validated ``TokenEndpointOutcome``. + The httpx edge: litellm's handler raises ``HTTPStatusError`` itself on a 4xx/5xx, and every + field the caller reads comes out of a validated ``TokenEndpointOutcome``. """ from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 # defer heavy handler import to call time - get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler is partially typed + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler factory params are coarsely typed ) from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 # deferred with the handler import try: client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) - response = await client.post( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # handler is partially typed + response: Final = await client.post( # pyright: ignore[reportUnknownMemberType] # handler params are coarsely typed url, headers={"Accept": "application/json", **headers}, data=form ) except httpx.HTTPStatusError as status_err: @@ -114,8 +113,6 @@ async def post_client_credentials_grant( return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}") except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable return TokenEndpointUnreachable(detail=str(exc)) - if not isinstance(response, httpx.Response): - return TokenEndpointUnreachable(detail="token endpoint returned no response") try: body: Final = _TOKEN_BODY_ADAPTER.validate_json(response.content) except ValidationError: From 80e0bc2dc6861222781f3dc87aa5ab72786dae8f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 10:58:10 -0700 Subject: [PATCH 10/16] docs(claude.md): require tests to check behavior, not code structure Claude-Session: https://claude.ai/code/session_017nZW6omb93ZuAfqqCzKSU5 --- CLAUDE.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 930825aeb89..d9e9e8f1586 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -23,6 +23,8 @@ When adding new features, add meaningful tests. Don't add tests that don't check Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression) +Never test structure of code only function of it + `tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` From d44d281d1d873fe3bf813e931bed27a8ac3ae7ee Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 29 Aug 2026 18:11:58 -0700 Subject: [PATCH 11/16] fix(proxy): emit timing headers and overhead for /v1/messages and /v1/responses (#38840) --- litellm/litellm_core_utils/litellm_logging.py | 14 ++ .../llm_response_utils/response_metadata.py | 98 ++++++----- litellm/llms/custom_httpx/llm_http_handler.py | 2 + litellm/proxy/common_request_processing.py | 28 +++- litellm/responses/streaming_iterator.py | 2 + .../test_response_metadata.py | 154 ++++++++++++++++++ .../test_litellm_logging.py | 107 ++++++++++++ .../custom_httpx/test_llm_http_handler.py | 79 +++++++++ .../proxy/test_common_request_processing.py | 146 +++++++++++++++++ .../responses/test_streaming_iterator.py | 21 +++ 10 files changed, 609 insertions(+), 42 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c27822d0479..97c4d038734 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -547,6 +547,9 @@ class Logging(LiteLLMLoggingBaseClass): # Init Caching related details self.caching_details: CachingDetails | None = None + # Timing for results that cannot carry ``_hidden_params`` (plain-dict /v1/messages + # responses and the bridge stream wrappers); see ``update_response_metadata``. + self.response_timing_metrics: Mapping[str, float] = {} # mutable-ok: kept deep-copyable # Passthrough endpoint guardrails config for field targeting self.passthrough_guardrails_config: dict[str, Any] | None = None @@ -566,6 +569,10 @@ class Logging(LiteLLMLoggingBaseClass): self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None + def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: + """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" + self.response_timing_metrics = dict(timing_metrics) # mutable-ok: kept deep-copyable + def process_dynamic_callbacks(self): """ Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -6005,6 +6012,13 @@ def get_standard_logging_object_payload( clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params) if clean_hidden_params["response_cost"] is None and raw_response_cost is not None: clean_hidden_params["response_cost"] = llm_response_cost + if clean_hidden_params["litellm_overhead_time_ms"] is None and status == "success": + # /v1/messages dict results and the bridge stream wrappers keep it on the logging object; + # failure payloads stay None like every response type that carries its own _hidden_params + timing_metrics: Final = ( + getattr(logging_obj, "response_timing_metrics", None) or {} # mutable-ok: empty fallback + ) + clean_hidden_params["litellm_overhead_time_ms"] = timing_metrics.get("litellm_overhead_time_ms") model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information( base_model=base_model, diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index a375560288f..ac8cf438a9b 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,4 +1,5 @@ import datetime +from collections.abc import Mapping from typing import Any, Final from litellm.constants import LITELLM_DETAILED_TIMING @@ -13,6 +14,39 @@ from litellm.types.utils import ( ) +def response_timing_metrics( + start_time: datetime.datetime, + end_time: datetime.datetime, + logging_obj: LiteLLMLoggingObject, + include_overhead: bool = True, +) -> Mapping[str, float]: + """``_response_ms`` for the whole call, plus ``litellm_overhead_time_ms`` when it can be derived. + + On a cache hit the overhead is the total minus the cache read; otherwise it is the total minus + the provider call (``llm_api_duration_ms``). It is omitted when neither duration was recorded, + and when ``include_overhead`` is False because the two durations cover different windows. + """ + total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000 + if not include_overhead: + return {"_response_ms": total_response_time_ms} # mutable-ok: read-only timing result + caching_details: Final = logging_obj.caching_details + cache_duration_ms: Final = ( + caching_details.get("cache_duration_ms") + if caching_details is not None and caching_details.get("cache_hit") is True + else None + ) + llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") + if cache_duration_ms is not None: + overhead_ms: float | None = total_response_time_ms - cache_duration_ms + elif llm_api_duration_ms is not None: + overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) + else: + overhead_ms = None + if overhead_ms is None: + return {"_response_ms": total_response_time_ms} + return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms} + + class ResponseMetadata: """ Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses @@ -52,7 +86,7 @@ class ResponseMetadata: } self._update_hidden_params(new_params) - def _update_hidden_params(self, new_params: dict) -> None: + def _update_hidden_params(self, new_params: Mapping[str, object]) -> None: """ Update hidden params - handles when self._hidden_params is a dict or HiddenParams object """ @@ -76,37 +110,24 @@ class ResponseMetadata: start_time: datetime.datetime, end_time: datetime.datetime, logging_obj: LiteLLMLoggingObject, + include_overhead: bool = True, ) -> None: """Set response timing metrics""" - total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000 + timing_metrics: Final = response_timing_metrics(start_time, end_time, logging_obj, include_overhead) + total_response_time_ms: Final = timing_metrics["_response_ms"] # Set total response time if supported if self.supports_response_time: self.result._response_ms = total_response_time_ms ######################################################### - # 1. Add _response_ms total duration + # 1. Add _response_ms total duration and the LiteLLM overhead within it + # (total minus the cache read on a cache hit, else total minus the provider call) ######################################################### - self._update_hidden_params( - { - "_response_ms": total_response_time_ms, - } - ) + self._update_hidden_params(timing_metrics) ######################################################### - # 2. Add LiteLLM overhead duration - ######################################################### - llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") - if llm_api_duration_ms is not None: - overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) - self._update_hidden_params( - { - "litellm_overhead_time_ms": overhead_ms, - } - ) - - ######################################################### - # 3. Add callback processing duration + # 2. Add callback processing duration ######################################################### callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None) if callback_duration_ms is not None: @@ -117,24 +138,9 @@ class ResponseMetadata: ) ######################################################### - # 4. Add duration for reading from cache - # In this case overhead from litellm is the difference between the cache read duration and the total response time - ######################################################### - if ( - logging_obj.caching_details is not None - and logging_obj.caching_details.get("cache_hit") is True - and (cache_duration_ms := logging_obj.caching_details.get("cache_duration_ms")) is not None - ): - overhead_ms = total_response_time_ms - cache_duration_ms - self._update_hidden_params( - { - "litellm_overhead_time_ms": overhead_ms, - } - ) - - ######################################################### - # 5. Detailed per-phase timing (opt-in via env var) + # 3. Detailed per-phase timing (opt-in via env var) ######################################################### + llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None: detailed: Final[dict] = { "timing_llm_api_ms": round(llm_api_duration_ms, 4), @@ -170,6 +176,7 @@ def update_response_metadata( kwargs: dict, start_time: datetime.datetime, end_time: datetime.datetime, + include_overhead: bool = True, ) -> None: """ Updates response metadata including hidden params and timing metrics @@ -177,11 +184,22 @@ def update_response_metadata( - response._hidden_params - response._hidden_params["litellm_overhead_time_ms"] - response.response_time_ms + A result that cannot hold ``_hidden_params`` gets its timing on ``logging_obj`` instead. + Callers whose ``end_time`` covers more than the recorded provider call (a stream read to + completion) pass ``include_overhead=False``, since the overhead cannot be derived there. """ - if result is None or not hasattr(result, "_hidden_params"): + if result is None: + return + if not hasattr(result, "_hidden_params"): + # /v1/messages returns a plain dict and the Anthropic / Responses bridge stream wrappers + # cannot hold ``_hidden_params``: keep only the timing on the logging object (no cost + # recompute) so the proxy headers and the standard logging payload can still read it. + logging_obj.set_response_timing_metrics( + response_timing_metrics(start_time, end_time, logging_obj, include_overhead) + ) return metadata: Final = ResponseMetadata(result) metadata.set_hidden_params(logging_obj, model, kwargs) - metadata.set_timing_metrics(start_time, end_time, logging_obj) + metadata.set_timing_metrics(start_time, end_time, logging_obj, include_overhead) metadata.apply() diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 573ba85416f..834f7d564a2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2872,6 +2872,7 @@ class BaseLLMHTTPHandler: headers=headers, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + logging_obj=logging_obj, **body_kwargs, ) @@ -2903,6 +2904,7 @@ class BaseLLMHTTPHandler: url=api_base, headers=headers, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + logging_obj=logging_obj, **body_kwargs, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a23669af462..c6a6639409a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1475,6 +1475,22 @@ async def _await_llm_call_cancelling_on_disconnect( monitor.cancel() +def _timing_values( + *, + hidden_params: Mapping[str, object], + logging_obj: LiteLLMLoggingObj | None, + use_logging_obj: bool, +) -> Mapping[str, object]: + """Both timing values from one source, so the two headers always describe the same window. + + /v1/messages returns a plain dict and the Anthropic / Responses bridge stream wrappers carry no + ``_hidden_params``, so ``update_response_metadata`` leaves their timing on the logging object. + """ + if hidden_params.get("_response_ms") is not None or not use_logging_obj or logging_obj is None: + return hidden_params + return getattr(logging_obj, "response_timing_metrics", None) or {} # mutable-ok: empty fallback + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data @@ -1495,10 +1511,16 @@ class ProxyBaseLLMRequestProcessing: request_data: dict | None = {}, timeout: float | httpx.Timeout | None = None, litellm_logging_obj: LiteLLMLoggingObj | None = None, + read_timing_from_logging_obj: bool = True, **kwargs, ) -> dict: exclude_values: Final = {"", None, "None"} hidden_params = hidden_params or {} + timing_values: Final = _timing_values( + hidden_params=hidden_params, + logging_obj=litellm_logging_obj, + use_logging_obj=read_timing_from_logging_obj, + ) cost_breakdown: Final = _get_cost_breakdown_from_logging_obj( litellm_logging_obj=litellm_logging_obj, response_cost=response_cost @@ -1566,8 +1588,8 @@ class ProxyBaseLLMRequestProcessing: "x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit), "x-litellm-key-max-budget": str(user_api_key_dict.max_budget), "x-litellm-key-spend": str(updated_spend), - "x-litellm-response-duration-ms": str(hidden_params.get("_response_ms", None)), - "x-litellm-overhead-duration-ms": str(hidden_params.get("litellm_overhead_time_ms", None)), + "x-litellm-response-duration-ms": str(timing_values.get("_response_ms")), + "x-litellm-overhead-duration-ms": str(timing_values.get("litellm_overhead_time_ms")), "x-litellm-callback-duration-ms": str(hidden_params.get("callback_duration_ms", None)), **( { @@ -3224,6 +3246,8 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, timeout=timeout, litellm_logging_obj=_litellm_logging_obj, + # a failed request reports no timing, matching /v1/chat/completions + read_timing_from_logging_obj=False, ) # Extract headers from exception - check both e.headers and e.response.headers headers = getattr(e, "headers", None) or {} diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 368fd481e63..63adf950142 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -681,6 +681,8 @@ class BaseResponsesAPIStreamingIterator: kwargs=request_payload, start_time=self.start_time, end_time=end_time, + # the provider call was timed to first byte, so the whole stream minus it is not overhead + include_overhead=False, ) except Exception: # Non-blocking diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py index 996530daa2e..a06b6bbf3cc 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -5,6 +5,7 @@ Covers the callback_duration_ms timing metric that flows from the Logging object through _hidden_params to the x-litellm-callback-duration-ms response header. """ +import asyncio import datetime from unittest.mock import MagicMock @@ -13,6 +14,7 @@ import litellm.proxy.common_request_processing as common_request_processing_mod from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ResponseMetadata, + response_timing_metrics, update_response_metadata, ) from litellm.proxy._types import UserAPIKeyAuth @@ -124,6 +126,158 @@ class TestDictResultsSkipMetadataUpdate: logging_obj._response_cost_calculator.assert_not_called() assert "_hidden_params" not in anthropic_response + def test_update_response_metadata_keeps_timing_on_logging_obj_for_dict_results(self): + """LIT-5466: the /v1/messages dict cannot carry _hidden_params, so its timing + (the input to x-litellm-overhead-duration-ms and the SLP litellm_overhead_time_ms) + lands on the logging object instead - still without recomputing cost.""" + anthropic_response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + logging_obj = MagicMock() + logging_obj.model_call_details = {"llm_api_duration_ms": 900.0} + logging_obj.caching_details = None + + update_response_metadata( + result=anthropic_response, + logging_obj=logging_obj, + model="openai/gpt-4o-mini", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + logging_obj.set_response_timing_metrics.assert_called_once_with( + {"_response_ms": 1000.0, "litellm_overhead_time_ms": 100.0} + ) + logging_obj._response_cost_calculator.assert_not_called() + assert "_hidden_params" not in anthropic_response + + def test_update_response_metadata_keeps_timing_for_stream_wrapper_without_hidden_params(self): + """The /v1/messages bridge streams a bare async generator, which cannot hold + _hidden_params either; when the provider duration is unknown only the total is kept.""" + + async def sse_stream(): + yield b"event: message_start\n\n" + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.caching_details = None + + async def drive(): + stream = sse_stream() + try: + update_response_metadata( + result=stream, + logging_obj=logging_obj, + model="openai/gpt-4o-mini", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 0, 250000), + ) + finally: + await stream.aclose() + + asyncio.run(drive()) + + logging_obj.set_response_timing_metrics.assert_called_once_with({"_response_ms": 250.0}) + logging_obj._response_cost_calculator.assert_not_called() + + def test_update_response_metadata_leaves_logging_obj_alone_for_objects_with_hidden_params(self): + """ModelResponse keeps carrying its own timing; the logging-object carrier is not written.""" + result = ModelResponse() + logging_obj = MagicMock() + logging_obj.model_call_details = {"llm_api_duration_ms": 900.0} + logging_obj.caching_details = None + logging_obj._response_cost_calculator = MagicMock(return_value=0.001) + logging_obj.litellm_call_id = "test-call-id" + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="gpt-4", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + logging_obj.set_response_timing_metrics.assert_not_called() + assert result._hidden_params["litellm_overhead_time_ms"] == 100.0 + + def test_update_response_metadata_omits_overhead_for_completed_stream(self): + """LIT-5466: the Responses streaming iterator finishes the whole stream before updating + metadata, and the provider call it recorded stopped at the first byte.""" + result = ModelResponse() + logging_obj = MagicMock() + logging_obj.model_call_details = {"llm_api_duration_ms": 200.0} + logging_obj.caching_details = None + logging_obj._response_cost_calculator = MagicMock(return_value=0.001) + logging_obj.litellm_call_id = "test-call-id" + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="gpt-4", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + include_overhead=False, + ) + + assert result._hidden_params["_response_ms"] == 1000.0 + assert "litellm_overhead_time_ms" not in result._hidden_params + + +class TestResponseTimingMetrics: + """response_timing_metrics() is the single source of _response_ms / litellm_overhead_time_ms.""" + + START = datetime.datetime(2025, 1, 1, 0, 0, 0) + END = datetime.datetime(2025, 1, 1, 0, 0, 1) + + def _make_logging_obj(self, llm_api_duration_ms=None, caching_details=None): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + if llm_api_duration_ms is not None: + logging_obj.model_call_details["llm_api_duration_ms"] = llm_api_duration_ms + logging_obj.caching_details = caching_details + return logging_obj + + def test_overhead_is_total_minus_provider_call(self): + logging_obj = self._make_logging_obj(llm_api_duration_ms=900.0) + assert response_timing_metrics(self.START, self.END, logging_obj) == { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 100.0, + } + + def test_overhead_omitted_when_no_provider_or_cache_duration_recorded(self): + logging_obj = self._make_logging_obj() + assert response_timing_metrics(self.START, self.END, logging_obj) == {"_response_ms": 1000.0} + + def test_cache_hit_overhead_is_total_minus_cache_read(self): + logging_obj = self._make_logging_obj( + llm_api_duration_ms=900.0, caching_details={"cache_hit": True, "cache_duration_ms": 250.0} + ) + assert response_timing_metrics(self.START, self.END, logging_obj) == { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 750.0, + } + + def test_cache_miss_ignores_cache_duration(self): + logging_obj = self._make_logging_obj(caching_details={"cache_hit": False, "cache_duration_ms": 250.0}) + assert response_timing_metrics(self.START, self.END, logging_obj) == {"_response_ms": 1000.0} + + def test_cache_hit_without_recorded_cache_duration_falls_back_to_provider_call(self): + logging_obj = self._make_logging_obj(llm_api_duration_ms=900.0, caching_details={"cache_hit": True}) + assert response_timing_metrics(self.START, self.END, logging_obj) == { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 100.0, + } + + def test_overhead_omitted_when_caller_measured_a_wider_window(self): + """A stream read to completion times the provider call to first byte, so the rest of the + stream is token generation, not LiteLLM overhead.""" + logging_obj = self._make_logging_obj(llm_api_duration_ms=200.0) + assert response_timing_metrics(self.START, self.END, logging_obj, include_overhead=False) == { + "_response_ms": 1000.0 + } + class TestCallbackDurationInCustomHeaders: """Test that callback_duration_ms flows into get_custom_headers.""" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index a45a6aaa2a3..947a55410ef 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5994,3 +5994,110 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o request_kwargs=untouched, ) assert "litellm_gateway_injected_cache" not in untouched["metadata"] + + +def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): + """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead + recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs).""" + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + logging_obj.set_response_timing_metrics({"_response_ms": 1000.0, "litellm_overhead_time_ms": 100.0}) + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "gpt-4o", "messages": []}, + init_response_obj={"id": "msg_1", "type": "message", "role": "assistant", "content": []}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["hidden_params"]["litellm_overhead_time_ms"] == 100.0 + + +def test_get_standard_logging_object_payload_survives_logging_obj_without_timing_metrics(logging_obj): + """The payload is built inside a blanket except that returns None, so a logging object without + the timing carrier (custom subclasses, older pickles) must not silently drop every spend log.""" + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + del logging_obj.response_timing_metrics + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "gpt-4o", "messages": []}, + init_response_obj={"id": "msg_1", "type": "message", "role": "assistant", "content": []}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["hidden_params"]["litellm_overhead_time_ms"] is None + + +def test_get_standard_logging_object_payload_failure_status_keeps_overhead_none(logging_obj): + """A post_call guardrail can fail the request after the upstream call succeeded; the failure + payload keeps litellm_overhead_time_ms None, matching responses that carry their own _hidden_params.""" + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + logging_obj.set_response_timing_metrics({"_response_ms": 1000.0, "litellm_overhead_time_ms": 100.0}) + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "gpt-4o", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="failure", + ) + + assert payload is not None + assert payload["hidden_params"]["litellm_overhead_time_ms"] is None + + +def test_get_standard_logging_object_payload_prefers_response_hidden_params_overhead(logging_obj): + """A response that carries its own litellm_overhead_time_ms (chat completions) wins over the logging object.""" + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + logging_obj.set_response_timing_metrics({"_response_ms": 1000.0, "litellm_overhead_time_ms": 100.0}) + response = ModelResponse() + response._hidden_params = {"litellm_overhead_time_ms": 5.0} + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "call-1", "model": "gpt-4o", "messages": []}, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["hidden_params"]["litellm_overhead_time_ms"] == 5.0 + + +def test_response_timing_metrics_survive_deepcopy(logging_obj): + """Proxy pre-call hooks deep-copy the logging object; the timing carrier must stay copyable.""" + import copy + + assert logging_obj.response_timing_metrics == {} + logging_obj.set_response_timing_metrics({"_response_ms": 12.5}) + + assert copy.deepcopy(logging_obj).response_timing_metrics == {"_response_ms": 12.5} diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index b37c0f466d2..26f841c1146 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -271,6 +271,85 @@ async def test_async_response_api_handler_streams_when_provider_transform_adds_s assert client.post.call_args.kwargs["json"]["stream"] is True +@pytest.mark.asyncio +async def test_async_response_api_handler_streaming_passes_logging_obj_to_post(): + """LIT-5466: @track_llm_api_timing only records llm_api_duration_ms when the POST + receives logging_obj; without it streaming /v1/responses never gets + x-litellm-overhead-duration-ms (the non-streaming site is pinned by + test_async_responses_records_llm_api_duration below).""" + handler = BaseLLMHTTPHandler() + config = Mock() + config.validate_environment.return_value = {} + config.get_complete_url.return_value = "https://chatgpt.example.com/responses" + config.transform_responses_api_request.return_value = {"model": "gpt-5", "input": "hi", "stream": True} + config.sign_request.return_value = ({}, None) + client = AsyncHTTPHandler() + client.post = AsyncMock( + return_value=httpx.Response( + 200, + request=httpx.Request("POST", "https://chatgpt.example.com/responses"), + ) + ) + logging_obj = Mock() + + await handler.async_response_api_handler( + model="gpt-5", + input="hi", + responses_api_provider_config=config, + response_api_optional_request_params={}, + custom_llm_provider="chatgpt", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + client=client, + ) + + assert client.post.call_args.kwargs["logging_obj"] is logging_obj + + +@pytest.mark.asyncio +async def test_async_responses_records_llm_api_duration(): + """aresponses must feed the httpx timing into the logging obj, so the proxy can emit + x-litellm-overhead-duration-ms on /v1/responses (mirrors the arerank regression test).""" + + def handle(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "model": "gpt-4o-mini", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.aresponses( + model="openai/gpt-4o-mini", + input="ping", + api_key="fake-key", + client=client, + ) + + assert response._hidden_params["litellm_overhead_time_ms"] is not None + assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"] + + def test_get_agentic_loop_settings_defaults_and_overrides(): handler = BaseLLMHTTPHandler() diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index dd5049b0f28..833c3754022 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2549,6 +2549,152 @@ class TestStreamingOverheadHeader: assert "x-litellm-overhead-duration-ms" in headers assert headers["x-litellm-overhead-duration-ms"] == "42.5" + @staticmethod + def _timing_logging_obj(timing_metrics): + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + logging_obj = LiteLLMLoggingObj( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=None, + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + logging_obj.set_response_timing_metrics(timing_metrics) + return logging_obj + + def test_get_custom_headers_reads_timing_from_logging_obj_when_response_has_no_hidden_params(self): + """ + LIT-5466: /v1/messages results and the bridge stream wrappers carry no + _hidden_params, so the timing headers come from the logging object. + """ + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + hidden_params={}, + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), + ) + + assert headers["x-litellm-response-duration-ms"] == "500.0" + assert headers["x-litellm-overhead-duration-ms"] == "42.5" + + def test_get_custom_headers_skips_logging_obj_timing_on_the_failure_path(self): + """LIT-5466: a failed request reports no timing, the same as /v1/chat/completions.""" + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + hidden_params={}, + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), + read_timing_from_logging_obj=False, + ) + + assert "x-litellm-response-duration-ms" not in headers + assert "x-litellm-overhead-duration-ms" not in headers + + def test_get_custom_headers_takes_both_timing_values_from_one_source(self): + """A response that timed itself but has no overhead (lazy provider streams) does not pick + up the logging object's overhead, which was measured over a different window.""" + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + hidden_params={"_response_ms": 300.0}, + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), + ) + + assert headers["x-litellm-response-duration-ms"] == "300.0" + assert "x-litellm-overhead-duration-ms" not in headers + + def test_get_custom_headers_survives_a_logging_object_without_timing_metrics(self): + """Duck-typed logging objects (older custom code, test doubles) must not break headers.""" + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + class _NoTimingLoggingObj: + litellm_call_id = "test-call-id" + litellm_params = {} + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + hidden_params={}, + litellm_logging_obj=_NoTimingLoggingObj(), + ) + + assert "x-litellm-overhead-duration-ms" not in headers + + def test_get_custom_headers_prefers_response_hidden_params_over_logging_obj_timing(self): + """A response that carries its own timing (chat completions) is not overridden.""" + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + hidden_params={"_response_ms": 300.0, "litellm_overhead_time_ms": 7.5}, + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), + ) + + assert headers["x-litellm-response-duration-ms"] == "300.0" + assert headers["x-litellm-overhead-duration-ms"] == "7.5" + + def test_get_custom_headers_omits_timing_when_no_source_has_it(self): + """No timing on the response and none on the logging object leaves both headers out.""" + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + hidden_params={}, + litellm_logging_obj=self._timing_logging_obj({}), + ) + + assert "x-litellm-response-duration-ms" not in headers + assert "x-litellm-overhead-duration-ms" not in headers + def test_get_custom_headers_omits_overhead_when_none(self): """ get_custom_headers() omits x-litellm-overhead-duration-ms diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 38407c94fe7..677faf7f655 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -305,3 +305,24 @@ def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypat asyncio.run(_short_lived_script()) assert len(writes) == 1 + + +def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): + """LIT-5466: the provider call is timed to first byte, so at stream completion the total minus + that duration is token generation, not LiteLLM overhead.""" + logging_obj = _logging_obj_stub() + logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 200.0} + logging_obj.caching_details = None + + class _CompletedEvent: + def __init__(self) -> None: + self._hidden_params: dict = {} + + iterator = _make_iterator(sse_events=[], logging_obj=logging_obj) + iterator.completed_response = _CompletedEvent() + iterator.start_time = datetime(2025, 1, 1, 0, 0, 0) + + iterator._run_post_success_hooks(datetime(2025, 1, 1, 0, 0, 10)) + + assert iterator.completed_response._hidden_params["_response_ms"] == 10000.0 + assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params From f62aa1b3a82eef2213e742f750a486ca0daae0e0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 29 Aug 2026 19:35:10 -0700 Subject: [PATCH 12/16] fix(tests): derive the no-cache-read-rate savings baseline from the model map --- .../proxy/spend_tracking/test_savings.py | 40 +++++++++++++++++-- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index 5fda4fb20b5..c3297ee6ae9 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -6,6 +6,7 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.proxy.spend_tracking.savings import ( _baseline_usage, + _resolve_model, compute_autorouter_savings, compute_savings_spend, marks_gateway_injection, @@ -754,24 +755,55 @@ def test_a_baseline_that_prices_caching_implicitly_still_pays_for_its_prompt(): assert reported > 0, "routing a cold first turn onto a cheaper model is a saving, not a loss" +def _priced_chat_model_without_cache_read_rate() -> tuple[str, str, str]: + """A chat model the bundled map prices per token for input and output but not for cache + reads, derived from the map itself: a hardcoded pick goes stale the moment the registry + prices that model's cache reads, which is exactly how this test's premise last broke. + Candidates go through the savings module's own resolver, so the pick is one the code + under test can actually price.""" + for key in sorted(litellm.model_cost): + entry = litellm.model_cost[key] + provider = entry.get("litellm_provider") + if not isinstance(provider, str) or not key.startswith(f"{provider}/"): + continue + if entry.get("mode") != "chat" or entry.get("cache_read_input_token_cost") is not None: + continue + if not entry.get("input_cost_per_token") or not entry.get("output_cost_per_token"): + continue + if _resolve_model(key, None) is None: + continue + priced = compute_autorouter_savings( + baseline_model=key, + selected_model="claude-haiku-4-5", + selected_provider="anthropic", + usage=_usage(fresh=1_000, cached=0, written=0, out=100), + conversation_continuing=True, + ) + if priced == 0.0: + continue + return key, key.removeprefix(f"{provider}/"), provider + raise AssertionError("the bundled map has no per-token chat model without a cache-read rate") + + def test_a_baseline_with_no_cache_read_rate_is_charged_its_input_rate(): """The same hole on the other bucket. A baseline whose entry has no `cache_read_input_token_cost` reads for 0.0, so a continuing turn priced the whole prompt at nothing and every switch away from it reported a loss. """ + baseline_key, baseline_name, baseline_provider = _priced_chat_model_without_cache_read_rate() continuing = _usage(fresh=0, cached=0, written=20_000, out=1_000) reported = compute_autorouter_savings( - baseline_model="xai/grok-4", + baseline_model=baseline_key, selected_model="claude-haiku-4-5", selected_provider="anthropic", usage=continuing, conversation_continuing=True, ) - grok = litellm.get_model_info("grok-4", "xai") - assert grok.get("cache_read_input_token_cost") is None, "pick a baseline with no cache-read rate" + baseline = litellm.get_model_info(baseline_name, baseline_provider) + assert baseline.get("cache_read_input_token_cost") is None, "pick a baseline with no cache-read rate" haiku = litellm.get_model_info("claude-haiku-4-5", "anthropic") - baseline_pays_input = 20_000 * grok["input_cost_per_token"] + 1_000 * grok["output_cost_per_token"] + baseline_pays_input = 20_000 * baseline["input_cost_per_token"] + 1_000 * baseline["output_cost_per_token"] actually_paid = 20_000 * haiku["cache_creation_input_token_cost"] + 1_000 * haiku["output_cost_per_token"] assert reported == pytest.approx(baseline_pays_input - actually_paid) From cafdfda8ba586429dcaf0c6c1dcf7e7c5caf70a7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 05:52:17 +0000 Subject: [PATCH 13/16] feat(cli): set ENABLE_TOOL_SEARCH=true for lite claude Claude Code turns tool search off when ANTHROPIC_BASE_URL is a proxy. lite claude, lite up, login --config-claude, and autoroute now force ENABLE_TOOL_SEARCH=true so MCP tools stay deferred through the proxy Co-authored-by: Mateo Wang --- litellm/proxy/client/cli/README.md | 6 +++--- litellm/proxy/client/cli/commands/agents.py | 7 ++++++- .../client/cli/commands/autoroute/settings.py | 3 +++ .../proxy/client/cli/commands/claude_settings.py | 15 +++++++++++---- .../proxy/client/cli/autoroute/test_commands.py | 1 + .../proxy/client/cli/autoroute/test_settings.py | 7 +++++++ .../test_litellm/proxy/client/cli/test_agents.py | 14 ++++++++++++++ .../proxy/client/cli/test_auth_commands.py | 1 + .../proxy/client/cli/test_claude_settings.py | 1 + .../proxy/client/cli/test_up_commands.py | 12 +++++++++++- 10 files changed, 58 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index fe417396317..0da219df700 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -489,7 +489,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). Options (these belong to the wrapper, so put them before the agent's own flags): @@ -505,7 +505,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true`, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. @@ -529,7 +529,7 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi lite --base-url https://your-proxy.example.com login --config-claude ``` -It writes the same two settings `lite up` does, `env.ANTHROPIC_BASE_URL` and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. +It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index e05e85ae483..4d3a441272e 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -13,6 +13,8 @@ from .auth import context_secret_vault, get_stored_api_key, login ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_ENV: Final = "ANTHROPIC_API_KEY" +ENABLE_TOOL_SEARCH_ENV: Final = "ENABLE_TOOL_SEARCH" +ENABLE_TOOL_SEARCH_VALUE: Final = "true" OPENAI_BASE_URL_ENV: Final = "OPENAI_BASE_URL" OPENAI_API_KEY_ENV: Final = "OPENAI_API_KEY" @@ -61,7 +63,9 @@ def build_agent_env( Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL, so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the /v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray - Anthropic key cannot win over the bearer token we set. + Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH is + forced on because Claude Code turns tool search off when ANTHROPIC_BASE_URL + is not a first-party Anthropic host. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -69,6 +73,7 @@ def build_agent_env( env[ANTHROPIC_BASE_URL_ENV] = root env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key env.pop(ANTHROPIC_API_KEY_ENV, None) + env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE if PROFILE_OPENAI in profiles: env[OPENAI_BASE_URL_ENV] = root + "/v1" env[OPENAI_API_KEY_ENV] = api_key diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py index 9fcb11a585b..8331d41ae78 100644 --- a/litellm/proxy/client/cli/commands/autoroute/settings.py +++ b/litellm/proxy/client/cli/commands/autoroute/settings.py @@ -9,6 +9,8 @@ API_KEY_HELPER_KEY: Final = "apiKeyHelper" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" +ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" +ENABLE_TOOL_SEARCH_VALUE: Final = "true" # Force every one of Claude Code's own model tiers to request the auto-router by name. # Router's auto-router registry is keyed by the literal requested model string # (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*" @@ -37,6 +39,7 @@ def merge_claude_settings_static_token( **base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), ANTHROPIC_AUTH_TOKEN_KEY: auth_token, + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, } env.pop(ANTHROPIC_API_KEY_KEY, None) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index e18e5b1b7ee..6d88c69bbd4 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -21,6 +21,8 @@ ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" +ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" +ENABLE_TOOL_SEARCH_VALUE: Final = "true" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" @@ -68,16 +70,19 @@ def merge_claude_settings( ) -> dict[str, JsonValue]: """Return a new settings dict wired to route Claude Code through the proxy. - Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a - stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued - token (same reasoning as build_agent_env in agents.py). Every other key is - preserved untouched. + Only env.ANTHROPIC_BASE_URL, env.ENABLE_TOOL_SEARCH, and the top-level + apiKeyHelper are overridden; a stray env.ANTHROPIC_API_KEY is dropped so it + cannot outrank the helper-issued token (same reasoning as build_agent_env + in agents.py). ENABLE_TOOL_SEARCH is forced on because Claude Code turns + tool search off when ANTHROPIC_BASE_URL is not a first-party Anthropic host. + Every other key is preserved untouched. """ raw_env: Final = settings.get(ENV_KEY, {}) base_env: Final = raw_env if isinstance(raw_env, dict) else {} env: Final = { **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, } return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} @@ -141,6 +146,8 @@ __all__ = ( "ANTHROPIC_API_KEY_KEY", "ANTHROPIC_BASE_URL_KEY", "API_KEY_HELPER_KEY", + "ENABLE_TOOL_SEARCH_KEY", + "ENABLE_TOOL_SEARCH_VALUE", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", "CLAUDE_SETTINGS_PATH", diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py index 74bf1c95777..4a3b3ef22c4 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -157,6 +157,7 @@ class TestUpCommand: assert captured["settings"]["theme"] == "dark" assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:5483" assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" + assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" assert "apiKeyHelper" not in captured["settings"] assert captured["settings_mode"] == 0o600 diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py index 40d3e7f2aee..ded8f0808a0 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py @@ -19,6 +19,13 @@ def test_sets_base_url_and_auth_token(): merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + + +def test_forces_tool_search_on(): + settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" def test_drops_stray_api_key(): diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index 32dfb8d521d..d03c5d9db2d 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -77,9 +77,19 @@ class TestBuildAgentEnv: ) assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert env["ENABLE_TOOL_SEARCH"] == "true" assert "OPENAI_BASE_URL" not in env assert "OPENAI_API_KEY" not in env + def test_anthropic_profile_forces_tool_search_on(self): + env = build_agent_env( + {"ENABLE_TOOL_SEARCH": "false"}, + "http://localhost:4000", + "sk-key", + frozenset({"anthropic"}), + ) + assert env["ENABLE_TOOL_SEARCH"] == "true" + def test_anthropic_profile_drops_existing_api_key(self): env = build_agent_env( {"ANTHROPIC_API_KEY": "real-key"}, @@ -96,6 +106,7 @@ class TestBuildAgentEnv: assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert env["OPENAI_API_KEY"] == "sk-key" assert "ANTHROPIC_BASE_URL" not in env + assert "ENABLE_TOOL_SEARCH" not in env def test_both_profiles_set_everything(self): env = build_agent_env( @@ -105,6 +116,7 @@ class TestBuildAgentEnv: assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" assert env["OPENAI_API_KEY"] == "sk-key" + assert env["ENABLE_TOOL_SEARCH"] == "true" def test_preserves_unrelated_env_and_does_not_mutate_input(self): base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} @@ -201,6 +213,7 @@ class TestRunAgent: env = calls["env"] assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert env["ENABLE_TOOL_SEARCH"] == "true" assert "ANTHROPIC_API_KEY" not in env assert "OPENAI_BASE_URL" not in env @@ -218,6 +231,7 @@ class TestRunAgent: assert calls["env"]["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert calls["env"]["OPENAI_API_KEY"] == "sk-key" assert "ANTHROPIC_BASE_URL" not in calls["env"] + assert "ENABLE_TOOL_SEARCH" not in calls["env"] def test_codex_injects_proxy_provider_args_before_user_args(self): calls = {} diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 85a4d90abf9..1d0a99b8e0a 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1373,6 +1373,7 @@ class TestLoginConfigClaude: assert result.exit_code == 0 written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com" + assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" assert "Configured Claude Code" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index 9010fb4c022..898f9ab1ed7 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -48,6 +48,7 @@ class TestWriteClaudeSettings: written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" + assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path): diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 9958286884b..a1aa9eeeb8a 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -55,8 +55,14 @@ class TestMergeClaudeSettings: } merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["apiKeyHelper"] == "new-helper" + def test_forces_tool_search_on(self): + settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} + merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} merged = merge_claude_settings(settings, "http://localhost:4000", "helper") @@ -64,7 +70,10 @@ class TestMergeClaudeSettings: def test_works_from_empty_settings(self): merged = merge_claude_settings({}, "http://localhost:4000", "helper") - assert merged["env"] == {"ANTHROPIC_BASE_URL": "http://localhost:4000"} + assert merged["env"] == { + "ANTHROPIC_BASE_URL": "http://localhost:4000", + "ENABLE_TOOL_SEARCH": "true", + } assert merged["apiKeyHelper"] == "helper" def test_does_not_mutate_input(self): @@ -486,6 +495,7 @@ class TestUpCommand: assert captured["backup_existed"] is True assert captured["settings"]["theme"] == "dark" assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" assert captured["settings"]["apiKeyHelper"] == "/usr/local/bin/lite auth print-token" assert json.loads(settings_path.read_text()) == original assert not backup_path.exists() From 096e016baf04fc833ce617025eace2f8937cefe2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 05:58:00 +0000 Subject: [PATCH 14/16] fix(cli): sort claude_settings __all__ for ruff RUF022 Co-authored-by: Mateo Wang --- litellm/proxy/client/cli/commands/claude_settings.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 6d88c69bbd4..41ba25447e1 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -146,11 +146,11 @@ __all__ = ( "ANTHROPIC_API_KEY_KEY", "ANTHROPIC_BASE_URL_KEY", "API_KEY_HELPER_KEY", - "ENABLE_TOOL_SEARCH_KEY", - "ENABLE_TOOL_SEARCH_VALUE", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", "CLAUDE_SETTINGS_PATH", + "ENABLE_TOOL_SEARCH_KEY", + "ENABLE_TOOL_SEARCH_VALUE", "ENV_KEY", "SETTINGS_FILE_OWNERS", "ClaudeSettingsError", From 43c64e24715cdd1aa7595544f1d4aba55f7c2512 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 06:00:45 +0000 Subject: [PATCH 15/16] fix(cli): keep an existing ENABLE_TOOL_SEARCH value Default remains true so lite claude turns tool search back on through a proxy. An explicit false or auto in the env or settings is left alone Co-authored-by: Mateo Wang --- litellm/proxy/client/cli/README.md | 4 ++-- litellm/proxy/client/cli/commands/agents.py | 10 ++++++---- .../client/cli/commands/autoroute/settings.py | 2 +- .../proxy/client/cli/commands/claude_settings.py | 14 +++++++------- .../proxy/client/cli/autoroute/test_settings.py | 4 ++-- tests/test_litellm/proxy/client/cli/test_agents.py | 4 ++-- .../proxy/client/cli/test_up_commands.py | 4 ++-- 7 files changed, 22 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 0da219df700..3ddce35b53d 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -489,7 +489,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). Options (these belong to the wrapper, so put them before the agent's own flags): @@ -505,7 +505,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true`, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` when that key is missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 4d3a441272e..45e05d353fb 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -63,9 +63,10 @@ def build_agent_env( Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL, so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the /v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray - Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH is - forced on because Claude Code turns tool search off when ANTHROPIC_BASE_URL - is not a first-party Anthropic host. + Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH + defaults to true because Claude Code turns tool search off when + ANTHROPIC_BASE_URL is not a first-party Anthropic host; a value already in + the environment is left alone. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -73,7 +74,8 @@ def build_agent_env( env[ANTHROPIC_BASE_URL_ENV] = root env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key env.pop(ANTHROPIC_API_KEY_ENV, None) - env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE + if ENABLE_TOOL_SEARCH_ENV not in env: + env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE if PROFILE_OPENAI in profiles: env[OPENAI_BASE_URL_ENV] = root + "/v1" env[OPENAI_API_KEY_ENV] = api_key diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py index 8331d41ae78..60729b5410d 100644 --- a/litellm/proxy/client/cli/commands/autoroute/settings.py +++ b/litellm/proxy/client/cli/commands/autoroute/settings.py @@ -36,10 +36,10 @@ def merge_claude_settings_static_token( raw_env: Final = settings.get(ENV_KEY, {}) base_env: Final = raw_env if isinstance(raw_env, dict) else {} env: Final[dict[str, JsonValue]] = { + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), ANTHROPIC_AUTH_TOKEN_KEY: auth_token, - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, } env.pop(ANTHROPIC_API_KEY_KEY, None) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 41ba25447e1..ea1fa019c83 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -70,19 +70,19 @@ def merge_claude_settings( ) -> dict[str, JsonValue]: """Return a new settings dict wired to route Claude Code through the proxy. - Only env.ANTHROPIC_BASE_URL, env.ENABLE_TOOL_SEARCH, and the top-level - apiKeyHelper are overridden; a stray env.ANTHROPIC_API_KEY is dropped so it - cannot outrank the helper-issued token (same reasoning as build_agent_env - in agents.py). ENABLE_TOOL_SEARCH is forced on because Claude Code turns - tool search off when ANTHROPIC_BASE_URL is not a first-party Anthropic host. - Every other key is preserved untouched. + Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a + stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued + token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH + defaults to true because Claude Code turns tool search off when + ANTHROPIC_BASE_URL is not a first-party Anthropic host; an existing value is + left alone. Every other key is preserved untouched. """ raw_env: Final = settings.get(ENV_KEY, {}) base_env: Final = raw_env if isinstance(raw_env, dict) else {} env: Final = { + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, } return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py index ded8f0808a0..87a33c79a79 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py @@ -22,10 +22,10 @@ def test_sets_base_url_and_auth_token(): assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" -def test_forces_tool_search_on(): +def test_preserves_existing_tool_search(): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(): diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index d03c5d9db2d..0191dad3d94 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -81,14 +81,14 @@ class TestBuildAgentEnv: assert "OPENAI_BASE_URL" not in env assert "OPENAI_API_KEY" not in env - def test_anthropic_profile_forces_tool_search_on(self): + def test_anthropic_profile_preserves_existing_tool_search(self): env = build_agent_env( {"ENABLE_TOOL_SEARCH": "false"}, "http://localhost:4000", "sk-key", frozenset({"anthropic"}), ) - assert env["ENABLE_TOOL_SEARCH"] == "true" + assert env["ENABLE_TOOL_SEARCH"] == "false" def test_anthropic_profile_drops_existing_api_key(self): env = build_agent_env( diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index a1aa9eeeb8a..053c90c36b4 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -58,10 +58,10 @@ class TestMergeClaudeSettings: assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["apiKeyHelper"] == "new-helper" - def test_forces_tool_search_on(self): + def test_preserves_existing_tool_search(self): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} merged = merge_claude_settings(settings, "http://localhost:4000", "helper") - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} From f079e4061bf986d6bb368da864f6c9dec22a0cac Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:22:36 -0700 Subject: [PATCH 16/16] fix(proxy): deliver budget alerts on webhook-only alerting and accept ALERTING_WEBHOOK_URL (#38441) * fix(proxy): deliver budget alerts on webhook-only alerting and accept ALERTING_WEBHOOK_URL ProxyLogging.budget_alerts forwarded to the alerting pipeline only when 'slack' was in general_settings.alerting, so alerting: ['webhook'] plus WEBHOOK_URL silently never delivered a budget alert (the config /health/services?service=webhook exists to test). Forward when 'webhook' is present too; SlackAlerting.send_alert already fans out per channel. Also accept a provider-neutral ALERTING_WEBHOOK_URL env fallback for the Slack-format channel (any Slack-compatible receiver works), mark it as a sensitive var, and de-brand the admin UI alerting copy. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): format settings.tsx with prettier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(ui): regenerate schema.d.ts for updated alerting description Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: retrigger checks after ALERTING_WEBHOOK_URL docs merged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + .../SlackAlerting/slack_alerting.py | 8 +-- litellm/proxy/_types.py | 2 +- litellm/proxy/proxy_server.py | 3 +- litellm/proxy/utils.py | 6 +- .../SlackAlerting/test_slack_alerting.py | 55 ++++++++++++++++++- .../test_slack_alerting_digest.py | 17 ++++++ .../utils/proxy_logging/test_alerting.py | 29 ++++++++++ .../src/components/settings.tsx | 5 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 10 files changed, 116 insertions(+), 12 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cc6db6c10cc..9872783bfab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1728,6 +1728,7 @@ SENTRY_DENYLIST: Final = [ "jwt_token", "private_key", "SLACK_WEBHOOK_URL", + "ALERTING_WEBHOOK_URL", "webhook_url", "LANGFUSE_SECRET_KEY", # Email Configuration diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index d7d06387d85..94d734546be 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1485,9 +1485,9 @@ Model Info: elif self.default_webhook_url is not None: _digest_webhook = self.default_webhook_url else: - _digest_webhook = os.getenv("SLACK_WEBHOOK_URL", None) + _digest_webhook = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL") if _digest_webhook is None: - raise ValueError("Missing SLACK_WEBHOOK_URL from environment") + raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment") digest_key: Final = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}" @@ -1516,10 +1516,10 @@ Model Info: elif self.default_webhook_url is not None: slack_webhook_url = self.default_webhook_url else: - slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL", None) + slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL") if slack_webhook_url is None: - raise ValueError("Missing SLACK_WEBHOOK_URL from environment") + raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment") payload: Final = {"text": formatted_message} headers: Final = {"Content-type": "application/json"} diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5ba5e8fa1aa..0f2d97b8b1c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2541,7 +2541,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): ) alerting: list | None = Field( None, - description="List of alerting integrations. Today, just slack - `alerting: ['slack']`", + description="List of alerting integrations - e.g. `alerting: ['slack', 'webhook', 'email']`. 'slack' posts Slack-format messages to any Slack-compatible webhook (Slack, Rocket.Chat, Mattermost); 'webhook' posts structured JSON budget alerts to WEBHOOK_URL", ) alert_types: list[AlertType] | None = Field( None, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 887716a383a..da2e09bd7f6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1363,7 +1363,7 @@ _OPENAPI_HTTP_METHODS: Final = { # the UI. Kept here at module scope to match the analogous descriptor # `is_secret` flags in litellm.proxy.config_resolvers and the # `_CACHE_SENSITIVE_FIELDS` constant in the cache endpoint file. -_ALERTING_SENSITIVE_VARS: Final[set[str]] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} +_ALERTING_SENSITIVE_VARS: Final[set[str]] = {"ALERTING_WEBHOOK_URL", "SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} def _strip_operation_id_method_suffix(operation_id: str) -> str: @@ -16566,6 +16566,7 @@ async def create_config_audit_log( _EXTRA_SECRET_CALLBACK_ENV_VARS: Final = frozenset( { + "ALERTING_WEBHOOK_URL", "GALILEO_USERNAME", "GENERIC_LOGGER_HEADERS", "OTEL_HEADERS", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cf56fc0b1dd..051d36c4d0f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -645,7 +645,7 @@ class ProxyLogging: self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache) self.max_budget_limiter = _PROXY_MaxBudgetLimiter() self.cache_control_check = _PROXY_CacheControlCheck() - self.alerting: list | None = None + self.alerting: list[str] | None = None self.alerting_threshold: float = 300 # default to 5 min. threshold self.alert_types: list[AlertType] = DEFAULT_ALERT_TYPES self.alert_to_webhook_url: dict | None = None @@ -2364,7 +2364,9 @@ class ProxyLogging: # do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails) return - if self.alerting is not None and ("slack" in self.alerting or "ms_teams" in self.alerting): + if self.alerting is not None and ( + "slack" in self.alerting or "ms_teams" in self.alerting or "webhook" in self.alerting + ): if self.slack_alerting_instance is not None: await self.slack_alerting_instance.budget_alerts( type=type, diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index cfbd3e76a88..55e2dcdc270 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -12,7 +12,7 @@ import litellm from litellm.caching.caching import DualCache from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.proxy._types import CallInfo, Litellm_EntityType -from litellm.types.integrations.slack_alerting import SlackAlertingCacheKeys +from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys class TestSlackAlerting(unittest.TestCase): @@ -366,3 +366,56 @@ async def test_scheduled_daily_report_threads_the_pod_lock_manager_through(): _, kwargs = slack_alerting._run_scheduler_helper.await_args assert kwargs["pod_lock_manager"] is pod_lock_manager + + +def _slack_alerting_with_env_resolution() -> SlackAlerting: + slack_alerting: Final = SlackAlerting(alerting=["slack"], internal_usage_cache=DualCache()) + slack_alerting.periodic_started = True + return slack_alerting + + +@pytest.mark.asyncio +async def test_send_alert_falls_back_to_alerting_webhook_url_env(monkeypatch): + monkeypatch.delenv("SLACK_WEBHOOK_URL", raising=False) + monkeypatch.setenv("ALERTING_WEBHOOK_URL", "https://chat.example.com/hooks/abc") + slack_alerting: Final = _slack_alerting_with_env_resolution() + + await slack_alerting.send_alert( + message="budget crossed", + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + assert slack_alerting.log_queue[0]["url"] == "https://chat.example.com/hooks/abc" + + +@pytest.mark.asyncio +async def test_send_alert_prefers_slack_webhook_url_over_fallback(monkeypatch): + monkeypatch.setenv("SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/T0/B0/X0") + monkeypatch.setenv("ALERTING_WEBHOOK_URL", "https://chat.example.com/hooks/abc") + slack_alerting: Final = _slack_alerting_with_env_resolution() + + await slack_alerting.send_alert( + message="budget crossed", + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + assert slack_alerting.log_queue[0]["url"] == "https://hooks.slack.com/services/T0/B0/X0" + + +@pytest.mark.asyncio +async def test_send_alert_raises_when_no_webhook_url_configured(monkeypatch): + monkeypatch.delenv("SLACK_WEBHOOK_URL", raising=False) + monkeypatch.delenv("ALERTING_WEBHOOK_URL", raising=False) + slack_alerting: Final = _slack_alerting_with_env_resolution() + + with pytest.raises(ValueError, match="SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL"): + await slack_alerting.send_alert( + message="budget crossed", + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py index edce5c5f3a2..d614823c0ef 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py @@ -79,6 +79,23 @@ class TestDigestMode(unittest.IsolatedAsyncioTestCase): self.assertEqual(len(self.slack_alerting.digest_buckets), 2) + async def test_digest_falls_back_to_alerting_webhook_url_env(self): + """With SLACK_WEBHOOK_URL unset, the digest entry resolves ALERTING_WEBHOOK_URL instead.""" + env = {k: v for k, v in os.environ.items() if k != "SLACK_WEBHOOK_URL"} + env["ALERTING_WEBHOOK_URL"] = "https://chat.example.com/hooks/abc" + with unittest.mock.patch.dict(os.environ, env, clear=True): + await self.slack_alerting.send_alert( + message="`Requests are hanging`", + level="Medium", + alert_type=AlertType.llm_requests_hanging, + alerting_metadata={}, + request_model="gemini-2.5-flash", + api_base="None", + ) + + bucket = list(self.slack_alerting.digest_buckets.values())[0] + self.assertEqual(bucket["webhook_url"], "https://chat.example.com/hooks/abc") + async def test_non_digest_alert_goes_to_queue(self): """Alert types without digest enabled should go straight to the log queue.""" message = "Budget exceeded" diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py b/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py index cede859cb38..77c0f71dbf9 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py @@ -115,6 +115,35 @@ async def test_budget_alerts_slack_when_slack_alerting(proxy_logging): assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"} +@pytest.mark.asyncio +async def test_budget_alerts_webhook_only_forwards_to_slack_alerting_instance(proxy_logging): + proxy_logging.alerting = ["webhook"] + captured: Dict[str, Any] = {} + + async def fake_alert(**kwargs): + captured.update(kwargs) + + proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=fake_alert) + proxy_logging.email_logging_instance = None + await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) + snapshot = { + "type": captured["type"], + "user_info_is_callinfo": isinstance(captured["user_info"], CallInfo), + "user_id": captured["user_info"].user_id, + } + assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"} + + +@pytest.mark.asyncio +async def test_budget_alerts_email_only_skips_slack_alerting_instance(proxy_logging): + proxy_logging.alerting = ["email"] + proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock()) + proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock()) + await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) + proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called() + proxy_logging.email_logging_instance.budget_alerts.assert_called_once() + + @pytest.mark.asyncio async def test_budget_alerts_soft_budget_with_alert_emails_bypasses_global(proxy_logging): proxy_logging.alerting = None diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 05e66985e6f..72da01d0918 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -522,7 +522,8 @@ const Settings: React.FC = ({ accessToken, userRole, userID,

- Alerts are only supported for Slack Webhook URLs. Get your webhook urls from{" "} + Alerts are sent to any Slack-compatible incoming webhook URL (Slack, Rocket.Chat, Mattermost, etc.). Get + Slack webhook urls from{" "} here @@ -532,7 +533,7 @@ const Settings: React.FC = ({ accessToken, userRole, userID, - Slack Webhook URL + Webhook URL (Slack-compatible) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 137c67e837c..8a564a07489 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25182,7 +25182,7 @@ export interface components { alert_types?: components["schemas"]["AlertType"][] | null; /** * Alerting - * @description List of alerting integrations. Today, just slack - `alerting: ['slack']` + * @description List of alerting integrations - e.g. `alerting: ['slack', 'webhook', 'email']`. 'slack' posts Slack-format messages to any Slack-compatible webhook (Slack, Rocket.Chat, Mattermost); 'webhook' posts structured JSON budget alerts to WEBHOOK_URL */ alerting?: unknown[] | null; /**