diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 7e4cc2d6100..521b4315e6e 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 23914 + "limit": 22947 }, "reportArgumentType": { - "limit": 2580 + "limit": 2579 }, "reportAssignmentType": { "limit": 323 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 7573 + "limit": 7312 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5719 + "limit": 5707 }, "reportMissingTypeArgument": { - "limit": 15657 + "limit": 15642 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,22 +99,22 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44832 + "limit": 44776 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39269 + "limit": 39237 }, "reportUnknownParameterType": { - "limit": 19988 + "limit": 19969 }, "reportUnknownVariableType": { - "limit": 30923 + "limit": 30881 }, "reportUnnecessaryCast": { - "limit": 118 + "limit": 117 }, "reportUnnecessaryComparison": { "limit": 699 diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 90a82e8fa28..a62a2b0c724 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -10,7 +10,7 @@ A2A Streaming Events (in order): 4. Status update (kind: "status-update") - Final status "completed" with final=true """ -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Callable, Coroutine, Mapping from typing import Any, Final import litellm @@ -54,7 +54,7 @@ class A2ACompletionBridgeHandler: agent_extra_headers: Mapping[str, str] | None, *, stream: bool, - ) -> Mapping[str, Any]: + ) -> Mapping[str, object]: # Extract message from params message: Final = params.get("message", {}) @@ -63,7 +63,7 @@ class A2ACompletionBridgeHandler: # Get completion params custom_llm_provider: Final = litellm_params.get("custom_llm_provider") - model: Final = litellm_params.get("model", "agent") + model: Final[str] = litellm_params.get("model", "agent") # Build full model string if provider specified # Skip prepending if model already starts with the provider prefix @@ -109,13 +109,16 @@ class A2ACompletionBridgeHandler: return completion_params @staticmethod - async def _acompletion(completion_params: Mapping[str, Any]) -> ModelResponse | CustomStreamWrapper: - return await litellm.acompletion(**completion_params) + async def _acompletion(completion_params: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper: + acompletion_fn: Final[Callable[..., Coroutine[object, object, ModelResponse | CustomStreamWrapper]]] = vars( + litellm + )["acompletion"] + return await acompletion_fn(**completion_params) @staticmethod async def handle_non_streaming( request_id: str, - params: dict[str, Any], + params: dict[str, object], litellm_params: dict[str, Any], api_base: str | None = None, agent_extra_headers: dict[str, str] | None = None, @@ -296,8 +299,8 @@ class A2ACompletionBridgeHandler: # Convenience functions that delegate to the class methods async def handle_a2a_completion( request_id: str, - params: dict[str, Any], - litellm_params: dict[str, Any], + params: dict[str, object], + litellm_params: dict[str, object], api_base: str | None = None, agent_extra_headers: dict[str, str] | None = None, ) -> dict[str, object]: @@ -313,8 +316,8 @@ async def handle_a2a_completion( async def handle_a2a_completion_streaming( request_id: str, - params: dict[str, Any], - litellm_params: dict[str, Any], + params: dict[str, object], + litellm_params: dict[str, object], api_base: str | None = None, agent_extra_headers: dict[str, str] | None = None, ) -> AsyncIterator[dict[str, object]]: diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 322393cd9c4..1c6ebf0b95c 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -12,7 +12,8 @@ Provides standalone functions with @client decorator for LiteLLM logging integra import asyncio import datetime import uuid -from collections.abc import AsyncIterator, Coroutine +from collections.abc import AsyncIterator, Coroutine, Mapping +from types import ModuleType from typing import TYPE_CHECKING, Any, Final, Optional, cast import litellm @@ -38,12 +39,15 @@ if TYPE_CHECKING: SendMessageResponse, SendStreamingMessageRequest, SendStreamingMessageResponse, + SendStreamingMessageSuccessResponse, Task, ) + from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest + from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse # Runtime imports — requires a2a-sdk>=1.1.0 A2A_SDK_AVAILABLE = False -_a2a_conversions: Any = None +_a2a_conversions: ModuleType | None = None try: from a2a.client import Client, ClientCallContext, ClientConfig, create_client @@ -128,7 +132,7 @@ _A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output def _set_litellm_params_on_logging_obj( kwargs: dict[str, Any], - litellm_params: dict[str, Any], + litellm_params: Mapping[str, object], ) -> None: """ Merge the agent's pricing params into model_call_details["litellm_params"] @@ -150,7 +154,7 @@ def _set_litellm_params_on_logging_obj( logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params} -def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str: +def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) -> str: """ Extract agent info and set model/custom_llm_provider for cost tracking. @@ -179,7 +183,7 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str: return agent_name -def _get_a2a_client_agent_card(a2a_client: Any) -> Optional["AgentCard"]: +def _get_a2a_client_agent_card(a2a_client: "A2AClientType") -> Optional["AgentCard"]: agent_card = cast(Optional["AgentCard"], getattr(a2a_client, "_litellm_agent_card", None)) if agent_card is not None: return agent_card @@ -191,9 +195,9 @@ def _get_a2a_client_agent_card(a2a_client: Any) -> Optional["AgentCard"]: async def _send_message_via_completion_bridge( request: "SendMessageRequest", - custom_llm_provider: str, + custom_llm_provider: object, api_base: str | None, - litellm_params: dict[str, Any], + litellm_params: dict[str, object], agent_extra_headers: dict[str, str] | None = None, ) -> LiteLLMSendMessageResponse: """ @@ -224,6 +228,20 @@ def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallConte return getattr(a2a_client, "_litellm_call_context", None) +def _to_core_send_message_request(request: "SendMessageRequest") -> "CoreSendMessageRequest": + from a2a.compat.v0_3 import conversions + + return conversions.to_core_send_message_request(request) + + +def _to_compat_stream_response( + event: "CoreStreamResponse", request_id: str | int +) -> "SendStreamingMessageSuccessResponse": + from a2a.compat.v0_3 import conversions + + return conversions.to_compat_stream_response(event, request_id=request_id) + + async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse": """Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response.""" if _a2a_conversions is None: @@ -231,17 +249,14 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - pb_request: Final = _a2a_conversions.to_core_send_message_request(request) + pb_request: Final = _to_core_send_message_request(request) last_event = None async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)): last_event = event if last_event is None: raise RuntimeError("A2A send_message failed: no response received from agent.") - stream_compat: Final = _a2a_conversions.to_compat_stream_response( - last_event, - request_id=request.id, - ) + stream_compat: Final = _to_compat_stream_response(last_event, request_id=request.id) result: Final = stream_compat.result if not isinstance(result, (Message, Task)): raise RuntimeError( @@ -306,12 +321,9 @@ async def _stream_messages( "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - pb_request: Final = _a2a_conversions.to_core_send_message_request(request) + pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request) async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)): - compat_chunk = _a2a_conversions.to_compat_stream_response( - event, - request_id=request.id, - ) + compat_chunk = _to_compat_stream_response(event, request_id=request.id) yield SendStreamingMessageResponse(root=compat_chunk) @@ -368,10 +380,10 @@ async def asend_message( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendMessageRequest"] = None, api_base: str | None = None, - litellm_params: dict[str, Any] | None = None, + litellm_params: dict[str, object] | None = None, agent_id: str | None = None, agent_extra_headers: dict[str, str] | None = None, - **kwargs: Any, + **kwargs: object, ) -> LiteLLMSendMessageResponse: """ Async: Send a message to an A2A agent. @@ -485,7 +497,7 @@ async def asend_message( response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id)) # Calculate token usage from request and response - response_dict: Final = a2a_response.model_dump(mode="json", exclude_none=True) + response_dict: Final[dict[str, object]] = a2a_response.model_dump(mode="json", exclude_none=True) ( prompt_tokens, completion_tokens, @@ -516,7 +528,7 @@ def send_message( a2a_client: "A2AClientType", request: "SendMessageRequest", **kwargs: Any, -) -> LiteLLMSendMessageResponse | Coroutine[Any, Any, LiteLLMSendMessageResponse]: +) -> LiteLLMSendMessageResponse | Coroutine[object, object, LiteLLMSendMessageResponse]: """ Sync: Send a message to an A2A agent. @@ -545,9 +557,9 @@ def _build_streaming_logging_obj( request: "SendStreamingMessageRequest", agent_name: str, agent_id: str | None, - litellm_params: dict[str, Any] | None, - metadata: dict[str, Any] | None, - proxy_server_request: dict[str, Any] | None, + litellm_params: dict[str, object] | None, + metadata: dict[str, object] | None, + proxy_server_request: dict[str, object] | None, ) -> Logging: """Build logging object for streaming A2A requests.""" start_time: Final = datetime.datetime.now() @@ -588,10 +600,10 @@ async def asend_message_streaming( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendStreamingMessageRequest"] = None, api_base: str | None = None, - litellm_params: dict[str, Any] | None = None, + litellm_params: dict[str, object] | None = None, agent_id: str | None = None, - metadata: dict[str, Any] | None = None, - proxy_server_request: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + proxy_server_request: dict[str, object] | None = None, agent_extra_headers: dict[str, str] | None = None, **kwargs: object, ) -> AsyncIterator[Any]: diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 4f127f476c3..e43e0dfd5f7 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,6 +1,8 @@ import json from collections.abc import AsyncIterator, Iterator -from typing import Any, Final, cast +from typing import Any, Final, TypedDict, cast + +from typing_extensions import ReadOnly from litellm import verbose_logger from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema @@ -28,6 +30,19 @@ from litellm.types.utils import ( ) +class _GenAITextPart(TypedDict, total=False): + text: ReadOnly[str] + + +class _GenAISystemInstruction(TypedDict, total=False): + parts: ReadOnly[list[_GenAITextPart]] + + +class _GenAIPart(TypedDict, total=False): + text: ReadOnly[str] + functionCall: ReadOnly[dict[str, object]] + + class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ Wrapper for streaming Google GenAI generate_content responses. @@ -36,9 +51,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): sent_first_chunk: bool = False # State tracking for accumulating partial tool calls - accumulated_tool_calls: dict[str, dict[str, Any]] + accumulated_tool_calls: dict[str, dict[str, str]] - def __init__(self, completion_stream: Any): + def __init__(self, completion_stream: object): self.sent_first_chunk = False self.accumulated_tool_calls = {} self._returned_response = False @@ -85,7 +100,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, @@ -94,7 +109,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): # 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 = { + function_call_part: _GenAIPart = { "functionCall": { "name": tool_call_data["name"] or "undefined_tool_name", "args": parsed_args, @@ -110,7 +125,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): tool_call_data["arguments"], ) if parts: - final_chunk: Final = { + final_chunk: Final[dict[str, object]] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -273,9 +288,9 @@ class GoogleGenAIAdapter: def _add_generic_litellm_params_to_request( self, - completion_request_dict: dict[str, Any], + completion_request_dict: dict[str, object], litellm_params: GenericLiteLLMParams | None = None, - ) -> dict: + ) -> dict[str, object]: """Add generic litellm params to request. e.g add api_base, api_key, api_version, etc. Args: @@ -295,7 +310,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) @@ -307,12 +322,12 @@ class GoogleGenAIAdapter: tools: list[dict[str, Any]], ) -> list[ChatCompletionToolParam]: """Transform Google GenAI tools to OpenAI tools format""" - openai_tools: Final[list[dict[str, Any]]] = [] + openai_tools: Final[list[dict[str, object]]] = [] for tool in tools: if "functionDeclarations" in tool: for func_decl in tool["functionDeclarations"]: - function_chunk: dict[str, Any] = { + function_chunk: dict[str, object] = { "name": func_decl.get("name", ""), } @@ -321,7 +336,7 @@ class GoogleGenAIAdapter: if "parametersJsonSchema" in func_decl: function_chunk["parameters"] = func_decl["parametersJsonSchema"] - openai_tool = {"type": "function", "function": function_chunk} + openai_tool: dict[str, object] = {"type": "function", "function": function_chunk} openai_tools.append(openai_tool) # normalize the tool schemas @@ -345,7 +360,7 @@ 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]] = [] @@ -461,7 +476,7 @@ class GoogleGenAIAdapter: def translate_completion_to_generate_content( self, response: ModelResponse, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform litellm completion response to Google GenAI generate_content format @@ -490,7 +505,7 @@ class GoogleGenAIAdapter: parts = [{"text": message_content}] if message_content else [] # Create Google GenAI format response - generate_content_response: Final[dict[str, Any]] = { + generate_content_response: Final[dict[str, object]] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -524,7 +539,7 @@ class GoogleGenAIAdapter: self, response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """ Transform streaming litellm completion chunk to Google GenAI generate_content format @@ -560,7 +575,7 @@ class GoogleGenAIAdapter: return None # Create Google GenAI streaming format response - streaming_chunk: Final[dict[str, Any]] = { + streaming_chunk: Final[dict[str, object]] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -597,9 +612,9 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, message: Any, - ) -> list[dict[str, Any]]: + ) -> list[_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: @@ -614,7 +629,7 @@ class GoogleGenAIAdapter: except json.JSONDecodeError: args = {} - function_call_part = { + function_call_part: _GenAIPart = { "functionCall": { "name": tool_call.function.name or "undefined_tool_name", "args": args, @@ -626,14 +641,14 @@ class GoogleGenAIAdapter: def _transform_openai_delta_to_google_genai_parts_with_accumulation( self, delta: Any, wrapper: GoogleGenAIStreamWrapper - ) -> list[dict[str, Any]]: + ) -> list[_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}) @@ -686,7 +701,7 @@ class GoogleGenAIAdapter: # 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 diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 97e831f5822..a474a11601d 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -6,12 +6,13 @@ 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 types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict import httpx +from typing_extensions import Never, ReadOnly from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -48,7 +49,20 @@ _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 _ServiceToolCall(TypedDict): + id: ReadOnly[str] + + +class _ServiceMessage(TypedDict, total=False): + content: ReadOnly[str] + tool_calls: ReadOnly[Sequence[_ServiceToolCall]] + + +class _ServiceChoice(TypedDict, total=False): + message: ReadOnly[_ServiceMessage] class _MalformedToolBlockingResponseError(Exception): @@ -143,7 +157,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 +205,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 +226,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 +267,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod async def _guarded( - coro: Any, + coro: Awaitable[GenericGuardrailAPIInputs], inputs: GenericGuardrailAPIInputs, label: str, ) -> GenericGuardrailAPIInputs: @@ -400,7 +414,7 @@ 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[object]) -> tuple[ChatCompletionMessageToolCall, ...]: """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) @@ -427,7 +441,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,14 +453,14 @@ 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[dict[str, object]] = { "role": "assistant", "content": content or None, } @@ -467,7 +481,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]: + def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]: """Collapse each message's content to a plain string for the webhook. litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, @@ -506,8 +520,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,7 +530,7 @@ 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]] = { + payload: Final[dict[str, object]] = { "model": inputs.get("model") or request_data.get("model") or "", "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), } @@ -540,8 +554,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _extract_request_data( call_details: Mapping[str, Any], - request_data: Mapping[str, Any] | None, - ) -> Mapping[str, Any]: + request_data: Mapping[str, object] | None, + ) -> Mapping[str, object]: """Extract original request data from model_call_details for the response moderation service envelope. @@ -576,7 +590,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: + def _sanitize_proxy_server_request(proxy_server_request: object) -> 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,17 +600,18 @@ 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: Mapping[str, object], call_details: Mapping[str, str]) -> str: """Get the model name for the ModifyResponseException.""" response: Final = request_data.get("response") if response and hasattr(response, "model"): - return response.model or "unknown" + response_model: Final[str | None] = getattr(response, "model", None) + return response_model or "unknown" return call_details.get("model", "unknown") # -- 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: Mapping[str, str], request_data: Mapping[str, str] | 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 +625,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: dict[str, object], source: Mapping[str, str]) -> 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 +645,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: dict[str, object], source: Mapping[str, object]) -> None: """Prepend ``source["system"]`` onto ``payload["messages"]``. Builds a NEW messages list rather than mutating ``payload["messages"]`` @@ -658,7 +673,9 @@ 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: Mapping[str, object], 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) @@ -697,7 +714,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: Mapping[str, object], event_type: str): try: payload: Final = await self._prepare_log_payload(kwargs, event_type) if payload is None: @@ -862,7 +879,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): base: Final = call_details.get("standard_logging_object") if base is not None: - payload: dict = safe_deep_copy(base) + payload: dict[str, object] = safe_deep_copy(base) else: verbose_logger.debug( "Rubrik: standard_logging_object not yet on model_call_details " @@ -908,7 +925,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): cls, call_details: Mapping[str, Any], user_api_key_dict: "UserAPIKeyAuth", - ) -> dict[str, Any]: + ) -> dict[str, object]: # 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") @@ -996,7 +1013,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) -> Mapping[str, Any]: """POST ``payload`` to a Rubrik webhook and return its dict response. Raises: @@ -1010,7 +1027,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): headers=self._headers, ) http_response.raise_for_status() - result: Final = http_response.json() + result: Final[object] = http_response.json() if not isinstance(result, dict): raise TypeError( f"{service_name} returned non-dict JSON " @@ -1021,8 +1038,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): async def _post_to_response_moderation_endpoint( self, - response_data: Mapping[str, Any], - request_data: Mapping[str, Any], + response_data: Mapping[str, object], + request_data: Mapping[str, object], ) -> Mapping[str, Any]: """Post the ``{request, response}`` envelope to the after_completion webhook and return its (possibly rewritten) response. @@ -1039,7 +1056,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]) -> Mapping[str, Any]: """Post a bare OpenAI request to the before_prompt webhook. Returns ``{}`` (passthrough) or a synthetic chat.completion (block). @@ -1054,7 +1071,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): chat.completion whose ``choices[0].message.content`` is the refusal explanation. """ - choices: Final = service_response.get("choices") + choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices") if not choices: return None message: Final = choices[0].get("message") or _EMPTY_MAPPING @@ -1086,7 +1103,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Expects service_response in OpenAI chat completion format: {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} """ - choices: Final = service_response.get("choices") or () + choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or () if not choices: raise _MalformedToolBlockingResponseError("Response moderation service returned empty response") diff --git a/litellm/interactions/litellm_responses_transformation/handler.py b/litellm/interactions/litellm_responses_transformation/handler.py index f826c980ac4..8fee0fcd1b5 100644 --- a/litellm/interactions/litellm_responses_transformation/handler.py +++ b/litellm/interactions/litellm_responses_transformation/handler.py @@ -2,8 +2,8 @@ Handler for transforming interactions API requests to litellm.responses requests. """ -from collections.abc import AsyncIterator, Coroutine, Iterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator +from typing import Any, Final import litellm from litellm.interactions.litellm_responses_transformation.streaming_iterator import ( @@ -37,7 +37,7 @@ class LiteLLMResponsesInteractionsHandler: ) -> ( InteractionsAPIResponse | Iterator[InteractionsAPIStreamingResponse] - | Coroutine[Any, Any, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]] + | Coroutine[object, object, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]] ): """ Handle Interactions API request by calling litellm.responses(). @@ -55,13 +55,15 @@ class LiteLLMResponsesInteractionsHandler: InteractionsAPIResponse or streaming iterator """ # Transform interactions request to responses request - responses_request = LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request( - model=model, - input=input, - optional_params=optional_params, - custom_llm_provider=custom_llm_provider, - stream=stream, - **kwargs, + responses_request: Final = ( + LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request( + model=model, + input=input, + optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + stream=stream, + **kwargs, + ) ) if _is_async: @@ -76,7 +78,10 @@ class LiteLLMResponsesInteractionsHandler: # Call litellm.responses() # Note: litellm.responses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] # but the type checker may see it as a coroutine in some contexts - responses_response: Final = litellm.responses( + responses_fn: Final[Callable[..., ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]] = vars(litellm)[ + "responses" + ] + responses_response: Final = responses_fn( **responses_request, ) @@ -92,8 +97,7 @@ class LiteLLMResponsesInteractionsHandler: ) # At this point, responses_response must be ResponsesAPIResponse (not streaming) - # Cast to satisfy type checker since we've already checked it's not a streaming iterator - responses_api_response: Final = cast(ResponsesAPIResponse, responses_response) + responses_api_response: Final = responses_response # Transform responses response to interactions response return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response( @@ -112,7 +116,10 @@ class LiteLLMResponsesInteractionsHandler: """Async handler for interactions API requests.""" # Call litellm.aresponses() # Note: litellm.aresponses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] - responses_response: Final = await litellm.aresponses( + aresponses_fn: Final[ + Callable[..., Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]] + ] = vars(litellm)["aresponses"] + responses_response: Final = await aresponses_fn( **responses_request, ) @@ -128,8 +135,7 @@ class LiteLLMResponsesInteractionsHandler: ) # At this point, responses_response must be ResponsesAPIResponse (not streaming) - # Cast to satisfy type checker since we've already checked it's not a streaming iterator - responses_api_response: Final = cast(ResponsesAPIResponse, responses_response) + responses_api_response: Final = responses_response # Transform responses response to interactions response return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_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..a7c462a8fb0 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 Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, Optional, TypedDict, Union, cast + +from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger @@ -29,9 +31,8 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth 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 @@ -534,7 +535,7 @@ def _augment_system_with_summary( return [{"type": "text", "text": prefix.rstrip()}, *system] -def _resolve_trigger_tokens(edit_spec: dict[str, object]) -> tuple[int, list[str]]: +def _resolve_trigger_tokens(edit_spec: Mapping[str, object]) -> tuple[int, list[str]]: """Validate and resolve ``trigger.value``. Raises ``AnthropicContextManagementError`` if the explicitly-supplied value @@ -568,7 +569,7 @@ def _resolve_trigger_tokens(edit_spec: dict[str, object]) -> tuple[int, list[str return value, warnings -def _build_summary_prompt(edit_spec: dict[str, object], tools: list[dict[str, object]] | None) -> str: +def _build_summary_prompt(edit_spec: Mapping[str, object], tools: Sequence[Mapping[str, object]] | None) -> str: custom: Final = edit_spec.get("instructions") if isinstance(custom, str) and custom.strip(): return custom @@ -623,7 +624,7 @@ def _count_effective_tokens( try: openai_shape = adapter.translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", messages_without_compaction, ) ) @@ -736,7 +737,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: +) -> dict[str, object] | None: """Translate Anthropic-shaped ``system`` to an OpenAI system message. Accepts a bare string or a list of Anthropic content blocks; returns @@ -773,7 +774,7 @@ def _build_summary_messages( try: openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", stripped, ) ) @@ -809,7 +810,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,10 +821,29 @@ 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[list[object]] = [*content, {"type": "text", "text": extra_text}] + return appended return [content, {"type": "text", "text": extra_text}] +class _SummaryCallUserKwarg(TypedDict, total=False): + user: ReadOnly[object] + + +class _SummaryCallRegionKwarg(TypedDict, total=False): + allowed_model_region: ReadOnly[str] + + +class _SummaryCallKwargs(TypedDict): + model: ReadOnly[str] + messages: ReadOnly[list[dict[str, object]]] + max_tokens: ReadOnly[int] + timeout: ReadOnly[float] + litellm_metadata: ReadOnly[Mapping[str, object]] + user: NotRequired[ReadOnly[object]] + allowed_model_region: NotRequired[ReadOnly[str]] + + async def _call_summary_model( *, summary_model: str, @@ -860,22 +880,24 @@ 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]] = { - "model": summary_model, - "messages": summary_messages, - "max_tokens": max_tokens, - "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, - "litellm_metadata": metadata, - } # The end-user id must also travel as the top-level ``user`` kwarg: legacy # limiter hooks and prometheus end-user tracking read it from there rather # 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: - call_kwargs["user"] = end_user_id - if allowed_model_region is not None: - call_kwargs["allowed_model_region"] = allowed_model_region + call_kwargs: Final[_SummaryCallKwargs] = { + "model": summary_model, + "messages": summary_messages, + "max_tokens": max_tokens, + "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, + "litellm_metadata": metadata, + **(_SummaryCallUserKwarg(user=end_user_id) if end_user_id else _SummaryCallUserKwarg()), + **( + _SummaryCallRegionKwarg(allowed_model_region=allowed_model_region) + if allowed_model_region is not None + else _SummaryCallRegionKwarg() + ), + } if llm_router is not None and hasattr(llm_router, "acompletion"): return await llm_router.acompletion(**call_kwargs) return await litellm.acompletion(**call_kwargs) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 1ce83e226e7..b77ba2f9460 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -2,11 +2,12 @@ import asyncio import hashlib import json import os -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Any, Final, Literal, NamedTuple, cast import httpx from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -23,6 +24,22 @@ from litellm.utils import _add_path_to_api_base azure_ad_cache: Final = DualCache() +class _AzureAdTokenJson(TypedDict, total=False): + access_token: ReadOnly[str] + expires_in: ReadOnly[int] + + +class _AzureV1ClientParams(TypedDict, total=False, extra_items=object): + base_url: ReadOnly[str] + + +class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object): + api_version: ReadOnly[str] + base_url: ReadOnly[str] + max_retries: ReadOnly[int] + timeout: ReadOnly[float | httpx.Timeout] + + class AzureOpenAIError(BaseLLMException): def __init__( self, @@ -220,7 +237,7 @@ def get_azure_ad_token_from_oidc( message=req_token.text, ) - azure_ad_token_json: Final = req_token.json() + azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json() azure_ad_token_access_token = azure_ad_token_json.get("access_token", None) azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None) @@ -486,7 +503,7 @@ class BaseAzureLLM(BaseOpenAILLM): v1_api_key = _async_v1_api_key - v1_params: Final[dict[str, Any]] = { + v1_params: Final[_AzureV1ClientParams] = { "api_key": v1_api_key, "base_url": f"{api_base}/openai/v1/", } @@ -643,7 +660,7 @@ class BaseAzureLLM(BaseOpenAILLM): api_base += "/" api_base += f"{model}" - azure_client_params: Final[dict[str, Any]] = { + azure_client_params: Final[_AzureGatewayClientParams] = { "api_version": api_version, "base_url": f"{api_base}", "http_client": litellm.client_session, @@ -702,7 +719,7 @@ class BaseAzureLLM(BaseOpenAILLM): @staticmethod def _get_base_azure_url( api_base: str | None, - litellm_params: GenericLiteLLMParams | dict[str, Any] | None, + litellm_params: GenericLiteLLMParams | Mapping[str, object] | None, route: Literal["/openai/responses", "/openai/vector_stores"] | str, default_api_version: str | Literal["latest", "preview"] | None = None, ) -> str: @@ -757,7 +774,9 @@ class BaseAzureLLM(BaseOpenAILLM): return False return api_version in {"preview", "latest", "v1"} - def _resolve_env_var(self, litellm_params: dict[str, Any], param_key: str, env_var_key: str) -> str | None: + def _resolve_env_var( + self, litellm_params: Mapping[str, str | None], param_key: str, env_var_key: str + ) -> str | None: """Resolve the environment variable for a given parameter key. The logic here is different from `params.get(key, os.getenv(env_var))` because diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 4ff7323c33f..b50a9ae04d1 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -2,17 +2,18 @@ import base64 import json import os import time -from collections.abc import Iterable, Mapping, MutableMapping +from collections.abc import Iterable, Mapping, MutableMapping, Sequence from functools import cache from itertools import chain from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, TypeAlias, TypedDict from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, ConfigDict, TypeAdapter +from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -63,10 +64,39 @@ from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resol S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers" -def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]: +def _frozen_mapping(items: Iterable[tuple[str, object]]) -> Mapping[str, object]: return MappingProxyType(dict(items)) +_EmbeddingBatchInput: TypeAlias = ( + str | int | float | Sequence[str] | Sequence[int] | Sequence[Sequence[int]] | Mapping[str, object] +) + + +class _OpenAIBatchRecordBody(TypedDict, total=False): + model: ReadOnly[str] + prompt: ReadOnly[str | Sequence[str] | Sequence[int] | Sequence[Sequence[int]]] + input: ReadOnly[_EmbeddingBatchInput] + metadata: ReadOnly[Mapping[str, object]] + + +class _OpenAIBatchRecord(TypedDict, total=False): + custom_id: ReadOnly[str] + url: ReadOnly[str] + body: ReadOnly[_OpenAIBatchRecordBody] + + +class _BedrockBatchRecord(TypedDict): + recordId: ReadOnly[str] + modelInput: ReadOnly[Mapping[str, object]] + + +class _S3UploadResponse(TypedDict, total=False): + Key: ReadOnly[str] + Bucket: ReadOnly[str] + ContentLength: ReadOnly[int] + + # JSONL batch records are untyped json, so the `/v1/responses` fields are # validated into their concrete Responses API types before being handed to the # Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't @@ -231,7 +261,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def _get_s3_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: Sequence[_OpenAIBatchRecord], ) -> str: """ Gets a unique S3 object name for the Bedrock batch processing job @@ -341,7 +371,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): OPENAI_RESPONSES_URL = "/v1/responses" @staticmethod - def _classify_batch_record(openai_jsonl_record: Mapping[str, Any]) -> BedrockBatchRecordKind: + def _classify_batch_record(openai_jsonl_record: _OpenAIBatchRecord) -> BedrockBatchRecordKind: """ Decide which OpenAI endpoint shape an OpenAI batch JSONL line carries. @@ -484,7 +514,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): return value if isinstance(value, str) and value else None @staticmethod - def _coerce_embedding_input_to_string(raw_input: Any, model: str = "") -> str: + def _coerce_embedding_input_to_string(raw_input: _EmbeddingBatchInput | None, model: str = "") -> str: """ Normalize an OpenAI /v1/embeddings `input` field into the single string that Bedrock Titan v2 InvokeModel expects in `inputText`. @@ -541,8 +571,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def _map_openai_embedding_to_bedrock_params( self, - openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + openai_request_body: _OpenAIBatchRecordBody, + ) -> dict[str, object]: """ Transform an OpenAI /v1/embeddings request body into the Bedrock InvokeModel `modelInput` for embedding models that AWS @@ -588,7 +618,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): return dict(titan_config._transform_request(input=input_text, inference_params=inference_params)) @staticmethod - def _transform_text_completion_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + def _transform_text_completion_body_to_chat_body( + openai_request_body: _OpenAIBatchRecordBody, + ) -> Mapping[str, object]: """ Rewrite an OpenAI `/v1/completions` batch body as a Chat Completions body. @@ -610,7 +642,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) @staticmethod - def _transform_responses_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + def _transform_responses_body_to_chat_body(openai_request_body: _OpenAIBatchRecordBody) -> Mapping[str, object]: """ Rewrite an OpenAI `/v1/responses` batch body as a Chat Completions body. @@ -631,23 +663,25 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): "Batch record for /v1/responses is missing required `input` field: " f"model={openai_request_body.get('model', '')}" ) - chat_body: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( - model=openai_request_body.get("model", ""), - input=_responses_input_adapter().validate_python(responses_input), - responses_api_request=_responses_request_adapter().validate_python( - _frozen_mapping( - (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") - ) - ), - metadata=openai_request_body.get("metadata"), + chat_body: Final[Mapping[str, object]] = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=openai_request_body.get("model", ""), + input=_responses_input_adapter().validate_python(responses_input), + responses_api_request=_responses_request_adapter().validate_python( + _frozen_mapping( + (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") + ) + ), + metadata=openai_request_body.get("metadata"), + ) ) return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) @staticmethod def _transform_batch_body_to_chat_body( - openai_request_body: Mapping[str, Any], + openai_request_body: _OpenAIBatchRecordBody, record_kind: BedrockBatchRecordKind, - ) -> Mapping[str, Any]: + ) -> Mapping[str, object]: """ Normalize a non-embedding batch body to the Chat Completions shape the per-provider Bedrock transformations expect. @@ -666,7 +700,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): self, openai_request_body: Mapping[str, Any], provider: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform OpenAI request body to Bedrock-compatible modelInput parameters using existing transformation logic. @@ -677,7 +711,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ from litellm.types.utils import LlmProviders - _model: Final = openai_request_body.get("model", "") + _model: Final[str] = openai_request_body.get("model", "") messages: Final = openai_request_body.get("messages", []) optional_params: Final = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} @@ -733,8 +767,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): } def _transform_openai_jsonl_content_to_bedrock_jsonl_content( - self, openai_jsonl_content: list[dict[str, Any]] - ) -> list[dict[str, Any]]: + self, openai_jsonl_content: Sequence[_OpenAIBatchRecord] + ) -> list[_BedrockBatchRecord]: """ Transforms OpenAI JSONL content to Bedrock batch format @@ -1026,7 +1060,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): response_headers: Final = raw_response.headers # Extract S3 object information from the response # S3 PUT object returns ETag and other metadata in headers - content_length: Final = response_headers.get("Content-Length", "0") + content_length: Final[str] = response_headers.get("Content-Length", "0") # Use the actual upload URL that was used for the S3 upload upload_url: Final = litellm_params.get("upload_url") @@ -1224,7 +1258,9 @@ class BedrockJsonlFilesTransformation: object_name: Final = self._get_s3_object_name(openai_jsonl_content=openai_jsonl_content) return bedrock_jsonl_string, object_name - def _transform_openai_jsonl_content_to_bedrock_jsonl_content(self, openai_jsonl_content: list[dict[str, Any]]): + def _transform_openai_jsonl_content_to_bedrock_jsonl_content( + self, openai_jsonl_content: Sequence[_OpenAIBatchRecord] + ): """ Delegate to the main BedrockFilesConfig transformation method """ @@ -1233,7 +1269,7 @@ class BedrockJsonlFilesTransformation: def _get_s3_object_name( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: Sequence[_OpenAIBatchRecord], ) -> str: """ Gets a unique S3 object name for the Bedrock batch processing job @@ -1285,7 +1321,7 @@ class BedrockJsonlFilesTransformation: return content def transform_s3_bucket_response_to_openai_file_object( - self, create_file_data: CreateFileRequest, s3_upload_response: dict[str, Any] + self, create_file_data: CreateFileRequest, s3_upload_response: _S3UploadResponse ) -> OpenAIFileObject: """ Transforms S3 Bucket upload file response to OpenAI FileObject diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 2e0ae30a192..b8e57fa7cc0 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 import httpx from httpx._types import RequestFiles +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION @@ -31,6 +33,29 @@ else: LiteLLMLoggingObj = Any +class _RunwayTaskResponse(TypedDict, total=False): + id: ReadOnly[str] + status: ReadOnly[str] + createdAt: ReadOnly[str] + completedAt: ReadOnly[str] + output: ReadOnly[Sequence[str] | str] + failureCode: ReadOnly[str] + failure: ReadOnly[str] + progress: ReadOnly[int] + + +class _VideoObjectData(TypedDict, extra_items=object): + id: ReadOnly[str] + object: ReadOnly[Literal["video"]] + status: ReadOnly[str] + created_at: ReadOnly[int] + + +def _parse_runway_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse: + response_data: Final[_RunwayTaskResponse] = raw_response.json() + return response_data + + class RunwayMLVideoConfig(BaseVideoConfig): """ Configuration class for RunwayML video generation. @@ -78,7 +103,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Handle input_reference parameter - map to promptImage if "input_reference" in video_create_optional_params: @@ -180,7 +205,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): } """ # Build the request data - request_data: Final[dict[str, Any]] = { + request_data: Final[dict[str, object]] = { "model": model, "promptText": prompt, } @@ -189,7 +214,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): 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[RequestFiles] = [] # Append the specific endpoint for video generation full_api_base: Final = f"{api_base}/image_to_video" @@ -216,10 +241,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): We map this to OpenAI VideoObject format. """ - response_data: Final = raw_response.json() + response_data: Final = _parse_runway_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_VideoObjectData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -326,7 +351,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Get task status to retrieve video URL url: Final = f"{api_base}/tasks/{encoded_video_id}" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, str]] = {} return url, params @@ -421,7 +446,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 +473,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,7 +509,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]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -494,7 +519,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): logging_obj: LiteLLMLoggingObj, ) -> VideoObject: """Transform the RunwayML video delete/cancel response.""" - response_data: Final = raw_response.json() + response_data: Final = _parse_runway_task_response(raw_response) video_obj: Final = VideoObject( id=response_data.get("id", ""), @@ -524,7 +549,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]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -537,10 +562,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ Transform the RunwayML video status retrieve response. """ - response_data: Final = raw_response.json() + response_data: Final = _parse_runway_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_VideoObjectData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -572,7 +597,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): 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..3db94211032 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 +from typing import Any, Final, TypedDict import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from typing_extensions import ReadOnly import litellm from litellm._uuid import uuid @@ -50,6 +51,7 @@ from litellm.types.llms.openai import ( HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, + OpenAIFilesPurpose, PathLike, ) from litellm.types.llms.vertex_ai import GcsBucketResponse @@ -62,6 +64,46 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +class _GcsObjectMetadataJson(TypedDict, total=False): + purpose: ReadOnly[OpenAIFilesPurpose] + + +class _GcsObjectJson(TypedDict, total=False): + id: ReadOnly[str] + name: ReadOnly[str] + size: ReadOnly[str] + timeCreated: ReadOnly[str] + metadata: ReadOnly[_GcsObjectMetadataJson] + + +class _VertexBatchRowRequest(TypedDict, total=False): + labels: ReadOnly[Mapping[str, object]] + + +class _VertexBatchRow(TypedDict, total=False): + request: ReadOnly[_VertexBatchRowRequest] + status: ReadOnly[str] + processed_time: ReadOnly[str] + + +class _OpenAIBatchOutputError(TypedDict): + code: ReadOnly[str] + message: ReadOnly[str] + + +class _OpenAIBatchOutputResponse(TypedDict): + status_code: ReadOnly[int] + request_id: ReadOnly[str] + body: ReadOnly[Mapping[str, object]] + + +class _OpenAIBatchOutputRow(TypedDict): + id: ReadOnly[str] + custom_id: ReadOnly[str] + response: ReadOnly[_OpenAIBatchOutputResponse | None] + error: ReadOnly[_OpenAIBatchOutputError | None] + + def _sanitize_gcp_label_value(value: str) -> str: """ Sanitize a string to meet GCP label value constraints. @@ -106,7 +148,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 _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: object) -> None: """ Store OpenAI batch custom_id for Vertex batch correlation. @@ -122,7 +164,7 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" raw: Final = labels.get("litellm_custom_id_raw") if raw: @@ -186,7 +228,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] @@ -246,6 +288,11 @@ def _iter_openai_jsonl_entries( yield json.loads(line) +def _parse_vertex_batch_output_row(line: str) -> _VertexBatchRow: + row: Final[_VertexBatchRow] = json.loads(line) + return row + + class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): """Streams an OpenAI batch JSONL upload as Vertex-wrapped JSONL one row at a time, so the transformed payload is never held in full. @@ -463,7 +510,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Transform VertexAI File upload response into OpenAI-style FileObject """ - response_json: Final = raw_response.json() + response_json: Final[GcsBucketResponse] = raw_response.json() try: response_object: Final = GcsBucketResponse(**response_json) @@ -523,7 +570,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - response_json: Final = raw_response.json() + response_json: Final[_GcsObjectJson] = raw_response.json() gcs_id = response_json.get("id", "") gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" return OpenAIFileObject( @@ -682,7 +729,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 +770,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,18 +789,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _transform_single_vertex_batch_output_to_openai( self, - vertex_output: dict[str, Any], + vertex_output: _VertexBatchRow, vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> dict[str, Any]: + ) -> _OpenAIBatchOutputRow: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. """ # Extract custom_id from request labels (prefer raw for OpenAI round-trip) request_data: Final = vertex_output.get("request", {}) - labels: Final = request_data.get("labels", {}) or {} + labels: Final[Mapping[str, object]] = request_data.get("labels", {}) or {} custom_id: Final = _get_litellm_batch_custom_id_from_labels(labels) # Check if there's an error diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index d28f5b5b120..16e72e3062d 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -7,10 +7,12 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer import base64 import time -from typing import TYPE_CHECKING, Any, Final, cast +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Final, TypedDict, cast import httpx from httpx._types import RequestFiles +from typing_extensions import ReadOnly from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils @@ -40,11 +42,37 @@ else: BaseLLMException = Any +class _VeoVideo(TypedDict, total=False): + gcsUri: ReadOnly[str] + bytesBase64Encoded: ReadOnly[str] + mimeType: ReadOnly[str] + + +class _VeoOperationResponse(TypedDict, total=False): + videos: ReadOnly[Sequence[_VeoVideo]] + + +class _VeoOperationMetadata(TypedDict, total=False): + createTime: ReadOnly[str] + + +class _VeoOperation(TypedDict, total=False): + name: ReadOnly[str] + done: ReadOnly[bool] + metadata: ReadOnly[_VeoOperationMetadata] + response: ReadOnly[_VeoOperationResponse] + + +def _parse_veo_operation(raw_response: httpx.Response) -> _VeoOperation: + operation: Final[_VeoOperation] = raw_response.json() + return operation + + def _build_vertex_video_usage_from_request_data( request_data: dict[str, Any] | None, -) -> dict[str, Any]: +) -> dict[str, float | str]: """Build usage metadata (duration, resolution) for video cost calculation.""" - usage_data: Final[dict[str, Any]] = {} + usage_data: Final[dict[str, float | str]] = {} if not request_data: return usage_data @@ -125,7 +153,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Map OpenAI-style parameters to Veo format. @@ -135,7 +163,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): - size → aspectRatio (e.g., "1280x720" → "16:9") - seconds → durationSeconds (defaults to 4 seconds if not provided) """ - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Map input_reference to image (will be processed in transform_video_create_request) if "input_reference" in video_create_optional_params: @@ -289,7 +317,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): } """ # Build instance with prompt - instance_dict: Final[dict[str, Any]] = {"prompt": prompt} + instance_dict: Final[dict[str, object]] = {"prompt": prompt} params_copy: Final = video_create_optional_request_params.copy() # Check if user wants to provide full instance dict @@ -324,13 +352,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # {"parameters": {"parameters": {...}}} ← wrong # {"parameters": {...}} ← correct nested_params: Final = params_copy.pop("parameters", None) - vertex_params: Final[dict[str, Any]] = {} + vertex_params: Final[dict[str, object]] = {} if isinstance(nested_params, dict): vertex_params.update(nested_params) vertex_params.update(params_copy) # Build request data directly (TypedDict doesn't have model_dump) - request_data: Final[dict[str, Any]] = {"instances": [instance_dict]} + request_data: Final[dict[str, object]] = {"instances": [instance_dict]} # Only add parameters if there are any if vertex_params: @@ -363,7 +391,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): - status: "processing" - usage: includes duration_seconds and optional video_resolution for cost calculation """ - response_data: Final = raw_response.json() + response_data: Final = _parse_veo_operation(raw_response) operation_name: Final = response_data.get("name") if not operation_name: @@ -441,7 +469,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): } } """ - response_data: Final = raw_response.json() + response_data: Final = _parse_veo_operation(raw_response) operation_name: Final = response_data.get("name", "") is_done: Final = response_data.get("done", False) @@ -513,7 +541,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): Extracts the base64 encoded video from the response and decodes it to bytes. """ - response_data: Final = raw_response.json() + response_data: Final = _parse_veo_operation(raw_response) if not response_data.get("done", False): raise ValueError( @@ -548,7 +576,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Video remix is not supported by Veo API. @@ -574,7 +602,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Video list is not supported by Veo API. @@ -615,7 +643,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): """Video delete is not supported.""" raise NotImplementedError("Video delete is not supported by Vertex AI Veo.") - 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 Vertex AI") def transform_video_create_character_response(self, raw_response, logging_obj): @@ -649,7 +677,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, prefetched_source_data: dict[str, Any] | None = None, ) -> tuple[str, dict]: """ @@ -667,12 +695,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): if not prefetched_source_data.get("done", False): raise ValueError("Source video generation is not complete yet. Check the video status before editing.") - videos: Final = prefetched_source_data.get("response", {}).get("videos", []) + source_response: Final[_VeoOperationResponse] = prefetched_source_data.get("response", {}) + videos: Final = source_response.get("videos", []) if not videos: raise ValueError("No videos found in the completed operation. Cannot edit.") source_video: Final = videos[0] - video_input: Final[dict[str, Any]] = {} + video_input: Final[dict[str, str]] = {} if "gcsUri" in source_video: video_input["gcsUri"] = source_video["gcsUri"] elif "bytesBase64Encoded" in source_video: @@ -684,13 +713,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): operation_name: Final = extract_original_video_id(video_id) model: Final = self.extract_model_from_operation_name(operation_name) or "" - instance_dict: Final[dict[str, Any]] = {"prompt": prompt, "video": video_input} - request_data: Final[dict[str, Any]] = {"instances": [instance_dict]} + instance_dict: Final[dict[str, object]] = {"prompt": prompt, "video": video_input} + request_data: Final[dict[str, object]] = {"instances": [instance_dict]} if extra_body: extra_body_copy: Final = dict(extra_body) nested_params: Final = extra_body_copy.pop("parameters", None) - vertex_params: Final[dict[str, Any]] = {} + vertex_params: Final[dict[str, object]] = {} if isinstance(nested_params, dict): vertex_params.update(nested_params) vertex_params.update(extra_body_copy) @@ -716,7 +745,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): usage includes duration_seconds and optional video_resolution from the edit request parameters for cost calculation. """ - response_data: Final = raw_response.json() + response_data: Final = _parse_veo_operation(raw_response) operation_name: Final = response_data.get("name") if not operation_name: diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index e61ede4478c..5534b8b5a3c 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -8,11 +8,33 @@ import json import os import re from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence -from typing import Any, Final +from collections.abc import Set as AbstractSet +from typing import Any, Final, Protocol from urllib.parse import quote from litellm.types.mcp_server.mcp_server_manager import MCPServer + +class _McpServerLike(Protocol): + @property + def server_id(self) -> str: ... + @property + def server_name(self) -> str | None: ... + @property + def alias(self) -> str | None: ... + @property + def short_prefix(self) -> str | None: ... + + +class McpServerPayloadLike(Protocol): + alias: str | None + + @property + def server_name(self) -> str | None: ... + @property + def tool_name_to_display_name(self) -> Mapping[str, str] | None: ... + + # Constants # # NOTE: The environment-backed values below are read once, when this module is @@ -102,7 +124,7 @@ def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: # at the end so the first emitted char comes from the high-order # bits of the digest (which is the position we constrain to be # alphabetic). - chars: Final = [] + chars: Final[list[str]] = [] for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH): is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1 alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET @@ -176,34 +198,34 @@ def lookup_mcp_server_auth_in_headers( MCP_TOOL_ALLOWLIST_ENFORCED_KEY: Final = "tool_allowlist_enforced" -def _parse_mcp_info_dict(mcp_info: Any) -> dict[str, Any] | None: +def _parse_mcp_info_dict(mcp_info: object) -> Mapping[str, object] | None: if mcp_info is None: return None if isinstance(mcp_info, dict): return mcp_info if isinstance(mcp_info, str): try: - parsed: Final = json.loads(mcp_info) + parsed: Final[object] = json.loads(mcp_info) except (ValueError, TypeError): return None return parsed if isinstance(parsed, dict) else None return None -def is_server_tool_allowlist_enforced(mcp_server: Any) -> bool: +def is_server_tool_allowlist_enforced(mcp_server: object) -> bool: mcp_info: Final = _parse_mcp_info_dict(getattr(mcp_server, "mcp_info", None)) if not mcp_info: return False return bool(mcp_info.get(MCP_TOOL_ALLOWLIST_ENFORCED_KEY)) -def server_applies_tool_allowlist(mcp_server: Any) -> bool: +def server_applies_tool_allowlist(mcp_server: object) -> bool: """Whether server-level allowed_tools whitelist filtering is active.""" - allowed_tools: Final = getattr(mcp_server, "allowed_tools", None) or [] + allowed_tools: Final[object] = getattr(mcp_server, "allowed_tools", None) or [] return is_server_tool_allowlist_enforced(mcp_server) or bool(allowed_tools) -def validate_and_normalize_mcp_server_payload(payload: Any) -> None: +def validate_and_normalize_mcp_server_payload(payload: McpServerPayloadLike) -> None: """ Validate and normalize MCP server payload fields (server_name, alias, and tool_name_to_display_name). @@ -233,8 +255,8 @@ def validate_and_normalize_mcp_server_payload(payload: Any) -> None: validate_tool_display_names(payload.tool_name_to_display_name) # Alias normalization and defaulting - alias = getattr(payload, "alias", None) - server_name: Final = getattr(payload, "server_name", None) + alias: str | None = getattr(payload, "alias", None) + server_name: Final[str | None] = getattr(payload, "server_name", None) if not alias and server_name: alias = normalize_server_name(server_name) @@ -257,7 +279,7 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str: ) -def get_server_prefix(server: Any) -> str: +def get_server_prefix(server: object) -> str: """Return the prefix for a server. When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``) @@ -270,23 +292,26 @@ def get_server_prefix(server: Any) -> str: alias if present, else server_name, else server_id. """ if is_short_mcp_tool_prefix_enabled(): - cached: Final = getattr(server, "short_prefix", None) + cached: Final[str | None] = getattr(server, "short_prefix", None) if cached: return cached - server_id: Final = getattr(server, "server_id", None) + server_id: Final[str | None] = getattr(server, "server_id", None) if server_id: return compute_short_server_prefix(server_id) - if hasattr(server, "alias") and server.alias: - return server.alias - if hasattr(server, "server_name") and server.server_name: - return server.server_name + alias: Final[str | None] = getattr(server, "alias", None) + if alias: + return alias + server_name: Final[str | None] = getattr(server, "server_name", None) + if server_name: + return server_name if hasattr(server, "server_id"): - return server.server_id + fallback_server_id: Final[str] = getattr(server, "server_id", "") + return fallback_server_id return "" -def iter_known_server_prefixes(server: Any) -> Iterator[str]: +def iter_known_server_prefixes(server: _McpServerLike) -> Iterator[str]: """Yield every prefix form that may appear in tool names for ``server``. Always includes the *current* prefix returned by ``get_server_prefix``. @@ -304,7 +329,7 @@ def iter_known_server_prefixes(server: Any) -> Iterator[str]: yield from _emit(get_server_prefix(server)) yield from _emit(getattr(server, "short_prefix", None)) - server_id: Final = getattr(server, "server_id", None) + server_id: Final[str | None] = getattr(server, "server_id", None) if server_id: try: yield from _emit(compute_short_server_prefix(server_id)) @@ -397,7 +422,7 @@ def match_known_server_prefix(name: str, known_prefixes: Iterable[str]) -> tuple return None -def strip_known_server_prefix(name: str, server: Any | None) -> str: +def strip_known_server_prefix(name: str, server: _McpServerLike | None) -> str: """Strip ``server``'s registered prefix from a prefixed tool/resource name. Unlike :func:`split_server_prefix_from_name`, which guesses the boundary at @@ -420,7 +445,7 @@ def strip_known_server_prefix(name: str, server: Any | None) -> str: def is_tool_name_prefixed( tool_name: str, - known_server_prefixes: set | None = None, + known_server_prefixes: AbstractSet[str] | None = None, ) -> bool: """ Check if tool name has a known MCP server prefix. @@ -640,7 +665,7 @@ def parse_admin_env_vars( if raw is None: continue if hasattr(raw, "model_dump"): - entry = raw.model_dump() + entry: Mapping[str, object] = raw.model_dump() elif isinstance(raw, dict): entry = raw else: diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index ca84ff47884..7d6fafe141f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -11,6 +11,7 @@ import requests from fastapi import HTTPException from httpx import HTTPStatusError from requests.auth import HTTPBasicAuth +from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -55,6 +56,26 @@ class _HiddenlayerResponse(TypedDict, total=False): modified_data: Mapping[str, _HiddenlayerModifiedSide] +class _LoggedCallMetadata(TypedDict, total=False): + headers: ReadOnly[Mapping[str, str]] + + +class _LoggedCallLitellmParams(TypedDict, total=False): + metadata: ReadOnly[_LoggedCallMetadata] + + +class _HiddenlayerOutputMessage(TypedDict, total=False): + content: ReadOnly[str | Sequence[Mapping[str, str]]] + + +class _HiddenlayerChoiceMessage(TypedDict, total=False): + content: ReadOnly[str] + + +class _HiddenlayerChoice(TypedDict, total=False): + message: ReadOnly[_HiddenlayerChoiceMessage] + + def is_saas(host: str) -> bool: """Checks whether the connection is to the SaaS platform""" @@ -155,7 +176,10 @@ class HiddenlayerGuardrail(CustomGuardrail): # from the logger object on the response from the model. headers = request_data.get("proxy_server_request", {}).get("headers", {}) if not headers and logging_obj and logging_obj.model_call_details: - headers = logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {}) + logged_litellm_params: Final[_LoggedCallLitellmParams] = logging_obj.model_call_details.get( + "litellm_params", {} + ) + headers = logged_litellm_params.get("metadata", {}).get("headers", {}) hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM" project_id: Final = headers.get("hl-project-id") @@ -408,7 +432,8 @@ class HiddenlayerGuardrailV2(CustomGuardrail): if input_type == "request": inputs["structured_messages"] = output - for message in output.get("messages", []): + modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = output.get("messages", []) + for message in modified_messages: content = message.get("content", "") if isinstance(content, list): text_parts = [ @@ -422,7 +447,8 @@ class HiddenlayerGuardrailV2(CustomGuardrail): inputs["texts"] = new_texts elif input_type == "response" and inputs.get("texts"): - inputs["texts"] = [output.get("choices", [{}])[-1].get("message", {}).get("content", "")] + redacted_choices: Final[Sequence[_HiddenlayerChoice]] = output.get("choices", [{}]) + inputs["texts"] = [redacted_choices[-1].get("message", {}).get("content", "")] elif input_type == "response" and inputs.get("tool_calls"): inputs["tool_calls"] = output diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 9f70ed63dcb..5f7374581a2 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -2,9 +2,10 @@ import importlib import os +from collections.abc import Callable, Iterator, Mapping from datetime import datetime, timezone from itertools import chain, count -from typing import Any, Final, Literal, Optional, cast +from typing import Any, Final, Literal, Optional, Protocol, cast from pydantic import ValidationError @@ -59,6 +60,13 @@ from .guardrail_initializers import ( initialize_tool_permission, ) + +class _GuardrailRowLike(Protocol): + @property + def guardrail_id(self) -> str: ... + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + guardrail_initializer_registry: Final = { SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock, SupportedGuardrailIntegrations.LAKERA.value: initialize_lakera, @@ -125,7 +133,9 @@ def get_guardrail_initializer_from_hooks(): # Check for guardrail_initializer_registry dictionary if hasattr(module, "guardrail_initializer_registry"): - registry = getattr(module, "guardrail_initializer_registry") + registry: Mapping[str, Callable[..., CustomGuardrail]] | None = getattr( + module, "guardrail_initializer_registry", None + ) if isinstance(registry, dict): discovered_initializers.update(registry) verbose_proxy_logger.debug( @@ -135,7 +145,7 @@ def get_guardrail_initializer_from_hooks(): # Check for standalone initialize_guardrail function (fallback for directory-based guardrails) elif hasattr(module, "initialize_guardrail"): # For directories with just initialize_guardrail, use the directory name as the key - initialize_fn = getattr(module, "initialize_guardrail") + initialize_fn: Callable[..., CustomGuardrail] | None = getattr(module, "initialize_guardrail", None) discovered_initializers[item] = initialize_fn verbose_proxy_logger.debug("Found initialize_guardrail function in %s", module_path) @@ -206,7 +216,9 @@ def get_guardrail_class_from_hooks(): # Check for guardrail_initializer_registry dictionary if hasattr(module, "guardrail_class_registry"): - registry = getattr(module, "guardrail_class_registry") + registry: Mapping[str, type[CustomGuardrail]] | None = getattr( + module, "guardrail_class_registry", None + ) if isinstance(registry, dict): discovered_classes.update(registry) @@ -275,7 +287,7 @@ class GuardrailRegistry: guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {})) # Create guardrail in DB - created_guardrail: Final = await GuardrailsRepository(prisma_client).table.create( + created_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.create( data={ "guardrail_name": guardrail_name, "litellm_params": litellm_params, @@ -321,7 +333,7 @@ class GuardrailRegistry: guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {})) # Update in DB - updated_guardrail: Final = await GuardrailsRepository(prisma_client).table.update( + updated_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.update( where={"guardrail_id": guardrail_id}, data={ "guardrail_name": guardrail_name, @@ -482,7 +494,7 @@ class InMemoryGuardrailHandler: custom_guardrail_callback = initializer(litellm_params, guardrail) elif isinstance(guardrail_type, str) and "." in guardrail_type: custom_guardrail_callback = self.initialize_custom_guardrail( - guardrail=cast(dict, guardrail), + guardrail=guardrail, guardrail_type=guardrail_type, litellm_params=litellm_params, config_file_path=config_file_path, @@ -512,7 +524,7 @@ class InMemoryGuardrailHandler: "skip_tool_message_in_guardrail are enabled together, which excludes every message from " "scanning, so no request content would ever be scanned. Remove one of the two." ) - configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None) + configured_run_in_parallel: Final[bool | None] = getattr(litellm_params, "run_in_parallel", None) if configured_run_in_parallel is not None: custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) @@ -532,7 +544,7 @@ class InMemoryGuardrailHandler: def initialize_custom_guardrail( self, - guardrail: dict, + guardrail: Guardrail, guardrail_type: str, litellm_params: LitellmParams, config_file_path: str | None = None, @@ -550,7 +562,9 @@ class InMemoryGuardrailHandler: guardrail_type, ) - _guardrail_class: Final = get_instance_fn(guardrail_type, config_file_path=config_file_path) + _guardrail_class: Final[Callable[..., CustomGuardrail]] = get_instance_fn( + guardrail_type, config_file_path=config_file_path + ) mode: Final = litellm_params.mode if mode is None: @@ -683,8 +697,8 @@ class InMemoryGuardrailHandler: @staticmethod def _normalize_litellm_params_for_comparison( - params: Any | None, - ) -> dict[str, Any] | None: + params: LitellmParams | Mapping[str, object] | None, + ) -> Mapping[str, object] | None: """ Render litellm_params to a canonical dict so an in-memory LitellmParams and the raw dict loaded from the DB compare equal when they describe the same diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index bfc70da46ea..6c25f096532 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -10,12 +10,19 @@ All /customer management endpoints """ #### END-USER/CUSTOMER MANAGEMENT #### +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta -from typing import Final +from typing import TYPE_CHECKING, Final, Protocol, TypeVar, overload import fastapi from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter + +if TYPE_CHECKING: + from prisma.models import LiteLLM_BudgetTable as PrismaBudgetRow + from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow + + from litellm.proxy.utils import PrismaClient import litellm from litellm._logging import verbose_proxy_logger @@ -41,6 +48,54 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import ( UnblockUsersResponse, ) +_RowT_co: Final = TypeVar("_RowT_co", covariant=True) +_STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) + +if TYPE_CHECKING: + + class _TableOps(Protocol[_RowT_co]): + async def find_first( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, + ) -> _RowT_co | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, + ) -> Sequence[_RowT_co]: ... + + async def create( + self, + data: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> _RowT_co: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> _RowT_co | None: ... + + async def upsert( + self, + where: Mapping[str, object], + data: Mapping[str, Mapping[str, object]], + ) -> _RowT_co: ... + + async def delete_many(self, where: Mapping[str, object]) -> int: ... + + +@overload +def _typed_table(repo: EndUserRepository) -> "_TableOps[PrismaEndUserRow]": ... +@overload +def _typed_table(repo: BudgetRepository) -> "_TableOps[PrismaBudgetRow]": ... +def _typed_table(repo: EndUserRepository | BudgetRepository) -> object: + return repo.table + + router: Final = APIRouter() @@ -89,7 +144,7 @@ async def block_user(data: BlockUsers): records: Final = [] if prisma_client is not None: for id in data.user_ids: - record = await EndUserRepository(prisma_client).table.upsert( + record = await _typed_table(EndUserRepository(prisma_client)).upsert( where={"user_id": id}, data={ "create": {"user_id": id, "blocked": True}, @@ -184,7 +239,7 @@ def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None: budget_kv_pairs[field_name] = value if budget_kv_pairs: - budget_request: Final = BudgetNewRequest(**budget_kv_pairs) + budget_request: Final = BudgetNewRequest.model_validate(budget_kv_pairs) validate_budget_duration(budget_request.budget_duration) if budget_request.budget_reset_at is None and budget_request.budget_duration is not None: budget_request.budget_reset_at = datetime.utcnow() + timedelta( @@ -195,10 +250,10 @@ def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None: async def _handle_customer_object_permission_update( - non_default_values: dict, + non_default_values: dict[str, object], end_user_table_data_typed: LiteLLM_EndUserTable | None, - update_end_user_table_data: dict, - prisma_client, + update_end_user_table_data: dict[str, object], + prisma_client: "PrismaClient", ) -> None: """ Handle object permission updates for customer endpoints. @@ -344,13 +399,13 @@ async def new_end_user( }, ) - new_end_user_obj: dict = {} + new_end_user_obj: dict[str, object] = {} ## CREATE BUDGET ## if set _new_budget: Final = new_budget_request(data) if _new_budget is not None: try: - budget_record: Final = await BudgetRepository(prisma_client).table.create( + budget_record: Final = await _typed_table(BudgetRepository(prisma_client)).create( data={ **_new_budget.model_dump(exclude_unset=True), "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -364,16 +419,18 @@ async def new_end_user( elif data.budget_id is not None: new_end_user_obj["budget_id"] = data.budget_id - _user_data: Final = data.dict(exclude_none=True) + _user_data: Final = _STR_OBJECT_DICT.validate_python(data.dict(exclude_none=True)) for k, v in _user_data.items(): if k not in BudgetNewRequest.model_fields: new_end_user_obj[k] = v ## Handle Object Permission - MCP Servers, Vector Stores etc. - new_end_user_obj = await _set_object_permission( - data_json=new_end_user_obj, - prisma_client=prisma_client, + new_end_user_obj = _STR_OBJECT_DICT.validate_python( + await _set_object_permission( + data_json=new_end_user_obj, + prisma_client=prisma_client, + ) ) # Ensure object_permission is not in the data being sent to create @@ -386,7 +443,7 @@ async def new_end_user( new_end_user_obj.pop("object_permission", None) ## WRITE TO DB ## - end_user_record: Final = await EndUserRepository(prisma_client).table.create( + end_user_record: Final = await _typed_table(EndUserRepository(prisma_client)).create( data=new_end_user_obj, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -442,7 +499,7 @@ async def end_user_info( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - user_info: Final = await EndUserRepository(prisma_client).table.find_first( + user_info: Final = await _typed_table(EndUserRepository(prisma_client)).find_first( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -535,13 +592,13 @@ async def update_end_user( from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client try: - data_json: Final[dict] = data.json() + data_json: Final = _STR_OBJECT_DICT.validate_python(data.json()) # get the row from db if prisma_client is None: raise Exception("Not connected to DB!") # get non default values for key - non_default_values: Final = {} + non_default_values: Final = dict[str, object]() for k, v in data_json.items(): if v is not None and v not in ( [], @@ -551,7 +608,7 @@ async def update_end_user( non_default_values[k] = v ## Get end user table data ## - end_user_table_data: Final = await EndUserRepository(prisma_client).table.find_first( + end_user_table_data: Final = await _typed_table(EndUserRepository(prisma_client)).find_first( where={"user_id": data.user_id}, include={"litellm_budget_table": True} ) @@ -563,14 +620,14 @@ async def update_end_user( param="user_id", ) - end_user_table_data_typed: Final = LiteLLM_EndUserTable(**end_user_table_data.model_dump()) + end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump()) ## Get budget table data ## end_user_budget_table: Final = end_user_table_data_typed.litellm_budget_table ## Get all params for budget table ## - budget_table_data: Final = {} - update_end_user_table_data: Final = {} + budget_table_data: Final = dict[str, object]() + update_end_user_table_data: Final = dict[str, object]() for k, v in non_default_values.items(): # budget_id is for linking to existing budget, not for creating new budget if k == "budget_id": @@ -593,7 +650,7 @@ async def update_end_user( if budget_table_data: if end_user_budget_table is None: ## Create new budget ## - budget_table_data_record = await BudgetRepository(prisma_client).table.create( + budget_table_data_record = await _typed_table(BudgetRepository(prisma_client)).create( data={ **budget_table_data, "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -605,7 +662,7 @@ async def update_end_user( update_end_user_table_data["budget_id"] = budget_table_data_record.budget_id else: ## Update existing budget ## - budget_table_data_record = await BudgetRepository(prisma_client).table.update( + budget_table_data_record = await _typed_table(BudgetRepository(prisma_client)).update( where={"budget_id": end_user_budget_table.budget_id}, data=budget_table_data, ) @@ -625,7 +682,7 @@ async def update_end_user( if data.user_id is not None and len(data.user_id) > 0: update_end_user_table_data["user_id"] = data.user_id verbose_proxy_logger.debug("In update customer, user_id condition block.") - response: Final = await EndUserRepository(prisma_client).table.update( + response: Final = await _typed_table(EndUserRepository(prisma_client)).update( where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True, "object_permission": True}, @@ -688,7 +745,7 @@ async def delete_end_user( verbose_proxy_logger.debug("/customer/delete: Received data = %s", data) if data.user_ids is not None and isinstance(data.user_ids, list) and len(data.user_ids) > 0: # First check if all users exist - existing_users: Final = await EndUserRepository(prisma_client).table.find_many( + existing_users: Final = await _typed_table(EndUserRepository(prisma_client)).find_many( where={"user_id": {"in": data.user_ids}} ) existing_user_ids: Final = {user.user_id for user in existing_users} @@ -703,7 +760,7 @@ async def delete_end_user( ) # All users exist, proceed with deletion - response: Final = await EndUserRepository(prisma_client).table.delete_many( + response: Final = await _typed_table(EndUserRepository(prisma_client)).delete_many( where={"user_id": {"in": data.user_ids}} ) verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) @@ -764,7 +821,7 @@ async def list_end_user( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - response: Final = await EndUserRepository(prisma_client).table.find_many( + response: Final = await _typed_table(EndUserRepository(prisma_client)).find_many( include={"litellm_budget_table": True, "object_permission": True} ) @@ -827,11 +884,10 @@ async def get_customer_daily_activity( exclude_end_user_ids_list = exclude_end_user_ids.split(",") if exclude_end_user_ids else None # Fetch organization aliases for metadata - where_condition: Final = {} + where_condition: Final = dict[str, object]() if end_user_ids_list: where_condition["user_id"] = {"in": list(end_user_ids_list)} - end_user_aliases: Final = await EndUserRepository(prisma_client).table.find_many(where=where_condition) - end_user_alias_metadata: Final = {e.user_id: {"alias": e.alias} for e in end_user_aliases} + end_user_aliases: Final = await _typed_table(EndUserRepository(prisma_client)).find_many(where=where_condition) # Query daily activity for organizations return await get_daily_activity( @@ -839,7 +895,7 @@ async def get_customer_daily_activity( table_name="litellm_dailyenduserspend", entity_id_field="end_user_id", entity_id=end_user_ids_list, - entity_metadata_field=end_user_alias_metadata, + entity_metadata_field={e.user_id: {"alias": e.alias} for e in end_user_aliases}, exclude_entity_ids=exclude_end_user_ids_list, start_date=start_date, end_date=end_date, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e156e5f0046..c4f5ceaba7a 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -50,6 +50,7 @@ from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, + McpServerPayloadLike, build_env_var_setup_url, collect_env_var_references, get_server_prefix, @@ -196,7 +197,7 @@ if MCP_AVAILABLE: server: MCPServer expires_at: datetime - def _validate_mcp_server_name_fields(payload: Any) -> None: + def _validate_mcp_server_name_fields(payload: McpServerPayloadLike) -> None: candidates: Final[list[tuple[str, str | None]]] = [] server_name: Final = getattr(payload, "server_name", None) @@ -223,7 +224,7 @@ if MCP_AVAILABLE: detail={"error": error_messages_text}, ) - def validate_and_normalize_mcp_server_payload(payload: Any) -> None: + def validate_and_normalize_mcp_server_payload(payload: McpServerPayloadLike) -> None: _base_validate_and_normalize_mcp_server_payload(payload) _validate_mcp_server_name_fields(payload) diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 0fdedafb2bf..b480d46f185 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -10,13 +10,21 @@ 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, TypeAlias, TypeVar, overload from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, Field, TypeAdapter if TYPE_CHECKING: + from prisma.models import LiteLLM_DailyToolSpend as PrismaDailyToolSpendRow + from prisma.models import LiteLLM_ObjectPermissionTable as PrismaObjectPermissionRow + from prisma.models import LiteLLM_SpendLogs as PrismaSpendLogRow + from prisma.models import LiteLLM_SpendLogToolIndex as PrismaSpendLogToolIndexRow + from prisma.models import LiteLLM_TeamTable as PrismaTeamRow + from prisma.models import LiteLLM_VerificationToken as PrismaVerificationTokenRow + from litellm.proxy.utils import PrismaClient from litellm._logging import verbose_proxy_logger @@ -49,6 +57,72 @@ from litellm.types.tool_management import ( ToolUsageLogsResponse, ) +_RowT_co: Final = TypeVar("_RowT_co", covariant=True) + +if TYPE_CHECKING: + + class _TableOps(Protocol[_RowT_co]): + 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[_RowT_co]: ... + + async def find_unique(self, where: Mapping[str, object]) -> _RowT_co | None: ... + + async def count(self, where: Mapping[str, object] | None = None) -> int: ... + + async def create(self, data: Mapping[str, object]) -> _RowT_co: ... + + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: ... + + async def delete(self, where: Mapping[str, object]) -> _RowT_co | None: ... + + async def group_by( + self, + by: Sequence[str], + sum: Mapping[str, bool] | None = None, + where: Mapping[str, object] | None = None, + order: Mapping[str, object] | None = None, + take: int | None = None, + ) -> Sequence[Mapping[str, object]]: ... + + class _SpendLogRow(Protocol): + @property + def messages(self) -> object: ... + @property + def proxy_server_request(self) -> str | Mapping[str, object] | None: ... + + +@overload +def _typed_table(repo: DailyToolSpendRepository) -> "_TableOps[PrismaDailyToolSpendRow]": ... +@overload +def _typed_table(repo: SpendLogToolIndexRepository) -> "_TableOps[PrismaSpendLogToolIndexRow]": ... +@overload +def _typed_table(repo: SpendLogsRepository) -> "_TableOps[PrismaSpendLogRow]": ... +@overload +def _typed_table(repo: VerificationTokenRepository) -> "_TableOps[PrismaVerificationTokenRow]": ... +@overload +def _typed_table(repo: TeamRepository) -> "_TableOps[PrismaTeamRow]": ... +@overload +def _typed_table(repo: ObjectPermissionRepository) -> "_TableOps[PrismaObjectPermissionRow]": ... +def _typed_table( + repo: DailyToolSpendRepository + | SpendLogToolIndexRepository + | SpendLogsRepository + | VerificationTokenRepository + | TeamRepository + | ObjectPermissionRepository, +) -> object: + return repo.table + + router: Final = APIRouter() TOOL_POLICY_OPTIONS: Final = ToolPolicyOptionsResponse( @@ -201,7 +275,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 = _typed_table(DailyToolSpendRepository(prisma_client)) top_tools: Final = _TOP_TOOL_ROWS.validate_python( await table.group_by( by=["tool_name"], @@ -222,7 +296,7 @@ async def get_tool_spend( for row in top_tools ] - daily_rows: Final = ( + daily_rows: Final[Sequence[PrismaDailyToolSpendRow]] = ( await table.find_many( where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}}, order=[{"date": "asc"}, {"spend": "desc"}], @@ -270,36 +344,43 @@ 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: +_ParsedJson: TypeAlias = dict[str, object] | list[object] | str | int | float | bool | None +_PARSED_JSON: Final[TypeAdapter[_ParsedJson]] = TypeAdapter(_ParsedJson) +_STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) + + +def _input_snippet_for_tool_log(sl: "_SpendLogRow | 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 = sl.messages if messages is not None: s = _snippet_str(messages, max_len) if s: return s - psr = getattr(sl, "proxy_server_request", None) + psr = sl.proxy_server_request 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): msgs = psr.get("messages") - if msgs is None and isinstance(psr.get("body"), dict): - msgs = psr["body"].get("messages") + if msgs is None: + body: Final = psr.get("body") + if isinstance(body, dict): + msgs = _STR_OBJECT_DICT.validate_python(body).get("messages") s = _snippet_str(msgs, max_len) if s: return s 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,7 +425,7 @@ 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} + where: Final[dict[str, object]] = {"tool_name": tool_name} if start_date or end_date: start_time_filter: datetime | None = None end_time_filter: datetime | None = None @@ -363,14 +444,14 @@ 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 + where["start_time"] = { + key: value + for key, value in (("gte", start_time_filter), ("lte", end_time_filter)) + if value is not None + } - total: Final = await SpendLogToolIndexRepository(prisma_client).table.count(where=where) - index_rows: Final = await SpendLogToolIndexRepository(prisma_client).table.find_many( + total: Final = await _typed_table(SpendLogToolIndexRepository(prisma_client)).count(where=where) + index_rows: Final = await _typed_table(SpendLogToolIndexRepository(prisma_client)).find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, @@ -380,7 +461,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 _typed_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,24 +532,24 @@ 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 _typed_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 = row.object_permission_id if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _typed_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 _typed_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}) - return getattr(row, "object_permission_id", None) if row else None + await _typed_table(ObjectPermissionRepository(prisma_client)).delete(where={"object_permission_id": new_id}) + row = await _typed_table(VerificationTokenRepository(prisma_client)).find_unique(where={"token": hashed}) + return row.object_permission_id if row else None return new_id @@ -478,24 +561,24 @@ 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 _typed_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 = row.object_permission_id if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _typed_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 _typed_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}) - return getattr(row, "object_permission_id", None) if row else None + await _typed_table(ObjectPermissionRepository(prisma_client)).delete(where={"object_permission_id": new_id}) + row = await _typed_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean}) + return row.object_permission_id if row else None return new_id diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 2e38abddd0f..9d5ddda017a 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -4,11 +4,11 @@ usage/spend data by querying the aggregated daily activity endpoints. """ import json -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import date -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, cast, overload -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -73,9 +73,36 @@ class SSEErrorEvent(TypedDict): SSEEvent = SSEStatusEvent | SSEToolCallEvent | SSEChunkEvent | SSEDoneEvent | SSEErrorEvent +class _EntityEntry(TypedDict, total=False): + metrics: ReadOnly[Mapping[str, float]] + metadata: ReadOnly[Mapping[str, str]] + + +class _DayDump(TypedDict, total=False): + breakdown: ReadOnly[Mapping[str, Mapping[str, _EntityEntry]]] + + +class _UsageDump(Protocol): + @overload + def get(self, key: Literal["metadata"], default: Mapping[str, float], /) -> Mapping[str, float]: ... + @overload + def get(self, key: Literal["results"], default: Sequence[_DayDump], /) -> Sequence[_DayDump]: ... + + +class _ToolFunctionDef(TypedDict): + name: ReadOnly[str] + description: ReadOnly[str] + parameters: ReadOnly[Mapping[str, object]] + + +class _ToolDef(TypedDict): + type: ReadOnly[str] + function: ReadOnly[_ToolFunctionDef] + + class ToolHandler(TypedDict): - fetch: Callable[..., Any] - summarise: Callable[[dict[str, Any]], str] + fetch: Callable[..., Awaitable[_UsageDump]] + summarise: Callable[[_UsageDump], str] label: str @@ -88,7 +115,7 @@ _DATE_PARAMS: Final = { "end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"}, } -_TOOL_USAGE: Final = { +_TOOL_USAGE: Final[_ToolDef] = { "type": "function", "function": { "name": "get_usage_data", @@ -111,7 +138,7 @@ _TOOL_USAGE: Final = { }, } -_TOOL_TEAM: Final = { +_TOOL_TEAM: Final[_ToolDef] = { "type": "function", "function": { "name": "get_team_usage_data", @@ -133,7 +160,7 @@ _TOOL_TEAM: Final = { }, } -_TOOL_TAG: Final = { +_TOOL_TAG: Final[_ToolDef] = { "type": "function", "function": { "name": "get_tag_usage_data", @@ -159,7 +186,7 @@ TOOLS_BASE: Final = [_TOOL_USAGE] TOOLS_ADMIN: Final = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG] -def get_tools_for_role(is_admin: bool) -> list[dict[str, Any]]: +def get_tools_for_role(is_admin: bool) -> list[_ToolDef]: """Return the tool list appropriate for the user's role.""" return TOOLS_ADMIN if is_admin else TOOLS_BASE @@ -254,7 +281,7 @@ async def _query_activity( ) -async def _fetch_usage_data(start_date: str, end_date: str, user_id: str | None = None) -> dict[str, Any]: +async def _fetch_usage_data(start_date: str, end_date: str, user_id: str | None = None) -> _UsageDump: resp: Final = await _query_activity( TABLE_DAILY_USER_SPEND, ENTITY_FIELD_USER, @@ -266,7 +293,7 @@ async def _fetch_usage_data(start_date: str, end_date: str, user_id: str | None return resp.model_dump(mode="json") -async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: str | None = None) -> dict[str, Any]: +async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: str | None = None) -> _UsageDump: resp: Final = await _query_activity( TABLE_DAILY_TEAM_SPEND, ENTITY_FIELD_TEAM, @@ -277,7 +304,7 @@ async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: str | return resp.model_dump(mode="json") -async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: str | None = None) -> dict[str, Any]: +async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: str | None = None) -> _UsageDump: resp: Final = await _query_activity( TABLE_DAILY_TAG_SPEND, ENTITY_FIELD_TAG, @@ -294,7 +321,7 @@ async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: str | None def _accumulate_breakdown( - results: list[dict[str, Any]], dimension: str, fields: list[str] + results: Sequence[_DayDump], dimension: str, fields: Sequence[str] ) -> dict[str, dict[str, float]]: """Aggregate a single breakdown dimension across days.""" totals: Final[dict[str, dict[str, float]]] = {} @@ -317,7 +344,7 @@ def _ranked_lines( return [fmt(name, vals) for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[:limit]] -def _summarise_usage_data(data: dict[str, Any]) -> str: +def _summarise_usage_data(data: _UsageDump) -> str: meta: Final = data.get("metadata", {}) results: Final = data.get("results", []) @@ -349,7 +376,7 @@ def _summarise_usage_data(data: dict[str, Any]) -> str: return "\n".join(sections) -def _summarise_entity_data(data: dict[str, Any], entity_label: str) -> str: +def _summarise_entity_data(data: _UsageDump, entity_label: str) -> str: """Summarise team/tag entity usage data.""" results: Final = data.get("results", []) if not results: @@ -409,16 +436,16 @@ def _sse(event: SSEEvent) -> str: def _resolve_fetch_kwargs( fn_name: str, - fn_args: dict[str, str], + fn_args: Mapping[str, str], user_id: str | None, is_admin: bool, -) -> dict[str, Any]: +) -> dict[str, str]: """Build keyword arguments for a tool's fetch function.""" start_date: Final = fn_args.get("start_date", "") end_date: Final = fn_args.get("end_date", "") if not start_date or not end_date: raise ValueError("Missing required start_date or end_date from tool arguments") - kwargs: Final[dict[str, Any]] = {"start_date": start_date, "end_date": end_date} + kwargs: Final[dict[str, str]] = {"start_date": start_date, "end_date": end_date} if fn_name == "get_usage_data": if not is_admin: if user_id is None: @@ -443,7 +470,7 @@ def _resolve_fetch_kwargs( async def _execute_tool_call( handler: ToolHandler, fn_name: str, - fn_args: dict[str, str], + fn_args: Mapping[str, str], user_id: str | None, is_admin: bool, ) -> str: @@ -455,13 +482,13 @@ async def _execute_tool_call( async def _process_tool_call( tc: Any, - chat_messages: list[dict[str, Any]], + chat_messages: list[Mapping[str, object]], user_id: str | None, is_admin: bool, ) -> AsyncIterator[str]: """Execute a single tool call, yielding SSE events for status.""" - fn_name: Final = tc.function.name - fn_args: Final = json.loads(tc.function.arguments) + fn_name: Final[str] = tc.function.name + fn_args: Final[Mapping[str, str]] = json.loads(tc.function.arguments) allowed_names: Final = {t["function"]["name"] for t in get_tools_for_role(is_admin)} handler: Final = TOOL_HANDLERS.get(fn_name) @@ -495,7 +522,7 @@ async def _process_tool_call( chat_messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_result}) -async def _stream_final_response(model: str, chat_messages: list[dict[str, Any]]) -> AsyncIterator[str]: +async def _stream_final_response(model: str, chat_messages: list[Mapping[str, object]]) -> AsyncIterator[str]: """Stream the final LLM response after tool results are appended.""" yield _sse({"type": "status", "message": "Analyzing results..."}) @@ -520,7 +547,7 @@ async def stream_usage_ai_chat( """Stream SSE events: status → tool_call → chunk → done.""" resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages - chat_messages: Final[list[dict[str, Any]]] = [ + chat_messages: Final[list[Mapping[str, object]]] = [ {"role": "system", "content": _build_system_prompt(is_admin)}, *truncated, ] diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index 09de170d542..0422c72cdb3 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -11,11 +11,19 @@ These endpoints use optimized single SQL queries with joins to efficiently calcu user metrics from tag activity data and return time series for dashboard visualization. """ +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta -from typing import Any, Final +from typing import TYPE_CHECKING, Final, Protocol, TypeVar, overload from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter + +if TYPE_CHECKING: + from prisma.models import LiteLLM_DailyTagSpend as PrismaDailyTagSpendRow + from prisma.models import LiteLLM_UserTable as PrismaUserRow + from prisma.models import LiteLLM_VerificationToken as PrismaVerificationTokenRow + + from litellm.proxy.utils import PrismaClient from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -103,6 +111,54 @@ class PerUserAnalyticsResponse(BaseModel): total_pages: int +class _DistinctTagRow(BaseModel): + tag: str + + +class _ActiveUsersRow(BaseModel): + tag: str + active_users: int + date: str + period_start: str | None = None + period_end: str | None = None + + +class _TagSummaryRow(BaseModel): + tag: str + unique_users: int | None = None + total_requests: float | int | str | None = None + successful_requests: float | int | str | None = None + failed_requests: float | int | str | None = None + total_tokens: float | int | str | None = None + total_spend: float | int | str | None = None + + +_DISTINCT_TAG_ROWS: Final = TypeAdapter(list[_DistinctTagRow]) +_ACTIVE_USERS_ROWS: Final = TypeAdapter(list[_ActiveUsersRow]) +_TAG_SUMMARY_ROWS: Final = TypeAdapter(list[_TagSummaryRow]) + +_RowT_co: Final = TypeVar("_RowT_co", covariant=True) + +if TYPE_CHECKING: + + class _TableOps(Protocol[_RowT_co]): + async def find_many(self, where: Mapping[str, object] | None = None) -> Sequence[_RowT_co]: ... + + +@overload +def _typed_table(repo: DailyTagSpendRepository) -> "_TableOps[PrismaDailyTagSpendRow]": ... +@overload +def _typed_table(repo: VerificationTokenRepository) -> "_TableOps[PrismaVerificationTokenRow]": ... +@overload +def _typed_table(repo: UserRepository) -> "_TableOps[PrismaUserRow]": ... +def _typed_table(repo: DailyTagSpendRepository | VerificationTokenRepository | UserRepository) -> object: + return repo.table + + +async def _query_raw(prisma_client: "PrismaClient", sql_query: str, *params: object) -> object: + return await prisma_client.db.query_raw(sql_query, *params) + + @router.get( "/tag/distinct", response_model=DistinctTagsResponse, @@ -141,9 +197,9 @@ async def get_distinct_user_agent_tags( LIMIT {MAX_TAGS} """ - db_response: Final = await prisma_client.db.query_raw(sql_query) + db_response: Final = _DISTINCT_TAG_ROWS.validate_python(await _query_raw(prisma_client, sql_query)) - results: Final = [DistinctTagResponse(tag=row["tag"]) for row in db_response] + results: Final = [DistinctTagResponse(tag=row.tag) for row in db_response] return DistinctTagsResponse(results=results) @@ -231,11 +287,10 @@ async def get_daily_active_users( ORDER BY dts.date DESC, active_users DESC """ - db_response: Final = await prisma_client.db.query_raw(sql_query, *params) + db_response: Final = _ACTIVE_USERS_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params)) results: Final = [ - TagActiveUsersResponse(tag=row["tag"], active_users=row["active_users"], date=row["date"]) - for row in db_response + TagActiveUsersResponse(tag=row.tag, active_users=row.active_users, date=row.date) for row in db_response ] return ActiveUsersAnalyticsResponse(results=results) @@ -346,15 +401,15 @@ async def get_weekly_active_users( ORDER BY week_offset DESC, active_users DESC """ - db_response: Final = await prisma_client.db.query_raw(sql_query, *params) + db_response: Final = _ACTIVE_USERS_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params)) results: Final = [ TagActiveUsersResponse( - tag=row["tag"], - active_users=row["active_users"], - date=row["date"], # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc. - period_start=row["period_start"], - period_end=row["period_end"], + tag=row.tag, + active_users=row.active_users, + date=row.date, # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc. + period_start=row.period_start, + period_end=row.period_end, ) for row in db_response ] @@ -467,15 +522,15 @@ async def get_monthly_active_users( ORDER BY month_offset DESC, active_users DESC """ - db_response: Final = await prisma_client.db.query_raw(sql_query, *params) + db_response: Final = _ACTIVE_USERS_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params)) results: Final = [ TagActiveUsersResponse( - tag=row["tag"], - active_users=row["active_users"], - date=row["date"], # This will be "Month 1 (Jan)", "Month 2 (Dec)", etc. - period_start=row["period_start"], - period_end=row["period_end"], + tag=row.tag, + active_users=row.active_users, + date=row.date, # This will be "Month 1 (Jan)", "Month 2 (Dec)", etc. + period_start=row.period_start, + period_end=row.period_end, ) for row in db_response ] @@ -565,17 +620,17 @@ async def get_tag_summary( ORDER BY total_requests DESC """ - db_response: Final = await prisma_client.db.query_raw(sql_query, *params) + db_response: Final = _TAG_SUMMARY_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params)) results: Final = [ TagSummaryMetrics( - tag=row["tag"], - unique_users=row["unique_users"] or 0, - total_requests=int(row["total_requests"] or 0), - successful_requests=int(row["successful_requests"] or 0), - failed_requests=int(row["failed_requests"] or 0), - total_tokens=int(row["total_tokens"] or 0), - total_spend=float(row["total_spend"] or 0.0), + tag=row.tag, + unique_users=row.unique_users or 0, + total_requests=int(row.total_requests or 0), + successful_requests=int(row.successful_requests or 0), + failed_requests=int(row.failed_requests or 0), + total_tokens=int(row.total_tokens or 0), + total_spend=float(row.total_spend or 0.0), ) for row in db_response ] @@ -648,7 +703,7 @@ async def get_per_user_analytics( start_date: Final = start_dt.strftime("%Y-%m-%d") # Build where clause with date range - where_clause: Final[dict[str, Any]] = {"date": {"gte": start_date, "lte": end_date}} + where_clause: Final[dict[str, object]] = {"date": {"gte": start_date, "lte": end_date}} # Add tag filtering if provided if tag_filters and len(tag_filters) > 0: @@ -657,7 +712,7 @@ async def get_per_user_analytics( where_clause["tag"] = {"contains": tag_filter} # Get all tag records in the date range with optional tag filtering - tag_records: Final = await DailyTagSpendRepository(prisma_client).table.find_many(where=where_clause) + tag_records: Final = await _typed_table(DailyTagSpendRepository(prisma_client)).find_many(where=where_clause) # Get unique api_keys api_keys: Final = set(record.api_key for record in tag_records if record.api_key) @@ -672,7 +727,7 @@ async def get_per_user_analytics( ) # Lookup user_id for each api_key - api_key_records: Final = await VerificationTokenRepository(prisma_client).table.find_many( + api_key_records: Final = await _typed_table(VerificationTokenRepository(prisma_client)).find_many( where={"token": {"in": list(api_keys)}} ) @@ -681,7 +736,9 @@ async def get_per_user_analytics( # Get user emails for the user_ids user_ids: Final = list(set(api_key_to_user_id.values())) - user_records: Final = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": user_ids}}) + user_records: Final = await _typed_table(UserRepository(prisma_client)).find_many( + where={"user_id": {"in": user_ids}} + ) # Create mapping from user_id to user_email user_id_to_email: Final = {record.user_id: record.user_email for record in user_records} diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index d8e9f8dfaee..1d71ea658e4 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -3,8 +3,10 @@ CRUD ENDPOINTS FOR PROMPTS """ import tempfile +from collections.abc import Awaitable, Mapping, Sequence +from datetime import datetime from pathlib import Path -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from fastapi import ( APIRouter, @@ -38,9 +40,68 @@ 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 + router: Final = APIRouter() +class _PromptRow(Protocol): + @property + def id(self) -> str: ... + @property + def prompt_id(self) -> str: ... + @property + def version(self) -> int: ... + @property + def environment(self) -> str: ... + @property + def created_by(self) -> str | None: ... + @property + def created_at(self) -> "datetime": ... + @property + def updated_at(self) -> "datetime": ... + @property + def litellm_params(self) -> str | Mapping[str, object]: ... + @property + def prompt_info(self) -> str | Mapping[str, object] | None: ... + + def model_dump(self) -> Mapping[str, object]: ... + + +class _PromptRowData(BaseModel): + prompt_id: str + version: int = 1 + environment: str = "development" + created_by: str | None = None + litellm_params: str | Mapping[str, object] | None = None + prompt_info: str | Mapping[str, object] | None = None + created_at: datetime | None = None + updated_at: datetime | None = None + + +class _PromptTableActions(Protocol): + def find_many( + self, + *, + where: Mapping[str, str | int], + order: Mapping[str, str] = ..., + take: int = ..., + distinct: Sequence[str] = ..., + ) -> Awaitable[Sequence[_PromptRow]]: ... + + def create(self, *, data: Mapping[str, str | int | None]) -> Awaitable[_PromptRow]: ... + + def update(self, *, where: Mapping[str, str | int], data: Mapping[str, str]) -> Awaitable[_PromptRow]: ... + + def delete_many(self, *, where: Mapping[str, str]) -> Awaitable[int]: ... + + +def _prompt_table(prisma_client: "PrismaClient") -> _PromptTableActions: + return PromptRepository(prisma_client).table + + def get_base_prompt_id(prompt_id: str) -> str: """ Extract the base prompt ID by stripping the version suffix if present. @@ -132,7 +193,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 +259,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 +273,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(prisma_client).find_many( where={"prompt_id": prompt_id, "environment": environment} ) @@ -221,7 +284,7 @@ async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment return 1 -def create_versioned_prompt_spec(db_prompt) -> PromptSpec: +def create_versioned_prompt_spec(db_prompt: _PromptRow) -> PromptSpec: """ Helper function to create a PromptSpec with versioned prompt_id from a DB prompt entry. @@ -235,38 +298,33 @@ def create_versioned_prompt_spec(db_prompt) -> PromptSpec: from litellm.types.prompts.init_prompts import PromptLiteLLMParams - prompt_dict: Final = db_prompt.model_dump() - base_prompt_id: Final = prompt_dict["prompt_id"] - version: Final = prompt_dict.get("version", 1) - environment: Final = prompt_dict.get("environment", "development") - created_by: Final = prompt_dict.get("created_by") + row: Final = _PromptRowData.model_validate(db_prompt.model_dump()) - # Parse litellm_params - litellm_params_data = prompt_dict.get("litellm_params") - if isinstance(litellm_params_data, str): - litellm_params_data = json.loads(litellm_params_data) - litellm_params: Final = PromptLiteLLMParams(**litellm_params_data) + litellm_params_data: Final = row.litellm_params + litellm_params_dict: Final[Mapping[str, object] | None] = ( + json.loads(litellm_params_data) if isinstance(litellm_params_data, str) else litellm_params_data + ) + litellm_params: Final = PromptLiteLLMParams.model_validate(litellm_params_dict) - # Parse prompt_info - prompt_info_data = prompt_dict.get("prompt_info") + prompt_info_data: Final = row.prompt_info if prompt_info_data: - if isinstance(prompt_info_data, str): - prompt_info_data = json.loads(prompt_info_data) - prompt_info = PromptInfo(**prompt_info_data) + prompt_info_dict: Final[Mapping[str, object]] = ( + json.loads(prompt_info_data) if isinstance(prompt_info_data, str) else prompt_info_data + ) + prompt_info = PromptInfo.model_validate(prompt_info_dict) else: prompt_info = PromptInfo(prompt_type="db") - # Create versioned prompt_id - versioned_prompt_id: Final = f"{base_prompt_id}.v{version}" + versioned_prompt_id: Final = f"{row.prompt_id}.v{row.version}" return PromptSpec( prompt_id=versioned_prompt_id, litellm_params=litellm_params, prompt_info=prompt_info, - created_at=prompt_dict.get("created_at"), - updated_at=prompt_dict.get("updated_at"), - environment=environment, - created_by=created_by, + created_at=row.created_at, + updated_at=row.updated_at, + environment=row.environment, + created_by=row.created_by, ) @@ -431,10 +489,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} + where_clause: Final[dict[str, str]] = {"prompt_id": base_prompt_id} if environment: where_clause["environment"] = environment - db_prompts: Final = await PromptRepository(prisma_client).table.find_many( + db_prompts: Final = await _prompt_table(prisma_client).find_many( where=where_clause, order={"version": "desc"}, ) @@ -590,7 +648,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(prisma_client).find_many( where={"prompt_id": base_prompt_id}, distinct=["environment"], ) @@ -602,13 +660,13 @@ 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]] = { + where_clause: Final[dict[str, str | int]] = { "prompt_id": base_prompt_id, "environment": environment, } 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(prisma_client).find_many( where=where_clause, order={"version": "desc"}, take=1, @@ -721,7 +779,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(prisma_client).create( data={ "prompt_id": request.prompt_id, "version": new_version, @@ -811,7 +869,7 @@ 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(prisma_client).find_many(where={"prompt_id": base_prompt_id}) if not existing_prompts: raise HTTPException( @@ -835,7 +893,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(prisma_client).create( data={ "prompt_id": base_prompt_id, "version": new_version, @@ -936,12 +994,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} + delete_where: Final[dict[str, str]] = {"prompt_id": base_prompt_id} if environment: delete_where["environment"] = environment # Delete versions from the database (scoped to environment if provided) - await PromptRepository(prisma_client).table.delete_many(where=delete_where) + await _prompt_table(prisma_client).delete_many(where=delete_where) # Remove matching prompts from memory — scope to environment if provided if environment: @@ -967,7 +1025,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 +1093,14 @@ 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]] = { + find_where: Final[dict[str, str | int]] = { "prompt_id": base_prompt_id, "environment": env, } 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(prisma_client).find_many( where=find_where, order={"version": "desc"}, take=1, @@ -1084,7 +1144,7 @@ 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]] = { + update_data: Final[dict[str, str]] = { "litellm_params": updated_litellm_params.model_dump_json(), "prompt_info": updated_prompt_info.model_dump_json(), } @@ -1092,7 +1152,7 @@ async def patch_prompt( 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(prisma_client).update( where={"id": target_row.id}, data=update_data, ) @@ -1216,7 +1276,7 @@ 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( + result: Final[object] = await base_llm_response_processor.base_process_llm_request( request=fastapi_request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 79791607b2e..47e30555a4f 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -1,10 +1,12 @@ import json import os import re +from collections.abc import Awaitable, Mapping, Sequence from importlib.resources import files -from typing import Any, Final +from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, HTTPException, Request +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -32,14 +34,66 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ( ) from litellm.types.utils import LlmProviders +if TYPE_CHECKING: + from datetime import datetime + router: Final = APIRouter() +class _ProviderSupportEntry(TypedDict, total=False): + display_name: ReadOnly[str] + endpoints: ReadOnly[Mapping[str, bool]] + + +class _ProvidersFile(TypedDict, total=False): + providers: ReadOnly[Mapping[str, _ProviderSupportEntry]] + + +class _EndpointProviderEntry(TypedDict): + slug: ReadOnly[str] + display_name: ReadOnly[str] + + +class _EndpointEntry(TypedDict): + key: ReadOnly[str] + label: ReadOnly[str] + endpoint: ReadOnly[str] + providers: ReadOnly[Sequence[_EndpointProviderEntry]] + + +class _PluginRow(Protocol): + @property + def id(self) -> str: ... + + @property + def name(self) -> str: ... + + @property + def enabled(self) -> bool: ... + + @property + def created_at(self) -> "datetime | None": ... + + @property + def updated_at(self) -> "datetime | None": ... + + @property + def manifest_json(self) -> str | None: ... + + +class _PluginTableActions(Protocol): + def find_many(self, *, where: Mapping[str, bool]) -> Awaitable[Sequence[_PluginRow]]: ... + + +def _plugin_table(prisma_client: object) -> _PluginTableActions: + return ClaudeCodePluginRepository(prisma_client).table + + # --------------------------------------------------------------------------- # /public/endpoints — helpers # --------------------------------------------------------------------------- -_ENDPOINT_METADATA: Final[dict[str, dict[str, str]]] = { +_ENDPOINT_METADATA: Final[Mapping[str, Mapping[str, str]]] = { "chat_completions": {"label": "Chat Completions", "endpoint": "/chat/completions"}, "messages": {"label": "Messages", "endpoint": "/messages"}, "responses": {"label": "Responses", "endpoint": "/responses"}, @@ -108,12 +162,12 @@ def _clean_display_name(raw: str) -> str: return _SLUG_SUFFIX_RE.sub("", raw).strip() -def _build_endpoints(raw: dict[str, Any]) -> list[dict[str, Any]]: +def _build_endpoints(raw: _ProvidersFile) -> list[_EndpointEntry]: """Transform raw provider_endpoints_support_backup.json into the response shape.""" - providers: Final[dict[str, Any]] = raw.get("providers", {}) + providers: Final = raw.get("providers", {}) # Collect endpoint keys in insertion order (union across all providers). - seen: Final[set] = set() + seen: Final[set[str]] = set() all_keys: Final[list[str]] = [] for provider_data in providers.values(): for key in provider_data.get("endpoints", {}): @@ -121,13 +175,13 @@ def _build_endpoints(raw: dict[str, Any]) -> list[dict[str, Any]]: seen.add(key) all_keys.append(key) - result: Final[list[dict[str, Any]]] = [] + result: Final[list[_EndpointEntry]] = [] for key in all_keys: meta = _ENDPOINT_METADATA.get(key) label = meta["label"] if meta else key.replace("_", " ").title() path = meta["endpoint"] if meta else "/" + key.replace("_", "/") - supporting: list[dict[str, str]] = [ + supporting: list[_EndpointProviderEntry] = [ { "slug": slug, "display_name": _clean_display_name(pd.get("display_name", slug)), @@ -140,8 +194,10 @@ def _build_endpoints(raw: dict[str, Any]) -> list[dict[str, Any]]: return result -def _load_endpoints() -> list[dict[str, Any]]: - raw = json.loads(files("litellm").joinpath("provider_endpoints_support_backup.json").read_text(encoding="utf-8")) +def _load_endpoints() -> list[_EndpointEntry]: + raw: Final[_ProvidersFile] = json.loads( + files("litellm").joinpath("provider_endpoints_support_backup.json").read_text(encoding="utf-8") + ) return _build_endpoints(raw) @@ -235,12 +291,7 @@ async def get_mcp_servers(): ) public_mcp_servers: Final = global_mcp_server_manager.get_public_mcp_servers() - return [ - MCPPublicServer( - **server.model_dump(), - ) - for server in public_mcp_servers - ] + return [MCPPublicServer.model_validate(server.model_dump()) for server in public_mcp_servers] @router.get( @@ -259,7 +310,7 @@ async def public_skill_hub(): try: prisma_client: Final = await _get_prisma_client() - plugins: Final = await ClaudeCodePluginRepository(prisma_client).table.find_many(where={"enabled": True}) + plugins: Final = await _plugin_table(prisma_client).find_many(where={"enabled": True}) items: Final = [] for plugin in plugins: raw = plugin.manifest_json or {} diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index e7ae6031e45..9e2b1c9d82d 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -7,7 +7,8 @@ Provides: """ import base64 -from typing import Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status @@ -31,6 +32,9 @@ from litellm.proxy.vector_store_endpoints.utils import ( ) from litellm.repositories.table_repositories import ManagedVectorStoresRepository +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + router: Final = APIRouter() @@ -58,7 +62,7 @@ def _append_payload_to_scan_stack( payload_stack.append((value, next_depth)) -def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]: +def _collect_vector_store_ids_from_payload(payload: object) -> set[str]: vector_store_ids: Final[set[str]] = set() payload_stack: Final = [(payload, 0)] @@ -95,7 +99,7 @@ def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]: async def _authorize_nested_vector_store_ids( - payload: Any, + payload: object, user_api_key_dict: UserAPIKeyAuth, ) -> None: for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload)): @@ -109,7 +113,7 @@ def _build_file_metadata_entry( response: Any, file_data: tuple[str, bytes, str] | None = None, file_url: str | None = None, -) -> dict[str, Any]: +) -> Mapping[str, str | int | None]: """ Build a file metadata entry for storing in vector_store_metadata. @@ -159,8 +163,8 @@ def _build_file_metadata_entry( async def _save_vector_store_to_db_from_rag_ingest( response: Any, - ingest_options: dict[str, Any], - prisma_client, + ingest_options: Mapping[str, dict[str, str | None]], + prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, file_data: tuple[str, bytes, str] | None = None, file_url: str | None = None, @@ -299,9 +303,9 @@ async def parse_rag_ingest_request( headers: Final = _safe_get_request_headers(request) content_type = headers.get("content-type", "") - file_data = None - file_url = None - file_id = None + file_data: tuple[str, bytes, str] | None = None + file_url: str | None = None + file_id: str | None = None ingest_options: dict[str, Any] = {} if "multipart/form-data" in content_type: @@ -315,7 +319,7 @@ async def parse_rag_ingest_request( file_data = (file_obj.filename, file_content, file_obj.content_type) # Parse JSON from 'request' form field (contains full request body as JSON) - request_json_str: Final = form_data.get("request") + request_json_str: Final[str | bytes | None] = form_data.get("request") if request_json_str: request_data: Final = orjson.loads(request_json_str) ingest_options = request_data.get("ingest_options", {}) @@ -382,7 +386,7 @@ async def parse_rag_ingest_request( "api_key", "api_base", } - vector_store_opts: Final = ingest_options.get("vector_store", {}) + vector_store_opts: Final[object] = ingest_options.get("vector_store", {}) if isinstance(vector_store_opts, dict): for field in _BLOCKED_VECTOR_STORE_CREDENTIAL_PARAMS: if field in vector_store_opts: @@ -658,7 +662,7 @@ async def rag_query( ) # Add litellm data - request_data: dict[str, Any] = {} + request_data: dict[str, object] = {} request_data = await add_litellm_data_to_request( data=request_data, request=request, diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 383ada5a1bc..31ab3596418 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,9 +10,10 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast from fastapi import Request, Response +from fastapi.responses import StreamingResponse from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -20,25 +21,30 @@ 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 + async def background_streaming_task( polling_id: str, - data: dict, + data, polling_handler: ResponsePollingHandler, request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, - general_settings: dict, - llm_router, - proxy_config, - proxy_logging_obj, + general_settings, + 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 +75,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, diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7854b17a06f..e9e7ae908a5 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, Sequence +from typing import TYPE_CHECKING, Any, Final, TypeAlias, 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.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: TypeAlias = 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: list[str]) -> dict[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 @@ -96,14 +99,14 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: def _replace_file_search_tools( tools: Iterable[ToolParam] | None, -) -> tuple[list[dict[str, Any]], list[str]]: +) -> tuple[list[object], list[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]]] = [] + non_file_search: Final[list[object]] = [] vector_store_ids: Final[list[str]] = [] for tool in tools or []: @@ -172,7 +175,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) -> Any: """Read a field from either a dict/TypedDict or an attribute-based object.""" if isinstance(result, dict): return result.get(key, default) @@ -211,7 +214,7 @@ def _format_search_results_as_tool_output( def _build_search_results_for_include( results: list[VectorStoreSearchResult], -) -> list[dict[str, Any]]: +) -> list[dict[str, object]]: """ Convert VectorStoreSearchResult objects to the format expected in file_search_call.search_results (mirrors OpenAI's include= format). @@ -220,7 +223,7 @@ 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]]] = [] + formatted: Final[list[dict[str, object]]] = [] for result in results: file_id = _get_field(result, "file_id") or "" content_items = _get_field(result, "content") or [] @@ -243,7 +246,7 @@ def _build_file_search_call_output( queries: list[str], results: list[VectorStoreSearchResult] | None = None, include_search_results: bool = False, -) -> dict[str, Any]: +) -> dict[str, object]: """Build the file_search_call output item (mirrors OpenAI's format). Args: @@ -268,14 +271,14 @@ def _build_file_search_call_output( def _build_file_citation_annotations( results: list[VectorStoreSearchResult], text: str, -) -> list[dict[str, Any]]: +) -> list[dict[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]]] = [] + annotations: Final[list[dict[str, object]]] = [] index: Final = len(text) # cite at end of text block - seen_file_ids: Final[set] = set() + seen_file_ids: Final[set[object]] = set() for result in results: file_id = _get_field(result, "file_id") @@ -298,7 +301,7 @@ def _build_file_citation_annotations( def _build_message_output( response_text: str, results: list[VectorStoreSearchResult], -) -> dict[str, Any]: +) -> dict[str, object]: """Build the message output item with optional file_citation annotations.""" annotations: Final = _build_file_citation_annotations(results, response_text) return { @@ -330,8 +333,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: dict[str, object], + message_output: dict[str, object], first_response: ResponsesAPIResponse | None = None, ) -> ResponsesAPIResponse: """ @@ -343,7 +346,7 @@ 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_output: Final[list[dict[str, object]]] = [file_search_call_output, message_output] synthesized: Final = ResponsesAPIResponse( id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"), object="response", @@ -383,12 +386,12 @@ 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]]: +) -> tuple[bool, dict[str, object]]: include_items: Final[list[str]] = list(kwargs.get("include") or []) include_search_results: Final = "file_search_call.results" in include_items original_stream: Final = kwargs.get("stream") - updated_kwargs = kwargs + updated_kwargs: dict[str, object] = kwargs if original_stream: verbose_logger.debug( "Streaming is not yet supported for emulated file_search. Disabling stream for this request." @@ -398,7 +401,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 +413,7 @@ 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]: +def _resolve_queries_from_args(args: dict[str, Any], input: object) -> list[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: @@ -423,13 +426,13 @@ def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: async def _execute_file_search_tool_calls( - file_search_calls: list[Any], + file_search_calls: Sequence[object], all_vs_ids: list[str], - input: Any, + input: object, file_search_call_id: str, -) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult]]: +) -> tuple[list[object], list[str], list[VectorStoreSearchResult]]: """Run the vector search for each file_search tool_call and collect results.""" - tool_results: Final[list[dict[str, Any]]] = [] + tool_results: Final[list[object]] = [] all_queries: Final[list[str]] = [] all_results: Final[list[VectorStoreSearchResult]] = [] @@ -465,17 +468,17 @@ async def _execute_file_search_tool_calls( def _build_follow_up_input( - input: Any, + input: object, first_response: ResponsesAPIResponse, - tool_results: list[dict[str, Any]], -) -> list[Any]: + tool_results: list[object], +) -> list[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 = ( + original_input_items: Final[list[object]] = ( list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}] ) first_response_output_items: Final[list[Any]] = [] @@ -491,7 +494,7 @@ def _build_follow_up_input( 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 diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 186852f91c2..022b9ece32e 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -68,7 +68,7 @@ async def create_mcp_list_tools_events( # Convert tools to dict format for the event _mcp_tools_dict: Final = [ tool.model_dump() - if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")) + if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump", None)) else tool.__dict__ if hasattr(tool, "__dict__") else {"name": getattr(tool, "name", str(tool))} @@ -356,7 +356,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers_obj) # Also check if headers are provided in tools array (from request body) - tools: Final = self.original_request_params.get("tools") + tools: Final[Sequence[object] | None] = self.original_request_params.get("tools") if tools: for tool in tools: if isinstance(tool, dict) and tool.get("type") == "mcp": @@ -395,7 +395,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): def _make_stream_error_event(self) -> ResponsesAPIStreamingResponse: err: Final = self._stream_error - status_code: Final = getattr(err, "status_code", None) + status_code: Final[object] = getattr(err, "status_code", None) return ErrorEvent( type=ResponsesAPIStreamEvents.ERROR, sequence_number=self._last_sequence_number + 1, @@ -515,7 +515,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Capture the response ID from the first event to ensure consistency if self._cached_response_id is None and hasattr(chunk, "response"): - response_obj = getattr(chunk, "response", None) + response_obj: ResponsesAPIResponse | None = getattr(chunk, "response", None) if response_obj and hasattr(response_obj, "id"): self._cached_response_id = response_obj.id verbose_logger.debug("Cached response ID: %s", self._cached_response_id) @@ -559,7 +559,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): """Check if this chunk indicates the response is completed""" from litellm.types.llms.openai import ResponsesAPIStreamEvents - return getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + chunk_type: Final[object] = getattr(chunk, "type", None) + return chunk_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: """ @@ -571,14 +572,14 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): chunk: Final = await cast(Any, self.base_iterator).__anext__() if self._cached_response_id is None and hasattr(chunk, "response"): - new_response: Final = getattr(chunk, "response", None) + new_response: Final[ResponsesAPIResponse | None] = getattr(chunk, "response", None) new_response_id: Final = getattr(new_response, "id", None) if new_response is not None else None if new_response_id: self._cached_response_id = new_response_id # Ensure response ID consistency - update chunk if needed if self._cached_response_id and hasattr(chunk, "response"): - response_obj = getattr(chunk, "response", None) + response_obj: ResponsesAPIResponse | None = getattr(chunk, "response", None) if response_obj and hasattr(response_obj, "id"): if response_obj.id != self._cached_response_id: verbose_logger.debug( @@ -605,7 +606,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): from litellm.responses.main import aresponses # Make the initial response API call - but avoid the MCP wrapper - params: Final = self.original_request_params.copy() + params: Final[dict[str, object]] = self.original_request_params.copy() params["stream"] = True # Ensure streaming # Use the pre-fetched all_tools from original_request_params (no re-processing needed) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d7f6ece5cd1..25e5fcb6976 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,7 +5,7 @@ import json import time import traceback import uuid -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -1035,7 +1035,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @runtime_checkable class _HasModelDump(Protocol): - def model_dump(self, *, exclude_none: bool = ...) -> Mapping[str, object]: ... + def model_dump(self, *, exclude_none: bool = ...) -> dict[str, object]: ... @runtime_checkable @@ -1043,8 +1043,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) -> dict[str, Any]: + if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): return obj @@ -1134,7 +1134,8 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []): + annotations_payload: Final[Sequence[dict[str, object]]] = part_payload.get("annotations", []) or [] + for annotation_index, annotation in enumerate(annotations_payload): events.append( openai_types.OutputTextAnnotationAddedEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, @@ -1200,7 +1201,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 +1216,8 @@ def _build_synthetic_response_events( ) if item_type == "message": - for content_index, part in enumerate(output_item_payload.get("content", []) or []): + content_parts: Sequence[object] = output_item_payload.get("content", []) or [] + for content_index, part in enumerate(content_parts): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( @@ -1261,7 +1264,8 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): + summaries: Sequence[object] = output_item_payload.get("summary", []) or [] + for summary_index, summary in enumerate(summaries): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1463,7 +1467,8 @@ class ResponsesWebSocketStreaming: # masked response.completed. if self.output_guardrail_callbacks: try: - _evt_type = json.loads(response_str).get("type") + _evt_payload: Mapping[str, object] = json.loads(response_str) + _evt_type = _evt_payload.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: @@ -1527,7 +1532,7 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = json.loads(message) + msg_obj: Final[dict[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return message @@ -1544,7 +1549,8 @@ class ResponsesWebSocketStreaming: self.request_data["metadata"] = {} modified = model_modified - for cb in self.guardrail_callbacks: + guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks) + for cb in guardrail_cbs: presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) # response.create carries client text in two shapes: # flat: {"type": "response.create", "input": ..., "instructions": ...} @@ -1655,7 +1661,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final[dict[str, object]] = json.loads(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -2012,7 +2018,7 @@ class ManagedResponsesWebSocketHandler: async def _parse_message(self, raw_message: str) -> dict[str, object] | 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[dict[str, object]] = json.loads(raw_message) except json.JSONDecodeError: await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None @@ -2293,11 +2299,10 @@ 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) - if requested_model is None or requested_model == self.model_group: - model = self.model - else: - model = requested_model + requested_model: Final[str | None] = call_kwargs.pop("model", None) + model: Final[str] = ( + self.model if requested_model is None or requested_model == self.model_group else requested_model + ) previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None) current_messages: Final = self._input_to_messages(call_kwargs.get("input")) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 1907b5aa447..4b5def790ed 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -93,9 +93,9 @@ class ResponsesAPIRequestUtils: @staticmethod def merge_client_forwarded_headers( - extra_headers: dict[str, Any] | None, + extra_headers: dict[str, object] | None, client_headers: dict[str, str] | None, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """ Merge headers forwarded by the proxy (`headers` kwarg, set when `forward_client_headers_to_llm_api` is enabled) into `extra_headers`. @@ -210,9 +210,9 @@ class ResponsesAPIRequestUtils: valid_keys: Final = get_type_hints(ResponsesAPIOptionalRequestParams).keys() custom_llm_provider: Final = params.pop("custom_llm_provider", None) - special_params: Final = params.pop("kwargs", {}) + special_params: Final[dict[str, object]] = params.pop("kwargs", {}) - additional_drop_params: Final = params.pop("additional_drop_params", None) + additional_drop_params: Final[list[str] | None] = params.pop("additional_drop_params", None) non_default_params: Final = PreProcessNonDefaultParams.base_pre_process_non_default_params( passed_params=params, special_params=special_params, @@ -401,9 +401,9 @@ class ResponsesAPIRequestUtils: @staticmethod def _update_encrypted_content_item_ids_in_response( - response: Union["ResponsesAPIResponse", dict[str, Any]], + response: Union["ResponsesAPIResponse", dict[str, object]], model_id: str | None, - ) -> Union["ResponsesAPIResponse", dict[str, Any]]: + ) -> Union["ResponsesAPIResponse", dict[str, object]]: """Rewrite item IDs for output items that contain ``encrypted_content``. Encodes ``model_id`` into the item ID so that follow-up requests can be @@ -415,7 +415,7 @@ class ResponsesAPIRequestUtils: if not model_id: return response - output: list | None = None + output: object = None if isinstance(response, dict): output = response.get("output") else: @@ -459,7 +459,7 @@ class ResponsesAPIRequestUtils: return response @staticmethod - def _restore_encrypted_content_item_ids_in_input(request_input: Any) -> Any: + def _restore_encrypted_content_item_ids_in_input(request_input: object) -> Any: """Decode litellm-encoded item IDs in request input back to original IDs. Called before forwarding the request to the upstream provider so the @@ -867,7 +867,7 @@ class ResponsesAPIRequestUtils: ) @staticmethod - def collect_container_ids_from_responses_response(response: Any) -> list[str]: + def collect_container_ids_from_responses_response(response: object) -> list[str]: """Return unique container IDs referenced in a Responses API payload.""" if response is None: return [] @@ -953,7 +953,7 @@ class ResponsesAPIRequestUtils: @staticmethod def extract_mcp_headers_from_request( secret_fields: dict[str, Any] | None, - tools: Iterable[Any] | None, + tools: Iterable[object] | None, ) -> tuple[ str | None, dict[str, dict[str, str]] | None, diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 1120323b4f9..e4ac45df4d5 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -8,9 +8,11 @@ Use this to route requests between Teams """ import re -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict + +from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY @@ -25,9 +27,39 @@ else: LitellmRouter = Any +class _TagRoutingLitellmParams(TypedDict, total=False): + tags: ReadOnly[Sequence[str] | None] + tag_regex: ReadOnly[Sequence[str] | None] + + +class _TagRoutingDeployment(TypedDict, total=False): + model_name: ReadOnly[str] + litellm_params: ReadOnly[_TagRoutingLitellmParams] + model_info: ReadOnly[Mapping[str, object] | None] + + +class _TagRoutingMatchStamp(TypedDict): + matched_deployment: ReadOnly[str | None] + matched_via: ReadOnly[str] + matched_value: ReadOnly[str] + request_tags: ReadOnly[Sequence[str]] + user_agent: ReadOnly[str] + + +class _TagRoutingMetadata(TypedDict, total=False): + tags: ReadOnly[Sequence[str] | None] + inherited_tags: ReadOnly[Sequence[str] | None] + user_agent: ReadOnly[str] + tag_routing: ReadOnly[_TagRoutingMatchStamp] + _consumed_request_tags: ReadOnly[object] + + +_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) + + 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. @@ -77,11 +109,11 @@ def is_valid_deployment_tag( def _match_deployment( - deployment: Any, - request_tags: list[str] | None, - header_strings: list[str], + deployment: _TagRoutingDeployment, + request_tags: Sequence[str] | None, + header_strings: Sequence[str], match_any: bool, -) -> dict[str, str] | None: +) -> Mapping[str, str] | None: """ Determine whether *deployment* matches the current request. @@ -94,8 +126,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: @@ -166,38 +198,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[ def _exclude_deployments( - deployments: Sequence[Any] | Mapping[Any, Any], + deployments: Iterable[_TagRoutingDeployment], excluded_set: frozenset[str], -) -> list[Any]: +) -> list[_TagRoutingDeployment]: 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: Iterable[_TagRoutingDeployment], required_set: frozenset[str], -) -> tuple[Any, ...]: +) -> tuple[_TagRoutingDeployment, ...]: 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: Iterable[_TagRoutingDeployment], +) -> tuple[_TagRoutingDeployment, ...]: 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: Iterable[_TagRoutingDeployment]) -> frozenset[str]: return frozenset( - tag for d in deployments for tag in (d.get("litellm_params", MappingProxyType({})).get("tags") or ()) + tag for d in deployments for tag in (d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ()) ) def _unknown_required_tag_hides_an_answer( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: Iterable[_TagRoutingDeployment], excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -221,23 +253,23 @@ def _unknown_required_tag_hides_an_answer( def _chain_allows_fail_open( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: Iterable[_TagRoutingDeployment], excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], ) -> bool: if _unknown_required_tag_hides_an_answer(healthy_deployments, excluded_set, required_set, routing_confirmed): return False - return any((d.get("model_info") or {}).get("allow_fail_open") is True for d in healthy_deployments) + return any((d.get("model_info") or _EMPTY_MODEL_INFO).get("allow_fail_open") is True for d in healthy_deployments) def _trusted_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: Iterable[_TagRoutingDeployment], excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, inherited_required_set: frozenset[str] | None, -) -> tuple[Any, ...]: +) -> tuple[_TagRoutingDeployment, ...]: # 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 @@ -264,8 +296,8 @@ def _trusted_only_pool( def _resolve_or_fail_open( - pool: Sequence[Any], - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + pool: Sequence[_TagRoutingDeployment], + healthy_deployments: Iterable[_TagRoutingDeployment], excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -273,7 +305,7 @@ def _resolve_or_fail_open( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_TagRoutingDeployment, ...]: if pool: return tuple(pool) if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed): @@ -293,7 +325,7 @@ def _resolve_or_fail_open( def _resolve_constraint_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: Iterable[_TagRoutingDeployment], excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -301,7 +333,7 @@ def _resolve_constraint_only_pool( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_TagRoutingDeployment, ...]: pool: Final = ( _require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set) if required_set @@ -323,8 +355,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: Iterable[_TagRoutingDeployment], +) -> Iterable[_TagRoutingDeployment]: 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 @@ -334,8 +366,8 @@ def _all_deployments_or_fallback( def _chain_tag_filtering_override( llm_router_instance: LitellmRouter, model: str, - healthy_deployments: Sequence[Any] | Mapping[Any, Any], -) -> bool | None: + healthy_deployments: Iterable[_TagRoutingDeployment], +) -> object: # Resolved from every deployment configured for this model group, not just the # ones that survived cooldown/health filtering (async_get_healthy_deployments # filters cooldowns before calling get_deployments_for_tag) -- otherwise the @@ -347,14 +379,14 @@ def _chain_tag_filtering_override( # than crashing the request. all_deployments: Final = _all_deployments_or_fallback(llm_router_instance, model, healthy_deployments) for d in all_deployments: - value = (d.get("model_info") or MappingProxyType({})).get("enable_tag_filtering") + value = (d.get("model_info") or _EMPTY_MODEL_INFO).get("enable_tag_filtering") if value is not None: return value return None def _inherited_constraint_sets( - inherited_tags: object, routing_prefix: str + inherited_tags: Sequence[str] | None, routing_prefix: str ) -> tuple[frozenset[str] | None, frozenset[str] | None]: # None means no origin information is available at all (e.g. this request # bypassed the proxy layer that populates metadata.inherited_tags, as direct @@ -385,15 +417,18 @@ def _tag_known_to_group( if tag_set & routing_confirmed: return True try: - all_deployments: Final = llm_router_instance._get_all_deployments(model_name=model) + all_deployments: Final[Sequence[_TagRoutingDeployment]] = llm_router_instance._get_all_deployments( + model_name=model + ) except Exception: # noqa: BLE001 # fail safe toward "unrecognized" so lookup errors preserve the existing silent-fallback behavior return False return any( - tag_set.intersection(d.get("litellm_params", MappingProxyType({})).get("tags") or ()) for d in all_deployments + tag_set.intersection(d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ()) + for d in all_deployments ) -def _request_tags_after_router_consumption(metadata: Mapping[Any, Any], model: str) -> Sequence[str] | None: +def _request_tags_after_router_consumption(metadata: _TagRoutingMetadata, model: str) -> Sequence[str] | None: # The pre-routing hook stamps which tags selected the router it rewrote the request # to: those tags already did their job and must not also constrain deployment choice # inside the routed group. The request's other tags still apply there, on top of the @@ -451,7 +486,8 @@ async def get_deployments_for_tag( verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name)) if metadata_variable_name in request_kwargs: - metadata: Final = request_kwargs[metadata_variable_name] + metadata: Final[_TagRoutingMetadata] = request_kwargs[metadata_variable_name] + stampable_metadata: Final[dict[str, object]] = request_kwargs[metadata_variable_name] request_tags: Final = _request_tags_after_router_consumption(metadata, model) match_any: Final = llm_router_instance.tag_filtering_match_any routing_prefix: Final = llm_router_instance.tag_routing_prefix or "" @@ -496,8 +532,8 @@ async def get_deployments_for_tag( request_tags, ) - new_healthy_deployments: Final[list[Any]] = [] - default_deployments: Final[list[Any]] = [] + new_healthy_deployments: Final[list[_TagRoutingDeployment]] = [] + default_deployments: Final[list[_TagRoutingDeployment]] = [] if has_positive_filter: verbose_logger.debug( @@ -523,7 +559,7 @@ async def get_deployments_for_tag( match_result["matched_value"], ) if "tag_routing" not in metadata: - metadata["tag_routing"] = { + stampable_metadata["tag_routing"] = { "matched_deployment": deployment.get("model_name"), "matched_via": match_result["matched_via"], "matched_value": match_result["matched_value"], @@ -568,7 +604,7 @@ 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 = [] + _default_deployments_with_tags: Final[list[_TagRoutingDeployment]] = [] for deployment in healthy_deployments: if "default" in deployment.get("litellm_params", {}).get("tags", []): _default_deployments_with_tags.append(deployment) @@ -603,7 +639,7 @@ def _tags_in_metadata(metadata: object) -> list[str]: def _get_tags_from_request_kwargs( - request_kwargs: Mapping[Any, Any] | None = None, + request_kwargs: Mapping[str, object] | None = None, metadata_variable_name: Literal["metadata", "litellm_metadata"] | None = None, ) -> list[str]: """ diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index dff010bfd30..17c8f02dfdd 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3058 + "limit": 3046 }, "ANN002": { "limit": 71 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1384 + "limit": 1342 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 505 }, "B009": { - "limit": 64 + "limit": 60 }, "B010": { "limit": 190 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1224 + "limit": 1220 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index a7286d9a89a..0104733bcf3 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23001 + "limit": 22943 }, "LIT002": { - "limit": 27146 + "limit": 27141 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1077 + "limit": 1074 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16731 + "limit": 16722 }, "LIT011": { "limit": 5596