diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 32b8eb3d4d0..8f42847bd2c 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 28842 + "limit": 27781 }, "reportArgumentType": { - "limit": 2634 + "limit": 2614 }, "reportAssignmentType": { "limit": 329 @@ -18,13 +18,13 @@ "limit": 40 }, "reportDeprecated": { - "limit": 215 + "limit": 214 }, "reportDuplicateImport": { "limit": 19 }, "reportExplicitAny": { - "limit": 9103 + "limit": 8969 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5843 + "limit": 5834 }, "reportMissingTypeArgument": { - "limit": 15816 + "limit": 15712 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1078 + "limit": 1069 }, "reportOptionalOperand": { "limit": 0 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 218 + "limit": 215 }, "reportTypedDictNotRequiredAccess": { "limit": 27 @@ -99,22 +99,22 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45098 + "limit": 45016 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39826 + "limit": 39641 }, "reportUnknownParameterType": { - "limit": 20237 + "limit": 20140 }, "reportUnknownVariableType": { - "limit": 31371 + "limit": 31198 }, "reportUnnecessaryCast": { - "limit": 122 + "limit": 119 }, "reportUnnecessaryComparison": { "limit": 701 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 864 + "limit": 860 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 339da998d56..31d0c6109a3 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,10 +6,13 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from collections.abc import AsyncIterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any, Final, Protocol, cast, overload, runtime_checkable from uuid import uuid4 +from pydantic import TypeAdapter +from typing_extensions import TypedDict + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -17,6 +20,47 @@ from litellm.llms.custom_httpx.http_handler import ( ) +class _A2AStatusView(TypedDict, total=False): + state: str + + +class _A2AMessageView(TypedDict, total=False): + role: str + parts: Sequence[Mapping[str, object]] + messageId: str + + +class _A2AArtifactView(TypedDict, total=False): + parts: Sequence[Mapping[str, object]] + + +class _A2ATaskResultView(TypedDict, total=False): + id: str + status: _A2AStatusView + history: Sequence[_A2AMessageView] + artifacts: Sequence[_A2AArtifactView] + message: _A2AMessageView + + +class _A2ATaskResponseView(TypedDict, total=False): + result: _A2ATaskResultView + + +@runtime_checkable +class _SupportsModelDump(Protocol): + def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ... + + +@runtime_checkable +class _SupportsDict(Protocol): + def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ... + + +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_TASK_VIEW_ADAPTER: Final = TypeAdapter(_A2ATaskResponseView) +_STR_ADAPTER: Final = TypeAdapter(str) + + class PydanticAITransformation: """ Transformation layer for Pydantic AI agents. @@ -27,8 +71,16 @@ class PydanticAITransformation: - Fake streaming by chunking non-streaming responses """ + @overload @staticmethod - def _remove_none_values(obj: Any) -> Any: + def _remove_none_values(obj: Mapping[str, object]) -> dict[str, object]: ... + + @overload + @staticmethod + def _remove_none_values(obj: object) -> object: ... + + @staticmethod + def _remove_none_values(obj: object) -> object: """ Recursively remove None values from a dict/list structure. @@ -42,14 +94,18 @@ class PydanticAITransformation: Cleaned object with None values removed """ if isinstance(obj, dict): - return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None} + mapping: Final[dict[str, object]] = obj + return {k: PydanticAITransformation._remove_none_values(v) for k, v in mapping.items() if v is not None} elif isinstance(obj, list): - return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None] + items: Final[list[object]] = obj + return [PydanticAITransformation._remove_none_values(item) for item in items if item is not None] else: return obj @staticmethod - def _params_to_dict(params: Any) -> dict[str, Any]: + def _params_to_dict( + params: "_SupportsModelDump | _SupportsDict | Mapping[str, object]", + ) -> Mapping[str, object]: """ Convert params to a dict, handling Pydantic models. @@ -59,10 +115,10 @@ class PydanticAITransformation: Returns: Dict representation of params """ - if hasattr(params, "model_dump"): + if isinstance(params, _SupportsModelDump): # Pydantic v2 model return params.model_dump(mode="python", exclude_none=True) - elif hasattr(params, "dict"): + elif isinstance(params, _SupportsDict): # Pydantic v1 model return params.dict(exclude_none=True) elif isinstance(params, dict): @@ -80,7 +136,7 @@ class PydanticAITransformation: max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -112,9 +168,10 @@ class PydanticAITransformation: }, ) response.raise_for_status() - poll_data = response.json() + poll_data = _JSON_OBJECT_ADAPTER.validate_python(response.json()) - result = poll_data.get("result", {}) + view = _TASK_VIEW_ADAPTER.validate_python(poll_data) + result = view.get("result", {}) status = result.get("status", {}) state = status.get("state", "") @@ -133,10 +190,10 @@ class PydanticAITransformation: async def _send_and_poll_raw( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -160,7 +217,9 @@ class PydanticAITransformation: # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI if "message" in params_dict: - params_dict["message"]["kind"] = "message" + message: Final = _JSON_OBJECT_ADAPTER.validate_python(params_dict["message"]) + message["kind"] = "message" + params_dict["message"] = message # Build A2A JSON-RPC request using message/send method for FastA2A compatibility a2a_request: Final = { @@ -189,10 +248,11 @@ class PydanticAITransformation: }, ) response.raise_for_status() - response_data = response.json() + response_data = _JSON_OBJECT_ADAPTER.validate_python(response.json()) # Check if task is already completed - result: Final = response_data.get("result", {}) + view: Final = _TASK_VIEW_ADAPTER.validate_python(response_data) + result: Final = view.get("result", {}) status: Final = result.get("status", {}) state: Final = status.get("state", "") @@ -217,10 +277,10 @@ class PydanticAITransformation: async def send_non_streaming_request( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -253,10 +313,10 @@ class PydanticAITransformation: async def send_and_get_raw_response( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -282,9 +342,9 @@ class PydanticAITransformation: @staticmethod def _transform_to_a2a_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform Pydantic AI task response to standard A2A non-streaming format. @@ -313,7 +373,7 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Build standard A2A message - a2a_message: Final = { + a2a_message: Final[Mapping[str, object]] = { "kind": "message", "role": "agent", "parts": parts if parts else [{"kind": "text", "text": full_text}], @@ -328,7 +388,7 @@ class PydanticAITransformation: } @staticmethod - def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]: + def _extract_response_text(response_data: Mapping[str, object]) -> tuple[str, str, Sequence[Mapping[str, object]]]: """ Extract response text from completed task response. @@ -342,7 +402,8 @@ class PydanticAITransformation: Returns: Tuple of (full_text, message_id, parts) """ - result: Final = response_data.get("result", {}) + view: Final = _TASK_VIEW_ADAPTER.validate_python(response_data) + result: Final = view.get("result", {}) # Try to extract from artifacts first (preferred for results) artifacts: Final = result.get("artifacts", []) @@ -351,7 +412,7 @@ class PydanticAITransformation: parts = artifact.get("parts", []) for part in parts: if part.get("kind") == "text": - text = part.get("text", "") + text = _STR_ADAPTER.validate_python(part.get("text", "")) if text: return text, str(uuid4()), parts @@ -364,7 +425,7 @@ class PydanticAITransformation: full_text = "" for part in parts: if part.get("kind") == "text": - full_text += part.get("text", "") + full_text += _STR_ADAPTER.validate_python(part.get("text", "")) if full_text: return full_text, message_id, parts @@ -376,18 +437,18 @@ class PydanticAITransformation: full_text = "" for part in parts: if part.get("kind") == "text": - full_text += part.get("text", "") + full_text += _STR_ADAPTER.validate_python(part.get("text", "")) return full_text, message_id, parts return "", str(uuid4()), [] @staticmethod async def fake_streaming_from_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Convert a non-streaming A2A response into fake streaming chunks. @@ -410,9 +471,10 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history - result: Final = response_data.get("result", {}) + view: Final = _TASK_VIEW_ADAPTER.validate_python(response_data) + result: Final = view.get("result", {}) history: Final = result.get("history", []) - input_message = {} + input_message: _A2AMessageView = {} for msg in history: if msg.get("role") == "user": input_message = msg @@ -426,7 +488,7 @@ class PydanticAITransformation: # 1. Emit initial task event (kind: "task", status: "submitted") # Format matches A2ACompletionBridgeTransformation.create_task_event - task_event: Final = { + task_event: Final[dict[str, object]] = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -452,7 +514,7 @@ class PydanticAITransformation: # 2. Emit status update (kind: "status-update", status: "working") # Format matches A2ACompletionBridgeTransformation.create_status_update_event - working_event: Final = { + working_event: Final[dict[str, object]] = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -478,7 +540,7 @@ class PydanticAITransformation: chunk_text = full_text[i : i + chunk_size] is_last_chunk = (i + chunk_size) >= len(full_text) - artifact_event = { + artifact_event: dict[str, object] = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -503,7 +565,7 @@ class PydanticAITransformation: await asyncio.sleep(delay_ms / 1000.0) # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event: Final = { + completed_event: Final[dict[str, object]] = { "jsonrpc": "2.0", "id": request_id, "result": { diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 858d10df53b..bf2cfc4716a 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,6 +1,7 @@ import asyncio import json -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from collections.abc import Mapping, Sized +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeGuard, cast import litellm from litellm._logging import verbose_logger @@ -20,15 +21,61 @@ from .litellm_logging import Logging as LiteLLMLogging if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.guardrails import GuardrailEventHooks + CLIENT_CONNECTION_CLASS = ClientConnection else: CLIENT_CONNECTION_CLASS = Any +def _is_str_keyed_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, dict) + + +def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, list) + + +def _output_modalities_from_beta(mods: object) -> list[str] | None: + mods_list: Final = mods if _is_object_list(mods) else [] + mods_set: Final = {m.lower() for m in mods_list if isinstance(m, str)} + if "audio" in mods_set: + return ["audio"] + if "text" in mods_set: + return ["text"] + return None + + +def _parse_json(payload: str | bytes) -> object: + return cast(object, json.loads(payload)) + + +def _json_loads_dict(payload: str | bytes) -> dict[str, object]: + parsed: Final = _parse_json(payload) + return parsed if _is_str_keyed_dict(parsed) else {} + + +def _is_two_item_pair(value: object) -> TypeGuard[tuple[object, object]]: # guard-ok: isinstance-checked 2-item list/tuple narrowed to a pair for unpacking # fmt: skip + return isinstance(value, Sized) and len(value) == 2 and isinstance(value, (list, tuple)) + + +class _ClientWebSocketExceptionTypes(Protocol): + ConnectionClosed: type[BaseException] + + +class ClientWebSocketInterface(Protocol): + exceptions: _ClientWebSocketExceptionTypes + + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... - def normalize(self, event: dict) -> dict: ... - def patch_outgoing_session(self, session: dict) -> dict: ... + def normalize(self, event: dict[str, object]) -> dict[str, object]: ... + def patch_outgoing_session(self, session: dict[str, object]) -> dict[str, object]: ... DefaultLoggedRealTimeEventTypes: Final = [ @@ -43,13 +90,13 @@ DefaultLoggedRealTimeEventTypes: Final = [ class RealTimeStreaming: def __init__( self, - websocket: Any, + websocket: ClientWebSocketInterface, backend_ws: CLIENT_CONNECTION_CLASS, logging_obj: LiteLLMLogging, provider_config: BaseRealtimeConfig | None = None, model: str = "", - user_api_key_dict: Any | None = None, - request_data: dict | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, + request_data: Mapping[str, object] | None = None, backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, @@ -57,11 +104,11 @@ class RealTimeStreaming: self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj - self.messages: list[OpenAIRealtimeEvents] = [] - self.input_message: dict = {} - self.input_messages: list[dict[str, str]] = [] - self.session_tools: list[dict] = [] - self.tool_calls: list[dict] = [] + self.messages: list[OpenAIRealtimeEvents | Mapping[str, object]] = [] + self.input_message: dict[str, object] = {} + self.input_messages: list[dict[str, object]] = [] + self.session_tools: list[object] = [] + self.tool_calls: list[dict[str, object]] = [] # Detect whether the client is explicitly opting into the beta protocol. self._client_wants_beta = self._detect_beta_header(websocket) @@ -84,7 +131,7 @@ class RealTimeStreaming: self.current_delta_type: ALL_DELTA_TYPES | None = None self.session_configuration_request: str | None = None self.user_api_key_dict = user_api_key_dict - self.request_data: dict = request_data or {} + self.request_data: Mapping[str, object] = request_data or {} # Violation counter for end_session_after_n_fails support self._violation_count: int = 0 # When a text message is blocked, hold the guardrail reason so the next @@ -127,7 +174,7 @@ class RealTimeStreaming: ] ) _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"]) - _AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = { + _AUDIO_FORMAT_MAP: Mapping[str, Mapping[str, str | int]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000}, @@ -149,16 +196,16 @@ class RealTimeStreaming: def _should_store_message( self, - message_obj: dict | OpenAIRealtimeEvents, + message_obj: Mapping[str, object], ) -> bool: _msg_type: Final = message_obj["type"] if "type" in message_obj else None if self.logged_real_time_event_types == "*": return True - if _msg_type and _msg_type in self.logged_real_time_event_types: + if _msg_type and isinstance(_msg_type, str) and _msg_type in self.logged_real_time_event_types: return True return False - def store_message(self, message: str | bytes | dict | OpenAIRealtimeEvents): + def store_message(self, message: str | bytes | dict[str, object] | OpenAIRealtimeEvents): """Store message in list""" if isinstance(message, bytes): message = message.decode("utf-8") @@ -183,34 +230,35 @@ class RealTimeStreaming: return self.messages.append(typed_obj) - def _collect_user_input_from_client_event(self, message: str | dict) -> None: + def _collect_user_input_from_client_event(self, message: str | dict[str, object]) -> None: """Extract user text content from client WebSocket events for spend logging.""" try: if isinstance(message, str): - msg_obj = json.loads(message) - elif isinstance(message, dict): - msg_obj = message + msg_obj: dict[str, object] = _json_loads_dict(message) else: - return + msg_obj = message msg_type: Final = msg_obj.get("type", "") if msg_type == "conversation.item.create": - item: Final = msg_obj.get("item", {}) + item_raw: Final = msg_obj.get("item", {}) + item: Final = item_raw if _is_str_keyed_dict(item_raw) else {} if item.get("role") == "user": - content_list: Final = item.get("content", []) + content_raw: Final = item.get("content", []) + content_list: Final = content_raw if _is_object_list(content_raw) else [] for content in content_list: - if isinstance(content, dict) and content.get("type") == "input_text": + if _is_str_keyed_dict(content) and content.get("type") == "input_text": text = content.get("text", "") if text: self.input_messages.append({"role": "user", "content": text}) elif msg_type == "session.update": - session: Final = msg_obj.get("session", {}) + session_raw: Final = msg_obj.get("session", {}) + session: Final = session_raw if _is_str_keyed_dict(session_raw) else {} instructions: Final = session.get("instructions", "") if instructions: self.input_messages.append({"role": "system", "content": instructions}) tools: Final = session.get("tools") - if tools and isinstance(tools, list): + if tools and _is_object_list(tools): self.session_tools = tools # GA: session.type is required; log it for traceability but no action needed verbose_logger.debug("Realtime session.type: %s", session.get("type")) @@ -219,18 +267,18 @@ class RealTimeStreaming: except (json.JSONDecodeError, AttributeError, TypeError): pass - def _collect_user_input_from_backend_event(self, event_obj: dict | OpenAIRealtimeEvents) -> None: + def _collect_user_input_from_backend_event(self, event_obj: Mapping[str, object]) -> None: """Extract user voice transcription from backend events for spend logging.""" try: event_type: Final = event_obj.get("type", "") if event_type == "conversation.item.input_audio_transcription.completed": - transcript: Final = cast(str, event_obj.get("transcript", "")) + transcript: Final = event_obj.get("transcript", "") if transcript: self.input_messages.append({"role": "user", "content": transcript}) except (AttributeError, TypeError): pass - def _detect_transcription_session_from_backend(self, event_obj: dict | OpenAIRealtimeEvents) -> None: + def _detect_transcription_session_from_backend(self, event_obj: Mapping[str, object]) -> None: """Flag transcription-only sessions from backend session events.""" try: event_type: Final = event_obj.get("type", "") @@ -240,13 +288,13 @@ class RealTimeStreaming: ): self._is_transcription_session = True elif event_type in ("session.created", "session.updated"): - session: Final = cast(dict, event_obj).get("session", {}) or {} - if session.get("type") == "transcription": + session_raw: Final = event_obj.get("session") + if _is_str_keyed_dict(session_raw) and session_raw.get("type") == "transcription": self._is_transcription_session = True except (AttributeError, TypeError): pass - def _capture_transcription_usage(self, event_obj: dict | OpenAIRealtimeEvents) -> None: + def _capture_transcription_usage(self, event_obj: Mapping[str, object]) -> None: """ Append a usage-only transcription completed event to the logged results so the cost calculator can bill it by audio duration. The default logged event @@ -264,24 +312,28 @@ class RealTimeStreaming: if self._should_store_message(event_obj): return self.messages.append( - cast( - OpenAIRealtimeEvents, - { - "type": "conversation.item.input_audio_transcription.completed", - "usage": usage, - }, - ) + { + "type": "conversation.item.input_audio_transcription.completed", + "usage": usage, + } ) except (AttributeError, TypeError): pass - def _collect_tool_calls_from_response_done(self, event_obj: dict | OpenAIRealtimeEvents) -> None: + def _collect_tool_calls_from_response_done(self, event_obj: Mapping[str, object]) -> None: """Extract function_call items from response.done events for spend logging.""" try: if event_obj.get("type") != "response.done": return - response: Final = cast(dict[str, Any], event_obj.get("response", {})) - for item in response.get("output", []): + response: Final = event_obj.get("response", {}) + if not _is_str_keyed_dict(response): + return + output: Final = response.get("output", []) + if not _is_object_list(output): + return + for item in output: + if not _is_str_keyed_dict(item): + return if item.get("type") == "function_call": self.tool_calls.append( { @@ -296,7 +348,7 @@ class RealTimeStreaming: except (AttributeError, TypeError): pass - def store_input(self, message: str | dict): + def store_input(self, message: str | dict[str, object]): """Store input message""" self.input_message = message if isinstance(message, dict) else {} self._collect_user_input_from_client_event(message) @@ -338,10 +390,10 @@ class RealTimeStreaming: sent = False for msg in transformed: try: - msg_obj = json.loads(msg) + msg_obj: object = _parse_json(msg) except (json.JSONDecodeError, TypeError): msg_obj = None - if isinstance(msg_obj, dict) and self.provider_config.is_setup_message(msg_obj): + if _is_str_keyed_dict(msg_obj) and self.provider_config.is_setup_message(msg_obj): if self._content_sent_after_setup: verbose_logger.debug("Dropping follow-up setup after content was already sent to backend") continue @@ -350,7 +402,9 @@ class RealTimeStreaming: self._cache_session_configuration_request(msg) sent = True else: - is_content_message = isinstance(msg_obj, dict) and self.provider_config.is_content_message(msg_obj) + is_content_message = _is_str_keyed_dict(msg_obj) and self.provider_config.is_content_message( + msg_obj + ) # Send first, then mutate state, so a failed send leaves both # ``session_configuration_request`` and # ``_content_sent_after_setup`` untouched. Caching or marking @@ -384,7 +438,7 @@ class RealTimeStreaming: return message try: - message_obj: Final = json.loads(message) + message_obj: Final = _json_loads_dict(message) except (json.JSONDecodeError, TypeError): return message @@ -395,7 +449,7 @@ class RealTimeStreaming: return message session: Final = message_obj.get("session") - if not isinstance(session, dict): + if not _is_str_keyed_dict(session): return message if session.get("type") == "transcription": @@ -405,7 +459,7 @@ class RealTimeStreaming: changed = False transcription: Final = session.get("input_audio_transcription") - if isinstance(transcription, dict) and transcription.get("model") != authorized_model: + if _is_str_keyed_dict(transcription) and transcription.get("model") != authorized_model: session["input_audio_transcription"] = { **transcription, "model": authorized_model, @@ -413,11 +467,11 @@ class RealTimeStreaming: changed = True audio: Final = session.get("audio") - if isinstance(audio, dict): + if _is_str_keyed_dict(audio): audio_input: Final = audio.get("input") - if isinstance(audio_input, dict): + if _is_str_keyed_dict(audio_input): nested_transcription: Final = audio_input.get("transcription") - if isinstance(nested_transcription, dict) and nested_transcription.get("model") != authorized_model: + if _is_str_keyed_dict(nested_transcription) and nested_transcription.get("model") != authorized_model: session["audio"] = { **audio, "input": { @@ -453,7 +507,7 @@ class RealTimeStreaming: for message in messages: try: - msg_type = json.loads(message).get("type") + msg_type = _json_loads_dict(message).get("type") except (json.JSONDecodeError, TypeError): collapsed.extend(pending_appends) pending_appends = [] @@ -487,14 +541,14 @@ class RealTimeStreaming: if self._backend_setup_complete and not self._flushing_pending_messages_until_setup: return False try: - msg_obj: Final = json.loads(message) + msg_obj: Final = _json_loads_dict(message) except (json.JSONDecodeError, TypeError): return False return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES def _buffer_pending_message_until_setup(self, message: str) -> None: try: - msg_type = json.loads(message).get("type") + msg_type: object = _json_loads_dict(message).get("type") except (json.JSONDecodeError, TypeError): msg_type = None @@ -546,22 +600,22 @@ class RealTimeStreaming: return self._event_normalizer.should_drop(event) return False - def _normalize_event_for_ga_client(self, event: dict) -> dict: + def _normalize_event_for_ga_client(self, event: dict[str, object]) -> dict[str, object]: """Apply per-provider GA normalization before forwarding to clients.""" if self._event_normalizer is not None: return self._event_normalizer.normalize(event) return event - def _event_to_client_json(self, event: dict) -> str: + def _event_to_client_json(self, event: dict[str, object]) -> str: return json.dumps(self._normalize_event_for_ga_client(event)) - async def _send_event_to_client(self, event: Any, event_str: str) -> bool: + async def _send_event_to_client(self, event: object, event_str: str) -> bool: if self._should_drop_event_from_client(event): return False - if isinstance(event, dict): + if _is_str_keyed_dict(event): event = self._normalize_event_for_ga_client(event) event_str = json.dumps(event) - if self._client_wants_beta and isinstance(event, dict): + if self._client_wants_beta and _is_str_keyed_dict(event): try: translated: Final = self._translate_event_to_beta(event) if translated is None: @@ -587,20 +641,20 @@ class RealTimeStreaming: ``return_new_content_delta_events`` modality lookup, ...). """ try: - message_obj: Final = json.loads(transformed_message) - if "setup" in message_obj: + message_obj: Final = _parse_json(transformed_message) + if _is_str_keyed_dict(message_obj) and "setup" in message_obj: self.session_configuration_request = transformed_message except (json.JSONDecodeError, TypeError): return def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, Any]] = { + turn_detection: Final[Mapping[str, str | bool]] = { "type": "server_vad", "create_response": False, } if self._backend_uses_beta_protocol: - session: dict[str, Any] = {"turn_detection": turn_detection} + session: Mapping[str, object] = {"turn_detection": turn_detection} else: session = { "type": "realtime", @@ -654,7 +708,7 @@ class RealTimeStreaming: def _has_realtime_guardrails_for_event_hooks( self, - event_hooks: list[Any], + event_hooks: "list[GuardrailEventHooks]", ) -> bool: """Return True if any callback would run for one of ``event_hooks``.""" from litellm.integrations.custom_guardrail import CustomGuardrail @@ -699,7 +753,7 @@ class RealTimeStreaming: transcript: str, item_id: str | None = None, pre_block_backend_message: str | None = None, - event_hooks: list[Any] | None = None, + event_hooks: "list[GuardrailEventHooks] | None" = None, ) -> bool: """ Run registered guardrails on realtime text (transcript, user message, tool output). @@ -724,8 +778,8 @@ class RealTimeStreaming: if event_hooks is None: event_hooks = [GuardrailEventHooks.realtime_input_transcription] _realtime_event_types: Final = event_hooks - _check_data: Final = {**self.request_data, "transcript": transcript} - _already_run: Final[set] = set() + _check_data: Final[dict[str, object]] = {**self.request_data, "transcript": transcript} + _already_run: Final[set[int]] = set() for callback in litellm.callbacks: if not isinstance(callback, CustomGuardrail): @@ -753,8 +807,8 @@ class RealTimeStreaming: raise # Extract the human-readable error from the detail dict (HTTPException) # or fall back to str(e) for plain ValueError. - detail = getattr(e, "detail", None) - if isinstance(detail, dict): + detail: object = getattr(e, "detail", None) + if _is_str_keyed_dict(detail): safe_msg = detail.get("error") or str(e) elif detail is not None: safe_msg = str(detail) @@ -762,7 +816,8 @@ class RealTimeStreaming: safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." # Use realtime_violation_message if configured; fall back to guardrail error text. - error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg + violation_message: object = getattr(callback, "realtime_violation_message", None) + error_msg = violation_message or safe_msg # Deliver any caller-supplied backend message FIRST so that # protocol contracts requiring a specific ordering (e.g. @@ -826,7 +881,7 @@ class RealTimeStreaming: return True return False - async def _handle_provider_config_message(self, raw_response) -> None: + async def _handle_provider_config_message(self, raw_response: str | bytes) -> None: """Process a backend message when a provider_config is set (transformed path).""" returned_object: Final = self.provider_config.transform_realtime_response( raw_response, @@ -855,7 +910,7 @@ class RealTimeStreaming: for event in events: if self._should_drop_event_from_client(event): continue - is_session_created_event = isinstance(event, dict) and event.get("type") == "session.created" + is_session_created_event = _is_str_keyed_dict(event) and event.get("type") == "session.created" if is_session_created_event: if self._uses_deferred_backend_setup() and not self._backend_setup_complete: self._backend_setup_complete = True @@ -893,14 +948,18 @@ class RealTimeStreaming: await self._maybe_send_guardrail_turn_detection_update() continue ## GUARDRAIL: run on transcription events in provider_config path too - if isinstance(event, dict) and event.get("type") == "conversation.item.input_audio_transcription.completed": + if ( + _is_str_keyed_dict(event) + and event.get("type") == "conversation.item.input_audio_transcription.completed" + ): transcript = event.get("transcript", "") - self._collect_user_input_from_backend_event(cast(dict, event)) + item_id_raw = event.get("item_id") + self._collect_user_input_from_backend_event(event) self.store_message(event_str) await self._send_event_to_client(event, event_str) blocked = await self.run_realtime_guardrails( - cast(str, transcript), - item_id=cast(str | None, event.get("item_id")), + transcript if isinstance(transcript, str) else "", + item_id=item_id_raw if isinstance(item_id_raw, str) else None, ) if not blocked: await self._send_to_backend(json.dumps({"type": "response.create"})) @@ -910,15 +969,15 @@ class RealTimeStreaming: await self._send_event_to_client(event, event_str) @staticmethod - def _parse_backend_event(raw_response: str) -> dict | None: + def _parse_backend_event(raw_response: str) -> dict[str, object] | None: """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" try: - event: Final = json.loads(raw_response) + event: Final = _parse_json(raw_response) except (json.JSONDecodeError, TypeError): return None - return event if isinstance(event, dict) else None + return event if _is_str_keyed_dict(event) else None - async def _handle_raw_backend_message(self, event_obj: dict, raw_response: str) -> bool: + async def _handle_raw_backend_message(self, event_obj: dict[str, object], raw_response: str) -> bool: """Process a backend message without provider_config (raw path). Returns True if the caller should skip the default store+forward (i.e. continue the loop). @@ -949,9 +1008,10 @@ class RealTimeStreaming: self._capture_transcription_usage(event_obj) return True + item_id_raw: Final = event_obj.get("item_id") blocked: Final = await self.run_realtime_guardrails( - transcript, - item_id=event_obj.get("item_id"), + transcript if isinstance(transcript, str) else "", + item_id=item_id_raw if isinstance(item_id_raw, str) else None, ) if not blocked: await self._send_to_backend(json.dumps({"type": "response.create"})) @@ -1013,27 +1073,39 @@ class RealTimeStreaming: await self.log_messages() @staticmethod - def _detect_beta_header(websocket: Any) -> bool: + def _detect_beta_header(websocket: object) -> bool: """Return True if the client sent 'OpenAI-Beta: realtime=v1'. Checks the raw ASGI scope headers so it works for both FastAPI WebSocket objects and any test doubles that expose a .scope dict. """ try: - headers: Final = websocket.scope.get("headers", []) - for name, value in headers: - if isinstance(name, bytes): - name = name.decode("latin-1") - if isinstance(value, bytes): - value = value.decode("latin-1") - if name.lower() == "openai-beta" and "realtime=v1" in value.lower(): + scope: Final[object] = getattr(websocket, "scope", None) + if not _is_str_keyed_dict(scope): + return False + headers: Final = scope.get("headers", []) + if not _is_object_list(headers): + return False + for header_pair in headers: + if not _is_two_item_pair(header_pair): + return False + raw_name, raw_value = header_pair + name = raw_name.decode("latin-1") if isinstance(raw_name, bytes) else raw_name + if not isinstance(name, str): + return False + if name.lower() != "openai-beta": + continue + value = raw_value.decode("latin-1") if isinstance(raw_value, bytes) else raw_value + if not isinstance(value, str): + return False + if "realtime=v1" in value.lower(): return True except Exception: pass return False @staticmethod - def _remap_beta_session_to_ga(session: dict) -> dict: + def _remap_beta_session_to_ga(session: dict[str, object]) -> dict[str, object]: """ Convert a beta-style session.update payload to the GA nested schema. @@ -1064,16 +1136,14 @@ class RealTimeStreaming: if "modalities" in session: mods: Final = session.pop("modalities") if "output_modalities" not in session: - mods_set: Final = {m.lower() for m in (mods or [])} - if "audio" in mods_set: - session["output_modalities"] = ["audio"] - elif "text" in mods_set: - session["output_modalities"] = ["text"] + normalized: Final = _output_modalities_from_beta(mods) + if normalized is not None: + session["output_modalities"] = normalized # 3-7. Lift flat audio fields into the nested audio object - audio: Final[dict[str, Any]] = {} - inp: Final[dict[str, Any]] = {} - out: Final[dict[str, Any]] = {} + audio: Final[dict[str, object]] = {} + inp: Final[dict[str, object]] = {} + out: Final[dict[str, object]] = {} # voice → audio.output.voice if "voice" in session: @@ -1105,10 +1175,12 @@ class RealTimeStreaming: if audio: # Merge with any existing GA-style `audio` block the client already set, # letting the remapped values take precedence within each sub-key. - existing: Final = session.get("audio") or {} + existing_candidate: Final = session.get("audio") + existing: Final[dict[str, object]] = existing_candidate if _is_str_keyed_dict(existing_candidate) else {} for sub_key, sub_val in audio.items(): - if sub_key in existing and isinstance(existing[sub_key], dict) and isinstance(sub_val, dict): - existing[sub_key] = {**existing[sub_key], **sub_val} + current = existing.get(sub_key) + if _is_str_keyed_dict(current) and _is_str_keyed_dict(sub_val): + existing[sub_key] = {**current, **sub_val} else: existing[sub_key] = sub_val session["audio"] = existing @@ -1116,7 +1188,7 @@ class RealTimeStreaming: return session @staticmethod - def _translate_event_to_beta(event: dict) -> dict | None: + def _translate_event_to_beta(event: dict[str, object]) -> dict[str, object] | None: """Translate a single GA event dict to its beta equivalent. Returns None when the event must be dropped (the GA-only @@ -1129,38 +1201,48 @@ class RealTimeStreaming: if event_type == "conversation.item.done": return None - renamed_type: Final = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES.get(event_type) - has_item: Final = isinstance(event.get("item"), dict) + renamed_type: Final = ( + RealTimeStreaming._GA_TO_BETA_EVENT_TYPES.get(event_type) if isinstance(event_type, str) else None + ) + item_raw: Final = event.get("item") + has_item: Final = isinstance(item_raw, dict) response: Final = event.get("response") - has_response_output: Final = isinstance(response, dict) and isinstance(response.get("output"), list) + output_raw: Final = response.get("output") if _is_str_keyed_dict(response) else None + has_response_output: Final = isinstance(output_raw, list) if renamed_type is None and not has_item and not has_response_output: return event translated: Final = dict(event) if renamed_type is not None: translated["type"] = renamed_type - if has_item: - translated["item"] = RealTimeStreaming._translate_item_content_types(dict(translated["item"])) - if has_response_output: - resp: Final = dict(translated["response"]) + if _is_str_keyed_dict(item_raw): + translated["item"] = RealTimeStreaming._translate_item_content_types(dict(item_raw)) + if _is_str_keyed_dict(response) and _is_object_list(output_raw): + resp: Final = dict(response) resp["output"] = [ - (RealTimeStreaming._translate_item_content_types(dict(o)) if isinstance(o, dict) else o) - for o in resp["output"] + (RealTimeStreaming._translate_item_content_types(dict(o)) if _is_str_keyed_dict(o) else o) + for o in output_raw ] translated["response"] = resp return translated @staticmethod - def _translate_item_content_types(item: dict) -> dict: + def _translate_item_content_types(item: dict[str, object]) -> dict[str, object]: """Replace GA content type names with beta names inside a single item.""" - if "content" not in item or not isinstance(item["content"], list): + content_raw: Final = item.get("content") if "content" in item else None + if not _is_object_list(content_raw): return item - new_content: Final = [] - for block in item["content"]: - if isinstance(block, dict) and block.get("type") in RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES: + new_content: Final[list[object]] = [] + for block in content_raw: + block_type = block.get("type") if _is_str_keyed_dict(block) else None + if ( + _is_str_keyed_dict(block) + and isinstance(block_type, str) + and block_type in RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES + ): block = dict(block) - block["type"] = RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES[block["type"]] + block["type"] = RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES[block_type] new_content.append(block) item["content"] = new_content return item @@ -1176,12 +1258,14 @@ class RealTimeStreaming: try: from litellm.types.guardrails import GuardrailEventHooks - msg_obj = json.loads(message) - msg_type = msg_obj.get("type") + msg_obj = _json_loads_dict(message) + msg_type_raw = msg_obj.get("type") + msg_type = msg_type_raw if isinstance(msg_type_raw, str) else None if msg_type == "conversation.item.create": # Check user text messages for prompt injection - item = msg_obj.get("item", {}) + item_raw = msg_obj.get("item", {}) + item = item_raw if _is_str_keyed_dict(item_raw) else {} # Check function_call_output first so a client cannot # bypass the tool-result guardrail by also setting # role="user" on a function_call_output item. @@ -1241,11 +1325,14 @@ class RealTimeStreaming: # interaction turn. continue elif item.get("role") == "user": - content_list = item.get("content", []) + content_raw = item.get("content", []) + content_list: list[object] = content_raw if _is_object_list(content_raw) else [] texts = [ - c.get("text", "") + text for c in content_list - if isinstance(c, dict) and c.get("type") == "input_text" + if _is_str_keyed_dict(c) + and c.get("type") == "input_text" + and isinstance(text := c.get("text", ""), str) ] combined_text = " ".join(texts) if combined_text: @@ -1281,10 +1368,11 @@ class RealTimeStreaming: and self._has_audio_transcription_guardrails() ): session = msg_obj.setdefault("session", {}) - if isinstance(session, dict): - existing_td = session.get("turn_detection") - if not isinstance(existing_td, dict): - existing_td = {} + if _is_str_keyed_dict(session): + existing_td_raw = session.get("turn_detection") + existing_td: dict[str, object] = ( + existing_td_raw if _is_str_keyed_dict(existing_td_raw) else {} + ) existing_td["create_response"] = False session["turn_detection"] = existing_td message = json.dumps(msg_obj) @@ -1308,27 +1396,27 @@ class RealTimeStreaming: and self._has_audio_transcription_guardrails() ): session = msg_obj.get("session") - if isinstance(session, dict): + if _is_str_keyed_dict(session): td_overridden = False - flat_td = session.get("turn_detection") - flat_td_present = flat_td is not None + flat_td_raw = session.get("turn_detection") + flat_td_present = flat_td_raw is not None if flat_td_present: - if not isinstance(flat_td, dict): - flat_td = {} + flat_td: dict[str, object] = flat_td_raw if _is_str_keyed_dict(flat_td_raw) else {} if flat_td.get("create_response") is not False: flat_td["create_response"] = False session["turn_detection"] = flat_td td_overridden = True nested_td_present = False audio = session.get("audio") - if isinstance(audio, dict): + if _is_str_keyed_dict(audio): audio_input = audio.get("input") - if isinstance(audio_input, dict): - nested_td = audio_input.get("turn_detection") - if nested_td is not None: + if _is_str_keyed_dict(audio_input): + nested_td_raw = audio_input.get("turn_detection") + if nested_td_raw is not None: nested_td_present = True - if not isinstance(nested_td, dict): - nested_td = {} + nested_td: dict[str, object] = ( + nested_td_raw if _is_str_keyed_dict(nested_td_raw) else {} + ) if nested_td.get("create_response") is not False: nested_td["create_response"] = False audio_input["turn_detection"] = nested_td @@ -1351,14 +1439,14 @@ class RealTimeStreaming: # session shape unchanged. if msg_type == "session.update" and not self._backend_uses_beta_protocol: session = msg_obj.get("session", {}) - if isinstance(session, dict): + if _is_str_keyed_dict(session): session = self._remap_beta_session_to_ga(session) msg_obj["session"] = session message = json.dumps(msg_obj) if msg_type == "session.update" and self._event_normalizer: session = msg_obj.get("session") - if isinstance(session, dict): + if _is_str_keyed_dict(session): msg_obj["session"] = self._event_normalizer.patch_outgoing_session(session) message = json.dumps(msg_obj) @@ -1420,6 +1508,6 @@ class RealTimeStreaming: pass -def client_sent_openai_beta_realtime_header(websocket: Any) -> bool: +def client_sent_openai_beta_realtime_header(websocket: object) -> bool: """True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``.""" return RealTimeStreaming._detect_beta_header(websocket) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..2c6dd1f647b 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,12 +3,11 @@ import time from collections.abc import Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, TypeGuard, Union from litellm._logging import verbose_logger from litellm.types.llms.openai import ( ChatCompletionAssistantContentValue, - ChatCompletionAudioDelta, ) from litellm.types.utils import ( CacheCreationTokenDetails, @@ -16,7 +15,6 @@ from litellm.types.utils import ( ChatCompletionCustomToolCallPayload, ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, - Choices, CompletionTokensDetails, CompletionTokensDetailsWrapper, Function, @@ -30,6 +28,7 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, ) @@ -39,13 +38,44 @@ if TYPE_CHECKING: ) +def _is_str_keyed_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, dict) + + +def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, list) + + +def _field(source: object, key: str) -> object: + if _is_str_keyed_dict(source): + return source.get(key) + value: Final[object] = getattr(source, key, None) + return value + + +def _chunk_hidden_params(chunk: object) -> dict[str, object]: + candidate: Final = _field(chunk, "_hidden_params") + if _is_str_keyed_dict(candidate): + return candidate + return {} + + +class _UsageChunkFields(TypedDict): + prompt_tokens: int | None + completion_tokens: int | None + cache_creation_input_tokens: int | None + cache_read_input_tokens: int | None + completion_tokens_details: CompletionTokensDetails | None + prompt_tokens_details: PromptTokensDetailsWrapper | None + cost: float | None + + def capture_cache_creation_token_details( prompt_tokens_details: PromptTokensDetailsWrapper | None, current: CacheCreationTokenDetails | None, ) -> CacheCreationTokenDetails | None: - incoming: Final = cast( - CacheCreationTokenDetails | None, - getattr(prompt_tokens_details, "cache_creation_token_details", None), + incoming: Final[CacheCreationTokenDetails | None] = getattr( + prompt_tokens_details, "cache_creation_token_details", None ) if incoming is not None: return incoming @@ -58,9 +88,8 @@ def attach_cache_creation_token_details( ) -> PromptTokensDetailsWrapper | None: if prompt_tokens_details is None or cache_creation_token_details is None: return prompt_tokens_details - existing: Final = cast( - CacheCreationTokenDetails | None, - getattr(prompt_tokens_details, "cache_creation_token_details", None), + existing: Final[CacheCreationTokenDetails | None] = getattr( + prompt_tokens_details, "cache_creation_token_details", None ) if existing is not None: return prompt_tokens_details @@ -71,63 +100,52 @@ class ChunkProcessor: def __init__(self, chunks: list, messages: list | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages - self.first_chunk = chunks[0] + self.first_chunk: object = chunks[0] def _sort_chunks(self, chunks: list) -> list: if not chunks: return [] - first_chunk: Final = chunks[0] - first_hidden_params: dict[str, Any] = {} - if isinstance(first_chunk, dict): - candidate = first_chunk.get("_hidden_params", {}) - if isinstance(candidate, dict): - first_hidden_params = candidate - else: - candidate = getattr(first_chunk, "_hidden_params", {}) - if isinstance(candidate, dict): - first_hidden_params = candidate + first_hidden_params: Final = _chunk_hidden_params(chunks[0]) if first_hidden_params.get("created_at"): - def _created_at(chunk: Any) -> int | float: - if isinstance(chunk, dict): - params = chunk.get("_hidden_params", {}) - else: - params = getattr(chunk, "_hidden_params", {}) - if isinstance(params, dict): - return cast(int | float, params.get("created_at", float("inf"))) + def _created_at(chunk: object) -> int | float: + created_at: Final = _chunk_hidden_params(chunk).get("created_at", float("inf")) + if isinstance(created_at, (int, float)): + return created_at return float("inf") return sorted(chunks, key=_created_at) return chunks def update_model_response_with_hidden_params( - self, model_response: ModelResponse, chunk: dict[str, Any] | None = None + self, model_response: ModelResponse, chunk: object = None ) -> ModelResponse: if chunk is None: return model_response # set hidden params from chunk to model_response if model_response is not None and hasattr(model_response, "_hidden_params"): - model_response._hidden_params = chunk.get("_hidden_params", {}) + hidden_params: Final = _field(chunk, "_hidden_params") + model_response._hidden_params = hidden_params if _is_str_keyed_dict(hidden_params) else {} return model_response @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, - chunks: list[Any], - logging_obj: Any | None = None, + chunks: list[object], + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> None: if not chunks: return - model: Final = getattr(response, "model", None) + model: Final[str | None] = getattr(response, "model", None) if not model: return - custom_llm_provider = None + custom_llm_provider: object = None if logging_obj is not None: - custom_llm_provider = logging_obj.model_call_details.get("custom_llm_provider") + custom_llm_provider = _field(logging_obj.model_call_details, "custom_llm_provider") try: from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -159,18 +177,19 @@ class ChunkProcessor: ) @staticmethod - def _get_chunk_id(chunks: list[dict[str, Any]]) -> str: + def _get_chunk_id(chunks: Sequence[Mapping[str, object]]) -> str: """ Chunks: [{"id": ""}, {"id": "1"}, {"id": "1"}] """ for chunk in chunks: - if chunk.get("id"): - return chunk["id"] + chunk_id = chunk.get("id") + if isinstance(chunk_id, str) and chunk_id: + return chunk_id return "" @staticmethod - def _get_model_from_chunks(chunks: list[dict[str, Any]], first_chunk_model: str) -> str: + def _get_model_from_chunks(chunks: Sequence[Mapping[str, object]], first_chunk_model: str) -> str: """ Get the actual model from chunks, preferring a model that differs from the first chunk. @@ -181,31 +200,37 @@ class ChunkProcessor: # Look for a model in chunks that differs from the first chunk's model for chunk in chunks: chunk_model = chunk.get("model") - if chunk_model and chunk_model != first_chunk_model: + if isinstance(chunk_model, str) and chunk_model and chunk_model != first_chunk_model: return chunk_model # Fall back to first chunk's model if no different model found return first_chunk_model - def build_base_response(self, chunks: list[dict[str, Any]]) -> ModelResponse: + def build_base_response(self, chunks: Sequence[Mapping[str, object]]) -> ModelResponse: chunk = self.first_chunk id: Final = ChunkProcessor._get_chunk_id(chunks) - object: Final = chunk["object"] - created: Final = chunk["created"] - first_chunk_model: Final = chunk["model"] + object: Final = _field(chunk, "object") + created: Final = _field(chunk, "created") + first_chunk_model_raw: Final = _field(chunk, "model") + first_chunk_model: Final = first_chunk_model_raw if isinstance(first_chunk_model_raw, str) else "" # Get the actual model - for Azure Model Router, this finds the real model from later chunks model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model) - system_fingerprint: Final = chunk.get("system_fingerprint", None) + system_fingerprint: Final = _field(chunk, "system_fingerprint") - first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), chunk) - role: Final = first_chunk_with_choices["choices"][0]["delta"]["role"] + first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), None) + if first_chunk_with_choices is not None: + role_choices = first_chunk_with_choices.get("choices") + else: + role_choices = _field(chunk, "choices") + role = None + if _is_object_list(role_choices): + first_role_choice: Final = role_choices[0] + role = _field(_field(first_role_choice, "delta"), "role") finish_reason = "stop" for chunk in chunks: - if "choices" in chunk and len(chunk["choices"]) > 0: - chunk_finish_reason = None - if hasattr(chunk["choices"][0], "finish_reason"): - chunk_finish_reason = chunk["choices"][0].finish_reason - elif "finish_reason" in chunk["choices"][0]: - chunk_finish_reason = chunk["choices"][0]["finish_reason"] + chunk_choices = chunk.get("choices") if "choices" in chunk else None + if _is_object_list(chunk_choices) and len(chunk_choices) > 0: + first_choice = chunk_choices[0] + chunk_finish_reason = _field(first_choice, "finish_reason") if chunk_finish_reason is not None: finish_reason = chunk_finish_reason @@ -237,35 +262,32 @@ class ChunkProcessor: @staticmethod def _iter_tool_call_fragments( - tool_call_chunks: Sequence[Mapping[str, Any]], + tool_call_chunks: Sequence[Mapping[str, object]], ) -> Iterator[tuple[int, str, str]]: for chunk in tool_call_chunks: - for choice in chunk["choices"]: - delta = choice.get("delta") + choices = chunk["choices"] + if not _is_object_list(choices): + continue + for choice in choices: + delta = _field(choice, "delta") if not delta: continue - for tool_call in delta.get("tool_calls", ()): + tool_calls = _field(delta, "tool_calls") + if not _is_object_list(tool_calls): + continue + for tool_call in tool_calls: if not tool_call: continue - if isinstance(tool_call, dict): - index = tool_call.get("index", 0) - function = tool_call.get("function") - if isinstance(function, dict): - if function.get("arguments"): - yield index, "arguments", function["arguments"] - elif getattr(function, "arguments", None): - yield index, "arguments", function.arguments - custom = tool_call.get("custom") - if isinstance(custom, dict) and custom.get("input"): - yield index, "custom_input", custom["input"] - else: - index = getattr(tool_call, "index", 0) - function = getattr(tool_call, "function", None) - if getattr(function, "arguments", None): - yield index, "arguments", function.arguments - custom = getattr(tool_call, "custom", None) - if getattr(custom, "input", None): - yield index, "custom_input", custom.input + index_raw = _field(tool_call, "index") + index = index_raw if isinstance(index_raw, int) else 0 + function = _field(tool_call, "function") + arguments = _field(function, "arguments") + if isinstance(arguments, str) and arguments: + yield index, "arguments", arguments + custom = _field(tool_call, "custom") + custom_input = _field(custom, "input") + if isinstance(custom_input, str) and custom_input: + yield index, "custom_input", custom_input @staticmethod def _join_fragments_by_index_and_field( @@ -282,20 +304,24 @@ class ChunkProcessor: ) def get_combined_tool_content( - self, tool_call_chunks: Sequence[Mapping[str, Any]] + self, tool_call_chunks: Sequence[Mapping[str, object]] ) -> list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ]: # mutable-ok: assigned verbatim to Message.tool_calls, a list field tool_calls_list: list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ] = [] # mutable-ok: see return type - tool_call_map: Final[dict[int, dict[str, Any]]] = {} # Map to store tool calls by index + tool_call_map: Final[dict[int, dict[str, object]]] = {} # Map to store tool calls by index for chunk in tool_call_chunks: choices = chunk["choices"] + if not _is_object_list(choices): + continue for choice in choices: - delta = choice.get("delta", {}) - tool_calls = delta.get("tool_calls", []) + delta = _field(choice, "delta") + tool_calls = _field(delta, "tool_calls") + if not _is_object_list(tool_calls): + continue for tool_call in tool_calls: # Handle both dict and object formats @@ -303,23 +329,15 @@ class ChunkProcessor: continue # Check if tool_call has function (either as attribute or dict key) - has_function = False - has_custom = False - if isinstance(tool_call, dict): - has_function = "function" in tool_call and tool_call["function"] is not None - has_custom = "custom" in tool_call and tool_call["custom"] is not None - else: - has_function = hasattr(tool_call, "function") and tool_call.function is not None - has_custom = getattr(tool_call, "custom", None) is not None + has_function = _field(tool_call, "function") is not None + has_custom = _field(tool_call, "custom") is not None if not has_function and not has_custom: continue # Get index (handle both dict and object) - if isinstance(tool_call, dict): - index = tool_call.get("index", 0) - else: - index = getattr(tool_call, "index", 0) + index_raw = _field(tool_call, "index") + index = index_raw if isinstance(index_raw, int) else 0 if index not in tool_call_map: tool_call_map[index] = { @@ -331,62 +349,35 @@ class ChunkProcessor: } # Extract id, type, and function data (handle both dict and object) - if isinstance(tool_call, dict): - if tool_call.get("id"): - tool_call_map[index]["id"] = tool_call["id"] - if tool_call.get("type"): - tool_call_map[index]["type"] = tool_call["type"] + tool_call_id = _field(tool_call, "id") + if tool_call_id: + tool_call_map[index]["id"] = tool_call_id + tool_call_type = _field(tool_call, "type") + if tool_call_type: + tool_call_map[index]["type"] = tool_call_type - function = tool_call.get("function", {}) - if isinstance(function, dict): - if function.get("name"): - tool_call_map[index]["name"] = function["name"] - else: - # function is an object - if hasattr(function, "name") and function.name: - tool_call_map[index]["name"] = function.name + function = _field(tool_call, "function") + function_name = _field(function, "name") + if function_name: + tool_call_map[index]["name"] = function_name - custom = tool_call.get("custom") - if isinstance(custom, dict): - if custom.get("name"): - tool_call_map[index]["custom_name"] = custom["name"] - else: - # tool_call is an object - if hasattr(tool_call, "id") and tool_call.id: - tool_call_map[index]["id"] = tool_call.id - if hasattr(tool_call, "type") and tool_call.type: - tool_call_map[index]["type"] = tool_call.type - if hasattr(tool_call, "function"): - if hasattr(tool_call.function, "name") and tool_call.function.name: - tool_call_map[index]["name"] = tool_call.function.name - - custom = getattr(tool_call, "custom", None) - if custom is not None: - if getattr(custom, "name", None): - tool_call_map[index]["custom_name"] = custom.name + custom = _field(tool_call, "custom") + custom_name = _field(custom, "name") + if custom_name: + tool_call_map[index]["custom_name"] = custom_name # Preserve provider_specific_fields from streaming chunks - provider_fields = None - if isinstance(tool_call, dict): - provider_fields = tool_call.get("provider_specific_fields") - if not provider_fields and isinstance(tool_call.get("function"), dict): - provider_fields = tool_call["function"].get("provider_specific_fields") - else: - if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields: - provider_fields = tool_call.provider_specific_fields - elif ( - hasattr(tool_call, "function") - and hasattr(tool_call.function, "provider_specific_fields") - and tool_call.function.provider_specific_fields - ): - provider_fields = tool_call.function.provider_specific_fields + provider_fields = _field(tool_call, "provider_specific_fields") + if not provider_fields: + provider_fields = _field(function, "provider_specific_fields") if provider_fields: # Merge provider_specific_fields if multiple chunks have them if tool_call_map[index]["provider_specific_fields"] is None: tool_call_map[index]["provider_specific_fields"] = {} - if isinstance(provider_fields, dict): - tool_call_map[index]["provider_specific_fields"].update(provider_fields) + existing_provider_fields = tool_call_map[index]["provider_specific_fields"] + if _is_str_keyed_dict(provider_fields) and _is_str_keyed_dict(existing_provider_fields): + existing_provider_fields.update(provider_fields) joined_fragments: Final = self._join_fragments_by_index_and_field( self._iter_tool_call_fragments(tool_call_chunks) @@ -395,58 +386,71 @@ class ChunkProcessor: # Convert the map to a list of tool calls for index in sorted(tool_call_map.keys()): tool_call_data = tool_call_map[index] - if tool_call_data["id"] and tool_call_data["custom_name"]: + data_id = tool_call_data["id"] + data_custom_name = tool_call_data["custom_name"] + data_name = tool_call_data["name"] + data_type = tool_call_data["type"] + if isinstance(data_id, str) and data_id and isinstance(data_custom_name, str) and data_custom_name: tool_calls_list.append( ChatCompletionMessageCustomToolCall( - id=tool_call_data["id"], + id=data_id, custom=ChatCompletionCustomToolCallPayload( - name=tool_call_data["custom_name"], + name=data_custom_name, input=joined_fragments.get((index, "custom_input"), ""), ), ) ) - elif tool_call_data["id"] and tool_call_data["name"]: + elif isinstance(data_id, str) and data_id and isinstance(data_name, str) and data_name: combined_arguments = joined_fragments.get((index, "arguments"), "") or "{}" # Build function - provider_specific_fields should be on tool_call level, not function level function = Function( arguments=combined_arguments, - name=tool_call_data["name"], + name=data_name, ) - # Prepare params for ChatCompletionMessageToolCall - tool_call_params = { - "id": tool_call_data["id"], - "function": function, - "type": tool_call_data["type"] or "function", - } + resolved_type = data_type if isinstance(data_type, str) and data_type else "function" # Add provider_specific_fields if present (for thought signatures in Gemini 3) - if tool_call_data.get("provider_specific_fields"): - tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"] - - tool_call = ChatCompletionMessageToolCall(**tool_call_params) + provider_specific_fields = tool_call_data.get("provider_specific_fields") + if _is_str_keyed_dict(provider_specific_fields) and provider_specific_fields: + tool_call = ChatCompletionMessageToolCall( + id=data_id, + function=function, + type=resolved_type, + provider_specific_fields=provider_specific_fields, + ) + else: + tool_call = ChatCompletionMessageToolCall( + id=data_id, + function=function, + type=resolved_type, + ) tool_calls_list.append(tool_call) return tool_calls_list - def get_combined_function_call_content(self, function_call_chunks: list[dict[str, Any]]) -> FunctionCall: - argument_list: Final = [] - delta = function_call_chunks[0]["choices"][0]["delta"] - function_call = delta.get("function_call", "") - function_call_name: Final = function_call.name + def get_combined_function_call_content(self, function_call_chunks: Sequence[Mapping[str, object]]) -> FunctionCall: + argument_list: Final[list[str]] = [] + first_choices: Final = function_call_chunks[0]["choices"] + first_choice: Final = first_choices[0] if _is_object_list(first_choices) else None + first_function_call: Final = _field(_field(first_choice, "delta"), "function_call") + function_call_name_raw: Final = _field(first_function_call, "name") + function_call_name: Final = function_call_name_raw if isinstance(function_call_name_raw, str) else None for chunk in function_call_chunks: choices = chunk["choices"] + if not _is_object_list(choices): + continue for choice in choices: - delta = choice.get("delta", {}) - function_call = delta.get("function_call", "") + delta = _field(choice, "delta") + function_call = _field(delta, "function_call") # Check if a function call is present if function_call: - # Now, function_call is expected to be a dictionary - arguments = function_call.arguments - argument_list.append(arguments) + arguments = _field(function_call, "arguments") + if isinstance(arguments, str): + argument_list.append(arguments) combined_arguments: Final = "".join(argument_list) @@ -456,17 +460,20 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: Sequence[Mapping[str, object]], delta_key: str = "content" ) -> ChatCompletionAssistantContentValue: content_list: Final[list[str]] = [] for chunk in chunks: choices = chunk["choices"] + if not _is_object_list(choices): + continue for choice in choices: - delta = choice.get("delta", {}) - content = delta.get(delta_key, "") + delta = _field(choice, "delta") + content = _field(delta, delta_key) if content is None: continue # openai v1.0.0 sets content = None for chunks - content_list.append(content) + if isinstance(content, str): + content_list.append(content) # Combine the "content" strings into a single string || combine the 'function' strings into a single string combined_content: Final = "".join(content_list) @@ -475,7 +482,7 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] + self, chunks: Sequence[Mapping[str, object]] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -501,16 +508,20 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] + if not _is_object_list(choices): + continue for choice in choices: - delta = choice.get("delta", {}) - thinking = delta.get("thinking_blocks", None) - if thinking and isinstance(thinking, list): + delta = _field(choice, "delta") + thinking = _field(delta, "thinking_blocks") + if thinking and _is_object_list(thinking): for thinking_block in thinking: + if not _is_str_keyed_dict(thinking_block): + continue thinking_type = thinking_block.get("type", None) if thinking_type and thinking_type == "redacted_thinking": _flush_thinking_block() redacted_data = thinking_block.get("data", None) - if redacted_data: + if isinstance(redacted_data, str) and redacted_data: thinking_blocks.append( ChatCompletionRedactedThinkingBlock( type="redacted_thinking", @@ -519,10 +530,10 @@ class ChunkProcessor: ) else: thinking_text = thinking_block.get("thinking", None) - if thinking_text: + if isinstance(thinking_text, str) and thinking_text: current_thinking_text_parts.append(thinking_text) signature = thinking_block.get("signature", None) - if signature: + if isinstance(signature, str) and signature: current_signature = signature _flush_thinking_block() @@ -532,10 +543,12 @@ class ChunkProcessor: return thinking_blocks return None - def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue: + def get_combined_reasoning_content( + self, chunks: Sequence[Mapping[str, object]] + ) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: Sequence[Mapping[str, object]]) -> ChatCompletionAudioResponse: base64_data_list: Final[list[str]] = [] transcript_list: Final[list[str]] = [] expires_at: int | None = None @@ -543,10 +556,12 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] + if not _is_object_list(choices): + continue for choice in choices: - delta = choice.get("delta") or {} - audio: ChatCompletionAudioDelta | None = delta.get("audio") - if audio is not None: + delta = _field(choice, "delta") + audio = _field(delta, "audio") + if _is_str_keyed_dict(audio): for k, v in audio.items(): if k == "data" and v is not None and isinstance(v, str): base64_data_list.append(v) @@ -565,7 +580,7 @@ class ChunkProcessor: id=id, ) - def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict: + def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> _UsageChunkFields: prompt_tokens = 0 completion_tokens = 0 ## anthropic prompt caching information ## @@ -609,14 +624,12 @@ class ChunkProcessor: def count_reasoning_tokens(self, response: ModelResponse) -> int | None: reasoning_tokens: int | None = None for choice in response.choices: - if ( - hasattr(cast(Choices, choice).message, "reasoning_content") - and cast(Choices, choice).message.reasoning_content is not None - ): + reasoning_content: object = getattr(choice.message, "reasoning_content", None) + if isinstance(reasoning_content, str): if reasoning_tokens is None: reasoning_tokens = 0 reasoning_tokens += token_counter( - text=cast(Choices, choice).message.reasoning_content, + text=reasoning_content, count_response_tokens=True, ) @@ -640,7 +653,7 @@ class ChunkProcessor: def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: list[dict[str, object] | ModelResponse], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -707,27 +720,13 @@ class ChunkProcessor: server_tool_use = usage_chunk.server_tool_use else: server_tool_use = ServerToolUse.model_validate(usage_chunk.server_tool_use) - if ( - usage_chunk_dict["prompt_tokens_details"] is not None - and getattr( - usage_chunk_dict["prompt_tokens_details"], - "web_search_requests", - None, - ) - is not None - ): - web_search_requests = getattr( - usage_chunk_dict["prompt_tokens_details"], - "web_search_requests", - ) + chunk_prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] + if chunk_prompt_tokens_details is not None: + web_search_value: int | None = getattr(chunk_prompt_tokens_details, "web_search_requests", None) + if web_search_value is not None: + web_search_requests = web_search_value - prompt_tokens_details = ( - cast( - PromptTokensDetailsWrapper | None, - usage_chunk_dict["prompt_tokens_details"], - ) - or prompt_tokens_details - ) + prompt_tokens_details = chunk_prompt_tokens_details or prompt_tokens_details cache_creation_token_details = capture_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details @@ -758,7 +757,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: list[dict[str, object] | ModelResponse], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -781,14 +780,14 @@ class ChunkProcessor: if saw_non_cursor_completion: return completion_tokens - custom_llm_provider: str | None = None + custom_llm_provider: object = None if chunks: first_chunk: Final = chunks[0] if isinstance(first_chunk, dict): - hp = first_chunk.get("_hidden_params") + hp: object = first_chunk.get("_hidden_params") else: hp = getattr(first_chunk, "_hidden_params", None) - if isinstance(hp, dict): + if _is_str_keyed_dict(hp): custom_llm_provider = hp.get("custom_llm_provider") if custom_llm_provider == "anthropic" and completion_tokens == 1: @@ -797,7 +796,7 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: list[dict[str, object] | ModelResponse], model: str, completion_output: str, messages: list | None = None, @@ -851,8 +850,8 @@ class ChunkProcessor: setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens) # for anthropic if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): - returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate( + completion_tokens_details.model_dump() ) else: returned_usage.completion_tokens_details = completion_tokens_details diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a9751489473..43e37021e13 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,7 +1,6 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping from typing import ( TYPE_CHECKING, - Any, Final, cast, ) @@ -27,14 +26,14 @@ from litellm.types.utils import ModelResponse from litellm.utils import get_model_info if TYPE_CHECKING: - from litellm.proxy._types import UserAPIKeyAuth + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.router import Router # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"}) -def _messages_have_compaction_block(messages: list[dict]) -> bool: +def _messages_have_compaction_block(messages: list[dict[str, object]]) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: content = msg.get("content") @@ -54,7 +53,7 @@ def _proxy_router_fallback() -> "Router | None": return _proxy_router -def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: +def _extract_proxy_litellm_metadata(kwargs: dict[str, object]) -> dict[str, object] | None: """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise. The proxy attaches its auth/spend-attribution fields (``user_api_key``, @@ -75,14 +74,14 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | async def _prepare_context_managed_request( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, - additional_drop_params: list[str] | None, - llm_router: "Router | None", - user_api_key_auth: "UserAPIKeyAuth | None" = None, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: str | list[dict[str, object]] | None, + context_management_spec: object, + litellm_metadata: dict[str, object] | None, + additional_drop_params: object, + llm_router: object, + user_api_key_auth: object = None, ) -> PolyfillResult | None: """Apply client compaction history, then optional context_management polyfill.""" from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import ( @@ -102,11 +101,11 @@ async def _prepare_context_managed_request( if polyfill_will_run: history_result: PolyfillResult | None = None - working_messages: list[dict] = messages - working_system: Any | None = system + working_messages = messages + working_system = system else: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) working_messages = history_result.messages if history_result is not None else messages @@ -136,7 +135,7 @@ async def _prepare_context_managed_request( # to non-Anthropic backends that would reject them. if polyfill_will_run and history_result is None: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) return history_result @@ -144,8 +143,8 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, - context_management_spec: Any, - additional_drop_params: list[str] | None, + context_management_spec: object, + additional_drop_params: object, ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. @@ -171,8 +170,8 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, - context_management_spec: Any, - additional_drop_params: list[str] | None, + context_management_spec: object, + additional_drop_params: object, ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -195,7 +194,7 @@ def _spec_has_non_compact_edits( return any(isinstance(edit.get("type"), str) and edit.get("type") != COMPACT_EDIT_TYPE for edit in edits) -def _context_management_explicitly_dropped(additional_drop_params: list[str] | None) -> bool: +def _context_management_explicitly_dropped(additional_drop_params: object) -> bool: """True when the caller opted out of context_management via ``additional_drop_params``. ``drop_params`` deliberately does NOT gate the polyfill: ``context_management`` @@ -209,9 +208,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N def _normalize_spec_edits( *, - context_management_spec: Any, - additional_drop_params: list[str] | None, -) -> list[dict[str, Any]] | None: + context_management_spec: object, + additional_drop_params: object, +) -> list[dict[str, object]] | None: """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run. Delegates spec-shape normalization to the dispatcher's ``_normalize_spec`` @@ -223,6 +222,9 @@ def _normalize_spec_edits( if _context_management_explicitly_dropped(additional_drop_params): return None + if not isinstance(context_management_spec, (dict, list)): + return None + from litellm.llms.anthropic.experimental_pass_through.context_management.dispatcher import ( _normalize_spec, ) @@ -236,14 +238,14 @@ def _normalize_spec_edits( async def _run_polyfill_if_enabled( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, - additional_drop_params: list[str] | None, - llm_router: "Router | None", - user_api_key_auth: "UserAPIKeyAuth | None" = None, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: str | list[dict[str, object]] | None, + context_management_spec: object, + litellm_metadata: dict[str, object] | None, + additional_drop_params: object, + llm_router: object, + user_api_key_auth: object = None, ) -> PolyfillResult | None: """Run the async context_management polyfill if a spec is present. @@ -260,6 +262,9 @@ async def _run_polyfill_if_enabled( if _context_management_explicitly_dropped(additional_drop_params): return None + if not isinstance(context_management_spec, (dict, list)): + return PolyfillResult(messages=messages, system=system) + try: return await apply_context_management( model=model, @@ -304,9 +309,9 @@ ANTHROPIC_ADAPTER: Final = AnthropicAdapter() class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def _route_openai_thinking_to_responses_api_if_needed( - completion_kwargs: dict[str, Any], + completion_kwargs: dict[str, object], *, - thinking: dict[str, Any] | None, + thinking: dict[str, object] | None, ) -> None: """ When users call `litellm.anthropic.messages.*` with a non-Anthropic model and @@ -319,8 +324,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: If the user provides a `summary` field in the thinking dict, it is passed through to the OpenAI reasoning params (opt-in per OpenAI spec). """ - custom_llm_provider = completion_kwargs.get("custom_llm_provider") - if custom_llm_provider is None: + raw_provider: Final = completion_kwargs.get("custom_llm_provider") + if isinstance(raw_provider, str): + custom_llm_provider = raw_provider + elif raw_provider is None: try: _, inferred_provider, _, _ = litellm.utils.get_llm_provider( model=cast(str, completion_kwargs.get("model")) @@ -328,6 +335,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: custom_llm_provider = inferred_provider except Exception: custom_llm_provider = None + else: + return if custom_llm_provider != "openai": return @@ -369,7 +378,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def _normalize_reasoning_effort( - completion_kwargs: dict[str, Any], + completion_kwargs: dict[str, object], ) -> None: """ Normalize reasoning_effort values based on target model capabilities. @@ -386,7 +395,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: return model: Final = cast(str, completion_kwargs.get("model", "")) - custom_llm_provider: Final = completion_kwargs.get("custom_llm_provider") + raw_provider: Final = completion_kwargs.get("custom_llm_provider") + custom_llm_provider: Final = raw_provider if isinstance(raw_provider, str) else None if isinstance(reasoning_effort, str): normalized = normalize_reasoning_effort_value( @@ -396,32 +406,35 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_kwargs["reasoning_effort"] = normalized elif isinstance(reasoning_effort, dict) and "effort" in reasoning_effort: effort: Final = reasoning_effort["effort"] - normalized = normalize_reasoning_effort_value(effort, model=model, custom_llm_provider=custom_llm_provider) - if normalized != effort: - completion_kwargs["reasoning_effort"] = { - **reasoning_effort, - "effort": normalized, - } + if isinstance(effort, str): + normalized = normalize_reasoning_effort_value( + effort, model=model, custom_llm_provider=custom_llm_provider + ) + if normalized != effort: + completion_kwargs["reasoning_effort"] = { + **reasoning_effort, + "effort": normalized, + } @staticmethod def _prepare_completion_kwargs( *, max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - metadata: dict | None = None, + metadata: dict[str, object] | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, - system: str | list[dict[str, Any]] | None = None, + system: str | list[dict[str, object]] | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, - ) -> tuple[dict[str, Any], dict[str, str]]: + output_format: dict[str, object] | None = None, + extra_kwargs: dict[str, object] | None = None, + ) -> tuple[dict[str, object], dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. Returns: @@ -433,7 +446,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -478,7 +491,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if openai_request is None: raise ValueError("Failed to translate request to OpenAI format") - completion_kwargs: Final[dict[str, Any]] = dict(openai_request) + completion_kwargs: Final[dict[str, object]] = dict(openai_request) if stream: completion_kwargs["stream"] = stream @@ -528,31 +541,29 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - metadata: dict | None = None, + metadata: dict[str, object] | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, - **kwargs, + output_format: dict[str, object] | None = None, + **kwargs: object, ) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]: """Handle non-Anthropic models asynchronously using the adapter""" context_management: Final = kwargs.pop("context_management", None) - additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) - requested_router: Final[Router | None] = kwargs.pop("litellm_router", None) - litellm_router: Final[Router | None] = ( - requested_router if requested_router is not None else _proxy_router_fallback() - ) + additional_drop_params: Final = kwargs.get("additional_drop_params", None) + requested_router: Final = kwargs.pop("litellm_router", None) + litellm_router: Final = requested_router if requested_router is not None else _proxy_router_fallback() proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( + user_api_key_auth: Final = ( proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None ) @@ -592,7 +603,13 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) - completion_response: Final = await litellm.acompletion(**completion_kwargs) + async def _invoke_acompletion( + call_kwargs: Mapping[str, object], + acompletion_fn: "Callable[..., Awaitable[ModelResponse | CustomStreamWrapper]]" = litellm.acompletion, + ) -> "ModelResponse | CustomStreamWrapper": + return await acompletion_fn(**call_kwargs) + + completion_response: Final = await _invoke_acompletion(completion_kwargs) if stream: transformed_stream: Final = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( @@ -618,21 +635,21 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - metadata: dict | None = None, + metadata: dict[str, object] | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: dict[str, object] | None = None, _is_async: bool = False, - **kwargs, + **kwargs: object, ) -> ( AnthropicMessagesResponse | Iterator[bytes] @@ -665,7 +682,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: # ``compact_20260112`` editor can ``await`` the summarization model); # bridge to it via ``run_async_function``. context_management: Final = kwargs.pop("context_management", None) - additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) + additional_drop_params: Final = kwargs.get("additional_drop_params", None) # Deliberately do NOT auto-attach the proxy ``llm_router`` here: # ``run_async_function`` spawns a new event loop in a worker thread # to bridge to the async dispatcher, but the proxy router's httpx @@ -677,7 +694,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: # ``llm_router`` is ``None``, which is safe to call from the bridged # loop. The async ``async_anthropic_messages_handler`` path is # unaffected because it ``await``s within the original event loop. - litellm_router: Final[Router | None] = kwargs.pop("litellm_router", None) + litellm_router: Final = kwargs.pop("litellm_router", None) # Skip the async bridge entirely when there is nothing for either the # polyfill or the client-history slice-only fallback to do. The vast @@ -689,9 +706,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: polyfill_result: PolyfillResult | None = None else: proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( + user_api_key_auth: Final = ( proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None ) + polyfill_result = run_async_function( _prepare_context_managed_request, model=model, @@ -729,7 +747,13 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) - completion_response: Final = litellm.completion(**completion_kwargs) + def _invoke_completion( + call_kwargs: Mapping[str, object], + completion_fn: "Callable[..., ModelResponse | CustomStreamWrapper]" = litellm.completion, + ) -> "ModelResponse | CustomStreamWrapper": + return completion_fn(**call_kwargs) + + completion_response: Final = _invoke_completion(completion_kwargs) if stream: transformed_stream: Final = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index c4b5cc628e2..e4eb0cd7bb4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -436,6 +436,15 @@ def anthropic_messages_handler( # This is needed by agentic hooks (e.g., websearch_interception) to make follow-up requests original_model: Final = model + messages_arg: Final = cast(list[dict[str, object]], messages) + metadata_arg: Final = cast(dict[str, object] | None, metadata) + system_arg: Final = cast(str | list[object] | None, system) + thinking_arg: Final = cast(dict[str, object] | None, thinking) + tool_choice_arg: Final = cast(dict[str, object] | None, tool_choice) + tools_arg: Final = cast(list[dict[str, object]] | None, tools) + container_arg: Final = cast(dict[str, object] | None, container) + kwargs_obj: Final = cast(dict[str, object], kwargs) + litellm_params: Final = GenericLiteLLMParams( **kwargs, api_key=api_key, @@ -493,24 +502,24 @@ def anthropic_messages_handler( if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): return anthropic_messages_with_mcp( max_tokens=max_tokens, - messages=messages, + messages=messages_arg, model=model, - metadata=metadata, + metadata=metadata_arg, stop_sequences=stop_sequences, stream=stream, - system=system, + system=system_arg, temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, + thinking=thinking_arg, + tool_choice=tool_choice_arg, + tools=tools_arg, top_k=top_k, top_p=top_p, - container=container, + container=container_arg, api_key=api_key, api_base=api_base, client=client, custom_llm_provider=custom_llm_provider, - **kwargs, + **kwargs_obj, ) anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None = None @@ -532,16 +541,16 @@ def anthropic_messages_handler( # Route to Responses API for OpenAI / Azure, chat/completions for everything else. _shared_kwargs: Final = dict( max_tokens=max_tokens, - messages=messages, + messages=messages_arg, model=model, - metadata=metadata, + metadata=metadata_arg, stop_sequences=stop_sequences, stream=stream, - system=system, + system=system_arg, temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, + thinking=thinking_arg, + tool_choice=tool_choice_arg, + tools=tools_arg, top_k=top_k, top_p=top_p, _is_async=is_async, @@ -549,7 +558,7 @@ def anthropic_messages_handler( api_base=api_base, client=client, custom_llm_provider=custom_llm_provider, - **kwargs, + **kwargs_obj, ) if _should_route_to_responses_api(custom_llm_provider): return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(**_shared_kwargs) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py index 6ba129f5a7d..1f6fdd49642 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -7,8 +7,8 @@ tool through a ``tool_use`` content block, and results are fed back as ``tool_result`` blocks in a user message. """ -from collections.abc import AsyncIterator, Mapping, Sequence -from typing import Any, Final +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence +from typing import Final, Protocol from litellm._logging import verbose_logger from litellm.responses.mcp.request_context import MCPRequestContext @@ -24,14 +24,44 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( MAX_MCP_TOOL_USE_ITERATIONS: Final = 10 -def _get_response_content(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]: +class _ResolvedMCPContext(Protocol): + @property + def user_api_key_auth(self) -> object: ... + + @property + def mcp_auth_header(self) -> str | None: ... + + @property + def mcp_server_auth_headers(self) -> Mapping[str, Mapping[str, str]] | None: ... + + @property + def oauth2_headers(self) -> Mapping[str, str] | None: ... + + @property + def raw_headers(self) -> Mapping[str, str] | None: ... + + @property + def request_tags(self) -> Sequence[str] | None: ... + + @property + def litellm_trace_id(self) -> str | None: ... + + @property + def litellm_call_id(self) -> str | None: ... + + +def _resolved_context(context: MCPRequestContext) -> _ResolvedMCPContext: + return context + + +def _get_response_content(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, object]]: content: Final = response.get("content") if not isinstance(content, list): return () return tuple(block for block in content if isinstance(block, dict)) -def _extract_tool_use_blocks(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]: +def _extract_tool_use_blocks(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, object]]: """Return the ``tool_use`` content blocks the model emitted.""" return tuple(block for block in _get_response_content(response) if block.get("type") == "tool_use") @@ -41,7 +71,7 @@ def _get_stop_reason(response: AnthropicMessagesResponse) -> str | None: return stop_reason if isinstance(stop_reason, str) else None -def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> AnthropicMessagesUserMessageParam: +def _build_tool_result_message(tool_results: Sequence[Mapping[str, object]]) -> AnthropicMessagesUserMessageParam: """Turn executed tool results into the user message Anthropic expects.""" return AnthropicMessagesUserMessageParam( role="user", @@ -58,11 +88,11 @@ def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> Ant async def anthropic_messages_with_mcp( max_tokens: int, - messages: Sequence[Mapping[str, Any]], + messages: Sequence[Mapping[str, object]], model: str, - tools: Sequence[Mapping[str, Any]] | None = None, - **kwargs: Any, # kwargs-ok: forwarded verbatim to litellm.anthropic_messages, which owns the param contract -) -> AnthropicMessagesResponse | AsyncIterator[Any]: + tools: Sequence[Mapping[str, object]] | None = None, + **kwargs: object, # kwargs-ok: forwarded verbatim to litellm.anthropic_messages, which owns the param contract +) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]: """ Expand litellm_proxy MCP references for `/v1/messages` and run the tool loop. @@ -78,19 +108,49 @@ async def anthropic_messages_with_mcp( LiteLLM_Proxy_MCP_Handler, ) + async def _forward_initial_messages_api( + call_kwargs: Mapping[str, object], + *, + max_tokens: int, + messages: Sequence[Mapping[str, object]], + model: str, + tools: Sequence[Mapping[str, object]] | None, + messages_fn: Callable[ + ..., Awaitable[AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]] + ] = litellm.anthropic_messages, + ) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]: + return await messages_fn( + max_tokens=max_tokens, + messages=messages, + model=model, + tools=tools, + _skip_mcp_handler=True, + **call_kwargs, + ) + + async def _forward_messages_api( + call_kwargs: Mapping[str, object], + *, + messages: Sequence[Mapping[str, object]], + stream: bool, + messages_fn: Callable[ + ..., Awaitable[AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]] + ] = litellm.anthropic_messages, + ) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]: + return await messages_fn(messages=messages, stream=stream, **call_kwargs) + mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) if not mcp_references: - return await litellm.anthropic_messages( + return await _forward_initial_messages_api( + kwargs, max_tokens=max_tokens, messages=list(messages), model=model, tools=list(tools) if tools else None, - _skip_mcp_handler=True, - **kwargs, ) - context: Final = MCPRequestContext.resolve(kwargs=dict(kwargs), tools=tools) + context: Final = _resolved_context(MCPRequestContext.resolve(kwargs=dict(kwargs), tools=tools)) ( deduplicated_mcp_tools, @@ -114,7 +174,7 @@ async def anthropic_messages_with_mcp( ) stream: Final = bool(kwargs.pop("stream", False)) - base_call_args: Final[Mapping[str, Any]] = { + base_call_args: Final[Mapping[str, object]] = { "max_tokens": max_tokens, "model": model, "tools": all_tools or None, @@ -123,11 +183,11 @@ async def anthropic_messages_with_mcp( } if not should_auto_execute: - return await litellm.anthropic_messages(messages=list(messages), stream=stream, **base_call_args) + return await _forward_messages_api(base_call_args, messages=list(messages), stream=stream) - working_messages: Sequence[Mapping[str, Any]] = tuple(messages) - response: AnthropicMessagesResponse = await litellm.anthropic_messages( - messages=list(working_messages), stream=False, **base_call_args + working_messages: Sequence[Mapping[str, object]] = tuple(messages) + response: AnthropicMessagesResponse = await _forward_messages_api( + base_call_args, messages=list(working_messages), stream=False ) for _ in range(MAX_MCP_TOOL_USE_ITERATIONS): @@ -161,7 +221,7 @@ async def anthropic_messages_with_mcp( {"role": "assistant", "content": list(_get_response_content(response))}, _build_tool_result_message(tool_results), ) - response = await litellm.anthropic_messages(messages=list(working_messages), stream=False, **base_call_args) + response = await _forward_messages_api(base_call_args, messages=list(working_messages), stream=False) else: verbose_logger.warning( "MCP tool loop hit its %s iteration cap for model %s; returning the last response", diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 5e05ebc3c63..ba1238bb94d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -4,11 +4,16 @@ Handler for the Anthropic v1/messages -> OpenAI Responses API path. Used when the target model is an OpenAI or Azure model. """ -from collections.abc import AsyncIterator, Coroutine -from typing import Any, Final +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Mapping +from typing import TYPE_CHECKING, Final import litellm -from litellm.types.llms.anthropic import AnthropicMessagesRequest +from litellm.types.llms.anthropic import ( + AllAnthropicToolsValues, + AnthropicMessagesRequest, + AnthropicOutputConfig, + AnthropicOutputSchema, +) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -17,34 +22,37 @@ from litellm.types.llms.openai import ResponsesAPIResponse from .streaming_iterator import AnthropicResponsesStreamWrapper from .transformation import LiteLLMAnthropicToResponsesAPIAdapter +if TYPE_CHECKING: + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator + _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter() def _build_responses_kwargs( *, max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, - output_config: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, -) -> dict[str, Any]: + output_format: AnthropicOutputSchema | None = None, + extra_kwargs: dict[str, object] | None = None, +) -> dict[str, object]: """ Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses(). """ # Build a typed AnthropicMessagesRequest for the adapter - request_data: Final[dict[str, Any]] = { + request_data: Final[AnthropicMessagesRequest] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -71,7 +79,7 @@ def _build_responses_kwargs( request_data["output_format"] = output_format anthropic_request: Final = AnthropicMessagesRequest(**request_data) - responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request) + responses_kwargs: Final[dict[str, object]] = _ADAPTER.translate_request(anthropic_request) # Normalize reasoning effort based on model capabilities # (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported) @@ -82,13 +90,15 @@ def _build_responses_kwargs( ) effort: Final = reasoning["effort"] - normalized: Final = normalize_reasoning_effort_value( - effort, - model=model, - custom_llm_provider=(extra_kwargs or {}).get("custom_llm_provider"), - ) - if normalized != effort: - responses_kwargs["reasoning"] = {**reasoning, "effort": normalized} + raw_provider: Final = (extra_kwargs or {}).get("custom_llm_provider") + if isinstance(effort, str): + normalized: Final = normalize_reasoning_effort_value( + effort, + model=model, + custom_llm_provider=raw_provider if isinstance(raw_provider, str) else None, + ) + if normalized != effort: + responses_kwargs["reasoning"] = {**reasoning, "effort": normalized} if stream: responses_kwargs["stream"] = True @@ -124,23 +134,23 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, - output_config: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, - **kwargs, - ) -> AnthropicMessagesResponse | AsyncIterator: + output_format: AnthropicOutputSchema | None = None, + **kwargs: object, + ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: responses_kwargs: Final = _build_responses_kwargs( max_tokens=max_tokens, messages=messages, @@ -161,7 +171,13 @@ class LiteLLMMessagesToResponsesAPIHandler: extra_kwargs=kwargs, ) - result: Final = await litellm.aresponses(**responses_kwargs) + async def _invoke_aresponses( + call_kwargs: Mapping[str, object], + aresponses_fn: "Callable[..., Awaitable[ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]]" = litellm.aresponses, + ) -> "ResponsesAPIResponse | BaseResponsesAPIStreamingIterator": + return await aresponses_fn(**call_kwargs) + + result: Final = await _invoke_aresponses(responses_kwargs) if stream: wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) @@ -175,27 +191,27 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, - output_config: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, + output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: AnthropicOutputSchema | None = None, _is_async: bool = False, - **kwargs, + **kwargs: object, ) -> ( AnthropicMessagesResponse - | AsyncIterator[Any] - | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any]] + | AsyncIterator[bytes] + | Coroutine[None, None, AnthropicMessagesResponse | AsyncIterator[bytes]] ): if _is_async: return LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( @@ -239,7 +255,13 @@ class LiteLLMMessagesToResponsesAPIHandler: extra_kwargs=kwargs, ) - result: Final = litellm.responses(**responses_kwargs) + def _invoke_responses( + call_kwargs: Mapping[str, object], + responses_fn: Callable[..., object] = litellm.responses, + ) -> object: + return responses_fn(**call_kwargs) + + result: Final = _invoke_responses(responses_kwargs) if stream: wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 671e4633af4..2fb495a0ce0 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -212,9 +212,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj # fmt: off @@ -301,9 +301,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj async def async_get_messages( @@ -443,7 +443,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = await openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) # fmt: off @@ -539,7 +539,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = azure_openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) async def async_get_thread( self, @@ -566,7 +566,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # fmt: off @@ -642,7 +642,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # def delete_thread(self): # pass @@ -655,7 +655,7 @@ class AzureAssistantsAPI(BaseAzureLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, stream: bool | None, tools: Iterable[AssistantToolParam] | None, @@ -698,7 +698,7 @@ class AzureAssistantsAPI(BaseAzureLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, tools: Iterable[AssistantToolParam] | None, event_handler: AssistantEventHandler | None, @@ -724,24 +724,33 @@ class AzureAssistantsAPI(BaseAzureLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, tools: Iterable[AssistantToolParam] | None, event_handler: AssistantEventHandler | None, litellm_params: dict | None = None, ) -> AssistantStreamManager[AssistantEventHandler]: - data: Final[dict[str, Any]] = { - "thread_id": thread_id, - "assistant_id": assistant_id, - "additional_instructions": additional_instructions, - "instructions": instructions, - "metadata": metadata, - "model": model, - "tools": tools, - } + stream_method: Final = client.beta.threads.runs.stream if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return stream_method( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + event_handler=event_handler, + ) + return stream_method( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + ) # fmt: off @@ -752,7 +761,7 @@ class AzureAssistantsAPI(BaseAzureLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, stream: bool | None, tools: Iterable[AssistantToolParam] | None, @@ -774,7 +783,7 @@ class AzureAssistantsAPI(BaseAzureLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, stream: bool | None, tools: Iterable[AssistantToolParam] | None, @@ -797,7 +806,7 @@ class AzureAssistantsAPI(BaseAzureLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, stream: bool | None, tools: Iterable[AssistantToolParam] | None, diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index e3e1ef8ecd5..4922da3b6c8 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -11,7 +11,7 @@ from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from ....litellm_core_utils.realtime_streaming import RealTimeStreaming +from ....litellm_core_utils.realtime_streaming import ClientWebSocketInterface, RealTimeStreaming from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion @@ -128,12 +128,12 @@ class AzureOpenAIRealtime(AzureChatCompletion): ssl=ssl_context, ) as backend_ws: realtime_streaming: Final = RealTimeStreaming( - websocket, + cast(ClientWebSocketInterface, websocket), cast(ClientConnection, backend_ws), logging_obj, model=model, user_api_key_dict=user_api_key_dict, - request_data={"litellm_metadata": litellm_metadata or {}}, + request_data=cast(dict[str, object], {"litellm_metadata": litellm_metadata or {}}), backend_uses_beta_protocol=backend_uses_beta_protocol, force_transcription_model=( model if (query_params or {}).get("intent") == "transcription" else None diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a58397c9184..ba300f4625d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -19,7 +19,7 @@ from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.litellm_core_utils.realtime_streaming import ClientWebSocketInterface, RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, @@ -5848,7 +5848,7 @@ class BaseLLMHTTPHandler: if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata realtime_streaming: Final = RealTimeStreaming( - websocket, + cast(ClientWebSocketInterface, websocket), cast(ClientConnection, backend_ws), logging_obj, provider_config, diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 0343f22e7d1..7520d3cb701 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -12,6 +12,7 @@ from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ....litellm_core_utils.realtime_streaming import ( + ClientWebSocketInterface, RealtimeEventNormalizer, RealTimeStreaming, client_sent_openai_beta_realtime_header, @@ -135,7 +136,8 @@ class OpenAIRealtime(OpenAIChatCompletion): # Get provider-specific SSL configuration ssl_config: Final = self._get_ssl_config(url) - openai_beta_realtime: Final = client_sent_openai_beta_realtime_header(websocket) + client_ws: Final = cast(ClientWebSocketInterface, websocket) + openai_beta_realtime: Final = client_sent_openai_beta_realtime_header(client_ws) if not openai_beta_realtime: verbose_logger.debug( "OpenAI Realtime: connecting with GA protocol (no OpenAI-Beta header). " @@ -161,12 +163,12 @@ class OpenAIRealtime(OpenAIChatCompletion): ssl=ssl_config, ) as backend_ws: realtime_streaming: Final = RealTimeStreaming( - websocket, + client_ws, cast(ClientConnection, backend_ws), logging_obj, model=model, user_api_key_dict=user_api_key_dict, - request_data={"litellm_metadata": litellm_metadata or {}}, + request_data=cast(dict[str, object], {"litellm_metadata": litellm_metadata or {}}), force_transcription_model=( model if (query_params or {}).get("intent") == "transcription" else None ), diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 76618e0f742..a80796a40a2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -2,7 +2,7 @@ import asyncio import importlib from collections.abc import Awaitable, Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status @@ -128,7 +128,7 @@ if MCP_AVAILABLE: logging_results: Final = await asyncio.gather( _fire_mcp_tool_call_logging( logging_obj, - result, + cast("CallToolResult", result), start_time, end_time, user_api_key_auth=user_api_key_auth, diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index 0896c344f05..08b69c6ed7b 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -11,7 +11,7 @@ MCP Spec Reference: """ import typing -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Any, Final, NamedTuple, Optional, Protocol, Union, runtime_checkable if typing.TYPE_CHECKING: @@ -439,9 +439,9 @@ def _convert_mcp_messages_to_openai( ) # Separate marker items from regular content parts - tool_call_markers = [] - tool_result_markers = [] - regular_parts = [] + tool_call_markers: list[dict[str, object]] = [] + tool_result_markers: list[dict[str, object]] = [] + regular_parts: list[Mapping[str, object]] = [] for part in converted_parts: marker = part.get("_marker_type") if isinstance(part, dict) else None if marker == "tool_use": @@ -520,7 +520,7 @@ def _extract_text_parts( ) -> str | None: """Extract text parts from mixed content.""" items: Final = content if isinstance(content, list) else [content] - texts: Final = [] + texts: Final[list[str]] = [] for item in items: if getattr(item, "type", None) == "text": texts.append(getattr(item, "text", "")) @@ -1176,15 +1176,15 @@ async def _build_completion_kwargs( ) -async def _run_guardrails_and_call_llm( - completion_kwargs: dict[str, Any], +async def _apply_pre_call_hook( + completion_kwargs: dict[str, object], user_api_key_auth: "UserAPIKeyAuth", -) -> Any: +) -> dict[str, object]: try: from litellm.proxy.proxy_server import proxy_logging_obj as _plo if _plo is not None: - completion_kwargs = await typing.cast("ProxyLogging", _plo).pre_call_hook( + return await typing.cast("ProxyLogging", _plo).pre_call_hook( user_api_key_dict=user_api_key_auth, data=completion_kwargs, call_type="acompletion", @@ -1198,16 +1198,32 @@ async def _run_guardrails_and_call_llm( ) raise + return completion_kwargs + + +async def _call_acompletion( + acompletion_fn: "Callable[..., Awaitable[object]]", + completion_kwargs: dict[str, object], +) -> object: + return await acompletion_fn(**completion_kwargs) + + +async def _run_guardrails_and_call_llm( + completion_kwargs: dict[str, object], + user_api_key_auth: "UserAPIKeyAuth", +) -> Any: + completion_kwargs = await _apply_pre_call_hook(completion_kwargs, user_api_key_auth) + import litellm try: from litellm.proxy.proxy_server import llm_router if llm_router is not None: - return await llm_router.acompletion(**completion_kwargs) - return await litellm.acompletion(**completion_kwargs) + return await _call_acompletion(llm_router.acompletion, completion_kwargs) + return await _call_acompletion(litellm.acompletion, completion_kwargs) except ImportError: - return await litellm.acompletion(**completion_kwargs) + return await _call_acompletion(litellm.acompletion, completion_kwargs) async def handle_sampling_create_message( @@ -1284,7 +1300,7 @@ async def handle_sampling_create_message( client_ip=client_ip, ) - openai_messages: Final[Sequence[Mapping[str, object]]] = completion_kwargs["messages"] + openai_messages: Final[Sequence[Mapping[str, object]]] = completion_kwargs["messages"] or [] openai_tools: Final = completion_kwargs.get("tools") verbose_logger.debug( "MCP sampling: calling litellm.acompletion with model=%s, num_messages=%d, has_tools=%s", diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1c6ad84ddb4..b2a5922eeef 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -15,7 +15,7 @@ import types import uuid from collections.abc import AsyncIterator, Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final import httpx from fastapi import FastAPI, HTTPException @@ -145,7 +145,7 @@ try: ) # Robust auth lookup keyed by session_object. - _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() + _session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() except ImportError as e: verbose_logger.debug("MCP module not found: %s", e) MCP_AVAILABLE = False @@ -493,14 +493,14 @@ if MCP_AVAILABLE: def _gateway_create_initialization_options( self, notification_options: NotificationOptions | None = None, - experimental_capabilities: dict[str, dict[str, Any]] | None = None, + experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: opts: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) - updates: Final[dict[str, Any]] = {} + updates: Final[dict[str, str]] = {} merged: Final = _mcp_gateway_initialize_instructions.get() if merged is not None: updates["instructions"] = merged @@ -778,7 +778,7 @@ if MCP_AVAILABLE: get_virtual_tool_definitions, ) - return [Tool(**d) for d in get_virtual_tool_definitions()] + return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -825,7 +825,7 @@ if MCP_AVAILABLE: if not (host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta): return None - host_token: Final = getattr(host_ctx.meta, "progressToken", None) + host_token: Final[object] = getattr(host_ctx.meta, "progressToken", None) if host_token is None or not (hasattr(host_ctx, "session") and host_ctx.session): return None host_session: Final = host_ctx.session @@ -847,7 +847,7 @@ if MCP_AVAILABLE: async def _build_virtual_call_logging_obj( name: str, - arguments: dict[str, Any], + arguments: Mapping[str, object], user_api_key_auth: UserAPIKeyAuth, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual @@ -885,7 +885,7 @@ if MCP_AVAILABLE: async def _dispatch_virtual_mcp_tool( name: str, - arguments: dict[str, Any] | None, + arguments: Mapping[str, object] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, mcp_servers: list[str] | None = None, @@ -957,7 +957,7 @@ if MCP_AVAILABLE: ) @server.call_tool() - async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult: + async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: @@ -1621,7 +1621,7 @@ if MCP_AVAILABLE: async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, - prefetched_creds: dict[str, dict[str, Any]] | None = None, + prefetched_creds: Mapping[str, "OAuthCredentialPayload"] | None = None, ) -> dict[str, str] | None: """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. @@ -1634,9 +1634,7 @@ if MCP_AVAILABLE: resolve_user_oauth_access_token, ) - token: Final = await resolve_user_oauth_access_token( - getattr(user_api_key_auth, "user_id", None), server, prefetched_creds - ) + token: Final = await resolve_user_oauth_access_token(user_api_key_auth.user_id, server, prefetched_creds) return {"Authorization": f"Bearer {token}"} if token else None async def _prefetch_oauth_creds_for_user( @@ -1646,7 +1644,7 @@ if MCP_AVAILABLE: Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ - user_id: Final = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + user_id: Final = user_api_key_auth.user_id if user_api_key_auth else None if not user_id: return {} try: @@ -1871,7 +1869,7 @@ if MCP_AVAILABLE: list_tools_start_time: Final = datetime.now() litellm_logging_obj: LiteLLMLoggingObj | None = None - list_tools_request_data: dict[str, Any] = {} + list_tools_request_data: dict[str, object] = {} if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook @@ -1879,7 +1877,7 @@ if MCP_AVAILABLE: list_tools_call_id: Final = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) - spend_logs_metadata: Final[dict[str, Any]] = { + spend_logs_metadata: Final[dict[str, str | list[str]]] = { "mcp_operation": "list_tools", } if isinstance(list_tools_log_source, str): @@ -1916,19 +1914,18 @@ if MCP_AVAILABLE: _metadata_variable_name="metadata", ) - user_identifier: Final = getattr(user_api_key_auth, "end_user_id", None) or getattr( - user_api_key_auth, "user_id", None - ) + user_identifier: Final = user_api_key_auth.end_user_id or user_api_key_auth.user_id if user_identifier: list_tools_request_data["user"] = user_identifier try: - litellm_logging_obj, _ = function_setup( + function_setup_result: Final[tuple[LiteLLMLoggingObj, dict[str, object]]] = function_setup( original_function="list_mcp_tools", rules_obj=rules_obj, start_time=list_tools_start_time, **list_tools_request_data, ) + litellm_logging_obj, _ = function_setup_result if litellm_logging_obj: litellm_logging_obj.call_type = CallTypes.list_mcp_tools.value litellm_logging_obj.model = "MCP: list_tools" @@ -2615,7 +2612,7 @@ if MCP_AVAILABLE: async def execute_mcp_tool( name: str, - arguments: dict[str, Any], + arguments: Mapping[str, object], allowed_mcp_servers: list[MCPServer], start_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -2816,8 +2813,9 @@ if MCP_AVAILABLE: ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. - if isinstance(hook_result, dict) and "arguments" in hook_result: - arguments = hook_result["arguments"] + arguments = ( + hook_result["arguments"] if isinstance(hook_result, dict) and "arguments" in hook_result else arguments + ) verbose_logger.debug("Executing local registry tool: %s", name) # For BYOK servers the credential must be injected via a ContextVar @@ -2882,7 +2880,7 @@ if MCP_AVAILABLE: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) _request_resolved_auth_headers.reset(_resolved_token) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) # Try managed MCP server tool (the name is bare; the prefix boundary was # already resolved above against this server's registered prefixes) @@ -2952,11 +2950,10 @@ if MCP_AVAILABLE: server=prefix_server, raw_headers=raw_headers, ) - if "arguments" in hook_result: - arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args + arguments = hook_result.get("arguments", arguments) local_content = await _handle_local_mcp_tool(original_tool_name, arguments) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) return await _run_post_mcp_call_guardrails( result=response, @@ -3003,7 +3000,7 @@ if MCP_AVAILABLE: async def _fire_mcp_tool_call_logging( logging_obj: LiteLLMLoggingObj, - result: Any, + result: CallToolResult, start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -3070,7 +3067,7 @@ if MCP_AVAILABLE: @client async def call_mcp_tool( name: str, - arguments: dict[str, Any] | None = None, + arguments: Mapping[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3161,7 +3158,7 @@ if MCP_AVAILABLE: async def mcp_get_prompt( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, str] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3262,7 +3259,7 @@ if MCP_AVAILABLE: def _get_standard_logging_mcp_tool_call( name: str, - arguments: dict[str, Any], + arguments: Mapping[str, object], server_name: str | None, session_id: str | None = None, ) -> StandardLoggingMCPToolCall: @@ -3291,13 +3288,13 @@ if MCP_AVAILABLE: async def _handle_managed_mcp_tool( server_name: str, name: str, - arguments: dict[str, Any], + arguments: Mapping[str, object], user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, - litellm_logging_obj: Any | None = None, + litellm_logging_obj: LiteLLMLoggingObj | None = None, host_progress_callback: Callable | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" @@ -3320,7 +3317,7 @@ if MCP_AVAILABLE: return call_tool_result async def _handle_local_mcp_tool( - name: str, arguments: dict[str, Any] + name: str, arguments: Mapping[str, object] ) -> list[TextContent | ImageContent | EmbeddedResource]: """ Handle tool execution for local registry tools @@ -3387,7 +3384,7 @@ if MCP_AVAILABLE: mcp_servers_from_path = [servers_and_path] return mcp_servers_from_path - async def extract_mcp_auth_context(scope, path): + async def extract_mcp_auth_context(scope: Scope, path: str): """ Extracts mcp_servers from the path and processes the MCP request for auth context. Returns: (user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers) @@ -3426,7 +3423,8 @@ if MCP_AVAILABLE: Extract mcp-session-id from ASGI scope headers. Returns None if not present. """ - for header_name, header_value in scope.get("headers", []): + session_header_pairs: Final[list[tuple[bytes | str, bytes | str]]] = scope.get("headers") or [] + for header_name, header_value in session_header_pairs: name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": return header_value.decode() if isinstance(header_value, bytes) else str(header_value) @@ -3460,7 +3458,7 @@ if MCP_AVAILABLE: is best-effort in that mode. """ - def _bytes_for_hash(value: Any) -> bytes | None: + def _bytes_for_hash(value: object) -> bytes | None: """Only hash str/bytes secrets; skip mocks and other unexpected types.""" if value is None: return None @@ -3471,11 +3469,11 @@ if MCP_AVAILABLE: return None if user_api_key_auth is not None: - key_material: Final = _bytes_for_hash(getattr(user_api_key_auth, "api_key", None)) + key_material: Final = _bytes_for_hash(user_api_key_auth.api_key) if key_material: api_key_hash: Final = hashlib.sha256(key_material).hexdigest() return f"key:{api_key_hash}" - uid_material: Final = _bytes_for_hash(getattr(user_api_key_auth, "user_id", None)) + uid_material: Final = _bytes_for_hash(user_api_key_auth.user_id) if uid_material: user_id_hash: Final = hashlib.sha256(uid_material).hexdigest() return f"user:{user_id_hash}" @@ -3496,7 +3494,7 @@ if MCP_AVAILABLE: if not body: return False try: - data: Final = json.loads(body) + data: Final[object] = json.loads(body) if body else None return isinstance(data, dict) and data.get("method") == "initialize" except (json.JSONDecodeError, TypeError): return False @@ -3528,7 +3526,7 @@ if MCP_AVAILABLE: if message.get("type") != "http.request": break - body = message.get("body", b"") or b"" + body: bytes = message.get("body", b"") or b"" if body: # Only retain up to the remaining peek budget for sniffing. # The full ``message`` is already in memory (delivered by @@ -3571,9 +3569,9 @@ if MCP_AVAILABLE: Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header: Final = b"mcp-session-id" - _headers: Final = scope.get("headers", []) + _headers: Final[list[tuple[bytes | str, bytes | str]]] = scope.get("headers") or [] - def _normalize_header_name(header_name: Any) -> bytes | None: + def _normalize_header_name(header_name: object) -> bytes | None: if isinstance(header_name, bytes): return header_name.lower() if isinstance(header_name, str): @@ -3610,7 +3608,8 @@ if MCP_AVAILABLE: return False # --- Session not in this worker's memory --- - method: Final = scope.get("method", "").upper() + method_raw: Final[str] = scope.get("method") or "" + method: Final = method_raw.upper() if method == "DELETE": _remove_stateful_session_tracking(_session_id) @@ -3902,7 +3901,8 @@ if MCP_AVAILABLE: def _get_authorization_header_from_scope(scope: Scope) -> str | None: """First ``Authorization`` header value in the ASGI scope, or None.""" - for key, value in scope.get("headers", []): + auth_header_pairs: Final[list[tuple[bytes, bytes]]] = scope.get("headers") or [] + for key, value in auth_header_pairs: if key.lower() == b"authorization": return value.decode("latin-1") return None @@ -3921,7 +3921,8 @@ if MCP_AVAILABLE: ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream would leak the proxy key to a third-party MCP server. """ - has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", [])) + forwarded_header_pairs: Final[list[tuple[bytes, bytes]]] = scope.get("headers") or [] + has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in forwarded_header_pairs) if not has_litellm_key_header: return None return _get_authorization_header_from_scope(scope) @@ -4115,7 +4116,7 @@ if MCP_AVAILABLE: async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path") or "" ( user_api_key_auth, mcp_auth_header, @@ -4135,7 +4136,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + incoming_header_pairs: Final[list[tuple[bytes, bytes]]] = scope.get("headers") or [] + scope["headers"] = [(k, v) for k, v in incoming_header_pairs if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar (set by # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). @@ -4293,7 +4295,7 @@ if MCP_AVAILABLE: request_method: Final = (scope.get("method") or "").upper() if body and request_method == "POST": try: - _peeked: Final = json.loads(body) + _peeked: Final[object] = json.loads(body) if body else None if ( isinstance(_peeked, dict) and _peeked.get("jsonrpc") == "2.0" @@ -4436,7 +4438,7 @@ if MCP_AVAILABLE: async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path") or "" ( user_api_key_auth, mcp_auth_header, @@ -4456,7 +4458,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + incoming_header_pairs: Final[list[tuple[bytes, bytes]]] = scope.get("headers") or [] + scope["headers"] = [(k, v) for k, v in incoming_header_pairs if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar so the # downstream probe list matches the fully-authorized server set @@ -4680,7 +4683,8 @@ if MCP_AVAILABLE: ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": - for key, value in message.get("headers", []): + response_header_pairs: Final[list[tuple[bytes | str, bytes | str]]] = message.get("headers") or [] + for key, value in response_header_pairs: header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": session_id = value.decode() if isinstance(value, bytes) else str(value) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 71407c89813..f2893377166 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,9 +13,9 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from json import JSONDecodeError -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -104,11 +104,104 @@ class UpdatePublicModelGroupsRequest(BaseModel): model_config = ConfigDict(extra="forbid") +class _ProxyModelTableLike(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[LiteLLM_ProxyModelTable | None]: ... + + def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[LiteLLM_ProxyModelTable]]: ... + + def create(self, *, data: Mapping[str, object]) -> Awaitable[LiteLLM_ProxyModelTable]: ... + + def update( + self, *, where: Mapping[str, object], data: Mapping[str, object] + ) -> Awaitable[LiteLLM_ProxyModelTable]: ... + + def delete(self, *, where: Mapping[str, object]) -> Awaitable[LiteLLM_ProxyModelTable | None]: ... + + def delete_many(self, *, where: Mapping[str, object]) -> Awaitable[int]: ... + + +class _ModelRepositoryLike(Protocol): + @property + def table(self) -> _ProxyModelTableLike: ... + + +class _TeamRowLike(Protocol): + @property + def models(self) -> Sequence[object]: ... + + def model_dump(self) -> Mapping[str, object]: ... + + +class _TeamTableLike(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRowLike | None]: ... + + def update( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, bool] | None = None, + ) -> Awaitable[LiteLLM_TeamTable]: ... + + +class _TeamRepositoryLike(Protocol): + @property + def table(self) -> _TeamTableLike: ... + + +class _TeamAliasRefLike(Protocol): + @property + def team_id(self) -> str: ... + + +class _ModelAliasRowLike(Protocol): + @property + def id(self) -> int: ... + + @property + def model_aliases(self) -> Mapping[str, str]: ... + + @property + def team(self) -> _TeamAliasRefLike | None: ... + + +class _ModelAliasTableLike(Protocol): + def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRowLike]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ModelAliasRowLike]: ... + + +class _ModelTableRepositoryLike(Protocol): + @property + def table(self) -> _ModelAliasTableLike: ... + + +def _model_table(repo: _ModelRepositoryLike) -> _ProxyModelTableLike: + return repo.table + + +def _tx_model_table(table: _ProxyModelTableLike) -> _ProxyModelTableLike: + return table + + +def _team_table(repo: _TeamRepositoryLike) -> _TeamTableLike: + return repo.table + + +def _db_team_table(table: _TeamTableLike) -> _TeamTableLike: + return table + + +def _model_alias_table(repo: _ModelTableRepositoryLike) -> _ModelAliasTableLike: + return repo.table + + +def _row_model_info(row: LiteLLM_ProxyModelTable) -> object: + return row.model_info + + async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None: - db_model: Final = cast( - BaseModel | None, - await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}), - ) + db_model: Final = await _model_table(ModelRepository(prisma_client)).find_unique(where={"model_id": model_id}) if not db_model: return None @@ -338,7 +431,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model: Final = await ModelRepository(prisma_client).table.update( + updated_model: Final = await _model_table(ModelRepository(prisma_client)).update( where={"model_id": model_id}, data=update_data, ) @@ -363,7 +456,7 @@ async def patch_model( raise_if_reload_degraded_serving( before=live_before_reload, - written_models=[(model_id, getattr(updated_model, "model_info", None))], + written_models=[(model_id, _row_model_info(updated_model))], action="update", still_desired=still_desired_ids, ) @@ -441,7 +534,7 @@ async def _set_model_blocked_status( param=None, ) - updated_model: Final = await ModelRepository(prisma_client).table.update( + updated_model: Final = await _model_table(ModelRepository(prisma_client)).update( where={"model_id": data.model_id}, data={ "blocked": blocked, @@ -460,9 +553,7 @@ async def _set_model_blocked_status( user_api_key_dict=user_api_key_dict, table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, before_value=db_model.model_dump_json(exclude_none=True), - after_value=( - updated_model.model_dump_json(exclude_none=True) if isinstance(updated_model, BaseModel) else None - ), + after_value=updated_model.model_dump_json(exclude_none=True), litellm_changed_by=litellm_changed_by, litellm_proxy_admin_name=litellm_proxy_admin_name, ) @@ -470,7 +561,7 @@ async def _set_model_blocked_status( raise_if_reload_degraded_serving( before=live_before_reload, - written_models=[(data.model_id, getattr(updated_model, "model_info", None))], + written_models=[(data.model_id, _row_model_info(updated_model))], action=action, still_desired=still_desired_ids, ) @@ -563,12 +654,12 @@ async def _add_model_to_db( should_create_model_in_db: bool = True, ) -> LiteLLM_ProxyModelTable | None: # encrypt litellm params # - _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) + _litellm_params_dict: Final = model_params.litellm_params.model_dump(exclude_none=True) _original_litellm_model_name: Final = model_params.litellm_params.model for k, v in _litellm_params_dict.items(): encrypted_value = encrypt_value_helper(value=v, new_encryption_key=new_encryption_key) model_params.litellm_params[k] = encrypted_value - _data: Final[dict] = { + _data: Final[dict[str, object]] = { "model_id": model_params.model_info.id, "model_name": model_params.model_name, "litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), @@ -579,9 +670,9 @@ async def _add_model_to_db( if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id if should_create_model_in_db: - model_response = await ModelRepository(prisma_client).table.create(data=_data) + model_response = await _model_table(ModelRepository(prisma_client)).create(data=_data) else: - model_response = LiteLLM_ProxyModelTable(**_data) + model_response = LiteLLM_ProxyModelTable.model_validate(_data) return model_response @@ -769,7 +860,7 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient, table: Any | None = None + team_id: str, prisma_client: PrismaClient, table: _ProxyModelTableLike | None = None ) -> list[LiteLLM_ProxyModelTable]: """ Fetch all deployments for a given team_id from the database. @@ -785,8 +876,8 @@ async def _get_team_deployments( existing transaction. """ prefix: Final = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table - response: Final = await table.find_many( + resolved_table: Final = table if table is not None else _model_table(ModelRepository(prisma_client)) + response: Final = await resolved_table.find_many( where={ "model_name": {"startswith": prefix}, } @@ -797,7 +888,7 @@ async def _get_team_deployments( # Confirm team_id in model_info (defensive check) result: Final = [] for row in response: - model_info = model_info_as_mapping(row.model_info) + model_info = model_info_as_mapping(_row_model_info(row)) if model_info is not None and model_info.get("team_id") == team_id: result.append(row) return result @@ -806,7 +897,7 @@ async def _get_team_deployments( async def delete_team_models( team_ids: list[str], prisma_client: PrismaClient, - llm_router: Any | None, + llm_router: Router | None, ) -> list[str]: """ Delete every BYOK model owned by the given teams, from the DB and the router. @@ -821,11 +912,12 @@ async def delete_team_models( """ deleted_model_ids: Final[list[str]] = [] async with prisma_client.db.tx() as tx: + tx_table: Final = _tx_model_table(tx.litellm_proxymodeltable) for team_id in team_ids: - rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) + rows = await _get_team_deployments(team_id, prisma_client, table=tx_table) model_ids = [row.model_id for row in rows] if model_ids: - await tx.litellm_proxymodeltable.delete_many(where={"model_id": {"in": model_ids}}) + await tx_table.delete_many(where={"model_id": {"in": model_ids}}) deleted_model_ids.extend(model_ids) if deleted_model_ids: @@ -852,10 +944,10 @@ async def _get_team_public_model_names( deployments: Final = await _get_team_deployments(team_id, prisma_client) public_names: Final[set[str]] = set() for row in deployments: - model_info = model_info_as_mapping(row.model_info) + model_info = model_info_as_mapping(_row_model_info(row)) if model_info is not None: public_name = model_info.get("team_public_model_name") - if public_name: + if isinstance(public_name, str) and public_name: public_names.add(public_name) return public_names @@ -920,11 +1012,12 @@ async def _remove_unbacked_team_models( if not names_to_remove: return - existing_team_row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + team_table: Final = _db_team_table(prisma_client.db.litellm_teamtable) + existing_team_row: Final = await team_table.find_unique(where={"team_id": team_id}) if existing_team_row is None: return - updated_team_row: Final[LiteLLM_TeamTable] = await prisma_client.db.litellm_teamtable.update( + updated_team_row: Final[LiteLLM_TeamTable] = await team_table.update( where={"team_id": team_id}, data={"models": [model for model in existing_team_row.models if model not in names_to_remove]}, include={"object_permission": True}, @@ -953,7 +1046,7 @@ async def _update_existing_team_model_assignment( """ def _get_team_public_model_name( - model_info: dict | str | None, + model_info: object, ) -> str | None: parsed: Final = model_info_as_mapping(model_info) if parsed is None: @@ -978,7 +1071,8 @@ async def _update_existing_team_model_assignment( other_deployments_with_old_name: Final = [ d for d in team_deployments - if d.model_name != db_model.model_name and _get_team_public_model_name(d.model_info) == old_public_name + if d.model_name != db_model.model_name + and _get_team_public_model_name(_row_model_info(d)) == old_public_name ] # Add new name first, then delete old name to prevent access loss on partial failure @@ -1062,7 +1156,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row: Final = await TeamRepository(prisma_client).table.find_unique( + _existing_team_row: Final = await _team_table(TeamRepository(prisma_client)).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1091,7 +1185,7 @@ class ModelManagementAuthChecks: ) -> Literal[True]: ## Check team model auth if model_params.model_info is not None and model_params.model_info.team_id is not None: - team_obj_row: Final = await TeamRepository(prisma_client).table.find_unique( + team_obj_row: Final = await _team_table(TeamRepository(prisma_client)).find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -1169,7 +1263,9 @@ async def delete_model( }, ) - model_in_db: Final = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_info.id}) + model_in_db: Final = await _model_table(ModelRepository(prisma_client)).find_unique( + where={"model_id": model_info.id} + ) if model_in_db is None: raise HTTPException( status_code=400, @@ -1192,7 +1288,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result: Final = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id}) + result: Final = await _model_table(ModelRepository(prisma_client)).delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1265,7 +1361,9 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases: Final = await ModelTableRepository(prisma_client).table.find_many(include={"team": True}) + team_model_aliases: Final = await _model_alias_table(ModelTableRepository(prisma_client)).find_many( + include={"team": True} + ) tasks: Final = [] removed_model_aliases: Final = [] for team_model_alias in team_model_aliases: @@ -1276,11 +1374,11 @@ async def delete_team_model_alias( key = list(model_aliases.keys())[list(model_aliases.values()).index(public_model_name)] if team_model_alias.team is not None: removed_model_aliases.append((team_model_alias.team.team_id, key)) - del model_aliases[key] + remaining_aliases = {alias: target for alias, target in model_aliases.items() if alias != key} tasks.append( - ModelTableRepository(prisma_client).table.update( + _model_alias_table(ModelTableRepository(prisma_client)).update( where={"id": id}, - data={"model_aliases": json.dumps(model_aliases)}, + data={"model_aliases": json.dumps(remaining_aliases)}, ) ) await asyncio.gather(*tasks) @@ -1423,9 +1521,7 @@ async def add_new_model( user_api_key_dict=user_api_key_dict, table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, before_value=None, - after_value=( - model_response.model_dump_json(exclude_none=True) if isinstance(model_response, BaseModel) else None - ), + after_value=model_response.model_dump_json(exclude_none=True), litellm_changed_by=user_api_key_dict.user_id, litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, ) @@ -1433,7 +1529,7 @@ async def add_new_model( raise_if_reload_degraded_serving( before=live_before_reload, - written_models=[(model_response.model_id, getattr(model_response, "model_info", None))], + written_models=[(model_response.model_id, _row_model_info(model_response))], action="create", still_desired=still_desired_ids, ) @@ -1501,7 +1597,9 @@ async def update_model( if _model_id is None: raise Exception("model_info.id not provided") - _existing_litellm_params = await ModelRepository(prisma_client).table.find_unique(where={"model_id": _model_id}) + _existing_litellm_params = await _model_table(ModelRepository(prisma_client)).find_unique( + where={"model_id": _model_id} + ) if _existing_litellm_params is None: if llm_router is not None and llm_router.get_deployment(model_id=_model_id) is not None: @@ -1551,11 +1649,11 @@ async def update_model( else: pass - _data: Final[dict] = { + _data: Final[dict[str, object]] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response: Final = await ModelRepository(prisma_client).table.update( + model_response: Final = await _model_table(ModelRepository(prisma_client)).update( where={"model_id": _model_id}, data=_data, ) @@ -1570,16 +1668,8 @@ async def update_model( action="updated", user_api_key_dict=user_api_key_dict, table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, - before_value=( - _existing_litellm_params.model_dump_json(exclude_none=True) - if isinstance(_existing_litellm_params, BaseModel) - else None - ), - after_value=( - model_response.model_dump_json(exclude_none=True) - if isinstance(model_response, BaseModel) - else None - ), + before_value=_existing_litellm_params.model_dump_json(exclude_none=True), + after_value=model_response.model_dump_json(exclude_none=True), litellm_changed_by=user_api_key_dict.user_id, litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, ) @@ -1587,7 +1677,7 @@ async def update_model( raise_if_reload_degraded_serving( before=live_before_reload, - written_models=[(_model_id, getattr(model_response, "model_info", None))], + written_models=[(_model_id, _row_model_info(model_response))], action="update", still_desired=still_desired_ids, ) @@ -1832,7 +1922,7 @@ async def get_auto_router_classifier_default_prompt( ) -def _deduplicate_litellm_router_models(models: list[dict]) -> list[dict]: +def _deduplicate_litellm_router_models(models: Sequence[dict[str, object]]) -> list[dict[str, object]]: """ Deduplicate models based on their model_info.id field. Returns a list of unique models keeping only the first occurrence of each model ID. @@ -1843,10 +1933,10 @@ def _deduplicate_litellm_router_models(models: list[dict]) -> list[dict]: Returns: List of deduplicated model dictionaries """ - seen_ids: Final = set() + seen_ids: Final[set[object]] = set() unique_models: Final = [] for model in models: - model_id = model.get("model_info", {}).get("id", None) + model_id = (model_info_as_mapping(model.get("model_info")) or {}).get("id") if model_id is not None and model_id not in seen_ids: unique_models.append(model) seen_ids.add(model_id) @@ -1997,18 +2087,18 @@ async def clear_cache() -> frozenset[str] | None: verbose_proxy_logger.debug("Clearing only DB models, preserving config models") # Get current models and filter out DB models - current_models: Final = llm_router.model_list.copy() + current_models: Final = cast(list[dict[str, object]], llm_router.model_list.copy()) config_models: Final = [] db_model_ids: Final = [] for model in current_models: - model_info = model.get("model_info", {}) - if model_info.get("db_model", False): - # This is a DB model, mark for deletion - db_model_ids.append(model_info.get("id")) - else: + model_info = model_info_as_mapping(model.get("model_info")) or {} + if not model_info.get("db_model", False): # This is a config model, preserve it config_models.append(model) + elif isinstance(candidate_id := model_info.get("id"), str): + # This is a DB model, mark for deletion + db_model_ids.append(candidate_id) # Clear only DB models for model_id in db_model_ids: @@ -2024,11 +2114,13 @@ async def clear_cache() -> frozenset[str] | None: # name from every router registry (no-op where absent); missing quality/adaptive # entries would otherwise make init raise "already exists" on reload and abort it. db_router_names: Final = { - model.get("model_name") + model_name for model in current_models - if model.get("model_name") is not None - and model.get("model_info", {}).get("db_model", False) - and str(model.get("litellm_params", {}).get("model", "")).startswith("auto_router/") + if isinstance(model_name := model.get("model_name"), str) + and (model_info_as_mapping(model.get("model_info")) or {}).get("db_model", False) + and str((model_info_as_mapping(model.get("litellm_params")) or {}).get("model", "")).startswith( + "auto_router/" + ) } for model_name in db_router_names: llm_router.auto_routers.pop(model_name, None) diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 987823d987f..e329412009b 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -18,13 +18,16 @@ Scoping: """ import json -from typing import Any, Final +from collections.abc import Awaitable, Mapping +from datetime import datetime +from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( CommonProxyErrors, + LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view, @@ -40,10 +43,95 @@ from litellm.types.memory_management import ( MemoryUpdateRequest, ) +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + router: Final = APIRouter() -def _serialize_metadata_for_prisma(metadata: Any) -> str: +class _MemoryRowLike(Protocol): + @property + def memory_id(self) -> str: ... + + @property + def key(self) -> str: ... + + @property + def value(self) -> str: ... + + @property + def metadata(self) -> object | None: ... + + @property + def user_id(self) -> str | None: ... + + @property + def team_id(self) -> str | None: ... + + @property + def created_at(self) -> datetime | None: ... + + @property + def created_by(self) -> str | None: ... + + @property + def updated_at(self) -> datetime | None: ... + + @property + def updated_by(self) -> str | None: ... + + +class _MemoryTableLike(Protocol): + def create(self, *, data: Mapping[str, str | None]) -> Awaitable[_MemoryRowLike]: ... + + def count(self, *, where: Mapping[str, object]) -> Awaitable[int]: ... + + def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str], + skip: int = 0, + take: int = 0, + ) -> Awaitable[list[_MemoryRowLike]]: ... + + def update(self, *, where: Mapping[str, str], data: Mapping[str, str | None]) -> Awaitable[_MemoryRowLike]: ... + + def delete(self, *, where: Mapping[str, str]) -> Awaitable[_MemoryRowLike | None]: ... + + +class _MemoryRepositoryLike(Protocol): + @property + def table(self) -> _MemoryTableLike: ... + + +class _TeamTableLike(Protocol): + def find_unique(self, *, where: Mapping[str, str]) -> Awaitable[LiteLLM_TeamTable | None]: ... + + +class _TeamRepositoryLike(Protocol): + @property + def table(self) -> _TeamTableLike: ... + + +class _HasMetadata(Protocol): + @property + def metadata(self) -> object | None: ... + + +def _memory_table(repo: _MemoryRepositoryLike) -> _MemoryTableLike: + return repo.table + + +def _team_table(repo: _TeamRepositoryLike) -> _TeamTableLike: + return repo.table + + +def _metadata_of(body: _HasMetadata) -> object | None: + return body.metadata + + +def _serialize_metadata_for_prisma(metadata: object) -> str: """ Encode a `metadata` payload for the `Json?` column. @@ -62,14 +150,14 @@ def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN -def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict | None: +def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict[str, object] | None: """ Prisma `where` fragment restricting rows to those the caller can see. Returns None for admins (no restriction). """ if user_api_key_has_admin_view(user_api_key_dict): return None - ors: Final[list[dict]] = [] + ors: Final[list[dict[str, str]]] = [] if user_api_key_dict.user_id: ors.append({"user_id": user_api_key_dict.user_id}) if user_api_key_dict.team_id: @@ -80,7 +168,7 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict | None: return {"OR": ors} -def _row_to_model(row: Any) -> LiteLLM_MemoryRow: +def _row_to_model(row: _MemoryRowLike) -> LiteLLM_MemoryRow: return LiteLLM_MemoryRow( memory_id=row.memory_id, key=row.key, @@ -95,7 +183,7 @@ def _row_to_model(row: Any) -> LiteLLM_MemoryRow: ) -def _require_prisma(): +def _require_prisma() -> "PrismaClient": from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -113,7 +201,9 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT return HTTPException(status_code=500, detail=default_detail) -async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth) -> None: +async def _assert_write_access( + prisma_client: "PrismaClient", row: _MemoryRowLike, user_api_key_dict: UserAPIKeyAuth +) -> None: """ Enforce ownership for mutations (PUT/DELETE). @@ -135,8 +225,8 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: """ if _is_admin(user_api_key_dict): return - row_user_id: Final = getattr(row, "user_id", None) - row_team_id: Final = getattr(row, "team_id", None) + row_user_id: Final = row.user_id + row_team_id: Final = row.team_id # Personal ownership. if row_user_id and row_user_id == user_api_key_dict.user_id: @@ -153,7 +243,7 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: ) -async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool: +async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool: """ True if the caller is a team admin of `team_id`, or an org admin for the team's organization. Mirrors the auth pattern used by team-management @@ -168,7 +258,7 @@ async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAu ) try: - team_obj: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + team_obj: Final = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) except Exception as e: verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e) return False @@ -269,7 +359,7 @@ async def create_memory( # `metadata` is a `Json?` column — prisma-client-python rejects raw # Python values, so JSON-encode any non-null payload and omit the field # entirely when None so the column defaults to SQL NULL. - create_data: Final[dict] = { + create_data: Final[dict[str, str | None]] = { "key": body.key, "value": body.value, "user_id": user_id, @@ -278,10 +368,10 @@ async def create_memory( "updated_by": user_api_key_dict.user_id, } if body.metadata is not None: - create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) + create_data["metadata"] = _serialize_metadata_for_prisma(_metadata_of(body)) try: - row: Final = await MemoryRepository(prisma_client).table.create(data=create_data) + row: Final = await _memory_table(MemoryRepository(prisma_client)).create(data=create_data) except Exception as e: # Key is globally unique. Any duplicate → 409. if _is_unique_violation(e): @@ -325,14 +415,14 @@ async def list_memory( # top-level "AND" — safer than `dict.update` since future visibility # filters could grow an "OR" key that would clobber this one if merged # by key. - key_filter: Final[dict] = {} + key_filter: Final[dict[str, str | dict[str, str]]] = {} if key_prefix is not None: key_filter["key"] = {"startsWith": key_prefix} elif key is not None: key_filter["key"] = key vis: Final = _visibility_filter(user_api_key_dict) - where: dict + where: Mapping[str, object] if vis is None: where = key_filter elif not key_filter: @@ -341,8 +431,8 @@ async def list_memory( where = {"AND": [key_filter, vis]} try: - total: Final = await MemoryRepository(prisma_client).table.count(where=where) - rows: Final = await MemoryRepository(prisma_client).table.find_many( + total: Final = await _memory_table(MemoryRepository(prisma_client)).count(where=where) + rows: Final = await _memory_table(MemoryRepository(prisma_client)).find_many( where=where, order={"updated_at": "desc"}, skip=(page - 1) * page_size, @@ -354,12 +444,16 @@ async def list_memory( return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total) -async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth) -> Any: +async def _find_memory_for_caller( + prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth +) -> _MemoryRowLike: """Look up a memory row by key, scoped to the caller's visibility.""" - key_filter: Final[dict] = {"key": key} + key_filter: Final[dict[str, str]] = {"key": key} vis: Final = _visibility_filter(user_api_key_dict) - where: Final[dict] = key_filter if vis is None else {"AND": [key_filter, vis]} - rows = await MemoryRepository(prisma_client).table.find_many(where=where, take=1, order={"updated_at": "desc"}) + where: Final[Mapping[str, object]] = key_filter if vis is None else {"AND": [key_filter, vis]} + rows: Final = await _memory_table(MemoryRepository(prisma_client)).find_many( + where=where, take=1, order={"updated_at": "desc"} + ) if not rows: raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found") return rows[0] @@ -415,11 +509,11 @@ async def upsert_memory( fields_sent: Final = body.model_fields_set metadata_in_payload: Final = "metadata" in fields_sent - data: Final[dict] = {} + data: Final[dict[str, str | None]] = {} if body.value is not None: data["value"] = body.value if metadata_in_payload: - data["metadata"] = _serialize_metadata_for_prisma(body.metadata) + data["metadata"] = _serialize_metadata_for_prisma(_metadata_of(body)) if not data: raise HTTPException( status_code=400, @@ -427,7 +521,7 @@ async def upsert_memory( ) data["updated_by"] = user_api_key_dict.user_id - async def _find_existing() -> Any: + async def _find_existing() -> _MemoryRowLike | None: """Return the caller-visible row for `key`, or None.""" try: return await _find_memory_for_caller(prisma_client, key, user_api_key_dict) @@ -444,7 +538,7 @@ async def upsert_memory( # their team) — otherwise a teammate could overwrite a personal # entry through the OR-based visibility filter. await _assert_write_access(prisma_client, existing, user_api_key_dict) - row = await MemoryRepository(prisma_client).table.update( + row = await _memory_table(MemoryRepository(prisma_client)).update( where={"memory_id": existing.memory_id}, data=data, ) @@ -459,7 +553,7 @@ async def upsert_memory( # Omit `metadata` when None so the column defaults to SQL NULL; # otherwise JSON-encode for Prisma — same pattern as # `create_memory` above. - create_data: Final[dict] = { + create_data: Final[dict[str, str | None]] = { "key": key, "value": body.value, "user_id": user_id, @@ -468,9 +562,9 @@ async def upsert_memory( "updated_by": user_api_key_dict.user_id, } if body.metadata is not None: - create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) + create_data["metadata"] = _serialize_metadata_for_prisma(_metadata_of(body)) try: - row = await MemoryRepository(prisma_client).table.create(data=create_data) + row = await _memory_table(MemoryRepository(prisma_client)).create(data=create_data) except Exception as e: # Race: a concurrent PUT/POST created the row after our check. # Re-read and fall back to an update so the PUT stays idempotent @@ -487,7 +581,7 @@ async def upsert_memory( ) # Same write-authorization check as the non-race path. await _assert_write_access(prisma_client, existing_after_race, user_api_key_dict) - row = await MemoryRepository(prisma_client).table.update( + row = await _memory_table(MemoryRepository(prisma_client)).update( where={"memory_id": existing_after_race.memory_id}, data=data, ) @@ -515,7 +609,7 @@ async def delete_memory( # Visibility != write authority — see the upsert handler for the rationale. await _assert_write_access(prisma_client, row, user_api_key_dict) try: - await MemoryRepository(prisma_client).table.delete(where={"memory_id": row.memory_id}) + await _memory_table(MemoryRepository(prisma_client)).delete(where={"memory_id": row.memory_id}) except Exception as e: raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 90daaaeae6b..167650c4531 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11452,7 +11452,7 @@ async def non_admin_all_models( ) # de-duplicate models. Only return unique model ids - unique_models: Final = _deduplicate_litellm_router_models(models=all_models) + unique_models: Final = _deduplicate_litellm_router_models(models=cast(list[dict[str, object]], all_models)) return unique_models diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 36f1e4cf480..f352b9745b8 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -17,7 +17,12 @@ from __future__ import annotations import hashlib import uuid -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import ConfigDict, TypeAdapter +from typing_extensions import TypeAliasType, TypedDict import litellm from litellm._logging import verbose_logger @@ -38,6 +43,35 @@ if TYPE_CHECKING: from litellm import Router from litellm.types.rag import RAGIngestOptions +_JSONValue = TypeAliasType( + "_JSONValue", + "Mapping[str, _JSONValue] | Sequence[_JSONValue] | str | int | float | bool | None", +) + + +class _S3Vector(TypedDict): + key: str + data: dict[str, list[float]] + metadata: dict[str, str] + + +class _S3QueryMatchView(TypedDict, total=False): + metadata: dict[str, _JSONValue] + + +class _S3QueryResultsView(TypedDict, total=False): + vectors: list[_S3QueryMatchView] + + +_RESPONSE_ADAPTER: Final = TypeAdapter(httpx.Response, config=ConfigDict(arbitrary_types_allowed=True)) +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, _JSONValue]) +_QUERY_RESULTS_VIEW_ADAPTER: Final = TypeAdapter(_S3QueryResultsView) +_STR_ADAPTER: Final = TypeAdapter(str) +_OPTIONAL_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) +_STR_LIST_ADAPTER: Final = TypeAdapter(list[str]) +_FLOAT_LIST_ADAPTER: Final = TypeAdapter(list[float]) +_INT_SOURCE_ADAPTER: Final[TypeAdapter[int | float | str]] = TypeAdapter(int | float | str) + class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): """ @@ -66,19 +100,23 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): BaseAWSLLM.__init__(self) # Extract config - self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] - self.index_name = self.vector_store_config.get("index_name") - self.distance_metric = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) - self.non_filterable_metadata_keys = self.vector_store_config.get( - "non_filterable_metadata_keys", - S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, + self.vector_bucket_name = _STR_ADAPTER.validate_python(self.vector_store_config["vector_bucket_name"]) + self.index_name = _OPTIONAL_STR_ADAPTER.validate_python(self.vector_store_config.get("index_name")) + self.distance_metric = _STR_ADAPTER.validate_python( + self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) + ) + self.non_filterable_metadata_keys = _STR_LIST_ADAPTER.validate_python( + self.vector_store_config.get( + "non_filterable_metadata_keys", + S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, + ) ) # Get dimension from config (will be auto-detected on first use if not provided) self.dimension = self._get_dimension_from_config() # Get AWS region using BaseAWSLLM method - _aws_region: Final = self.vector_store_config.get("aws_region_name") + _aws_region: Final = _OPTIONAL_STR_ADAPTER.validate_python(self.vector_store_config.get("aws_region_name")) self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( aws_region_name=str(_aws_region) if _aws_region else None ) @@ -135,7 +173,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): Returns None if dimension should be auto-detected. """ if "dimension" in self.vector_store_config: - return int(self.vector_store_config["dimension"]) + return int(_INT_SOURCE_ADAPTER.validate_python(self.vector_store_config["dimension"])) return None async def _ensure_config_initialized(self): @@ -166,7 +204,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): url: str, data: str | None = None, headers: dict[str, str] | None = None, - ) -> Any: + ) -> httpx.Response: """ Helper to sign and execute AWS API requests using httpx + SigV4. @@ -233,7 +271,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): else: raise ValueError(f"Unsupported HTTP method: {method}") - return response + return _RESPONSE_ADAPTER.validate_python(response) async def _ensure_vector_bucket_exists(self): """Create vector bucket if it doesn't exist using GetVectorBucket and CreateVectorBucket APIs.""" @@ -311,7 +349,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) # Prepare index configuration per AWS API docs - index_config: Final = { + index_config: Final[dict[str, str | int | dict[str, list[str]] | None]] = { "vectorBucketName": self.vector_bucket_name, "indexName": self.index_name, "dataType": "float32", @@ -336,7 +374,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): verbose_logger.exception("Error creating vector index: %s", e) raise - async def _put_vectors(self, vectors: list[dict[str, Any]]): + async def _put_vectors(self, vectors: list[_S3Vector]): """ Call PutVectors API to store vectors in S3 Vectors. @@ -442,7 +480,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): raise ValueError(error_msg) # Prepare vectors for PutVectors API - vectors: Final = [] + vectors: Final[list[_S3Vector]] = [] for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): # Build metadata dict metadata: dict[str, str] = { @@ -453,7 +491,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): if filename: metadata["filename"] = filename # Filterable - vector_obj = { + vector_obj: _S3Vector = { "key": f"{filename}_{i}" if filename else f"chunk_{i}", "data": {"float32": embedding}, "metadata": metadata, @@ -468,7 +506,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): vector_store_id: Final = f"{self.vector_bucket_name}:{self.index_name}" return vector_store_id, filename - async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> dict[str, Any] | None: + async def query_vector_store( + self, vector_store_id: str, query: str, top_k: int = 5 + ) -> dict[str, _JSONValue] | None: """ Query S3 Vectors using QueryVectors API. @@ -488,8 +528,8 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): embedding_model: Final = self.embedding_config.get("model", "text-embedding-3-small") - response = await litellm.aembedding(model=embedding_model, input=[query]) - query_embedding: Final = response.data[0]["embedding"] + embedding_response = await litellm.aembedding(model=embedding_model, input=[query]) + query_embedding: Final = _FLOAT_LIST_ADAPTER.validate_python(embedding_response.data[0]["embedding"]) # Call QueryVectors API url: Final = f"https://s3vectors.{self.aws_region_name}.api.aws/QueryVectors" @@ -507,14 +547,15 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body)) if response.status_code == 200: - results: Final = response.json() - verbose_logger.debug("Query returned %s results", len(results.get("vectors", []))) + results: Final = _JSON_OBJECT_ADAPTER.validate_python(response.json()) + vectors: Final = _QUERY_RESULTS_VIEW_ADAPTER.validate_python(results).get("vectors", []) + verbose_logger.debug("Query returned %s results", len(vectors)) # Check if query terms appear in results - if results.get("vectors"): - for result in results["vectors"]: + if vectors: + for result in vectors: metadata = result.get("metadata", {}) - source_text = metadata.get("source_text", "") + source_text = _STR_ADAPTER.validate_python(metadata.get("source_text", "")) if query.lower() in source_text.lower(): return results diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 820839fc6bf..4e5c4996f9d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,11 +5,11 @@ import json import time import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Mapping, MutableMapping from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable import httpx from openai._streaming import SSEDecoder @@ -33,6 +33,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, + ResponseAPIUsage, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -41,6 +42,7 @@ from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook if TYPE_CHECKING: + from litellm.caching.base_cache import BaseCache from litellm.proxy._types import UserAPIKeyAuth from litellm.types.responses.streaming_websocket import ( PresidioGuardrailCallback, @@ -68,6 +70,58 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verif return _is_json_object(value) and all(isinstance(item, str) for item in value.values()) +def _parse_json(payload: str | bytes) -> object: + return cast(object, json.loads(payload)) + + +def _get_attr(source: object, name: str) -> object: + value: Final[object] = getattr(source, name, None) + return value + + +def _mapping_get(source: Mapping[str, object], key: str, default: object = None) -> object: + return source.get(key, default) + + +def _pop_kwarg(kwargs: MutableMapping[str, object], key: str) -> object: + return kwargs.pop(key, None) + + +@runtime_checkable +class _SupportsModelDump(Protocol): + def model_dump(self, *, exclude_none: bool = ...) -> dict[str, object]: ... + + +@runtime_checkable +class _SupportsModelDumpJson(Protocol): + def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... + + +@runtime_checkable +class _StreamingDeploymentHookCallback(Protocol): + async def async_post_call_streaming_deployment_hook( + self, + request_data: dict[str, object], + response_chunk: ResponsesAPIStreamingResponse, + call_type: CallTypes | None, + ) -> ResponsesAPIStreamingResponse | None: ... + + +@runtime_checkable +class _StreamingHookInvoker(Protocol): + def __call__(self, chunk: ResponsesAPIStreamingResponse) -> Awaitable[ResponsesAPIStreamingResponse]: ... + + +@runtime_checkable +class _PiiUnmasker(Protocol): + def __call__(self, text: str, pii_tokens: dict[str, str]) -> str: ... + + +@runtime_checkable +class _ShouldStoreChecker(Protocol): + def __call__(self, *, original_function: object, kwargs: dict[str, object]) -> bool: ... + + def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None: model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None model_id: Final = model_info.get("id") if _is_json_object(model_info) else None @@ -112,7 +166,7 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]: - if isinstance(error_obj, dict): + if _is_json_object(error_obj): raw_message = error_obj.get("message") raw_type = error_obj.get("type") raw_code = error_obj.get("code") @@ -182,9 +236,10 @@ class BaseResponsesAPIStreamingIterator: # set hidden params for response headers (e.g., x-litellm-model-id) # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py + _litellm_params: Final = _mapping_get(self.logging_obj.model_call_details, "litellm_params", {}) _api_base: Final = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=_litellm_params if _is_json_object(_litellm_params) else {}, ) self._hidden_params: dict[str, object] = { "model_id": _model_id_from_metadata(litellm_metadata), @@ -227,10 +282,10 @@ class BaseResponsesAPIStreamingIterator: try: # Parse the JSON chunk - parsed_chunk: Final = json.loads(chunk) + parsed_chunk: Final = _parse_json(chunk) # Format as ResponsesAPIStreamingResponse - if isinstance(parsed_chunk, dict): + if _is_json_object(parsed_chunk): if self.responses_api_provider_config is None: raise ValueError("responses_api_provider_config is required to process live streaming chunks") openai_responses_api_chunk: Final = self.responses_api_provider_config.transform_streaming_response( @@ -288,8 +343,9 @@ class BaseResponsesAPIStreamingIterator: _stream_model_id, ) else: + _part_annotations: Final[object] = getattr(_part, "annotations", None) ResponsesAPIRequestUtils._encode_container_ids_in_annotations( - getattr(_part, "annotations", None), + _part_annotations, self.custom_llm_provider, _stream_model_id, ) @@ -302,9 +358,9 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - item: Final = getattr(openai_responses_api_chunk, "item", None) + item: Final[object] = getattr(openai_responses_api_chunk, "item", None) if item: - encrypted_content: Final = getattr(item, "encrypted_content", None) + encrypted_content: Final[object] = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): model_id: Final = _model_id_from_metadata(self.litellm_metadata) if model_id: @@ -414,8 +470,8 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None - error_info: Final = getattr(response_obj, "error", None) if response_obj else None + response_obj: Final = _get_attr(self.completed_response, "response") if self.completed_response else None + error_info: Final = _get_attr(response_obj, "error") if response_obj else None error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) exception: Final = litellm.APIError( @@ -426,10 +482,10 @@ class BaseResponsesAPIStreamingIterator: ) self._handle_failure(exception) - def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: + def _record_failed_response_usage(self, response_obj: object) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is None: return try: @@ -501,15 +557,15 @@ class BaseResponsesAPIStreamingIterator: if response_obj is None: return - caching_handler: Final = getattr(self.logging_obj, "_llm_caching_handler", None) + caching_handler: Final[object] = getattr(self.logging_obj, "_llm_caching_handler", None) if caching_handler is None: return - request_kwargs = getattr(caching_handler, "request_kwargs", None) - if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: + request_kwargs_raw: Final[object] = getattr(caching_handler, "request_kwargs", None) + if not _is_json_object(request_kwargs_raw) or request_kwargs_raw.get("stream") is not True: return - request_kwargs = request_kwargs.copy() - preset_cache_key = getattr(caching_handler, "preset_cache_key", None) + request_kwargs: Final = request_kwargs_raw.copy() + preset_cache_key: object = getattr(caching_handler, "preset_cache_key", None) request_cache_key: Final = request_kwargs.pop("cache_key", None) if preset_cache_key is None: preset_cache_key = request_cache_key @@ -519,8 +575,12 @@ class BaseResponsesAPIStreamingIterator: if preset_cache_key is not None: request_kwargs["cache_key"] = preset_cache_key - if not caching_handler._should_store_result_in_cache( - original_function=caching_handler.original_function, + should_store_fn: Final[object] = getattr(caching_handler, "_should_store_result_in_cache", None) + original_function: Final[object] = getattr(caching_handler, "original_function", None) + if not isinstance(should_store_fn, _ShouldStoreChecker): + return + if not should_store_fn( + original_function=original_function, kwargs=request_kwargs, ): return @@ -528,12 +588,13 @@ class BaseResponsesAPIStreamingIterator: if litellm.cache is None: return + dual_cache: Final[BaseCache | None] = getattr(caching_handler, "dual_cache", None) cached_response: Final = response_obj.model_dump_json() if is_async: cache_write_task: Final = asyncio.create_task( litellm.cache.async_add_cache( cached_response, - dynamic_cache_object=getattr(caching_handler, "dual_cache", None), + dynamic_cache_object=dual_cache, **request_kwargs, ) ) @@ -546,7 +607,7 @@ class BaseResponsesAPIStreamingIterator: else: litellm.cache.add_cache( cached_response, - dynamic_cache_object=getattr(caching_handler, "dual_cache", None), + dynamic_cache_object=dual_cache, **request_kwargs, ) @@ -572,11 +633,11 @@ class BaseResponsesAPIStreamingIterator: except Exception: typed_call_type = None - request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {}) - callbacks: Final = getattr(litellm, "callbacks", None) or [] + request_data: Final[dict[str, object]] = self.request_data or self.logging_obj.model_call_details + callbacks: Final = litellm.callbacks or [] hooks_ran = False for callback in callbacks: - if hasattr(callback, "async_post_call_streaming_deployment_hook"): + if isinstance(callback, _StreamingDeploymentHookCallback): hooks_ran = True result = await callback.async_post_call_streaming_deployment_hook( request_data=request_data, @@ -606,7 +667,7 @@ class BaseResponsesAPIStreamingIterator: if self.completed_response is None: return - request_payload: Final[dict[str, Any]] = {} + request_payload: Final[dict[str, object]] = {} if isinstance(self.request_data, dict): request_payload.update(self.request_data) try: @@ -616,8 +677,8 @@ class BaseResponsesAPIStreamingIterator: pass if "litellm_params" not in request_payload: try: - request_payload["litellm_params"] = getattr(self.logging_obj, "model_call_details", {}).get( - "litellm_params", {} + request_payload["litellm_params"] = _mapping_get( + self.logging_obj.model_call_details, "litellm_params", {} ) except Exception: request_payload["litellm_params"] = {} @@ -695,12 +756,14 @@ class BaseResponsesAPIStreamingIterator: pass -async def call_post_streaming_hooks_for_testing(iterator, chunk): +async def call_post_streaming_hooks_for_testing( + iterator: object, chunk: ResponsesAPIStreamingResponse +) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ - hook_fn: Final = getattr(iterator, "_call_post_streaming_deployment_hook", None) - if hook_fn is None: + hook_fn: Final[object] = getattr(iterator, "_call_post_streaming_deployment_hook", None) + if hook_fn is None or not isinstance(hook_fn, _StreamingHookInvoker): return chunk return await hook_fn(chunk) @@ -1016,18 +1079,18 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): return evt -def _dump_response_object(obj: Any) -> dict[str, Any]: - if hasattr(obj, "model_dump"): +def _dump_response_object(obj: object) -> dict[str, object]: + if isinstance(obj, _SupportsModelDump): return obj.model_dump() - if isinstance(obj, dict): + if _is_json_object(obj): return obj return {} def _build_response_status_event( event_type: Literal[ - "response.created", - "response.in_progress", + ResponsesAPIStreamEvents.RESPONSE_CREATED, + ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, ], transformed: ResponsesAPIResponse, ) -> ResponsesAPIStreamingResponse: @@ -1090,7 +1153,7 @@ def _add_text_like_part_events( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], + part_payload: dict[str, object], chunk_size: int, ) -> None: openai_types: Final = _get_openai_response_types() @@ -1107,7 +1170,10 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []): + annotations_raw: Final = part_payload.get("annotations") + for annotation_index, annotation in enumerate(annotations_raw if _is_json_array(annotations_raw) else []): + if not _is_json_object(annotation): + continue events.append( openai_types.OutputTextAnnotationAddedEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, @@ -1173,7 +1239,8 @@ def _build_synthetic_response_events( ] sequence_number = 0 - for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): + transformed_output: Final = _get_attr(transformed, "output") + for output_index, output_item in enumerate(transformed_output if _is_json_array(transformed_output) else []): 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") @@ -1187,7 +1254,8 @@ def _build_synthetic_response_events( ) if item_type == "message": - for content_index, part in enumerate(output_item_payload.get("content", []) or []): + item_content = output_item_payload.get("content") + for content_index, part in enumerate(item_content if _is_json_array(item_content) else []): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( @@ -1234,7 +1302,8 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): + item_summary = output_item_payload.get("summary") + for summary_index, summary in enumerate(item_summary if _is_json_array(item_summary) else []): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1327,7 +1396,7 @@ class ResponsesWebSocketStreaming: user_api_key_dict: UserAPIKeyAuth | None = None, request_data: dict[str, object] | None = None, first_message: str | None = None, - guardrail_callbacks: list[Any] | None = None, + guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, authorized_model: str | None = None, ): @@ -1339,7 +1408,7 @@ class ResponsesWebSocketStreaming: self.messages: list[dict[str, object]] = [] self.input_messages: list[dict[str, object]] = [] self.first_message = first_message - self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] + self.guardrail_callbacks: list[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. @@ -1353,9 +1422,12 @@ class ResponsesWebSocketStreaming: event = event.decode("utf-8") if isinstance(event, str): try: - event_obj = json.loads(event) + parsed_event = _parse_json(event) except (json.JSONDecodeError, TypeError): return + if not _is_json_object(parsed_event): + return + event_obj = parsed_event else: event_obj = event @@ -1366,7 +1438,10 @@ class ResponsesWebSocketStreaming: """Extract user input content from response.create for logging.""" try: if isinstance(message, str): - msg_obj = json.loads(message) + parsed_message = _parse_json(message) + if not _is_json_object(parsed_message): + return + msg_obj = parsed_message elif _is_json_object(message): msg_obj = message else: @@ -1436,9 +1511,10 @@ class ResponsesWebSocketStreaming: # masked response.completed. if self.output_guardrail_callbacks: try: - _evt_type = json.loads(response_str).get("type") + _evt_parsed = _parse_json(response_str) except (json.JSONDecodeError, TypeError): - _evt_type = None + _evt_parsed = None + _evt_type = _evt_parsed.get("type") if _is_json_object(_evt_parsed) else None if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES: continue @@ -1500,10 +1576,14 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = json.loads(message) + parsed_message: Final = _parse_json(message) except (json.JSONDecodeError, TypeError): return message + if not _is_json_object(parsed_message): + return message + msg_obj: Final = parsed_message + if msg_obj.get("type") != "response.create": return message @@ -1628,11 +1708,18 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + parsed_event: Final = _parse_json(response_str) except (json.JSONDecodeError, TypeError): return response_str + if not _is_json_object(parsed_event): + return response_str + evt_obj: Final = parsed_event + cb: Final = self.guardrail_callbacks[0] + unmask_fn: Final[object] = getattr(cb, "_unmask_pii_text", None) + if not isinstance(unmask_fn, _PiiUnmasker): + return response_str event_type: Final = evt_obj.get("type") if event_type == "response.completed": @@ -1652,7 +1739,7 @@ class ResponsesWebSocketStreaming: continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text(text, pii_tokens) + unmasked = unmask_fn(text, pii_tokens) if unmasked != text: content_block["text"] = unmasked modified = True @@ -1661,7 +1748,7 @@ class ResponsesWebSocketStreaming: if event_type in self._DELTA_EVENT_TYPES: delta: Final = evt_obj.get("delta") if isinstance(delta, str): - unmasked = cb._unmask_pii_text(delta, pii_tokens) + unmasked = unmask_fn(delta, pii_tokens) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) @@ -1684,10 +1771,14 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + parsed_event: Final = _parse_json(response_str) except (json.JSONDecodeError, TypeError): return response_str + if not _is_json_object(parsed_event): + return response_str + evt_obj: Final = parsed_event + if evt_obj.get("type") != "response.completed": return response_str @@ -1832,7 +1923,7 @@ class ManagedResponsesWebSocketHandler: model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: dict[str, object] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, @@ -1844,10 +1935,11 @@ class ManagedResponsesWebSocketHandler: self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.litellm_metadata: dict[str, Any] = litellm_metadata or {} - self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( + self.litellm_metadata: dict[str, object] = litellm_metadata or {} + model_group_raw: Final = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( "deployment_model_name" ) + self.model_group: str | None = model_group_raw if isinstance(model_group_raw, str) else None self.api_key = api_key self.api_base = api_base self.timeout = timeout @@ -1867,14 +1959,14 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ @staticmethod - def _serialize_chunk(chunk: Any) -> str | None: + def _serialize_chunk(chunk: object) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: - if hasattr(chunk, "model_dump_json"): + if isinstance(chunk, _SupportsModelDumpJson): return chunk.model_dump_json(exclude_none=True) - if hasattr(chunk, "model_dump"): + if isinstance(chunk, _SupportsModelDump): return json.dumps(chunk.model_dump(exclude_none=True), default=str) - if isinstance(chunk, dict): + if _is_json_object(chunk): return json.dumps(chunk, default=str) return json.dumps(str(chunk)) except Exception as exc: @@ -1925,27 +2017,31 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, Any], + completed_event: dict[str, object], ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into Responses API message dicts suitable for the next turn's ``input``. """ resp_obj: Final = completed_event.get("response", {}) - if not isinstance(resp_obj, dict): + if not _is_json_object(resp_obj): return [] messages: Final[list[dict[str, object]]] = [] - for item in resp_obj.get("output", []) or []: - if not isinstance(item, dict): + output_items: Final = resp_obj.get("output") + for item in output_items if _is_json_array(output_items) else []: + if not _is_json_object(item): continue item_type = item.get("type") role = item.get("role", "assistant") if item_type == "message": - content_parts = item.get("content") or [] + content_raw = item.get("content") + content_parts = content_raw if _is_json_array(content_raw) else [] text_parts = [ - p.get("text", "") + part_text for p in content_parts - if isinstance(p, dict) and p.get("type") in ("output_text", "text") + if _is_json_object(p) + and p.get("type") in ("output_text", "text") + and isinstance(part_text := p.get("text", ""), str) ] text = "".join(text_parts) if text: @@ -1985,10 +2081,12 @@ 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 = _parse_json(raw_message) except json.JSONDecodeError: await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None + if not _is_json_object(msg_obj): + return None if msg_obj.get("type") != "response.create": # Silently ignore non-response.create messages (e.g. warmup pings) return None @@ -2076,7 +2174,7 @@ class ManagedResponsesWebSocketHandler: def _apply_history( self, - call_kwargs: dict[str, Any], + call_kwargs: dict[str, object], previous_response_id: str | None, current_messages: list[dict[str, object]], prior_history: list[dict[str, object]], @@ -2129,7 +2227,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None: + def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key @@ -2148,24 +2246,33 @@ class ManagedResponsesWebSocketHandler: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) @staticmethod - def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: + def _update_proxy_request(call_kwargs: dict[str, object], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" - proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get("proxy_server_request") or {} - if not isinstance(proxy_server_request, dict): + metadata_raw: Final = call_kwargs.get("litellm_metadata") + proxy_server_request_raw: Final = ( + metadata_raw.get("proxy_server_request") if _is_json_object(metadata_raw) else None + ) + if proxy_server_request_raw and not _is_json_object(proxy_server_request_raw): return - body: Final = dict(proxy_server_request.get("body") or {}) + proxy_server_request_base: Final[dict[str, object]] = ( + proxy_server_request_raw if _is_json_object(proxy_server_request_raw) else {} + ) + body_raw: Final = proxy_server_request_base.get("body") + body: Final[dict[str, object]] = dict(body_raw) if _is_json_object(body_raw) else {} body["input"] = call_kwargs.get("input") body["store"] = call_kwargs.get("store") body["model"] = model for k in ("tools", "tool_choice", "instructions", "metadata"): if k in call_kwargs and call_kwargs[k] is not None: body[k] = call_kwargs[k] - proxy_server_request = {**proxy_server_request, "body": body} - if "litellm_metadata" not in call_kwargs: - call_kwargs["litellm_metadata"] = {} - call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request - call_kwargs.setdefault("litellm_params", {}) - call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request + proxy_server_request: Final[dict[str, object]] = {**proxy_server_request_base, "body": body} + metadata: Final[dict[str, object]] = metadata_raw if _is_json_object(metadata_raw) else {} + metadata["proxy_server_request"] = proxy_server_request + call_kwargs["litellm_metadata"] = metadata + litellm_params_raw: Final = call_kwargs.setdefault("litellm_params", {}) + litellm_params: Final[dict[str, object]] = litellm_params_raw if _is_json_object(litellm_params_raw) else {} + litellm_params["proxy_server_request"] = proxy_server_request + call_kwargs["litellm_params"] = litellm_params async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None: """ @@ -2189,7 +2296,9 @@ class ManagedResponsesWebSocketHandler: continue if chunk_type == "response.completed" and completed_event is None: try: - completed_event = json.loads(serialized) + parsed_completed = _parse_json(serialized) + if _is_json_object(parsed_completed): + completed_event = parsed_completed except Exception: pass try: @@ -2266,14 +2375,16 @@ class ManagedResponsesWebSocketHandler: # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. - requested_model: Final = call_kwargs.pop("model", None) + requested_model_raw: Final = _pop_kwarg(call_kwargs, "model") + requested_model: Final = requested_model_raw if isinstance(requested_model_raw, str) else None if requested_model is None or requested_model == self.model_group: model = self.model else: model = requested_model - previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None) - current_messages: Final = self._input_to_messages(call_kwargs.get("input")) + previous_response_id_raw: Final = _pop_kwarg(call_kwargs, "previous_response_id") + previous_response_id: Final = previous_response_id_raw if isinstance(previous_response_id_raw, str) else None + current_messages: Final = self._input_to_messages(_mapping_get(call_kwargs, "input")) # Fetch history once; reused in both _apply_history and _save_turn_history prior_history: Final = self._get_history_messages(previous_response_id) if previous_response_id else [] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a8af4eabb3f..29627c06e52 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,18 +1,18 @@ { "ANN001": { - "limit": 3121 + "limit": 3111 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 834 + "limit": 826 }, "ANN201": { - "limit": 2032 + "limit": 2030 }, "ANN202": { - "limit": 865 + "limit": 863 }, "ANN204": { "limit": 713 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1630 + "limit": 1554 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 505 }, "B009": { - "limit": 81 + "limit": 79 }, "B010": { "limit": 190 @@ -228,7 +228,7 @@ "limit": 0 }, "RUF012": { - "limit": 241 + "limit": 239 }, "RUF015": { "limit": 8 @@ -264,7 +264,7 @@ "limit": 58 }, "SIM102": { - "limit": 322 + "limit": 316 }, "SIM103": { "limit": 119 @@ -306,7 +306,7 @@ "limit": 0 }, "TID251": { - "limit": 1240 + "limit": 1232 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0a0cfe9a617..f61dfcfe900 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23235 }, "LIT002": { - "limit": 27176 + "limit": 27163 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1091 + "limit": 1090 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16769 + "limit": 16762 }, "LIT011": { - "limit": 5598 + "limit": 5595 } }