From 0db81a60ad71bb1ed6b3341b36fee3d9926a9a14 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 3 Aug 2026 03:42:04 +0000 Subject: [PATCH] chore(typing): clear basedpyright Any errors in responses, proxy hooks, and management endpoints Replaces Any seams with real types across the eight production files carrying the highest combined reportAny + reportExplicitAny counts. Payload dicts become TypedDicts, record shapes become Protocols, and db/provider helpers get precise return types, so the values carry their real shapes instead of laundering them through Any. Typing only; no runtime behavior changes. No cast(), no type: ignore, and no pyright or noqa suppressions were added. reportAny drops 19234 to 18638 and reportExplicitAny 6518 to 6289, with every other rule at or below its baseline repo-wide. The three lint budgets are ratcheted down to match. --- basedpyright-code-budget.json | 26 +- .../streaming_chunk_builder_utils.py | 167 +++++----- .../adapters/handler.py | 110 +++---- litellm/models/model.py | 2 +- .../mcp_server/sampling_handler.py | 28 +- .../hooks/parallel_request_limiter_v3.py | 210 +++++++------ .../model_management_endpoints.py | 160 ++++++---- .../managed_id_rewriter.py | 286 +++++++++++------- .../transformation.py | 252 ++++++++------- litellm/responses/streaming_iterator.py | 236 ++++++++++----- .../streaming_chunk_builder_utils.py | 113 ++++++- .../passthrough_endpoints/managed_ids.py | 75 +++++ .../model_management_endpoints.py | 11 +- litellm/types/responses/main.py | 14 +- ruff-strict-budget.json | 10 +- type-discipline-budget.json | 8 +- 16 files changed, 1123 insertions(+), 585 deletions(-) create mode 100644 litellm/types/passthrough_endpoints/managed_ids.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index f6dd90077b1..6e5e818d4d6 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 29813 + "limit": 29217 }, "reportArgumentType": { - "limit": 2645 + "limit": 2632 }, "reportAssignmentType": { "limit": 329 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 9473 + "limit": 9244 }, "reportFunctionMemberAccess": { "limit": 11 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5855 + "limit": 5850 }, "reportMissingTypeArgument": { - "limit": 15852 + "limit": 15810 }, "reportMissingTypeStubs": { "limit": 41 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1079 + "limit": 1071 }, "reportOptionalOperand": { "limit": 0 @@ -90,7 +90,7 @@ "limit": 12 }, "reportReturnType": { - "limit": 219 + "limit": 218 }, "reportTypedDictNotRequiredAccess": { "limit": 27 @@ -99,22 +99,22 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45324 + "limit": 45296 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40452 + "limit": 40362 }, "reportUnknownParameterType": { - "limit": 20309 + "limit": 20268 }, "reportUnknownVariableType": { - "limit": 31978 + "limit": 31896 }, "reportUnnecessaryCast": { - "limit": 177 + "limit": 173 }, "reportUnnecessaryComparison": { "limit": 1021 @@ -123,7 +123,7 @@ "limit": 7 }, "reportUnnecessaryIsInstance": { - "limit": 1204 + "limit": 1203 }, "reportUntypedBaseClass": { "limit": 165 diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 2e62a151f98..8a63ff8418c 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -11,7 +11,6 @@ from litellm.types.utils import ( CacheCreationTokenDetails, ChatCompletionAudioResponse, ChatCompletionMessageToolCall, - Choices, CompletionTokensDetails, CompletionTokensDetailsWrapper, Function, @@ -25,7 +24,13 @@ 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 ( + StreamingChunkDict, + StreamingUsageDict, + ToolCallAccumulator, + ToolCallParams, + UsageChunkCalculation, UsagePerChunk, ) from litellm.types.llms.openai import ( @@ -34,6 +39,25 @@ if TYPE_CHECKING: ) +def _as_str_keyed_dict(value: object) -> dict[str, object] | None: + if isinstance(value, dict): + return value + return None + + +def _extract_hidden_params(chunk: object) -> dict[str, object]: + chunk_dict = _as_str_keyed_dict(chunk) + params: object + if chunk_dict is not None: + params = chunk_dict.get("_hidden_params", {}) + else: + params = getattr(chunk, "_hidden_params", {}) + params_dict = _as_str_keyed_dict(params) + if params_dict is not None: + return params_dict + return {} + + class ChunkProcessor: def __init__(self, chunks: list, messages: list | None = None): self.chunks = self._sort_chunks(chunks) @@ -44,33 +68,19 @@ class ChunkProcessor: if not chunks: return [] - first_chunk = 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 = _extract_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"))) - return float("inf") + def _created_at(chunk: object) -> int | float: + params = _extract_hidden_params(chunk) + return cast(int | float, params.get("created_at", 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: "StreamingChunkDict | None" = None ) -> ModelResponse: if chunk is None: return model_response @@ -82,19 +92,20 @@ class ChunkProcessor: @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 = getattr(response, "model", None) + model = response.model 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") + model_call_details: dict[str, object] = logging_obj.model_call_details + custom_llm_provider = model_call_details.get("custom_llm_provider") try: from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -126,7 +137,7 @@ class ChunkProcessor: ) @staticmethod - def _get_chunk_id(chunks: list[dict[str, Any]]) -> str: + def _get_chunk_id(chunks: list["StreamingChunkDict"]) -> str: """ Chunks: [{"id": ""}, {"id": "1"}, {"id": "1"}] @@ -137,7 +148,7 @@ class ChunkProcessor: return "" @staticmethod - def _get_model_from_chunks(chunks: list[dict[str, Any]], first_chunk_model: str) -> str: + def _get_model_from_chunks(chunks: list["StreamingChunkDict"], first_chunk_model: str) -> str: """ Get the actual model from chunks, preferring a model that differs from the first chunk. @@ -153,7 +164,7 @@ class ChunkProcessor: # 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: list["StreamingChunkDict"]) -> ModelResponse: chunk = self.first_chunk id = ChunkProcessor._get_chunk_id(chunks) object = chunk["object"] @@ -168,11 +179,12 @@ class ChunkProcessor: 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"] + first_choice = chunk["choices"][0] + chunk_finish_reason: str | None = None + if hasattr(first_choice, "finish_reason"): + chunk_finish_reason = first_choice.finish_reason + elif "finish_reason" in first_choice: + chunk_finish_reason = first_choice["finish_reason"] if chunk_finish_reason is not None: finish_reason = chunk_finish_reason @@ -202,9 +214,11 @@ class ChunkProcessor: response = self.update_model_response_with_hidden_params(model_response=response, chunk=chunk) return response - def get_combined_tool_content(self, tool_call_chunks: list[dict[str, Any]]) -> list[ChatCompletionMessageToolCall]: + def get_combined_tool_content( + self, tool_call_chunks: list["StreamingChunkDict"] + ) -> list[ChatCompletionMessageToolCall]: tool_calls_list: list[ChatCompletionMessageToolCall] = [] - tool_call_map: dict[int, dict[str, Any]] = {} # Map to store tool calls by index + tool_call_map: dict[int, ToolCallAccumulator] = {} # Map to store tool calls by index for chunk in tool_call_chunks: choices = chunk["choices"] @@ -244,18 +258,22 @@ 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"] + dict_tool_call_id = tool_call.get("id") + if dict_tool_call_id: + tool_call_map[index]["id"] = dict_tool_call_id + dict_tool_call_type = tool_call.get("type") + if dict_tool_call_type: + tool_call_map[index]["type"] = dict_tool_call_type function = tool_call.get("function", {}) if isinstance(function, dict): - if function.get("name"): - tool_call_map[index]["name"] = function["name"] - if function.get("arguments"): - tool_call_map[index]["arguments"].append(function["arguments"]) - else: + dict_function_name = function.get("name") + if dict_function_name: + tool_call_map[index]["name"] = dict_function_name + dict_function_arguments = function.get("arguments") + if dict_function_arguments: + tool_call_map[index]["arguments"].append(dict_function_arguments) + elif function is not None: # function is an object if hasattr(function, "name") and function.name: tool_call_map[index]["name"] = function.name @@ -267,7 +285,7 @@ class ChunkProcessor: 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") and tool_call.function is not None: if hasattr(tool_call.function, "name") and tool_call.function.name: tool_call_map[index]["name"] = tool_call.function.name if hasattr(tool_call.function, "arguments") and tool_call.function.arguments: @@ -277,13 +295,15 @@ class ChunkProcessor: 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") + dict_tool_call_function = tool_call.get("function") + if not provider_fields and isinstance(dict_tool_call_function, dict): + provider_fields = dict_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 tool_call.function is not None and hasattr(tool_call.function, "provider_specific_fields") and tool_call.function.provider_specific_fields ): @@ -291,10 +311,12 @@ class ChunkProcessor: 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"] = {} + merged_provider_fields = tool_call_map[index]["provider_specific_fields"] + if merged_provider_fields is None: + merged_provider_fields = {} + tool_call_map[index]["provider_specific_fields"] = merged_provider_fields if isinstance(provider_fields, dict): - tool_call_map[index]["provider_specific_fields"].update(provider_fields) + merged_provider_fields.update(provider_fields) # Convert the map to a list of tool calls for index in sorted(tool_call_map.keys()): @@ -309,15 +331,16 @@ class ChunkProcessor: ) # Prepare params for ChatCompletionMessageToolCall - tool_call_params = { + tool_call_params: ToolCallParams = { "id": tool_call_data["id"], "function": function, "type": tool_call_data["type"] or "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"] + accumulated_provider_fields = tool_call_data["provider_specific_fields"] + if accumulated_provider_fields: + tool_call_params["provider_specific_fields"] = accumulated_provider_fields tool_call = ChatCompletionMessageToolCall(**tool_call_params) tool_calls_list.append(tool_call) @@ -350,7 +373,7 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: list["StreamingChunkDict"], delta_key: str = "content" ) -> ChatCompletionAssistantContentValue: content_list: list[str] = [] for chunk in chunks: @@ -369,7 +392,7 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] + self, chunks: list["StreamingChunkDict"] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -426,10 +449,10 @@ 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: list["StreamingChunkDict"]) -> 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: list["StreamingChunkDict"]) -> ChatCompletionAudioResponse: base64_data_list: list[str] = [] transcript_list: list[str] = [] expires_at: int | None = None @@ -459,9 +482,9 @@ class ChunkProcessor: id=id, ) - def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict: - prompt_tokens = 0 - completion_tokens = 0 + def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "UsageChunkCalculation": + prompt_tokens: int = 0 + completion_tokens: int = 0 ## anthropic prompt caching information ## cache_creation_input_tokens: int | None = None cache_read_input_tokens: int | None = None @@ -503,22 +526,19 @@ 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 - ): + if hasattr(choice.message, "reasoning_content") and choice.message.reasoning_content is not None: if reasoning_tokens is None: reasoning_tokens = 0 reasoning_tokens += token_counter( - text=cast(Choices, choice).message.reasoning_content, + text=choice.message.reasoning_content, count_response_tokens=True, ) return reasoning_tokens @staticmethod - def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None: - usage_chunk: Usage | dict[str, Any] | None = None + def _extract_usage_chunk(chunk: "StreamingChunkDict | ModelResponse | ModelResponseStream") -> Usage | None: + usage_chunk: Usage | StreamingUsageDict | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage elif "usage" in chunk: @@ -534,7 +554,7 @@ class ChunkProcessor: def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: list["StreamingChunkDict | ModelResponse"], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -615,10 +635,7 @@ class ChunkProcessor: "web_search_requests", ) - prompt_tokens_details = cast( - PromptTokensDetailsWrapper | None, - usage_chunk_dict["prompt_tokens_details"], - ) + prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] cache_creation_token_details = self._capture_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details @@ -679,7 +696,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: list["StreamingChunkDict | ModelResponse"], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -702,7 +719,7 @@ 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 = chunks[0] if isinstance(first_chunk, dict): @@ -718,7 +735,7 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: list["StreamingChunkDict | ModelResponse"], model: str, completion_output: str, messages: list | None = None, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a5aa1509969..6557beb1c62 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,6 +1,7 @@ from collections.abc import AsyncIterator, Coroutine, Iterator from typing import ( Any, + TypeAlias, cast, ) @@ -27,8 +28,11 @@ from litellm.utils import get_model_info # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. ANTHROPIC_ONLY_REQUEST_KEYS: frozenset[str] = frozenset({"output_config"}) +ContextManagementSpec: TypeAlias = dict[str, object] | list[dict[str, object]] | None +AnthropicSystemPrompt: TypeAlias = str | list[dict[str, object]] | None -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") @@ -40,7 +44,7 @@ def _messages_have_compaction_block(messages: list[dict]) -> bool: return False -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``, @@ -61,14 +65,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, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: AnthropicSystemPrompt, + context_management_spec: ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, - llm_router: Any, - user_api_key_auth: Any = None, + 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 ( @@ -88,11 +92,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: list[dict[str, object]] = messages + working_system: AnthropicSystemPrompt = 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 @@ -122,7 +126,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 @@ -130,7 +134,7 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, - context_management_spec: Any, + context_management_spec: ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. @@ -157,7 +161,7 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, - context_management_spec: Any, + context_management_spec: ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -198,9 +202,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N def _normalize_spec_edits( *, - context_management_spec: Any, + context_management_spec: ContextManagementSpec, additional_drop_params: list[str] | None, -) -> list[dict[str, Any]] | None: +) -> 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`` @@ -225,14 +229,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, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: AnthropicSystemPrompt, + context_management_spec: ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, - llm_router: Any, - user_api_key_auth: Any = None, + llm_router: object, + user_api_key_auth: object = None, ) -> PolyfillResult | None: """Run the async context_management polyfill if a spec is present. @@ -293,9 +297,9 @@ ANTHROPIC_ADAPTER = 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 @@ -318,7 +322,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: except Exception: custom_llm_provider = None - if custom_llm_provider != "openai": + if not isinstance(custom_llm_provider, str) or custom_llm_provider != "openai": return if not isinstance(thinking, dict) or thinking.get("type") != "enabled": @@ -342,7 +346,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: reasoning_effort = completion_kwargs.get("reasoning_effort") summary = thinking.get("summary") if isinstance(reasoning_effort, str) and reasoning_effort: - reasoning_dict: dict[str, Any] = {"effort": reasoning_effort} + reasoning_dict: dict[str, object] = {"effort": reasoning_effort} if summary: reasoning_dict["summary"] = summary elif auto_summary: @@ -396,20 +400,20 @@ class LiteLLMMessagesToCompletionTransformationHandler: 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: AnthropicSystemPrompt = 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, + output_format: dict[str, object] | None = None, + extra_kwargs: dict[str, object] | None = None, ) -> tuple[dict[str, Any], dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. @@ -422,7 +426,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data = { + request_data: dict[str, object] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -467,7 +471,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if openai_request is None: raise ValueError("Failed to translate request to OpenAI format") - completion_kwargs: dict[str, Any] = dict(openai_request) + completion_kwargs: dict[str, object] = dict(openai_request) if stream: completion_kwargs["stream"] = stream @@ -517,21 +521,21 @@ 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, + output_format: dict[str, object] | None = None, **kwargs, - ) -> AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]: + ) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]: """Handle non-Anthropic models asynchronously using the adapter""" context_management = kwargs.pop("context_management", None) additional_drop_params: list[str] | None = kwargs.get("additional_drop_params", None) @@ -611,26 +615,26 @@ 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, ) -> ( AnthropicMessagesResponse | Iterator[bytes] - | AsyncIterator[Any] - | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]] + | AsyncIterator[bytes] + | Coroutine[None, None, AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]] ): """Handle non-Anthropic models using the adapter.""" if _is_async is True: diff --git a/litellm/models/model.py b/litellm/models/model.py index 209f26d4837..d94e42d6c5f 100644 --- a/litellm/models/model.py +++ b/litellm/models/model.py @@ -16,7 +16,7 @@ from litellm.types.llms.base import LiteLLMPydanticObjectBase class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): model_id: str model_name: str - litellm_params: dict + litellm_params: dict[str, object] model_info: dict | None = None blocked: bool = False created_at: datetime | None = None diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index 779cc5861d4..78da1a17aa1 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -295,8 +295,12 @@ def _convert_mcp_content_to_openai( return _convert_single_content(content) +def _required_attr(obj: object, name: str) -> object: + return getattr(obj, name) + + def _convert_single_content( - content: Any, + content: object, ) -> "dict[str, object] | list[dict[str, object]]": """Convert a single MCP content item to OpenAI format. @@ -310,7 +314,7 @@ def _convert_single_content( content_type = getattr(content, "type", None) if content_type == "text": - return {"type": "text", "text": content.text} + return {"type": "text", "text": _required_attr(content, "text")} elif content_type == "image": data = getattr(content, "data", "") mime_type = getattr(content, "mimeType", "image/png") @@ -581,12 +585,28 @@ def _convert_mcp_tool_choice_to_openai( return "auto" +class _SamplingToolCallFunction(Protocol): + @property + def name(self) -> str: ... + + @property + def arguments(self) -> str | dict[str, object]: ... + + +class _SamplingToolCall(Protocol): + @property + def id(self) -> str: ... + + @property + def function(self) -> _SamplingToolCallFunction: ... + + class _SamplingResponseMessage(Protocol): @property def content(self) -> str | None: ... @property - def tool_calls(self) -> Sequence[object] | None: ... + def tool_calls(self) -> Sequence[_SamplingToolCall] | None: ... class _SamplingResponseChoice(Protocol): @@ -641,7 +661,7 @@ def _convert_openai_response_to_mcp_result( stop_reason = "endTurn" actual_model: str = getattr(response, "model", model_name) or model_name # Check if response has tool calls - tool_calls = getattr(message, "tool_calls", None) + tool_calls: Sequence[_SamplingToolCall] | None = getattr(message, "tool_calls", None) if tool_calls: # Build ToolUseContent items content_parts: list[SamplingMessageContentBlock] = [] diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7253a684b3c..6c4f2be8ff9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -8,7 +8,7 @@ import asyncio import binascii import os import uuid -from collections.abc import Callable +from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime @@ -16,9 +16,10 @@ from typing import ( TYPE_CHECKING, Any, Literal, + Protocol, + TypeAlias, TypedDict, Union, - cast, ) from litellm import DualCache @@ -44,6 +45,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject from litellm.types.utils import ( CallTypes, + CallTypesLiteral, EmbeddingResponse, ModelResponse, TextCompletionResponse, @@ -54,6 +56,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + from litellm.types.agents import AgentResponse from litellm.types.caching import RedisPipelineIncrementOperation Span = Union[_Span, Any] @@ -300,6 +303,15 @@ _TPM_FLOOR_FRACTION = 4 PARALLEL_REQUEST_SLOT_TTL_SECONDS = 3600 +LuaReplyValue: TypeAlias = int | str | bytes +CachedCounterValue: TypeAlias = int | str | bytes | None +CachedGaugeValue: TypeAlias = dict[str, object] | int | str | bytes | None + + +class AsyncLuaScript(Protocol): + def __call__(self, *, keys: Sequence[str], args: Sequence[float | str]) -> Awaitable[list[LuaReplyValue]]: ... + + class RateLimitDescriptorRateLimitObject(TypedDict, total=False): requests_per_unit: int | None tokens_per_unit: int | None @@ -342,6 +354,29 @@ class RateLimitResponseWithDescriptors(TypedDict): response: RateLimitResponse +class WindowedKeyMetadata(TypedDict): + requests_limit: int | None + tokens_limit: int | None + window_size: int + descriptor_key: str + + +class AtomicCounterMeta(TypedDict): + descriptor_key: str + current_limit: int + rate_limit_type: Literal["requests", "tokens"] + window_key: str + counter_key: str + increment: int + ttl: int + window_size: int + + +class AtomicCounterState(TypedDict): + window_expired: bool + current: int + + @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -423,25 +458,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ): self.internal_usage_cache = internal_usage_cache self._time_provider = time_provider or datetime.now - if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + redis_cache = self.internal_usage_cache.dual_cache.redis_cache + if redis_cache is not None: + self.batch_rate_limiter_script: AsyncLuaScript | None = redis_cache.async_register_script( BATCH_RATE_LIMITER_SCRIPT ) - self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + self.token_increment_script: AsyncLuaScript | None = redis_cache.async_register_script( TOKEN_INCREMENT_SCRIPT ) - self.check_and_increment_by_n_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT) + self.check_and_increment_by_n_script: AsyncLuaScript | None = redis_cache.async_register_script( + CHECK_AND_INCREMENT_BY_N_SCRIPT ) - self.parallel_acquire_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + self.parallel_acquire_script: AsyncLuaScript | None = redis_cache.async_register_script( PARALLEL_ACQUIRE_SCRIPT ) - self.parallel_release_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + self.parallel_release_script: AsyncLuaScript | None = redis_cache.async_register_script( PARALLEL_RELEASE_SCRIPT ) - self.parallel_count_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - PARALLEL_COUNT_SCRIPT - ) + self.parallel_count_script: AsyncLuaScript | None = redis_cache.async_register_script(PARALLEL_COUNT_SCRIPT) else: self.batch_rate_limiter_script = None self.token_increment_script = None @@ -459,7 +493,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.tpm_reservation_enabled = os.getenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "true").lower() == "true" # Batch rate limiter (lazy loaded) - self._batch_rate_limiter: Any | None = None + self._batch_rate_limiter: CustomLogger | None = None # Serializes multi-phase check+increment sequences (batch + dynamic # limiters) within this process to close the TOCTOU window between @@ -477,7 +511,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # one round-trip. self._check_and_increment_lock = asyncio.Lock() - def _get_batch_rate_limiter(self) -> Any | None: + def _get_batch_rate_limiter(self) -> CustomLogger | None: """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: try: @@ -604,12 +638,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): keys: list[str], now_int: int, window_size: int, - ) -> list[Any]: + ) -> list[CachedCounterValue]: """ Implement sliding window rate limiting logic using in-memory cache operations. This follows the same logic as the Redis Lua script but uses async cache operations. """ - results: list[Any] = [] + results: list[CachedCounterValue] = [] # Process each window/counter pair for i in range(0, len(keys), 2): @@ -618,7 +652,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): increment_value = 1 # Get the window start time - window_start = await self.internal_usage_cache.async_get_cache( + window_start: CachedCounterValue = await self.internal_usage_cache.async_get_cache( key=window_key, litellm_parent_otel_span=None, local_only=True, @@ -645,7 +679,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): results.append(increment_value) # counter else: # Increment the counter - current_counter = await self.internal_usage_cache.async_get_cache( + current_counter: CachedCounterValue = await self.internal_usage_cache.async_get_cache( key=counter_key, litellm_parent_otel_span=None, local_only=True, @@ -679,8 +713,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def is_cache_list_over_limit( self, keys_to_fetch: list[str], - cache_values: list[Any], - key_metadata: dict[str, Any], + cache_values: Sequence[CachedCounterValue], + key_metadata: dict[str, WindowedKeyMetadata], ) -> RateLimitResponse: """ Check if the cache values are over the limit. @@ -783,7 +817,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, keys_to_fetch: list[str], now_int: int, - ) -> list[Any]: + ) -> list[CachedCounterValue]: """ Execute Redis operations grouped by hash tag for cluster compatibility. @@ -792,13 +826,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): now_int: int - Current timestamp Returns: - List[Any] - List of cache values + list[CachedCounterValue] - List of cache values """ if self.batch_rate_limiter_script is None: return [] key_groups = self._group_keys_by_hash_tag(keys_to_fetch) - all_cache_values = [] + all_cache_values: list[CachedCounterValue] = [] for hash_tag, group_keys in key_groups.items(): try: @@ -866,7 +900,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): windowed_response = RateLimitResponse(overall_code="OK", statuses=[]) if keys_to_fetch: ## CHECK IN-MEMORY CACHE - cache_values = await self.internal_usage_cache.async_batch_get_cache( + cache_values: Sequence[CachedCounterValue] | None = await self.internal_usage_cache.async_batch_get_cache( keys=keys_to_fetch, parent_otel_span=parent_otel_span, local_only=True, @@ -888,9 +922,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # For keys that don't exist yet, set them to 0 if cache_values is None: - cache_values = [] - for _ in keys_to_fetch: - cache_values.append(str(now_int) if _.endswith(":window") else 0) + cache_values = [str(now_int) if key.endswith(":window") else 0 for key in keys_to_fetch] elif self.batch_rate_limiter_script is not None: # NORMAL MODE: Increment counters in Redis # Group keys by hash tag for Redis cluster compatibility @@ -949,14 +981,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, descriptors: list[RateLimitDescriptor], skip_tpm_check: bool, - ) -> tuple[list[str], dict[str, dict[str, Any]], list[ParallelRequestGauge]]: + ) -> tuple[list[str], dict[str, WindowedKeyMetadata], list[ParallelRequestGauge]]: """ Split descriptors into the windowed (window_key, counter_key) fetch list with its per-window metadata, and the concurrency gauges for descriptors carrying a max_parallel_requests limit. """ keys_to_fetch: list[str] = [] - key_metadata: dict[str, dict[str, Any]] = {} + key_metadata: dict[str, WindowedKeyMetadata] = {} gauges: list[ParallelRequestGauge] = [] for descriptor in descriptors: descriptor_key = descriptor["key"] @@ -995,12 +1027,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not rate_limit_set: continue - key_metadata[window_key] = { - "requests_limit": (int(requests_limit) if requests_limit is not None else None), - "tokens_limit": int(tokens_limit) if tokens_limit is not None else None, - "window_size": int(window_size), - "descriptor_key": descriptor_key, - } + key_metadata[window_key] = WindowedKeyMetadata( + requests_limit=(int(requests_limit) if requests_limit is not None else None), + tokens_limit=int(tokens_limit) if tokens_limit is not None else None, + window_size=int(window_size), + descriptor_key=descriptor_key, + ) return keys_to_fetch, key_metadata, gauges def _gauge_status(self, gauge: ParallelRequestGauge, in_flight: int, code: str) -> RateLimitStatus: @@ -1012,7 +1044,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptor_key=gauge["descriptor_key"], ) - def _gauge_in_flight_from_cache_value(self, raw_value: Any) -> int: + def _gauge_in_flight_from_cache_value(self, raw_value: CachedGaugeValue) -> int: """ In-flight count from a cached gauge value: a dict of slot_id -> acquire timestamp when the in-memory registry is authoritative, or @@ -1114,7 +1146,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): gauge_keys: list[str], parent_otel_span: Span | None = None, ) -> list[int]: - values = await self.internal_usage_cache.async_batch_get_cache( + values: Sequence[CachedGaugeValue] | None = await self.internal_usage_cache.async_batch_get_cache( keys=gauge_keys, parent_otel_span=parent_otel_span, local_only=True, @@ -1143,7 +1175,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): cutoff = now - PARALLEL_REQUEST_SLOT_TTL_SECONDS states: list[tuple[dict[str, float] | None, int]] = [] for gauge in gauges: - raw_value = await self.internal_usage_cache.async_get_cache( + raw_value: CachedGaugeValue = await self.internal_usage_cache.async_get_cache( key=gauge["counter_key"], litellm_parent_otel_span=parent_otel_span, local_only=True, @@ -1218,7 +1250,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async with self._check_and_increment_lock: for counter_key in counter_keys: - raw_value = await self.internal_usage_cache.async_get_cache( + raw_value: CachedGaugeValue = await self.internal_usage_cache.async_get_cache( key=counter_key, litellm_parent_otel_span=parent_otel_span, local_only=True, @@ -1226,7 +1258,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if isinstance(raw_value, dict): if slot_id not in raw_value: continue - new_value: dict[str, float] | int = {key: ts for key, ts in raw_value.items() if key != slot_id} + new_value: dict[str, object] | int = {key: ts for key, ts in raw_value.items() if key != slot_id} elif raw_value is None: continue else: @@ -1277,7 +1309,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Build per-descriptor (keys, args, meta) groups. All keys within a # group share the descriptor's {key:value} hash tag, so a single Lua # call per group never triggers CROSSSLOT on Redis Cluster. - descriptor_groups: list[tuple[list[str], list[Any], list[dict[str, Any]]]] = [] + descriptor_groups: list[tuple[list[str], list[int], list[AtomicCounterMeta]]] = [] for descriptor, increment_amounts in zip(descriptors, increments): keys, args, meta = self._build_descriptor_atomic_payload( descriptor=descriptor, @@ -1300,7 +1332,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=parent_otel_span, ) - flat_meta: list[dict[str, Any]] = [m for _keys, _args, group_meta in descriptor_groups for m in group_meta] + flat_meta: list[AtomicCounterMeta] = [m for _keys, _args, group_meta in descriptor_groups for m in group_meta] async with self._check_and_increment_lock: return await self._atomic_check_and_increment_in_memory( per_counter_meta=flat_meta, @@ -1311,7 +1343,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, descriptor: RateLimitDescriptor, increment_amounts: dict[Literal["requests", "tokens"], int], - ) -> tuple[list[str], list[Any], list[dict[str, Any]]]: + ) -> tuple[list[str], list[int], list[AtomicCounterMeta]]: """ Build (KEYS, ARGV, per-counter meta) for a single descriptor's Lua call. All keys returned share the descriptor's {key:value} hash tag. @@ -1325,11 +1357,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): window_key = f"{{{descriptor_key}:{descriptor_value}}}:window" keys: list[str] = [] - args: list[Any] = [] - meta: list[dict[str, Any]] = [] + args: list[int] = [] + meta: list[AtomicCounterMeta] = [] - for rate_limit_type in ("requests", "tokens"): - rlt: Literal["requests", "tokens"] = cast(Literal["requests", "tokens"], rate_limit_type) + rate_limit_types: tuple[Literal["requests", "tokens"], ...] = ("requests", "tokens") + for rlt in rate_limit_types: if rlt == "requests": limit_value = rate_limit.get("requests_per_unit") inc_amount = int(increment_amounts.get("requests", 0) or 0) @@ -1350,22 +1382,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # [limit, increment, ttl_seconds, window_size_seconds]. args.extend([int(limit_value), inc_amount, ttl_seconds, window_size_seconds]) meta.append( - { - "descriptor_key": descriptor_key, - "current_limit": int(limit_value), - "rate_limit_type": rlt, - "window_key": window_key, - "counter_key": counter_key, - "increment": inc_amount, - "ttl": ttl_seconds, - "window_size": window_size_seconds, - } + AtomicCounterMeta( + descriptor_key=descriptor_key, + current_limit=int(limit_value), + rate_limit_type=rlt, + window_key=window_key, + counter_key=counter_key, + increment=inc_amount, + ttl=ttl_seconds, + window_size=window_size_seconds, + ) ) return keys, args, meta async def _atomic_lua_per_descriptor( self, - descriptor_groups: list[tuple[list[str], list[Any], list[dict[str, Any]]]], + descriptor_groups: list[tuple[list[str], list[int], list[AtomicCounterMeta]]], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """ @@ -1374,7 +1406,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptor i, refund descriptors 0..i-1's increments. On Lua failure mid-loop, refund applied increments and fall back to in-memory. """ - applied: list[list[dict[str, Any]]] = [] + applied: list[list[AtomicCounterMeta]] = [] statuses: list[RateLimitStatus] = [] for _idx, (keys, args, meta) in enumerate(descriptor_groups): @@ -1397,7 +1429,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): f"{self.window_size}s)." ) await self._refund_applied_descriptor_groups(applied) - flat_meta: list[dict[str, Any]] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta] + flat_meta: list[AtomicCounterMeta] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta] async with self._check_and_increment_lock: return await self._atomic_check_and_increment_in_memory( per_counter_meta=flat_meta, @@ -1415,7 +1447,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _refund_applied_descriptor_groups( self, - applied: list[list[dict[str, Any]]], + applied: list[list[AtomicCounterMeta]], ) -> None: """ Decrement counters for descriptor groups already applied via Lua. @@ -1441,8 +1473,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_atomic_response( self, - raw: list[Any], - per_counter_meta: list[dict[str, Any]], + raw: list[LuaReplyValue], + per_counter_meta: list[AtomicCounterMeta], ) -> RateLimitResponse: """Convert Lua script return value to RateLimitResponse. @@ -1493,7 +1525,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _atomic_check_and_increment_in_memory( self, - per_counter_meta: list[dict[str, Any]], + per_counter_meta: list[AtomicCounterMeta], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """In-memory all-or-nothing check-and-increment. Caller holds lock. @@ -1508,10 +1540,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): now_int = int(self._get_current_time().timestamp()) # Pass 1: read state, validate. - descriptor_state: list[dict[str, Any]] = [] + descriptor_state: list[AtomicCounterState] = [] for meta in per_counter_meta: window_size = meta["window_size"] - window_start = await self.internal_usage_cache.async_get_cache( + window_start: CachedCounterValue = await self.internal_usage_cache.async_get_cache( key=meta["window_key"], litellm_parent_otel_span=parent_otel_span, local_only=True, @@ -1547,7 +1579,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ], ) - descriptor_state.append({"window_expired": window_expired, "current": current_counter}) + descriptor_state.append(AtomicCounterState(window_expired=window_expired, current=current_counter)) # Pass 2: apply increments. statuses: list[RateLimitStatus] = [] @@ -1920,7 +1952,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ return rpm_limit_type == "dynamic" or tpm_limit_type == "dynamic" - def _get_agent_from_registry(self, agent_id: str) -> Any | None: + def _get_agent_from_registry(self, agent_id: str) -> "AgentResponse | None": """Look up an agent from the in-memory registry by ID.""" from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry @@ -1969,8 +2001,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if agent is None: return descriptors - agent_rpm = getattr(agent, "rpm_limit", None) - agent_tpm = getattr(agent, "tpm_limit", None) + agent_rpm = agent.rpm_limit + agent_tpm = agent.tpm_limit if agent_rpm is not None or agent_tpm is not None: descriptors.append( RateLimitDescriptor( @@ -1984,8 +2016,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - session_rpm = getattr(agent, "session_rpm_limit", None) - session_tpm = getattr(agent, "session_tpm_limit", None) + session_rpm = agent.session_rpm_limit + session_tpm = agent.session_tpm_limit if session_rpm is not None or session_tpm is not None: session_id = self._get_session_id_from_data(data) if session_id is not None: @@ -2244,7 +2276,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Fail safe: enforce limits if we can't check return True - def get_rate_limiter_for_call_type(self, call_type: str) -> Any | None: + def get_rate_limiter_for_call_type(self, call_type: str) -> CustomLogger | None: """Get the rate limiter for the call type.""" if call_type == "acreate_batch": batch_limiter = self._get_batch_rate_limiter() @@ -2372,8 +2404,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, - call_type: str, - ): + call_type: CallTypesLiteral, + ) -> Exception | str | dict[str, object] | None: """ Pre-call hook to check rate limits before making the API call. Supports dynamic rate limiting based on deployment health. @@ -2698,8 +2730,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get operations for this hash tag group group_operations = [op for op in pipeline_operations if op["key"] in group_keys] - keys = [] - args = [] + keys: list[str] = [] + args: list[float] = [] for op in group_operations: # Convert None TTL to 0 for Lua script @@ -2767,15 +2799,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @staticmethod def _merge_ratelimit_statuses_into_additional_headers( - additional_headers: dict[str, Any], + additional_headers: dict[str, object], statuses: list[RateLimitStatus], - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Return ``additional_headers`` extended with ``x-ratelimit-{descriptor_key}-{remaining|limit}-{rate_limit_type}`` entries. Non-mutating so callers pick their own target dict. """ - merged: dict[str, Any] = dict(additional_headers) + merged: dict[str, object] = dict(additional_headers) for status in statuses: prefix = f"x-ratelimit-{status['descriptor_key']}" merged[f"{prefix}-remaining-{status['rate_limit_type']}"] = status["limit_remaining"] @@ -2785,7 +2817,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _collect_tpm_scope_targets( self, standard_logging_metadata: dict[str, Any], - kwargs: Any, + kwargs: object, model_group: str | None, ) -> list[tuple[str, str]]: """ @@ -2882,8 +2914,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_success_event_pipeline_operations( self, - kwargs: Any, - response_obj: Any, + kwargs: dict[str, object], + response_obj: object, rate_limit_type: Literal["output", "input", "total"], ) -> list[RedisPipelineIncrementOperation]: """Build Redis pipeline increment ops for TPM / parallel-request counters.""" @@ -2894,7 +2926,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get metadata from standard_logging_object - this correctly handles both # 'metadata' and 'litellm_metadata' fields from litellm_params standard_logging_object = kwargs.get("standard_logging_object") or {} - standard_logging_metadata = standard_logging_object.get("metadata") or {} + standard_logging_metadata = ( + standard_logging_object.get("metadata") or {} if isinstance(standard_logging_object, dict) else {} + ) model_group = get_model_group_from_litellm_kwargs(kwargs) @@ -3008,9 +3042,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def async_logging_hook( self, kwargs: dict, - result: Any, + result: object, call_type: str, - ) -> tuple[dict, Any]: + ) -> tuple[dict, object]: """ Mirror the pre-call rate-limit snapshot into the SLP so streaming success callbacks see the same ``x-ratelimit-*`` headers the @@ -3027,8 +3061,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _mirror_ratelimit_response_into_logging_payload( self, - kwargs: Any, - response_obj: Any, + kwargs: object, + response_obj: object, ) -> None: """ Copy the stashed ``RateLimitResponse`` into the SLP's diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 35b64963d4d..03c72da9abb 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,8 +13,8 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Mapping, Sequence -from typing import Any, Literal, cast +from collections.abc import Awaitable, Mapping, Sequence +from typing import Literal, Protocol, TypeVar, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field @@ -55,6 +55,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.base_repository import BaseRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository @@ -64,6 +65,7 @@ from litellm.router_utils.auto_router_model_naming import ( validate_strategy_router_model_write, ) from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + PrismaCompatibleCreateDBModel, UpdateUsefulLinksRequest, ) from litellm.types.router import ( @@ -78,6 +80,67 @@ from litellm.utils import get_utc_datetime router = APIRouter() +_PrismaRowT = TypeVar("_PrismaRowT") +_RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel) + + +class _PrismaTableActions(Protocol[_PrismaRowT]): + async def find_unique( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaRowT | None: ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + include: Mapping[str, object] | None = None, + ) -> list[_PrismaRowT]: ... + + async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ... + + async def update( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _PrismaRowT: ... + + async def delete(self, *, where: Mapping[str, object]) -> _PrismaRowT | None: ... + + async def delete_many(self, *, where: Mapping[str, object]) -> int: ... + + +class _TeamModelAliasTeam(Protocol): + team_id: str + + +class _TeamModelAliasRow(Protocol): + id: int + model_aliases: dict[str, str] + team: _TeamModelAliasTeam | None + + +def _prisma_table( + repository: BaseRepository[_RepositoryModelT], +) -> _PrismaTableActions[_RepositoryModelT]: + return repository.table # any-ok: repository.table is the untyped Prisma client wrapper + + +def _model_alias_table(prisma_client: PrismaClient) -> _PrismaTableActions[_TeamModelAliasRow]: + return ModelTableRepository(prisma_client).table # any-ok: untyped Prisma client wrapper + + +def _transaction_model_table(transaction: object) -> _PrismaTableActions[LiteLLM_ProxyModelTable]: + return getattr(transaction, "litellm_proxymodeltable") # any-ok: untyped Prisma transaction client + + +def _audit_log_json(row: object) -> str | None: + return row.model_dump_json(exclude_none=True) if isinstance(row, BaseModel) else None + async def update_team(*args, **kwargs): """ @@ -96,10 +159,7 @@ class UpdatePublicModelGroupsRequest(BaseModel): async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None: - db_model = cast( - BaseModel | None, - await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}), - ) + db_model = await _prisma_table(ModelRepository(prisma_client)).find_unique(where={"model_id": model_id}) if not db_model: return None @@ -322,7 +382,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model = await ModelRepository(prisma_client).table.update( + updated_model = await _prisma_table(ModelRepository(prisma_client)).update( where={"model_id": model_id}, data=update_data, ) @@ -425,7 +485,7 @@ async def _set_model_blocked_status( param=None, ) - updated_model = await ModelRepository(prisma_client).table.update( + updated_model = await _prisma_table(ModelRepository(prisma_client)).update( where={"model_id": data.model_id}, data={ "blocked": blocked, @@ -444,9 +504,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=_audit_log_json(updated_model), litellm_changed_by=litellm_changed_by, litellm_proxy_admin_name=litellm_proxy_admin_name, ) @@ -552,24 +610,20 @@ async def _add_model_to_db( 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: dict = { + _data: PrismaCompatibleCreateDBModel = { "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), # type: ignore - "model_info": model_params.model_info.model_dump_json( # type: ignore - exclude_none=True - ), + "litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), + "model_info": model_params.model_info.model_dump_json(exclude_none=True), "created_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } 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 # type: ignore - ) + model_response = await _prisma_table(ModelRepository(prisma_client)).create(data=_data) else: - model_response = LiteLLM_ProxyModelTable(**_data) + model_response = LiteLLM_ProxyModelTable.model_validate(dict(_data)) return model_response @@ -757,7 +811,9 @@ 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: _PrismaTableActions[LiteLLM_ProxyModelTable] | None = None, ) -> list[LiteLLM_ProxyModelTable]: """ Fetch all deployments for a given team_id from the database. @@ -773,8 +829,8 @@ async def _get_team_deployments( existing transaction. """ prefix = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table - response = await table.find_many( + model_table = table or _prisma_table(ModelRepository(prisma_client)) + response = await model_table.find_many( where={ "model_name": {"startswith": prefix}, } @@ -783,7 +839,7 @@ async def _get_team_deployments( return [] # Confirm team_id in model_info (defensive check) - result = [] + result: list[LiteLLM_ProxyModelTable] = [] for row in response: model_info = model_info_as_mapping(row.model_info) if model_info is not None and model_info.get("team_id") == team_id: @@ -794,7 +850,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. @@ -809,11 +865,12 @@ async def delete_team_models( """ deleted_model_ids: list[str] = [] async with prisma_client.db.tx() as tx: + tx_model_table = _transaction_model_table(tx) 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_model_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_model_table.delete_many(where={"model_id": {"in": model_ids}}) deleted_model_ids.extend(model_ids) if deleted_model_ids: @@ -908,14 +965,15 @@ async def _remove_unbacked_team_models( if not names_to_remove: return - existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + team_table = _prisma_table(TeamRepository(prisma_client)) + existing_team_row = await team_table.find_unique(where={"team_id": team_id}) if existing_team_row is None: return - updated_team_row: LiteLLM_TeamTable = await prisma_client.db.litellm_teamtable.update( + updated_team_row: 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}, # type: ignore + include={"object_permission": True}, ) await _refresh_cached_team( team_row=updated_team_row, @@ -1050,7 +1108,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique( + _existing_team_row = await _prisma_table(TeamRepository(prisma_client)).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1079,7 +1137,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 = await TeamRepository(prisma_client).table.find_unique( + team_obj_row = await _prisma_table(TeamRepository(prisma_client)).find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -1157,7 +1215,7 @@ async def delete_model( }, ) - model_in_db = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_info.id}) + model_in_db = await _prisma_table(ModelRepository(prisma_client)).find_unique(where={"model_id": model_info.id}) if model_in_db is None: raise HTTPException( status_code=400, @@ -1180,7 +1238,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id}) + result = await _prisma_table(ModelRepository(prisma_client)).delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1253,9 +1311,9 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases = await ModelTableRepository(prisma_client).table.find_many(include={"team": True}) - tasks = [] - removed_model_aliases = [] + team_model_aliases = await _model_alias_table(prisma_client).find_many(include={"team": True}) + tasks: list[Awaitable[_TeamModelAliasRow]] = [] + removed_model_aliases: list[tuple[str, str]] = [] for team_model_alias in team_model_aliases: model_aliases = team_model_alias.model_aliases # {"alias": "public model name"} id = team_model_alias.id @@ -1266,7 +1324,7 @@ async def delete_team_model_alias( removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] tasks.append( - ModelTableRepository(prisma_client).table.update( + _model_alias_table(prisma_client).update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1411,9 +1469,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=_audit_log_json(model_response), litellm_changed_by=user_api_key_dict.user_id, litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, ) @@ -1481,7 +1537,7 @@ async def update_model( ) _model_id = None - _model_info = getattr(model_params, "model_info", None) + _model_info = model_params.model_info if _model_info is None: raise Exception("model_info not provided") @@ -1489,7 +1545,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 _prisma_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: @@ -1543,7 +1601,7 @@ async def update_model( "litellm_params": json.dumps(merged_dictionary), # type: ignore "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response = await ModelRepository(prisma_client).table.update( + model_response = await _prisma_table(ModelRepository(prisma_client)).update( where={"model_id": _model_id}, data=_data, # type: ignore ) @@ -1558,16 +1616,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=_audit_log_json(_existing_litellm_params), + after_value=_audit_log_json(model_response), litellm_changed_by=user_api_key_dict.user_id, litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, ) @@ -1787,7 +1837,7 @@ def model_info_as_mapping(model_info: object) -> Mapping[str, object] | None: if not isinstance(model_info, str): return None try: - parsed = json.loads(model_info) + parsed: object = json.loads(model_info) except (TypeError, ValueError): return None return parsed if isinstance(parsed, Mapping) else None diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 78fa732a67b..949a352258c 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -32,7 +32,9 @@ from __future__ import annotations import json import re -from typing import Any +from collections.abc import Sequence +from operator import attrgetter +from typing import TYPE_CHECKING, Protocol, TypeVar from urllib.parse import quote, unquote from fastapi import HTTPException @@ -46,11 +48,24 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.repositories.table_repositories import ( ManagedFileRepository, ManagedObjectRepository, + PrismaTableRepository, ) from litellm.types.llms.openai import OpenAIFileObject +from litellm.types.passthrough_endpoints.managed_ids import ( + ManagedResourceOwner, + ManagedResourceRow, + ManagedResourceTable, + PassthroughListResponse, + PrismaWhere, +) from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +_KeyT = TypeVar("_KeyT") + # --------------------------------------------------------------------------- # Field map # --------------------------------------------------------------------------- @@ -172,7 +187,7 @@ class _RawIdGuardBudget: def __init__(self, limit: int = _MAX_RAW_ID_GUARD_LOOKUPS) -> None: self._remaining = limit - self._seen: set = set() + self._seen: set[str] = set() def reserve(self, raw_id: str) -> bool: """Return True when a guard lookup for *raw_id* should run. Returns @@ -254,6 +269,76 @@ def _canonical_path(route: str) -> str: return stripped +def _as_str_keyed_dict(value: object) -> dict[str, object] | None: + """Narrow a decoded-JSON value to a mapping, or ``None`` when it is not one.""" + if isinstance(value, dict): + return value + return None + + +class _ManagedFileLookup(Protocol): + """The enterprise managed-files hook's cached file lookup.""" + + async def __call__(self, file_id: str, litellm_parent_otel_span: None = None) -> ManagedResourceOwner | None: ... + + +class _ManagedFileStore(Protocol): + """The enterprise managed-files hook's file writer.""" + + async def __call__( + self, + *, + file_id: str, + file_object: OpenAIFileObject | None, + litellm_parent_otel_span: None, + model_mappings: dict[str, str], + user_api_key_dict: UserAPIKeyAuth, + ) -> None: ... + + +async def _hook_get_unified_file_id(managed_files_hook: object, managed_id: str) -> ManagedResourceOwner | None: + """Look a managed file up in the enterprise hook's cache. + + The proxy hook registry hands the hook back as a bare ``CustomLogger``, so + the managed-files API is resolved by name; a hook that does not implement it + raises ``AttributeError`` exactly as a direct attribute access would. + """ + lookup: _ManagedFileLookup = attrgetter("get_unified_file_id")(managed_files_hook) # any-ok: untyped hook registry + return await lookup(managed_id, litellm_parent_otel_span=None) + + +async def _hook_store_unified_file_id( + managed_files_hook: object, + file_id: str, + file_object: OpenAIFileObject | None, + model_mappings: dict[str, str], + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Persist a managed file row through the enterprise hook, resolved by name + for the same reason as ``_hook_get_unified_file_id``.""" + store: _ManagedFileStore = attrgetter("store_unified_file_id")(managed_files_hook) # any-ok: untyped hook registry + await store( + file_id=file_id, + file_object=file_object, + litellm_parent_otel_span=None, + model_mappings=model_mappings, + user_api_key_dict=user_api_key_dict, + ) + + +def _managed_table(repository: PrismaTableRepository) -> ManagedResourceTable: + """Typed view of a repository's Prisma table actions.""" + return repository.table # any-ok: Prisma actions are reached through the untyped client wrapper + + +def _managed_file_table(prisma_client: PrismaClient) -> ManagedResourceTable: + return _managed_table(ManagedFileRepository(prisma_client)) + + +def _managed_object_table(prisma_client: PrismaClient) -> ManagedResourceTable: + return _managed_table(ManagedObjectRepository(prisma_client)) + + # --------------------------------------------------------------------------- # Shared resolver — used by all INPUT path extractors # --------------------------------------------------------------------------- @@ -263,8 +348,8 @@ async def _resolve_one( managed_id: str, provider: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, + prisma_client: PrismaClient | None, + managed_files_hook: object, ) -> str: """ Resolve a single value that may be a passthrough managed ID. @@ -307,10 +392,7 @@ async def _resolve_one( # File table — use hook's internal cache for speed when available if managed_files_hook is not None: try: - file_row = await managed_files_hook.get_unified_file_id( - managed_id, - litellm_parent_otel_span=None, - ) + file_row = await _hook_get_unified_file_id(managed_files_hook, managed_id) if file_row is not None: row_created_by = file_row.created_by row_team_id = file_row.team_id @@ -322,9 +404,7 @@ async def _resolve_one( ) if not found and prisma_client is not None: try: - db_row = await ManagedFileRepository(prisma_client).table.find_first( - where={"unified_file_id": managed_id} - ) + db_row = await _managed_file_table(prisma_client).find_first(where={"unified_file_id": managed_id}) if db_row is not None: row_created_by = db_row.created_by row_team_id = db_row.team_id @@ -338,9 +418,7 @@ async def _resolve_one( # Object table (batches, responses) if prisma_client is not None: try: - obj_row = await ManagedObjectRepository(prisma_client).table.find_first( - where={"unified_object_id": managed_id} - ) + obj_row = await _managed_object_table(prisma_client).find_first(where={"unified_object_id": managed_id}) if obj_row is not None: row_created_by = obj_row.created_by row_team_id = obj_row.team_id @@ -372,7 +450,7 @@ async def _guard_raw_provider_id( raw_id: str, provider: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: PrismaClient | None, budget: _RawIdGuardBudget | None = None, ) -> None: """Deny a raw provider ID that maps to a managed resource the caller does @@ -398,7 +476,7 @@ async def _guard_raw_provider_id( # id and scope to the current provider in the application layer (same as # _mint_or_reuse_file's dedup). try: - candidates = await ManagedFileRepository(prisma_client).table.find_many( + candidates = await _managed_file_table(prisma_client).find_many( where={"flat_model_file_ids": {"has": raw_id}}, ) except Exception: @@ -419,7 +497,7 @@ async def _guard_raw_provider_id( # Object rows store model_object_id as "passthrough:{provider}:{raw}", so # the lookup is exact and already provider-scoped. try: - existing = await ManagedObjectRepository(prisma_client).table.find_first( + existing = await _managed_object_table(prisma_client).find_first( where={"model_object_id": f"passthrough:{provider}:{raw_id}"} ) except Exception: @@ -434,7 +512,7 @@ async def _guard_raw_provider_id( # --------------------------------------------------------------------------- -def _build_managed_file_object(snapshot: dict[str, Any] | None, managed_id: str) -> OpenAIFileObject | None: +def _build_managed_file_object(snapshot: dict[str, object] | None, managed_id: str) -> OpenAIFileObject | None: """Build an ``OpenAIFileObject`` (with the managed ID swapped in) from an upstream file response so the DB-served list returns the same metadata as a direct file GET. Returns ``None`` when no usable snapshot is available, in @@ -442,7 +520,7 @@ def _build_managed_file_object(snapshot: dict[str, Any] | None, managed_id: str) if not snapshot: return None try: - return OpenAIFileObject(**{**snapshot, "id": managed_id}) + return OpenAIFileObject.model_validate({**snapshot, "id": managed_id}) except Exception: verbose_proxy_logger.debug( "managed_id_rewriter: file object snapshot incomplete; storing file row without list metadata", @@ -455,9 +533,9 @@ async def _mint_or_reuse_file( raw_id: str, provider: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, - file_object_snapshot: dict[str, Any] | None = None, + prisma_client: PrismaClient | None, + managed_files_hook: object, + file_object_snapshot: dict[str, object] | None = None, is_create_route: bool = True, ) -> str: """Return an existing managed file ID or mint + store a new one.""" @@ -479,7 +557,7 @@ async def _mint_or_reuse_file( # reuse a stable row instead of minting duplicate rows on every call. if prisma_client is not None: try: - candidates = await ManagedFileRepository(prisma_client).table.find_many( + candidates = await _managed_file_table(prisma_client).find_many( where={"flat_model_file_ids": {"has": raw_id}}, order={"created_at": "asc"}, ) @@ -525,10 +603,10 @@ async def _mint_or_reuse_file( ) if managed_files_hook is not None: try: - await managed_files_hook.store_unified_file_id( + await _hook_store_unified_file_id( + managed_files_hook, file_id=managed_id, file_object=_build_managed_file_object(file_object_snapshot, managed_id), - litellm_parent_otel_span=None, model_mappings={ _passthrough_sentinel_model_id(provider): raw_id, _PASSTHROUGH_PROVIDER_MARKER_KEY: _passthrough_provider_marker(provider), @@ -551,9 +629,9 @@ async def _mint_or_reuse_object( raw_id: str, provider: str, file_purpose: str, - body_snapshot: dict, + body_snapshot: dict[str, object], user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: PrismaClient | None, is_create_route: bool, ) -> str: """Return an existing managed object ID (batch/response) or mint + store one.""" @@ -569,7 +647,7 @@ async def _mint_or_reuse_object( # f"{purpose}:{provider}:{raw_id}" for the same reason. namespaced_model_object_id = f"passthrough:{provider}:{raw_id}" - async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str: + async def _reuse_existing(existing: ManagedResourceRow, refresh_snapshot: bool) -> str: """Resolve an already-persisted namespaced row: enforce the access check, optionally refresh the snapshot, and return its managed ID.""" if not can_access_resource(user_api_key_dict, existing.created_by, existing.team_id): @@ -598,7 +676,7 @@ async def _mint_or_reuse_object( # the batch's latest state (e.g. output_file_id / error_file_id that # were null at creation but populated once the batch completed). try: - await ManagedObjectRepository(prisma_client).table.update( + await _managed_object_table(prisma_client).update( where={"unified_object_id": existing.unified_object_id}, data={ "file_object": json.dumps(body_snapshot), @@ -618,7 +696,7 @@ async def _mint_or_reuse_object( # Dedup: look up by the namespaced key — guaranteed unique per provider. try: - existing = await ManagedObjectRepository(prisma_client).table.find_first( + existing = await _managed_object_table(prisma_client).find_first( where={"model_object_id": namespaced_model_object_id} ) except Exception: @@ -635,7 +713,7 @@ async def _mint_or_reuse_object( raw_id.split("_", 1)[0], ) try: - await ManagedObjectRepository(prisma_client).table.upsert( + await _managed_object_table(prisma_client).upsert( where={"unified_object_id": managed_id}, data={ "create": { @@ -659,7 +737,7 @@ async def _mint_or_reuse_object( # the winner's managed ID so both callers converge on one ID instead of # the loser silently keeping the raw id. try: - raced = await ManagedObjectRepository(prisma_client).table.find_first( + raced = await _managed_object_table(prisma_client).find_first( where={"model_object_id": namespaced_model_object_id} ) except Exception: @@ -681,11 +759,11 @@ async def rewrite_response_ids( provider: str, method: str, route: str, - body: dict, + body: dict[str, object], user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, -) -> dict: + prisma_client: PrismaClient | None, + managed_files_hook: object, +) -> dict[str, object]: """ Mint managed IDs for raw provider values listed in ``BUILTIN_OUTPUT_ID_FIELD_MAP`` and swap them into *body*. @@ -795,7 +873,7 @@ def is_passthrough_list_route(provider: str, method: str, route: str) -> bool: return (provider, canonical) in _LIST_ROUTE_TABLE -def _parse_file_object(file_object: Any) -> Any: +def _parse_file_object(file_object: object) -> object: """Prisma may return ``Json`` columns as either a parsed dict or the raw JSON string (depending on driver / row source). Mirror the handling used elsewhere (see ``openai_files_endpoints/common_utils.py``) so callers can @@ -803,13 +881,13 @@ def _parse_file_object(file_object: Any) -> Any: """ if isinstance(file_object, str): try: - return json.loads(file_object) + return json.loads(file_object) # any-ok: json.loads -> Any except (TypeError, ValueError): return None return file_object -def _empty_list_response() -> dict[str, Any]: +def _empty_list_response() -> PassthroughListResponse: return { "object": "list", "data": [], @@ -819,7 +897,7 @@ def _empty_list_response() -> dict[str, Any]: } -def _parse_list_limit(query_params: dict[str, Any] | None) -> tuple[int, int]: +def _parse_list_limit(query_params: dict[str, str] | None) -> tuple[int, int]: params = query_params or {} try: raw_limit = int(params.get("limit", 20)) @@ -830,17 +908,17 @@ def _parse_list_limit(query_params: dict[str, Any] | None) -> tuple[int, int]: async def _build_list_where_with_cursor( - prisma_client: Any, + prisma_client: PrismaClient, resource_kind: str, provider: str, - owner_filter: dict[str, Any], - query_params: dict[str, Any] | None, -) -> tuple[dict[str, Any], str]: + owner_filter: PrismaWhere, + query_params: dict[str, str] | None, +) -> tuple[PrismaWhere, str]: """Return a Prisma ``where`` clause and fetch order for a list query.""" params = query_params or {} after_id: str | None = params.get("after") before_id: str | None = params.get("before") - where: dict[str, Any] = dict(owner_filter) + where: PrismaWhere = dict(owner_filter) fetch_order = "desc" cursor_id = after_id or before_id @@ -851,9 +929,7 @@ async def _build_list_where_with_cursor( return where, fetch_order cursor_table = ( - ManagedFileRepository(prisma_client).table - if resource_kind == "files" - else ManagedObjectRepository(prisma_client).table + _managed_file_table(prisma_client) if resource_kind == "files" else _managed_object_table(prisma_client) ) cursor_field = "unified_file_id" if resource_kind == "files" else "unified_object_id" try: @@ -867,7 +943,7 @@ async def _build_list_where_with_cursor( # created_at is not unique, so the boundary must also compare the # unique id (the secondary sort key) to avoid skipping or repeating # rows that share the cursor row's timestamp across a page boundary. - boundary = { + boundary: PrismaWhere = { "OR": [ {"created_at": {op: cursor_row.created_at}}, { @@ -885,23 +961,23 @@ async def _build_list_where_with_cursor( async def _fetch_list_rows( - prisma_client: Any, + prisma_client: PrismaClient, resource_kind: str, - where: dict[str, Any], + where: PrismaWhere, fetch_order: str, fetch_limit: int, -) -> list[Any] | None: +) -> Sequence[ManagedResourceRow] | None: # created_at is not unique, so a second sort on the unique id column gives a # total order, keeping the limit+1 page boundary and cursor deterministic # across rows that share a created_at timestamp. try: if resource_kind == "files": - return await ManagedFileRepository(prisma_client).table.find_many( + return await _managed_file_table(prisma_client).find_many( where=where, order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}], take=fetch_limit, ) - return await ManagedObjectRepository(prisma_client).table.find_many( + return await _managed_object_table(prisma_client).find_many( where={**where, "file_purpose": "batch"}, order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}], take=fetch_limit, @@ -912,14 +988,14 @@ async def _fetch_list_rows( async def _fetch_provider_scoped_list_rows( - prisma_client: Any, + prisma_client: PrismaClient, resource_kind: str, provider: str, - where: dict[str, Any], + where: PrismaWhere, fetch_order: str, raw_limit: int, fetch_limit: int, -) -> tuple[list[Any], bool]: +) -> tuple[Sequence[ManagedResourceRow], bool]: """Fetch one page of list rows scoped to *provider* at the DB level. Both resource kinds carry a provider-distinguishing value that the query @@ -933,7 +1009,7 @@ async def _fetch_provider_scoped_list_rows( A DB failure returns an empty page (fail closed) so the caller never falls through to the upstream provider. """ - scoped_where = dict(where) + scoped_where: PrismaWhere = dict(where) if resource_kind == "files": scoped_where["flat_model_file_ids"] = {"has": _passthrough_provider_marker(provider)} else: @@ -951,43 +1027,44 @@ async def _fetch_provider_scoped_list_rows( return page, has_more -def _serialize_file_list_item(row: Any) -> dict[str, Any]: - item: dict[str, Any] = { +def _serialize_file_list_item(row: ManagedResourceRow) -> dict[str, object]: + item: dict[str, object] = { "id": row.unified_file_id, "object": "file", "created_at": int(row.created_at.timestamp()) if row.created_at else None, } - file_object = _parse_file_object(row.file_object) - if isinstance(file_object, dict): + file_object = _as_str_keyed_dict(_parse_file_object(row.file_object)) + if file_object is not None: item.update(file_object) item["id"] = row.unified_file_id # managed ID always wins over stored raw id return item -def _serialize_batch_list_item(row: Any) -> dict[str, Any]: - item: dict[str, Any] = {} - file_object = _parse_file_object(row.file_object) - if isinstance(file_object, dict): +def _serialize_batch_list_item(row: ManagedResourceRow) -> dict[str, object]: + item: dict[str, object] = {} + file_object = _as_str_keyed_dict(_parse_file_object(row.file_object)) + if file_object is not None: item.update(file_object) item["id"] = row.unified_object_id # managed ID always wins item["object"] = "batch" return item -def _list_boundary_ids(rows: list[Any], resource_kind: str) -> tuple[str | None, str | None]: +def _list_boundary_ids(rows: Sequence[ManagedResourceRow], resource_kind: str) -> tuple[str | None, str | None]: if not rows: return None, None - id_attr = "unified_file_id" if resource_kind == "files" else "unified_object_id" - return getattr(rows[0], id_attr), getattr(rows[-1], id_attr) + if resource_kind == "files": + return rows[0].unified_file_id, rows[-1].unified_file_id + return rows[0].unified_object_id, rows[-1].unified_object_id async def list_passthrough_ids_from_db( provider: str, route: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - query_params: dict[str, Any] | None = None, -) -> dict[str, Any] | None: + prisma_client: PrismaClient | None, + query_params: dict[str, str] | None = None, +) -> PassthroughListResponse | None: """Query the DB for managed IDs the caller owns and return an OpenAI-style paginated list response. @@ -1060,8 +1137,8 @@ async def rewrite_path_ids( path: str, provider: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, + prisma_client: PrismaClient | None, + managed_files_hook: object, ) -> str: """ Walk URL path segments and resolve any passthrough managed IDs to raw @@ -1092,12 +1169,12 @@ async def rewrite_path_ids( async def rewrite_query_ids( - params: dict[str, Any] | None, + params: dict[str, object] | None, provider: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, -) -> dict[str, Any] | None: + prisma_client: PrismaClient | None, + managed_files_hook: object, +) -> dict[str, object] | None: """ Walk query param values and resolve any passthrough managed IDs. Returns *params* unchanged (same object) when nothing is resolved. @@ -1124,12 +1201,12 @@ async def rewrite_query_ids( async def rewrite_body_ids( - body: dict[str, Any] | None, + body: dict[str, object] | None, provider: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, -) -> dict[str, Any] | None: + prisma_client: PrismaClient | None, + managed_files_hook: object, +) -> dict[str, object] | None: """ Recursively walk a request body dict/list and resolve any passthrough managed IDs. Skips litellm internal keys (``litellm_*``). @@ -1140,27 +1217,34 @@ async def rewrite_body_ids( budget = _RawIdGuardBudget() - async def _walk(node: Any, depth: int) -> Any: + async def _walk_dict(node: dict[_KeyT, object], depth: int) -> dict[_KeyT, object]: + result: dict[_KeyT, object] = {} + changed_inner = False + for k, v in node.items(): + # Skip litellm internal injection keys (e.g. litellm_logging_obj) + if isinstance(k, str) and k.startswith("litellm_"): + result[k] = v + continue + new_v = await _walk(v, depth + 1) + result[k] = new_v + if new_v is not v: + changed_inner = True + return result if changed_inner else node + + async def _walk_list(node: list[object], depth: int) -> list[object]: + new_list = [await _walk(item, depth + 1) for item in node] + if any(n is not o for n, o in zip(new_list, node)): + return new_list + return node + + async def _walk(node: object, depth: int) -> object: if depth >= _MAX_BODY_REWRITE_DEPTH: return node - if isinstance(node, dict): - result: dict[str, Any] = {} - changed_inner = False - for k, v in node.items(): - # Skip litellm internal injection keys (e.g. litellm_logging_obj) - if isinstance(k, str) and k.startswith("litellm_"): - result[k] = v - continue - new_v = await _walk(v, depth + 1) - result[k] = new_v - if new_v is not v: - changed_inner = True - return result if changed_inner else node + node_dict = _as_str_keyed_dict(node) + if node_dict is not None: + return await _walk_dict(node_dict, depth) elif isinstance(node, list): - new_list = [await _walk(item, depth + 1) for item in node] - if any(n is not o for n, o in zip(new_list, node)): - return new_list - return node + return await _walk_list(node, depth) elif isinstance(node, str): if is_managed(node): return await _resolve_one(node, provider, user_api_key_dict, prisma_client, managed_files_hook) @@ -1168,7 +1252,7 @@ async def rewrite_body_ids( return node return node - rewritten = await _walk(body, 0) + rewritten = await _walk_dict(body, 0) if rewritten is not body: verbose_proxy_logger.debug("managed_id_rewriter: body ids rewritten provider=%s", provider) return rewritten diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 6b1ca3564e3..235155ab7a7 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,7 +4,7 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Literal, cast from openai.types.responses import ResponseFunctionToolCall @@ -33,6 +33,8 @@ from litellm.types.llms.openai import ( ChatCompletionToolParam, ChatCompletionUserMessage, GenericChatCompletionMessage, + ImageURLListItem, + ImageURLObject, InputTokensDetails, OpenAIMcpServerTool, OpenAIWebSearchOptions, @@ -47,6 +49,7 @@ from litellm.types.llms.openai import ( ValidChatCompletionMessageContentTypesLiteral, ) from litellm.types.responses.main import ( + ApplyPatchToolCallLike, CustomToolCallOutputItem, GenericResponseOutputItem, GenericResponseOutputItemContentAnnotation, @@ -189,7 +192,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None = None, stream: bool | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, **kwargs, ) -> dict: """ @@ -446,7 +449,9 @@ class LiteLLMCompletionResponsesConfig: if not chat_completion_messages: continue - deduped_in_place: list[Any] = [] + deduped_in_place: list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage + ] = [] for m in chat_completion_messages: role = "" if isinstance(m, dict): @@ -456,7 +461,7 @@ class LiteLLMCompletionResponsesConfig: # Drop assistant tool_calls wrappers if we already have this call_id if role == "assistant": - tool_calls: Any = ( + tool_calls: object = ( m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None) ) call_id = "" @@ -518,7 +523,7 @@ class LiteLLMCompletionResponsesConfig: call_id = "" if role == "assistant": - tool_calls: Any = None + tool_calls: object = None if isinstance(tool_call_message, dict): tool_calls = tool_call_message.get("tool_calls") else: @@ -562,7 +567,16 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None: + def _find_previous_assistant_idx( + messages: Sequence[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message + ], + current_idx: int, + ) -> int | None: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": @@ -570,17 +584,19 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str: + def _recover_tool_call_id_from_assistant(assistant_message: object, message: object) -> str: """Try to recover empty tool_call_id from assistant message's tool_calls.""" tool_calls_raw = ( assistant_message.get("tool_calls") if isinstance(assistant_message, dict) else getattr(assistant_message, "tool_calls", None) ) - if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0: - first_tool_call = tool_calls_raw[0] - if isinstance(first_tool_call, dict): - tool_call_id_raw = first_tool_call.get("id", "") + tool_calls_items = LiteLLMCompletionResponsesConfig._as_object_list(tool_calls_raw) + if tool_calls_raw and tool_calls_items is not None and len(tool_calls_items) > 0: + first_tool_call = tool_calls_items[0] + first_tool_call_mapping = LiteLLMCompletionResponsesConfig._as_object_mapping(first_tool_call) + if first_tool_call_mapping is not None: + tool_call_id_raw: object = first_tool_call_mapping.get("id", "") return str(tool_call_id_raw) if tool_call_id_raw is not None else "" elif hasattr(first_tool_call, "id"): tool_call_id_raw = getattr(first_tool_call, "id", None) @@ -588,7 +604,7 @@ class LiteLLMCompletionResponsesConfig: return "" @staticmethod - def _get_tool_calls_list(assistant_message: Any) -> list[Any]: + def _get_tool_calls_list(assistant_message: object) -> list[object]: """Extract tool_calls as a list from assistant message.""" tool_calls_raw = ( assistant_message.get("tool_calls") @@ -597,14 +613,15 @@ class LiteLLMCompletionResponsesConfig: ) if tool_calls_raw is None: return [] - if isinstance(tool_calls_raw, list): - return tool_calls_raw + tool_calls_items = LiteLLMCompletionResponsesConfig._as_object_list(tool_calls_raw) + if tool_calls_items is not None: + return tool_calls_items if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)): return list(tool_calls_raw) return [] @staticmethod - def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool: + def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool: """Check if a tool_call with the given ID exists in the list.""" for tool_call in tool_calls: tool_call_id_to_check: str | None = None @@ -617,7 +634,7 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None: + def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None: """Reconstruct a minimal tool_call definition from tools list.""" for tool in tools: if isinstance(tool, dict): @@ -635,7 +652,7 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: + def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object: """ Safely read a field from dict-like or attribute-based objects. """ @@ -656,13 +673,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _create_tool_call_chunk( - tool_use_definition: dict[str, Any], tool_call_id: str, index: int + tool_use_definition: Mapping[str, object], tool_call_id: str, index: int ) -> ChatCompletionToolCallChunk: """Create a ChatCompletionToolCallChunk from tool_use_definition.""" function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function") function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments") - function: dict[str, Any] = { + function: dict[str, object] = { "name": function_name_raw or "", "arguments": function_arguments_raw or "{}", } @@ -681,7 +698,7 @@ class LiteLLMCompletionResponsesConfig: ) @staticmethod - def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None: + def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[str, object] | None: """ Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. """ @@ -689,7 +706,7 @@ class LiteLLMCompletionResponsesConfig: return None if isinstance(tool_use_definition, dict): - normalized_definition: dict[str, Any] = dict(tool_use_definition) + normalized_definition: dict[str, object] = dict(tool_use_definition) else: tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id") tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type") @@ -722,20 +739,42 @@ class LiteLLMCompletionResponsesConfig: return normalized_definition @staticmethod - def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None: + def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None: """Add a tool_call to an assistant message.""" if isinstance(assistant_message, dict): - prev_assistant_dict = cast(dict[str, Any], assistant_message) + prev_assistant_dict = cast(dict[str, object], assistant_message) if "tool_calls" not in prev_assistant_dict: prev_assistant_dict["tool_calls"] = [] - tool_calls_list = prev_assistant_dict["tool_calls"] - if isinstance(tool_calls_list, list): + tool_calls_list = LiteLLMCompletionResponsesConfig._as_object_list(prev_assistant_dict["tool_calls"]) + if tool_calls_list is not None: tool_calls_list.append(tool_call_chunk) elif hasattr(assistant_message, "tool_calls"): if assistant_message.tool_calls is None: assistant_message.tool_calls = [] - if isinstance(assistant_message.tool_calls, list): - assistant_message.tool_calls.append(tool_call_chunk) + existing_tool_calls: list[object] | None = LiteLLMCompletionResponsesConfig._as_object_list( + getattr(assistant_message, "tool_calls", None) + ) + if existing_tool_calls is not None: + existing_tool_calls.append(tool_call_chunk) + + @staticmethod + def _as_object_list(value: object) -> list[object] | None: + """Return ``value`` typed as a list of objects when it is a list, else ``None``.""" + if isinstance(value, list): + return value + return None + + @staticmethod + def _require_attr(obj: object, name: str) -> object: + """Read a required attribute, preserving the ``AttributeError`` raised when it is absent.""" + return getattr(obj, name) + + @staticmethod + def _as_object_mapping(value: object) -> dict[str, object] | None: + """Return ``value`` typed as a string-keyed mapping when it is a dict, else ``None``.""" + if isinstance(value, dict): + return value + return None @staticmethod def _ensure_tool_results_have_corresponding_tool_calls( @@ -746,7 +785,7 @@ class LiteLLMCompletionResponsesConfig: | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: Sequence[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -810,7 +849,7 @@ class LiteLLMCompletionResponsesConfig: # Type-safe way to set tool_call_id on tool message if isinstance(message, dict): # Cast to dict to allow setting tool_call_id - message_dict = cast(dict[str, Any], message) + message_dict = cast(dict[str, object], message) message_dict["tool_call_id"] = tool_call_id elif hasattr(message, "tool_call_id"): setattr(message, "tool_call_id", tool_call_id) @@ -862,7 +901,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_input_item_to_chat_completion_message( - input_item: Any, + input_item: Mapping[str, object], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API input item into a Chat Completion message @@ -897,9 +936,10 @@ class LiteLLMCompletionResponsesConfig: # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: return [] + role_value = input_item.get("role") return [ GenericChatCompletionMessage( - role=input_item.get("role") or "user", + role=str(role_value) if role_value else "user", content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( content ), @@ -907,7 +947,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_tool_call_output(input_item: Any) -> bool: + def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a tool call output """ @@ -920,7 +960,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_function_call(input_item: Any) -> bool: + def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a function call or custom tool call. Both need to be reconstructed as assistant tool_calls for Chat @@ -930,7 +970,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_tool_call_output_to_chat_completion_message( - tool_call_output: dict[str, Any], + tool_call_output: Mapping[str, object], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ ChatCompletionToolMessage is used to indicate the output from a tool call @@ -942,7 +982,7 @@ class LiteLLMCompletionResponsesConfig: return [] def _normalize_function_call_output_to_tool_content( - output: Any, + output: object, ) -> Any: """ Normalize Responses API function_call_output.output into a shape that downstream @@ -964,22 +1004,25 @@ class LiteLLMCompletionResponsesConfig: return output # Some adapters represent tool output as a list of "input_*" parts - if isinstance(output, list): - normalized_blocks: list[dict[str, Any]] = [] + output_parts = LiteLLMCompletionResponsesConfig._as_object_list(output) + if output_parts is not None: + normalized_blocks: list[dict[str, object]] = [] text_acc: list[str] = [] - for part in output: - if not isinstance(part, dict): + for part in output_parts: + part_mapping = LiteLLMCompletionResponsesConfig._as_object_mapping(part) + if part_mapping is None: continue - part_type = part.get("type") + part_type = part_mapping.get("type") if part_type in ("input_text", "output_text", "text"): - txt = part.get("text") + txt = part_mapping.get("text") if isinstance(txt, str) and txt: text_acc.append(txt) normalized_blocks.append({"type": "text", "text": txt}) elif part_type in ("input_image", "image_url"): - image_url_val = part.get("image_url") or part.get("url") - if isinstance(image_url_val, dict): - url = image_url_val.get("url") + image_url_val = part_mapping.get("image_url") or part_mapping.get("url") + image_url_mapping = LiteLLMCompletionResponsesConfig._as_object_mapping(image_url_val) + if image_url_mapping is not None: + url = image_url_mapping.get("url") if isinstance(url, str) and url: normalized_blocks.append({"type": "image_url", "image_url": {"url": url}}) elif isinstance(image_url_val, str) and image_url_val: @@ -1066,7 +1109,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_function_call_to_chat_completion_message( - function_call: dict[str, Any], + function_call: Mapping[str, object], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API function_call into a Chat Completion message with tool calls @@ -1091,11 +1134,13 @@ class LiteLLMCompletionResponsesConfig: if not raw_arguments and function_call.get("type") == "custom_tool_call": raw_input = function_call.get("input") or "" raw_arguments = json.dumps({"content": raw_input}) if raw_input else "" + call_id_value = function_call.get("call_id") or function_call.get("id") or "" + name_value = function_call.get("name") or "" tool_call = ChatCompletionToolCallChunk( - id=function_call.get("call_id") or function_call.get("id") or "", + id=str(call_id_value), type="function", function=ChatCompletionToolCallFunctionChunk( - name=function_call.get("name") or "", + name=str(name_value), arguments=str(raw_arguments or ""), ), index=0, @@ -1111,7 +1156,7 @@ class LiteLLMCompletionResponsesConfig: return [chat_completion_response_message] @staticmethod - def _resolve_file_id(item: dict[str, Any]) -> str | None: + def _resolve_file_id(item: Mapping[str, object]) -> object: """ Return the effective file_id for a Responses API input_file item. Explicit file_id takes precedence; file_url is used as fallback so @@ -1120,7 +1165,7 @@ class LiteLLMCompletionResponsesConfig: return item.get("file_id") or item.get("file_url") or None @staticmethod - def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]: + def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]: """ Transform a Responses API input_file item to a Chat Completion file item @@ -1130,77 +1175,83 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary with transformed file structure for Chat Completion """ - file_dict: dict[str, Any] = {} + file_dict: dict[str, object] = {} file_id = LiteLLMCompletionResponsesConfig._resolve_file_id(item) if file_id: file_dict["file_id"] = file_id if item.get("file_data"): file_dict["file_data"] = item["file_data"] - new_item: dict[str, Any] = {"type": "file", "file": file_dict} + new_item: dict[str, object] = {"type": "file", "file": file_dict} if "cache_control" in item: new_item["cache_control"] = item["cache_control"] return new_item @staticmethod def _transform_input_image_item_to_image_item( - item: dict[str, Any], + item: Mapping[str, object], ) -> ChatCompletionImageObject: """ Transform a Responses API input_image item to a Chat Completion image item """ + image_url_value = item.get("image_url") + detail_value = item.get("detail") image_url_obj = ChatCompletionImageUrlObject( - url=item.get("image_url") or "", detail=item.get("detail") or "auto" + url=str(image_url_value) if image_url_value else "", + detail=str(detail_value) if detail_value else "auto", ) return ChatCompletionImageObject(type="image_url", image_url=image_url_obj) @staticmethod def _transform_responses_api_content_to_chat_completion_content( - content: Any, - ) -> str | list[str | dict[str, Any]]: + content: object, + ) -> str | list[str | dict[str, object]]: """ Transform a Responses API content into a Chat Completion content Note: This function should not be called with None content. Callers should check for None before calling this function. """ + content_items = LiteLLMCompletionResponsesConfig._as_object_list(content) if content is None: # Defensive check: should not happen if callers check first # Return empty string as fallback to avoid type errors return "" elif isinstance(content, str): return content - elif isinstance(content, list): - content_list: list[str | dict[str, Any]] = [] - for item in content: + elif content_items is not None: + content_list: list[str | dict[str, object]] = [] + for item in content_items: + item_mapping = LiteLLMCompletionResponsesConfig._as_object_mapping(item) if isinstance(item, str): content_list.append(item) - elif isinstance(item, dict): - if item.get("type") == "input_file": + elif item_mapping is not None: + if item_mapping.get("type") == "input_file": content_list.append( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) + LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item_mapping) ) - elif item.get("type") == "input_image": - image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) + elif item_mapping.get("type") == "input_image": + image_block: dict[str, object] = dict( + LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item_mapping) ) - if "cache_control" in item: - image_block["cache_control"] = item["cache_control"] + if "cache_control" in item_mapping: + image_block["cache_control"] = item_mapping["cache_control"] content_list.append(image_block) else: # Skip text blocks with None text to avoid downstream errors - text_value = item.get("text") + text_value = item_mapping.get("text") if text_value is None: continue - content_block: dict[str, Any] = { + item_type = item_mapping.get("type") + content_block: dict[str, object] = { "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( - item.get("type") or "text" + str(item_type) if item_type else "text" ), "text": text_value, } - if "cache_control" in item: - content_block["cache_control"] = item["cache_control"] + if "cache_control" in item_mapping: + content_block["cache_control"] = item_mapping["cache_control"] content_list.append(content_block) return content_list else: @@ -1283,7 +1334,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(typed_tool.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - chat_completion_tool: dict[str, Any] = { + chat_completion_tool: dict[str, object] = { "type": "function", "function": { "name": typed_tool.get("name") or "", @@ -1324,7 +1375,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to Responses API request tool format. Inverse of @@ -1332,17 +1383,18 @@ class LiteLLMCompletionResponsesConfig: """ if chat_completion_tools is None or not chat_completion_tools: return [] - result: list[dict[str, Any]] = [] + result: list[dict[str, object]] = [] for tool in chat_completion_tools: if not isinstance(tool, dict): result.append(tool) # type: ignore continue if tool.get("type") == "function": - fn = cast(dict[str, Any], tool.get("function") or {}) - parameters = dict(fn.get("parameters", {}) or {}) + fn = cast(dict[str, object], tool.get("function") or {}) + fn_parameters = LiteLLMCompletionResponsesConfig._as_object_mapping(fn.get("parameters")) + parameters: dict[str, object] = dict(fn_parameters) if fn_parameters else {} if not parameters or "type" not in parameters: parameters["type"] = "object" - responses_tool: dict[str, Any] = { + responses_tool: dict[str, object] = { "type": "function", "name": fn.get("name") or "", "description": fn.get("description") or "", @@ -1496,9 +1548,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item: Any, + tool_call_item: object, index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format. @@ -1511,32 +1563,29 @@ class LiteLLMCompletionResponsesConfig: """ # Extract provider_specific_fields if present provider_specific_fields = getattr(tool_call_item, "provider_specific_fields", None) + item_getter = getattr(tool_call_item, "get", None) if provider_specific_fields and not isinstance(provider_specific_fields, dict): provider_specific_fields = ( dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} ) - elif hasattr(tool_call_item, "get") and callable(tool_call_item.get): # type: ignore - provider_fields = tool_call_item.get("provider_specific_fields") # type: ignore + elif callable(item_getter): + provider_fields = item_getter("provider_specific_fields") if provider_fields: provider_specific_fields = ( provider_fields if isinstance(provider_fields, dict) - else ( - dict(provider_fields) # type: ignore - if hasattr(provider_fields, "__dict__") - else {} - ) + else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {}) ) - function_dict: dict[str, Any] = { - "name": tool_call_item.name, - "arguments": tool_call_item.arguments, + function_dict: dict[str, object] = { + "name": LiteLLMCompletionResponsesConfig._require_attr(tool_call_item, "name"), + "arguments": LiteLLMCompletionResponsesConfig._require_attr(tool_call_item, "arguments"), } if provider_specific_fields: function_dict["provider_specific_fields"] = provider_specific_fields - tool_call_dict: dict[str, Any] = { + tool_call_dict: dict[str, object] = { "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item( getattr(tool_call_item, "id", None), getattr(tool_call_item, "call_id", None), @@ -1553,9 +1602,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_apply_patch_tool_call_to_chat_completion_tool_call( - tool_call_item: Any, + tool_call_item: ApplyPatchToolCallLike, index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. @@ -1573,7 +1622,7 @@ class LiteLLMCompletionResponsesConfig: import json operation_dict = tool_call_item.operation.model_dump() - tool_call_dict: dict[str, Any] = { + tool_call_dict: dict[str, object] = { "id": tool_call_item.call_id, "function": { "name": "apply_patch", @@ -1711,7 +1760,7 @@ class LiteLLMCompletionResponsesConfig: """ output_items: list = [] for choice in chat_completion_response.choices or []: - message = getattr(choice, "message", None) + message: object = getattr(choice, "message", None) if not message: continue psf = getattr(message, "provider_specific_fields", None) @@ -1783,13 +1832,13 @@ class LiteLLMCompletionResponsesConfig: """ image_generation_items: list[OutputImageGenerationCall] = [] - images = getattr(choice.message, "images", []) + images: list[ImageURLListItem] | None = getattr(choice.message, "images", []) if not images: return image_generation_items for idx, image_item in enumerate(images): # Extract base64 from data URL - image_url = image_item.get("image_url", {}).get("url", "") + image_url = image_item.get("image_url", ImageURLObject(url="")).get("url", "") base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: @@ -2034,8 +2083,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_text_format_to_response_format( - text_param: dict[str, Any] | Any, - ) -> dict[str, Any] | None: + text_param: object, + ) -> dict[str, object] | None: """ Transform Responses API text.format parameter to Chat Completion response_format parameter. @@ -2062,9 +2111,10 @@ class LiteLLMCompletionResponsesConfig: if not text_param: return None - if isinstance(text_param, dict): - format_param = text_param.get("format") - if format_param and isinstance(format_param, dict): + text_mapping = LiteLLMCompletionResponsesConfig._as_object_mapping(text_param) + if text_mapping is not None: + format_param = LiteLLMCompletionResponsesConfig._as_object_mapping(text_mapping.get("format")) + if format_param: format_type = format_param.get("type") if format_type == "json_schema": diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 3bcc19822a6..487c51f3f22 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 Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Protocol import httpx from openai._streaming import SSEDecoder @@ -34,6 +34,18 @@ from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook +if TYPE_CHECKING: + from litellm.types.guardrails import PresidioPerRequestConfig + from litellm.types.llms.openai import ( + PART_UNION_TYPES, + ContentPartDoneEvent, + ResponseAPIUsage, + ResponseCreatedEvent, + ResponseInProgressEvent, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, + ) + @lru_cache(maxsize=1) def _get_openai_response_types(): @@ -42,7 +54,7 @@ def _get_openai_response_types(): return openai_types -def _log_background_task_failure(task: asyncio.Task[Any], *, task_name: str) -> None: +def _log_background_task_failure(task: asyncio.Task[None], *, task_name: str) -> None: if task.cancelled(): return exception = task.exception() @@ -98,6 +110,20 @@ def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None] return message, error_type, code +class _StreamCachingHandler(Protocol): + """Caching-handler surface attached to the logging object, used to cache completed streamed responses.""" + + original_function: Callable[..., object] + + def _should_store_result_in_cache( + self, original_function: Callable[..., object], kwargs: dict[str, object] + ) -> bool: ... + + @staticmethod + def should_store(handler: _StreamCachingHandler, kwargs: dict[str, object]) -> bool: + return handler._should_store_result_in_cache(original_function=handler.original_function, kwargs=kwargs) + + def _status_code_for_error_fields(error_type: str | None, error_code: str | None) -> int: fields = tuple(field for field in (error_code, error_type) if field is not None) if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields): @@ -131,7 +157,7 @@ class BaseResponsesAPIStreamingIterator: self.logging_obj = logging_obj self.finished = False self.responses_api_provider_config = responses_api_provider_config - self.completed_response: Any | None = None + self.completed_response: ResponsesAPIStreamingResponse | None = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called self._yielded_first_chunk = False @@ -154,7 +180,7 @@ class BaseResponsesAPIStreamingIterator: model=model or "", optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) - _model_info: dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} + _model_info: dict[str, object] = litellm_metadata.get("model_info", {}) if litellm_metadata else {} self._hidden_params = { "model_id": _model_info.get("id", None), "api_base": _api_base, @@ -176,7 +202,7 @@ class BaseResponsesAPIStreamingIterator: llm_provider=self.custom_llm_provider or "", ) - def _process_chunk(self, chunk) -> Any | None: + def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None: """Process a single chunk of data from the stream""" if not chunk: return None @@ -196,7 +222,7 @@ class BaseResponsesAPIStreamingIterator: try: # Parse the JSON chunk - parsed_chunk = json.loads(chunk) + parsed_chunk: object = json.loads(chunk) # Format as ResponsesAPIStreamingResponse if isinstance(parsed_chunk, dict): @@ -212,7 +238,7 @@ class BaseResponsesAPIStreamingIterator: # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a # truthy child Mock for any attribute, which breaks tests and is wrong on stream. if "response" in parsed_chunk: - response_object = getattr(openai_responses_api_chunk, "response", None) + response_object: ResponsesAPIResponse | None = getattr(openai_responses_api_chunk, "response", None) if response_object is not None: response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=response_object, @@ -250,7 +276,7 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: - _part = getattr(openai_responses_api_chunk, "part", None) + _part: object = getattr(openai_responses_api_chunk, "part", None) if _part is not None: if isinstance(_part, dict): ResponsesAPIRequestUtils._encode_container_ids_in_annotations( @@ -273,7 +299,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - item = getattr(openai_responses_api_chunk, "item", None) + item: object = getattr(openai_responses_api_chunk, "item", None) if item: encrypted_content = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): @@ -299,9 +325,11 @@ class BaseResponsesAPIStreamingIterator: self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Any | None = getattr(openai_responses_api_chunk, "response", None) + response_obj: ResponsesAPIResponse | None = getattr( + openai_responses_api_chunk, "response", None + ) if response_obj: - usage_obj: Any | None = getattr(response_obj, "usage", None) + usage_obj: object | None = getattr(response_obj, "usage", None) if usage_obj is not None: try: cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) @@ -377,10 +405,10 @@ class BaseResponsesAPIStreamingIterator: ) self._run_post_success_hooks(end_time=end_time) - def _handle_logging_completed_response(self): + def _handle_logging_completed_response(self) -> None: """Base implementation - should be overridden by subclasses""" - def _handle_logging_failed_response(self): + def _handle_logging_failed_response(self) -> None: """ Handle logging for RESPONSE_FAILED events by routing to failure handlers. @@ -389,8 +417,10 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj = getattr(self.completed_response, "response", None) if self.completed_response else None - error_info = getattr(response_obj, "error", None) if response_obj else None + response_obj: ResponsesAPIResponse | None = ( + getattr(self.completed_response, "response", None) if self.completed_response else None + ) + error_info: object | None = getattr(response_obj, "error", None) if response_obj else None error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) exception = litellm.APIError( @@ -401,10 +431,10 @@ class BaseResponsesAPIStreamingIterator: ) self._handle_failure(exception) - def _record_failed_response_usage(self, response_obj: Any | None) -> None: + def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj = getattr(response_obj, "usage", None) + usage_obj: ResponseAPIUsage | None = getattr(response_obj, "usage", None) if usage_obj is None: return try: @@ -451,7 +481,7 @@ class BaseResponsesAPIStreamingIterator: is_pre_first_chunk=not self._yielded_first_chunk, ) - def _get_completed_response_object(self) -> Any | None: + def _get_completed_response_object(self) -> ResponsesAPIResponse | None: openai_types = _get_openai_response_types() completed_response = self.completed_response if isinstance(completed_response, openai_types.ResponsesAPIResponse): @@ -476,15 +506,15 @@ class BaseResponsesAPIStreamingIterator: if response_obj is None: return - caching_handler = getattr(self.logging_obj, "_llm_caching_handler", None) + caching_handler: _StreamCachingHandler | None = getattr(self.logging_obj, "_llm_caching_handler", None) if caching_handler is None: return - request_kwargs = getattr(caching_handler, "request_kwargs", None) + request_kwargs: object = getattr(caching_handler, "request_kwargs", None) if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: return request_kwargs = request_kwargs.copy() - preset_cache_key = getattr(caching_handler, "preset_cache_key", None) + preset_cache_key: object = getattr(caching_handler, "preset_cache_key", None) request_cache_key = request_kwargs.pop("cache_key", None) if preset_cache_key is None: preset_cache_key = request_cache_key @@ -494,10 +524,7 @@ 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, - kwargs=request_kwargs, - ): + if not _StreamCachingHandler.should_store(caching_handler, request_kwargs): return if litellm.cache is None: @@ -527,7 +554,9 @@ class BaseResponsesAPIStreamingIterator: self._completed_response_cached = True - async def _call_post_streaming_deployment_hook(self, chunk): + async def _call_post_streaming_deployment_hook( + self, chunk: ResponsesAPIStreamingResponse + ) -> ResponsesAPIStreamingResponse: """ Allow callbacks to modify streaming chunks before returning (parity with chat). """ @@ -564,7 +593,9 @@ class BaseResponsesAPIStreamingIterator: except Exception: return chunk - async def call_post_streaming_hooks_for_testing(self, chunk): + async def call_post_streaming_hooks_for_testing( + self, chunk: ResponsesAPIStreamingResponse + ) -> ResponsesAPIStreamingResponse: """ Helper to invoke streaming deployment hooks explicitly (used in tests). """ @@ -666,7 +697,9 @@ class BaseResponsesAPIStreamingIterator: pass -async def call_post_streaming_hooks_for_testing(iterator, chunk): +async def call_post_streaming_hooks_for_testing( + iterator: BaseResponsesAPIStreamingIterator, chunk: ResponsesAPIStreamingResponse +) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ @@ -707,7 +740,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __aiter__(self): return self - async def __anext__(self) -> Any: + async def __anext__(self) -> ResponsesAPIStreamingResponse: try: self._check_max_streaming_duration() while True: @@ -753,7 +786,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): self._handle_failure(e) raise e - def _handle_logging_completed_response(self): + def _handle_logging_completed_response(self) -> None: """Handle logging for completed responses in async context""" self._log_completed_response(is_async=True) @@ -789,7 +822,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __iter__(self): return self - def __next__(self): + def __next__(self) -> ResponsesAPIStreamingResponse: try: self._check_max_streaming_duration() while True: @@ -835,7 +868,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): self._handle_failure(e) raise e - def _handle_logging_completed_response(self): + def _handle_logging_completed_response(self) -> None: """Handle logging for completed responses in sync context""" self._log_completed_response(is_async=False) @@ -880,7 +913,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _set_events_from_response( self, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events = _build_synthetic_response_events( @@ -894,7 +927,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __aiter__(self): return self - async def __anext__(self) -> Any: + async def __anext__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopAsyncIteration evt = self._events[self._idx] @@ -908,7 +941,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __iter__(self): return self - def __next__(self) -> Any: + def __next__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopIteration evt = self._events[self._idx] @@ -923,7 +956,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __init__( self, - response: Any, + response: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, request_data: dict[str, Any] | None = None, call_type: str | None = None, @@ -941,13 +974,13 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): ) self._completed_response_cache_hit = True self._persist_completed_response_before_logging = False - self._events: list[Any] = [] + self._events: list[ResponsesAPIStreamingResponse] = [] self._idx = 0 self._set_events_from_response(transformed=response, logging_obj=logging_obj) def _set_events_from_response( self, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events = _build_synthetic_response_events( @@ -961,7 +994,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __aiter__(self): return self - async def __anext__(self) -> Any: + async def __anext__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopAsyncIteration evt = self._events[self._idx] @@ -975,7 +1008,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __iter__(self): return self - def __next__(self) -> Any: + def __next__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopIteration evt = self._events[self._idx] @@ -997,11 +1030,11 @@ def _dump_response_object(obj: Any) -> dict[str, Any]: def _build_response_status_event( event_type: Literal[ - "response.created", - "response.in_progress", + ResponsesAPIStreamEvents.RESPONSE_CREATED, + ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, ], - transformed: Any, -) -> Any: + transformed: ResponsesAPIResponse, +) -> ResponseCreatedEvent | ResponseInProgressEvent: openai_types = _get_openai_response_types() in_progress_response = transformed.model_copy( deep=True, @@ -1018,10 +1051,10 @@ def _build_content_part_done_event( output_index: int, content_index: int, part_payload: dict[str, Any], -) -> Any | None: +) -> ContentPartDoneEvent | None: openai_types = _get_openai_response_types() part_type = part_payload.get("type") - part: Any + part: PART_UNION_TYPES if part_type == "output_text": annotations = [ openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) @@ -1057,7 +1090,7 @@ def _build_content_part_done_event( def _add_text_like_part_events( *, - events: list[Any], + events: list[ResponsesAPIStreamingResponse], item_id: str, output_index: int, content_index: int, @@ -1123,13 +1156,13 @@ def _add_text_like_part_events( def _build_synthetic_response_events( *, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, chunk_size: int, -) -> list[Any]: +) -> list[ResponsesAPIStreamingResponse]: openai_types = _get_openai_response_types() if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Any | None = getattr(transformed, "usage", None) + usage_obj: object | None = getattr(transformed, "usage", None) if usage_obj is not None: try: cost: float | None = logging_obj._response_cost_calculator(result=transformed) @@ -1138,13 +1171,14 @@ def _build_synthetic_response_events( except Exception: pass - events: list[Any] = [ + events: list[ResponsesAPIStreamingResponse] = [ _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), ] sequence_number = 0 - for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): + output_items: Sequence[object] = getattr(transformed, "output", []) or [] + for output_index, output_item in enumerate(output_items): output_item_payload = _dump_response_object(output_item) item_id = str(output_item_payload.get("id") or transformed.id) item_type = output_item_payload.get("type") @@ -1158,7 +1192,8 @@ def _build_synthetic_response_events( ) if item_type == "message": - for content_index, part in enumerate(output_item_payload.get("content", []) or []): + content_payload: Sequence[object] = output_item_payload.get("content", []) or [] + for content_index, part in enumerate(content_payload): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( @@ -1205,7 +1240,8 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): + summary_payload_items: Sequence[object] = output_item_payload.get("summary", []) or [] + for summary_index, summary in enumerate(summary_payload_items): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1277,6 +1313,46 @@ RESPONSES_WS_LOGGED_EVENT_TYPES = [ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES = frozenset({"input_text", "output_text", "text"}) +class _ClientWebSocket(Protocol): + """Client-facing socket surface used by the Responses WebSocket handlers.""" + + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + +class _BackendWebSocket(Protocol): + """Upstream provider socket surface used by the Responses WebSocket handlers.""" + + async def recv(self, decode: bool = ...) -> str | bytes: ... + + async def send(self, message: str) -> None: ... + + async def close(self) -> None: ... + + +class _PIIMaskingGuardrail(Protocol): + """Guardrail surface used for Presidio PII masking/unmasking of WebSocket frames.""" + + def get_presidio_settings_from_request_data( + self, data: Mapping[str, object] + ) -> PresidioPerRequestConfig | None: ... + + async def check_pii( + self, + text: str, + output_parse_pii: bool, + presidio_config: PresidioPerRequestConfig | None, + request_data: Mapping[str, object], + ) -> str: ... + + def _unmask_pii_text(self, text: str, pii_tokens: Mapping[str, str]) -> str: ... + + @staticmethod + def unmask(guardrail: _PIIMaskingGuardrail, text: str, pii_tokens: Mapping[str, str]) -> str: + return guardrail._unmask_pii_text(text, pii_tokens) + + class ResponsesWebSocketStreaming: """ Manages bidirectional WebSocket forwarding for the Responses API @@ -1292,31 +1368,31 @@ class ResponsesWebSocketStreaming: def __init__( self, - websocket: Any, - backend_ws: Any, + websocket: _ClientWebSocket, + backend_ws: _BackendWebSocket, logging_obj: LiteLLMLoggingObj, - user_api_key_dict: Any | None = None, - request_data: dict | None = None, + user_api_key_dict: object | None = None, + request_data: dict[str, object] | None = None, first_message: str | None = None, - guardrail_callbacks: list[Any] | None = None, - output_guardrail_callbacks: list[Any] | None = None, + guardrail_callbacks: list[_PIIMaskingGuardrail] | None = None, + output_guardrail_callbacks: list[_PIIMaskingGuardrail] | None = None, authorized_model: str | None = None, ): self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.request_data: dict = request_data or {} - self.messages: list[dict] = [] + self.request_data: dict[str, object] = request_data or {} + self.messages: list[dict[str, object]] = [] self.input_messages: list[dict[str, str]] = [] self.first_message = first_message - self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] - self.output_guardrail_callbacks: list[Any] = output_guardrail_callbacks or [] + self.guardrail_callbacks: list[_PIIMaskingGuardrail] = guardrail_callbacks or [] + self.output_guardrail_callbacks: list[_PIIMaskingGuardrail] = output_guardrail_callbacks or [] # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model - def _should_store_event(self, event_obj: dict) -> bool: + def _should_store_event(self, event_obj: dict[str, object]) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES def _store_event(self, event: Any) -> None: @@ -1333,7 +1409,7 @@ class ResponsesWebSocketStreaming: if self._should_store_event(event_obj): self.messages.append(event_obj) - def _collect_input_from_client_event(self, message: Any) -> None: + def _collect_input_from_client_event(self, message: object) -> None: """Extract user input content from response.create for logging.""" try: if isinstance(message, str): @@ -1368,7 +1444,7 @@ class ResponsesWebSocketStreaming: except (json.JSONDecodeError, AttributeError, TypeError): pass - def _store_input(self, message: Any) -> None: + def _store_input(self, message: object) -> None: self._collect_input_from_client_event(message) if self.logging_obj: self.logging_obj.pre_call(input=message, api_key="") @@ -1429,7 +1505,7 @@ class ResponsesWebSocketStreaming: finally: await self._log_messages() - def _enforce_authorized_model(self, msg_obj: dict) -> bool: + def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool: """ Overwrite any ``model`` field in a ``response.create`` frame with the connection-authorized model to prevent deployment-substitution attacks. @@ -1495,8 +1571,9 @@ class ResponsesWebSocketStreaming: # nested: {"type": "response.create", "response": {"input": ..., "instructions": ...}} # Mask "input" and "instructions" in both shapes so PII is never # forwarded unmasked regardless of where the client places it. - nested_response = msg_obj.get("response") if isinstance(msg_obj.get("response"), dict) else None - text_containers: list[tuple[dict, str]] = [] + raw_nested_response = msg_obj.get("response") + nested_response = raw_nested_response if isinstance(raw_nested_response, dict) else None + text_containers: list[tuple[dict[str, object], str]] = [] for container in (msg_obj, nested_response): if container is None: continue @@ -1590,7 +1667,8 @@ class ResponsesWebSocketStreaming: if not self.guardrail_callbacks: return response_str - pii_tokens: dict[str, str] = (self.request_data.get("metadata") or {}).get("pii_tokens", {}) + metadata = self.request_data.get("metadata") + pii_tokens: dict[str, str] = metadata.get("pii_tokens", {}) if isinstance(metadata, dict) else {} if not pii_tokens: return response_str @@ -1618,7 +1696,7 @@ class ResponsesWebSocketStreaming: continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text(text, pii_tokens) + unmasked = _PIIMaskingGuardrail.unmask(cb, text, pii_tokens) if unmasked != text: content_block["text"] = unmasked modified = True @@ -1627,7 +1705,7 @@ class ResponsesWebSocketStreaming: if event_type in self._DELTA_EVENT_TYPES: delta = evt_obj.get("delta") if isinstance(delta, str): - unmasked = cb._unmask_pii_text(delta, pii_tokens) + unmasked = _PIIMaskingGuardrail.unmask(cb, delta, pii_tokens) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) @@ -1650,7 +1728,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj = json.loads(response_str) + evt_obj: Mapping[str, object] = json.loads(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1756,12 +1834,12 @@ class ResponsesWebSocketStreaming: # Managed WebSocket mode (HTTP-backed, provider-agnostic) # --------------------------------------------------------------------------- -_RESPONSE_CREATE_PARAMS: frozenset = ( +_RESPONSE_CREATE_PARAMS: frozenset[str] = ( _get_openai_response_types().ResponsesAPIRequestParams.__required_keys__ | _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__ ) -_MANAGED_WS_SKIP_KWARGS: frozenset = frozenset( +_MANAGED_WS_SKIP_KWARGS: frozenset[str] = frozenset( { "litellm_logging_obj", "litellm_call_id", @@ -1793,10 +1871,10 @@ class ManagedResponsesWebSocketHandler: def __init__( self, - websocket: Any, + websocket: _ClientWebSocket, model: str, logging_obj: LiteLLMLoggingObj, - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, litellm_metadata: dict[str, Any] | None = None, api_key: str | None = None, api_base: str | None = None, diff --git a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py index c9f9d4e6baa..fec05621fa3 100644 --- a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,8 +1,115 @@ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Mapping, Optional, Protocol, Sequence, Union -from typing_extensions import TypedDict +from typing_extensions import Required, TypedDict -from ..utils import CompletionTokensDetails, PromptTokensDetailsWrapper, ServerToolUse +from ..llms.openai import ChatCompletionAudioDelta +from ..utils import ( + CompletionTokensDetails, + CompletionTokensDetailsWrapper, + Function, + PromptTokensDetailsWrapper, + ServerToolUse, + Usage, +) + + +class StreamingUsageDict(TypedDict, total=False): + prompt_tokens: Optional[int] + completion_tokens: Optional[int] + total_tokens: Optional[int] + reasoning_tokens: Optional[int] + prompt_tokens_details: Union[PromptTokensDetailsWrapper, dict[str, object], None] + completion_tokens_details: Union[CompletionTokensDetailsWrapper, dict[str, object], None] + server_tool_use: Union[ServerToolUse, dict[str, object], None] + cost: Optional[float] + cache_creation_input_tokens: Optional[int] + cache_read_input_tokens: Optional[int] + + +class StreamingToolCallFunctionDict(TypedDict, total=False): + name: Optional[str] + arguments: Optional[str] + provider_specific_fields: Optional[Mapping[str, object]] + + +class StreamingToolCallFunctionLike(Protocol): + name: Optional[str] + arguments: Optional[str] + provider_specific_fields: Optional[Mapping[str, object]] + + +class StreamingToolCallDict(TypedDict, total=False): + id: Optional[str] + type: Optional[str] + index: int + function: Union[StreamingToolCallFunctionDict, StreamingToolCallFunctionLike, None] + provider_specific_fields: Optional[Mapping[str, object]] + + +class StreamingToolCallLike(Protocol): + id: Optional[str] + type: Optional[str] + index: int + function: Optional[StreamingToolCallFunctionLike] + provider_specific_fields: Optional[Mapping[str, object]] + + +class StreamingThinkingBlockDict(TypedDict, total=False): + type: Optional[str] + thinking: Optional[str] + data: Optional[str] + signature: Optional[str] + + +class StreamingChunkDelta(TypedDict, total=False): + role: Required[Optional[str]] + content: Optional[str] + reasoning_content: Optional[str] + tool_calls: Sequence[Union[StreamingToolCallDict, StreamingToolCallLike]] + thinking_blocks: Optional[Sequence[StreamingThinkingBlockDict]] + audio: Optional[ChatCompletionAudioDelta] + + +class StreamingChunkChoice(TypedDict, total=False): + index: int + finish_reason: Optional[str] + delta: Required[StreamingChunkDelta] + + +class StreamingChunkDict(TypedDict, total=False): + _hidden_params: dict[str, object] + id: Required[str] + object: Required[str] + created: Required[int] + model: Required[str] + system_fingerprint: Optional[str] + choices: Required[list[StreamingChunkChoice]] + usage: Union[Usage, StreamingUsageDict, None] + + +class UsageChunkCalculation(TypedDict): + prompt_tokens: Optional[int] + completion_tokens: Optional[int] + cache_creation_input_tokens: Optional[int] + cache_read_input_tokens: Optional[int] + completion_tokens_details: Optional[CompletionTokensDetails] + prompt_tokens_details: Optional[PromptTokensDetailsWrapper] + cost: Optional[float] + + +class ToolCallAccumulator(TypedDict): + id: Optional[str] + name: Optional[str] + type: Optional[str] + arguments: list[str] + provider_specific_fields: Optional[dict[str, object]] + + +class ToolCallParams(TypedDict, total=False): + id: str + function: Function + type: str + provider_specific_fields: dict[str, object] class UsagePerChunk(TypedDict): diff --git a/litellm/types/passthrough_endpoints/managed_ids.py b/litellm/types/passthrough_endpoints/managed_ids.py new file mode 100644 index 00000000000..2a3845c1705 --- /dev/null +++ b/litellm/types/passthrough_endpoints/managed_ids.py @@ -0,0 +1,75 @@ +""" +Types for passthrough managed-ID rewriting. + +The Prisma client is an untyped runtime wrapper, so the managed-file / +managed-object rows and the table actions read by +``litellm.proxy.pass_through_endpoints.managed_id_rewriter`` are described +structurally here instead of being imported from generated stubs. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Protocol, TypedDict + +PrismaWhere = dict[str, object] +PrismaOrder = dict[str, str] | list[dict[str, str]] + + +class ManagedResourceOwner(Protocol): + """Ownership columns carried by every managed-resource row.""" + + @property + def created_by(self) -> str | None: ... + + @property + def team_id(self) -> str | None: ... + + +class ManagedResourceRow(ManagedResourceOwner, Protocol): + """A ``LiteLLM_ManagedFileTable`` or ``LiteLLM_ManagedObjectTable`` row. + + ``unified_file_id`` exists only on file rows and ``unified_object_id`` only + on object rows; each is read exclusively off the table it belongs to. + """ + + @property + def created_at(self) -> datetime | None: ... + + @property + def file_object(self) -> object: ... + + @property + def unified_file_id(self) -> str: ... + + @property + def unified_object_id(self) -> str: ... + + +class ManagedResourceTable(Protocol): + """The Prisma table actions the managed-ID rewriter issues.""" + + async def find_first(self, *, where: PrismaWhere) -> ManagedResourceRow | None: ... + + async def find_many( + self, + *, + where: PrismaWhere, + order: PrismaOrder | None = None, + take: int | None = None, + ) -> Sequence[ManagedResourceRow]: ... + + async def update(self, *, where: PrismaWhere, data: PrismaWhere) -> ManagedResourceRow | None: ... + + async def upsert(self, *, where: PrismaWhere, data: Mapping[str, PrismaWhere]) -> ManagedResourceRow: ... + + +class PassthroughListResponse(TypedDict): + """OpenAI-style paginated list body served from the managed-ID tables.""" + + object: str + data: Sequence[Mapping[str, object]] + first_id: str | None + last_id: str | None + has_more: bool diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index db0c75e26ab..711718924df 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,10 +1,19 @@ -from typing import Dict, List, Union, Any, Optional +from typing import Dict, List, TypedDict, Union, Any, Optional from pydantic import BaseModel, Field from ...router import ModelGroupInfo +class PrismaCompatibleCreateDBModel(TypedDict): + model_id: Optional[str] + model_name: str + litellm_params: str + model_info: str + created_by: str + updated_by: str + + class ModelGroupInfoProxy(ModelGroupInfo): is_public_model_group: bool = Field(default=False) health_status: Optional[str] = Field(default=None) diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 32e07f9e52f..8b8de7b51f5 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -1,7 +1,7 @@ -from typing import List, Literal, Optional, Union +from typing import List, Literal, Optional, Protocol, Union from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from pydantic import PrivateAttr +from pydantic import BaseModel, PrivateAttr from typing_extensions import Any, List, Optional, TypedDict from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject @@ -140,3 +140,13 @@ class DecodedResponseId(TypedDict, total=False): custom_llm_provider: Optional[str] model_id: Optional[str] response_id: str + + +class ApplyPatchToolCallLike(Protocol): + """Structural view of an openai ``ResponseApplyPatchToolCall`` used by the completion bridge.""" + + @property + def call_id(self) -> str: ... + + @property + def operation(self) -> BaseModel: ... diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index d3ef01940bb..8427210e633 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3104 + "limit": 3099 }, "ANN002": { "limit": 69 @@ -9,13 +9,13 @@ "limit": 831 }, "ANN201": { - "limit": 2137 + "limit": 2135 }, "ANN202": { - "limit": 941 + "limit": 935 }, "ANN204": { - "limit": 724 + "limit": 723 }, "ANN205": { "limit": 127 @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 1851 + "limit": 1762 }, "ASYNC230": { "limit": 14 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f071c381916..f76775fa261 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23350 + "limit": 23334 }, "LIT002": { - "limit": 27256 + "limit": 27251 }, "LIT003": { "limit": 292 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1105 + "limit": 1097 }, "LIT007": { "limit": 0 @@ -24,6 +24,6 @@ "limit": 1004 }, "LIT009": { - "limit": 2465 + "limit": 2458 } }