From a38ce9e58599c328a40a77efebc39eb785981554 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 4 Aug 2026 04:30:21 +0000 Subject: [PATCH] chore(typing): clear basedpyright Any errors in streaming, passthrough, and proxy hot paths --- basedpyright-code-budget.json | 34 +- .../streaming_chunk_builder_utils.py | 452 ++++++---- .../litellm_core_utils/streaming_handler.py | 640 ++++++++------ .../streaming_handler_types.py | 233 +++++ .../adapters/handler.py | 177 ++-- litellm/proxy/common_request_processing.py | 277 ++++-- .../hooks/parallel_request_limiter_v3.py | 455 ++++++---- .../model_management_endpoints.py | 216 +++-- litellm/proxy/management_endpoints/ui_sso.py | 2 +- .../llm_passthrough_endpoints.py | 2 +- .../managed_id_rewriter.py | 355 ++++---- .../managed_id_rewriter_types.py | 148 ++++ .../pass_through_endpoints.py | 365 +++++--- .../session_handler.py | 2 +- .../transformation.py | 523 ++++++----- .../responses/mcp/mcp_streaming_iterator.py | 6 +- litellm/responses/streaming_iterator.py | 834 ++++++++++-------- ruff-strict-budget.json | 28 +- .../test_model_management_endpoints.py | 3 +- type-discipline-budget.json | 8 +- 20 files changed, 2961 insertions(+), 1799 deletions(-) create mode 100644 litellm/litellm_core_utils/streaming_handler_types.py create mode 100644 litellm/proxy/pass_through_endpoints/managed_id_rewriter_types.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index f6dd90077b1..6a61dec8506 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,15 +1,15 @@ { "reportAny": { - "limit": 29813 + "limit": 28296 }, "reportArgumentType": { - "limit": 2645 + "limit": 2634 }, "reportAssignmentType": { "limit": 329 }, "reportAttributeAccessIssue": { - "limit": 516 + "limit": 511 }, "reportCallIssue": { "limit": 123 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 9473 + "limit": 9093 }, "reportFunctionMemberAccess": { "limit": 11 @@ -54,13 +54,13 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5855 + "limit": 5829 }, "reportMissingTypeArgument": { - "limit": 15852 + "limit": 15785 }, "reportMissingTypeStubs": { - "limit": 41 + "limit": 40 }, "reportOperatorIssue": { "limit": 0 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1079 + "limit": 1071 }, "reportOptionalOperand": { "limit": 0 @@ -84,13 +84,13 @@ "limit": 77 }, "reportPrivateUsage": { - "limit": 2437 + "limit": 2436 }, "reportRedeclaration": { "limit": 12 }, "reportReturnType": { - "limit": 219 + "limit": 218 }, "reportTypedDictNotRequiredAccess": { "limit": 27 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45324 + "limit": 45178 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40452 + "limit": 40239 }, "reportUnknownParameterType": { - "limit": 20309 + "limit": 20229 }, "reportUnknownVariableType": { - "limit": 31978 + "limit": 31822 }, "reportUnnecessaryCast": { - "limit": 177 + "limit": 172 }, "reportUnnecessaryComparison": { - "limit": 1021 + "limit": 1019 }, "reportUnnecessaryContains": { "limit": 7 }, "reportUnnecessaryIsInstance": { - "limit": 1204 + "limit": 1199 }, "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..114f98bc7c8 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,17 +1,24 @@ import base64 import time -from typing import TYPE_CHECKING, Any, Union, cast +from collections.abc import Sequence +from typing import TYPE_CHECKING, Literal + +from pydantic import TypeAdapter +from typing_extensions import Required, TypedDict from litellm._logging import verbose_logger from litellm.types.llms.openai import ( + AllMessageValues, ChatCompletionAssistantContentValue, ChatCompletionAudioDelta, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, ) from litellm.types.utils import ( CacheCreationTokenDetails, ChatCompletionAudioResponse, + ChatCompletionDeltaToolCall, ChatCompletionMessageToolCall, - Choices, CompletionTokensDetails, CompletionTokensDetailsWrapper, Function, @@ -20,57 +27,158 @@ from litellm.types.utils import ( ModelResponseStream, PromptTokensDetailsWrapper, ServerToolUse, + StreamingChoices, Usage, ) from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, ) - from litellm.types.llms.openai import ( - ChatCompletionRedactedThinkingBlock, - ChatCompletionThinkingBlock, - ) + + +RawHiddenParamsDict = dict[str, object] + + +class ChunkHiddenParamsDict(TypedDict, total=False): + created_at: int | float + custom_llm_provider: str | None + + +class UsageDict(TypedDict, total=False): + prompt_tokens: int + completion_tokens: int + total_tokens: int + reasoning_tokens: int | None + cache_creation_input_tokens: int | None + cache_read_input_tokens: int | None + cost: float | None + prompt_tokens_details: dict[str, object] | None + completion_tokens_details: dict[str, object] | None + server_tool_use: dict[str, object] | None + + +class ToolCallFunctionDict(TypedDict, total=False): + name: str | None + arguments: str + provider_specific_fields: dict[str, object] | None + + +class ToolCallDict(TypedDict, total=False): + id: str | None + type: str | None + index: int + function: ToolCallFunctionDict | Function | None + provider_specific_fields: dict[str, object] | None + + +class DeltaDict(TypedDict, total=False): + role: str | None + content: str | None + reasoning_content: str | None + function_call: FunctionCall | None + tool_calls: list[ToolCallDict | ChatCompletionDeltaToolCall] | None + audio: ChatCompletionAudioDelta | None + thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None + + +class ChoiceDict(TypedDict, total=False): + index: int + delta: Required[DeltaDict] + finish_reason: str | None + + +class ChunkDict(TypedDict, total=False): + id: Required[str] + object: Required[str] + created: Required[int] + model: Required[str] + choices: Required[list[ChoiceDict | StreamingChoices]] + system_fingerprint: str | None + usage: Usage | UsageDict | None + _hidden_params: RawHiddenParamsDict + + +class AccumulatedToolCallDict(TypedDict): + id: str | None + name: str | None + type: str | None + arguments: list[str] + provider_specific_fields: dict[str, object] | None + + +class UsageChunkCalculationDict(TypedDict): + prompt_tokens: int | None + completion_tokens: int | None + cache_creation_input_tokens: int | None + cache_read_input_tokens: int | None + completion_tokens_details: CompletionTokensDetails | None + prompt_tokens_details: PromptTokensDetailsWrapper | None + cost: float | None + + +class CompletionTokensDetailsDumpDict(TypedDict, total=False): + accepted_prediction_tokens: int | None + audio_tokens: int | None + reasoning_tokens: int | None + rejected_prediction_tokens: int | None + text_tokens: int | None + image_tokens: int | None + video_tokens: int | None + + +class UsageDumpDict(TypedDict, total=False): + prompt_tokens: int + completion_tokens: int + total_tokens: int + completion_tokens_details: dict[str, object] | None + prompt_tokens_details: dict[str, object] | None + server_tool_use: dict[str, object] | None + cost: float | None + cache_creation_input_tokens: int + cache_read_input_tokens: int + + +_CHUNK_HIDDEN_PARAMS_ADAPTER: TypeAdapter[ChunkHiddenParamsDict] = TypeAdapter(ChunkHiddenParamsDict) +_AUDIO_DELTA_ADAPTER: TypeAdapter[ChatCompletionAudioDelta | None] = TypeAdapter(ChatCompletionAudioDelta | None) +_COMPLETION_TOKENS_DETAILS_DUMP_ADAPTER: TypeAdapter[CompletionTokensDetailsDumpDict] = TypeAdapter( + CompletionTokensDetailsDumpDict +) +_USAGE_DUMP_ADAPTER: TypeAdapter[UsageDumpDict] = TypeAdapter(UsageDumpDict) class ChunkProcessor: - def __init__(self, chunks: list, messages: list | None = None): + def __init__(self, chunks: list[ChunkDict], messages: list[AllMessageValues] | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages self.first_chunk = chunks[0] - def _sort_chunks(self, chunks: list) -> list: + @staticmethod + def _chunk_hidden_params(chunk: "ChunkDict | ModelResponseStream") -> ChunkHiddenParamsDict: + candidate: object = ( + chunk.get("_hidden_params", {}) if isinstance(chunk, dict) else getattr(chunk, "_hidden_params", {}) + ) + if isinstance(candidate, dict): + return _CHUNK_HIDDEN_PARAMS_ADAPTER.validate_python(candidate) + return {} + + @staticmethod + def _chunk_created_at(chunk: "ChunkDict | ModelResponseStream") -> int | float: + return ChunkProcessor._chunk_hidden_params(chunk).get("created_at", float("inf")) + + def _sort_chunks(self, chunks: list[ChunkDict]) -> list[ChunkDict]: 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 = self._chunk_hidden_params(chunks[0]) if first_hidden_params.get("created_at"): - - def _created_at(chunk: Any) -> int | float: - if isinstance(chunk, dict): - params = chunk.get("_hidden_params", {}) - else: - params = getattr(chunk, "_hidden_params", {}) - if isinstance(params, dict): - return cast(int | float, params.get("created_at", float("inf"))) - return float("inf") - - return sorted(chunks, key=_created_at) + return sorted(chunks, key=self._chunk_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: ChunkDict | None = None ) -> ModelResponse: if chunk is None: return model_response @@ -82,17 +190,17 @@ class ChunkProcessor: @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, - chunks: list[Any], - logging_obj: Any | None = None, + chunks: list[ChunkDict], + logging_obj: "Logging | None" = None, ) -> None: if not chunks: return - model = getattr(response, "model", None) + model: str | None = getattr(response, "model", None) if not model: return - custom_llm_provider = None + custom_llm_provider: str | None = None if logging_obj is not None: custom_llm_provider = logging_obj.model_call_details.get("custom_llm_provider") @@ -126,7 +234,7 @@ class ChunkProcessor: ) @staticmethod - def _get_chunk_id(chunks: list[dict[str, Any]]) -> str: + def _get_chunk_id(chunks: Sequence[ChunkDict]) -> str: """ Chunks: [{"id": ""}, {"id": "1"}, {"id": "1"}] @@ -137,7 +245,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: Sequence[ChunkDict], first_chunk_model: str) -> str: """ Get the actual model from chunks, preferring a model that differs from the first chunk. @@ -153,7 +261,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: Sequence[ChunkDict]) -> ModelResponse: chunk = self.first_chunk id = ChunkProcessor._get_chunk_id(chunks) object = chunk["object"] @@ -164,15 +272,19 @@ class ChunkProcessor: system_fingerprint = chunk.get("system_fingerprint", None) first_chunk_with_choices = next((c for c in chunks if c.get("choices")), chunk) - role = first_chunk_with_choices["choices"][0]["delta"]["role"] + first_choice = first_chunk_with_choices["choices"][0] + if isinstance(first_choice, dict): + role = first_choice["delta"].get("role") + else: + role = first_choice.delta.role finish_reason = "stop" for chunk in chunks: if "choices" in chunk and len(chunk["choices"]) > 0: - chunk_finish_reason = None - if hasattr(chunk["choices"][0], "finish_reason"): - chunk_finish_reason = chunk["choices"][0].finish_reason - elif "finish_reason" in chunk["choices"][0]: - chunk_finish_reason = chunk["choices"][0]["finish_reason"] + finish_reason_choice = chunk["choices"][0] + if isinstance(finish_reason_choice, dict): + chunk_finish_reason = finish_reason_choice.get("finish_reason") + else: + chunk_finish_reason = finish_reason_choice.finish_reason if chunk_finish_reason is not None: finish_reason = chunk_finish_reason @@ -202,15 +314,18 @@ 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: Sequence[ChunkDict]) -> 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, AccumulatedToolCallDict] = {} # Map to store tool calls by index for chunk in tool_call_chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta", {}) - tool_calls = delta.get("tool_calls", []) + tool_calls: Sequence[ToolCallDict | ChatCompletionDeltaToolCall] + if isinstance(choice, dict): + tool_calls = choice.get("delta", {}).get("tool_calls") or [] + else: + tool_calls = choice.delta.tool_calls or [] for tool_call in tool_calls: # Handle both dict and object formats @@ -222,7 +337,7 @@ class ChunkProcessor: if isinstance(tool_call, dict): has_function = "function" in tool_call and tool_call["function"] is not None else: - has_function = hasattr(tool_call, "function") and tool_call.function is not None + has_function = getattr(tool_call, "function", None) is not None if not has_function: continue @@ -244,22 +359,26 @@ class ChunkProcessor: # Extract id, type, and function data (handle both dict and object) if isinstance(tool_call, dict): - if tool_call.get("id"): - tool_call_map[index]["id"] = tool_call["id"] - if tool_call.get("type"): - tool_call_map[index]["type"] = tool_call["type"] + tool_call_id = tool_call.get("id") + if tool_call_id: + tool_call_map[index]["id"] = tool_call_id + tool_call_type = tool_call.get("type") + if tool_call_type: + tool_call_map[index]["type"] = tool_call_type function = tool_call.get("function", {}) if isinstance(function, dict): - if function.get("name"): - tool_call_map[index]["name"] = function["name"] - if function.get("arguments"): - tool_call_map[index]["arguments"].append(function["arguments"]) + function_name = function.get("name") + if function_name: + tool_call_map[index]["name"] = function_name + function_arguments = function.get("arguments") + if function_arguments: + tool_call_map[index]["arguments"].append(function_arguments) else: # function is an object - if hasattr(function, "name") and function.name: + if function is not None and hasattr(function, "name") and function.name: tool_call_map[index]["name"] = function.name - if hasattr(function, "arguments") and function.arguments: + if function is not None and hasattr(function, "arguments") and function.arguments: tool_call_map[index]["arguments"].append(function.arguments) else: # tool_call is an object @@ -274,27 +393,24 @@ class ChunkProcessor: tool_call_map[index]["arguments"].append(tool_call.function.arguments) # Preserve provider_specific_fields from streaming chunks - provider_fields = None + provider_fields: dict[str, object] | None = 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") + if not provider_fields: + function_value = tool_call.get("function") + if isinstance(function_value, dict): + provider_fields = function_value.get("provider_specific_fields") else: - if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields: - provider_fields = tool_call.provider_specific_fields - elif ( - hasattr(tool_call, "function") - and hasattr(tool_call.function, "provider_specific_fields") - and tool_call.function.provider_specific_fields - ): - provider_fields = tool_call.function.provider_specific_fields + object_provider_fields = getattr(tool_call, "provider_specific_fields", None) or getattr( + getattr(tool_call, "function", None), "provider_specific_fields", None + ) + if isinstance(object_provider_fields, dict) and object_provider_fields: + provider_fields = object_provider_fields - if provider_fields: + if isinstance(provider_fields, dict) and provider_fields: # Merge provider_specific_fields if multiple chunks have them - if tool_call_map[index]["provider_specific_fields"] is None: - tool_call_map[index]["provider_specific_fields"] = {} - if isinstance(provider_fields, dict): - tool_call_map[index]["provider_specific_fields"].update(provider_fields) + existing_fields = tool_call_map[index]["provider_specific_fields"] or {} + tool_call_map[index]["provider_specific_fields"] = {**existing_fields, **provider_fields} # Convert the map to a list of tool calls for index in sorted(tool_call_map.keys()): @@ -308,39 +424,48 @@ class ChunkProcessor: name=tool_call_data["name"], ) - # Prepare params for ChatCompletionMessageToolCall - tool_call_params = { - "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"] - - tool_call = ChatCompletionMessageToolCall(**tool_call_params) - tool_calls_list.append(tool_call) + provider_specific_fields = tool_call_data["provider_specific_fields"] + if provider_specific_fields: + tool_calls_list.append( + ChatCompletionMessageToolCall( + id=tool_call_data["id"], + function=function, + type=tool_call_data["type"] or "function", + provider_specific_fields=provider_specific_fields, + ) + ) + else: + tool_calls_list.append( + ChatCompletionMessageToolCall( + id=tool_call_data["id"], + function=function, + type=tool_call_data["type"] or "function", + ) + ) return tool_calls_list - def get_combined_function_call_content(self, function_call_chunks: list[dict[str, Any]]) -> FunctionCall: - argument_list = [] - delta = function_call_chunks[0]["choices"][0]["delta"] - function_call = delta.get("function_call", "") - function_call_name = function_call.name + def get_combined_function_call_content(self, function_call_chunks: Sequence[ChunkDict]) -> FunctionCall: + argument_list: list[str] = [] + first_choice = function_call_chunks[0]["choices"][0] + if isinstance(first_choice, dict): + first_function_call = first_choice["delta"].get("function_call") + else: + first_function_call = getattr(first_choice.delta, "function_call", None) + function_call_name = first_function_call.name for chunk in function_call_chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta", {}) - function_call = delta.get("function_call", "") + if isinstance(choice, dict): + function_call = choice.get("delta", {}).get("function_call") + else: + function_call = getattr(choice.delta, "function_call", None) # Check if a function call is present if function_call: - # Now, function_call is expected to be a dictionary - arguments = function_call.arguments - argument_list.append(arguments) + argument_list.append(function_call.arguments) combined_arguments = "".join(argument_list) @@ -350,14 +475,16 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: Sequence[ChunkDict], delta_key: Literal["content", "reasoning_content"] = "content" ) -> ChatCompletionAssistantContentValue: content_list: list[str] = [] for chunk in chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta", {}) - content = delta.get(delta_key, "") + if isinstance(choice, dict): + content = choice.get("delta", {}).get(delta_key, "") + else: + content = getattr(choice.delta, delta_key, "") if content is None: continue # openai v1.0.0 sets content = None for chunks content_list.append(content) @@ -369,13 +496,8 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] - ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: - from litellm.types.llms.openai import ( - ChatCompletionRedactedThinkingBlock, - ChatCompletionThinkingBlock, - ) - + self, chunks: Sequence[ChunkDict] + ) -> list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None: thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = [] current_thinking_text_parts: list[str] = [] current_signature: str | None = None @@ -396,12 +518,13 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta", {}) - thinking = delta.get("thinking_blocks", None) - if thinking and isinstance(thinking, list): + if isinstance(choice, dict): + thinking = choice.get("delta", {}).get("thinking_blocks", None) + else: + thinking = getattr(choice.delta, "thinking_blocks", None) + if isinstance(thinking, list) and thinking: for thinking_block in thinking: - thinking_type = thinking_block.get("type", None) - if thinking_type and thinking_type == "redacted_thinking": + if thinking_block.get("type", None) == "redacted_thinking": _flush_thinking_block() redacted_data = thinking_block.get("data", None) if redacted_data: @@ -426,10 +549,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: Sequence[ChunkDict]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: Sequence[ChunkDict]) -> ChatCompletionAudioResponse: base64_data_list: list[str] = [] transcript_list: list[str] = [] expires_at: int | None = None @@ -438,8 +561,12 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta") or {} - audio: ChatCompletionAudioDelta | None = delta.get("audio") + audio: ChatCompletionAudioDelta | None + if isinstance(choice, dict): + delta: DeltaDict = choice.get("delta") or {} + audio = delta.get("audio") + else: + audio = _AUDIO_DELTA_ADAPTER.validate_python(getattr(choice.delta, "audio", None)) if audio is not None: for k, v in audio.items(): if k == "data" and v is not None and isinstance(v, str): @@ -459,9 +586,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) -> UsageChunkCalculationDict: + 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,30 +630,30 @@ 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 - if hasattr(chunk, "usage") and chunk.usage is not None: - usage_chunk = chunk.usage - elif "usage" in chunk: - usage_chunk = chunk["usage"] - elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr( - chunk, "_hidden_params" - ): - usage_chunk = chunk._hidden_params.get("usage", None) + def _extract_usage_chunk(chunk: ChunkDict | ModelResponse | ModelResponseStream) -> Usage | None: + usage_chunk: Usage | UsageDict | None = None + if isinstance(chunk, dict): + usage_chunk = chunk.get("usage") + elif hasattr(chunk, "usage"): + usage_chunk = getattr(chunk, "usage", None) + elif hasattr(chunk, "_hidden_params"): + hidden_usage: object = chunk._hidden_params.get("usage", None) + if isinstance(hidden_usage, Usage): + return hidden_usage + if isinstance(hidden_usage, dict): + return Usage(**hidden_usage) + return None if isinstance(usage_chunk, dict): return Usage(**usage_chunk) @@ -534,7 +661,7 @@ class ChunkProcessor: def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence[ChunkDict], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -601,24 +728,15 @@ class ChunkProcessor: server_tool_use = usage_chunk.server_tool_use else: server_tool_use = ServerToolUse.model_validate(usage_chunk.server_tool_use) - if ( - usage_chunk_dict["prompt_tokens_details"] is not None - and getattr( - usage_chunk_dict["prompt_tokens_details"], - "web_search_requests", - None, - ) - is not None - ): - web_search_requests = getattr( - usage_chunk_dict["prompt_tokens_details"], - "web_search_requests", - ) - - prompt_tokens_details = cast( - PromptTokensDetailsWrapper | None, + web_search_requests_value: int | None = getattr( usage_chunk_dict["prompt_tokens_details"], + "web_search_requests", + None, ) + if usage_chunk_dict["prompt_tokens_details"] is not None and web_search_requests_value is not None: + web_search_requests = web_search_requests_value + + 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 @@ -654,9 +772,10 @@ class ChunkProcessor: prompt_tokens_details: PromptTokensDetailsWrapper | None, current: CacheCreationTokenDetails | None, ) -> CacheCreationTokenDetails | None: - incoming = cast( - CacheCreationTokenDetails | None, - getattr(prompt_tokens_details, "cache_creation_token_details", None), + if prompt_tokens_details is None: + return current + incoming: CacheCreationTokenDetails | None = getattr( + prompt_tokens_details, "cache_creation_token_details", None ) if incoming is not None: return incoming @@ -669,9 +788,8 @@ class ChunkProcessor: ) -> PromptTokensDetailsWrapper | None: if prompt_tokens_details is None or cache_creation_token_details is None: return prompt_tokens_details - existing = cast( - CacheCreationTokenDetails | None, - getattr(prompt_tokens_details, "cache_creation_token_details", None), + existing: CacheCreationTokenDetails | None = getattr( + prompt_tokens_details, "cache_creation_token_details", None ) if existing is not None: return prompt_tokens_details @@ -679,7 +797,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence[ChunkDict], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -704,13 +822,7 @@ class ChunkProcessor: custom_llm_provider: str | None = None if chunks: - first_chunk = chunks[0] - if isinstance(first_chunk, dict): - hp = first_chunk.get("_hidden_params") - else: - hp = getattr(first_chunk, "_hidden_params", None) - if isinstance(hp, dict): - custom_llm_provider = hp.get("custom_llm_provider") + custom_llm_provider = ChunkProcessor._chunk_hidden_params(chunks[0]).get("custom_llm_provider") if custom_llm_provider == "anthropic" and completion_tokens == 1: return 0 @@ -718,10 +830,10 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence[ChunkDict], model: str, completion_output: str, - messages: list | None = None, + messages: list[AllMessageValues] | None = None, reasoning_tokens: int | None = None, ) -> Usage: """ @@ -773,7 +885,7 @@ class ChunkProcessor: if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() + **_COMPLETION_TOKENS_DETAILS_DUMP_ADAPTER.validate_python(completion_tokens_details.model_dump()) ) else: returned_usage.completion_tokens_details = completion_tokens_details @@ -806,12 +918,12 @@ class ChunkProcessor: # Return a new usage object with the new values - returned_usage = Usage(**returned_usage.model_dump()) + returned_usage = Usage(**_USAGE_DUMP_ADAPTER.validate_python(returned_usage.model_dump())) return returned_usage -def concatenate_base64_list(base64_strings: list[str]) -> str: +def concatenate_base64_list(base64_strings: Sequence[str]) -> str: """ Concatenates a list of base64-encoded strings. diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fb7d06bee93..0232db246b6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -6,7 +6,7 @@ import logging import threading import time import traceback -from collections.abc import AsyncIterator, Callable, Iterator +from collections.abc import AsyncIterator, Callable, Iterator, Sequence from dataclasses import dataclass from typing import ( Any, @@ -17,7 +17,8 @@ from typing import ( import anyio import httpx -from pydantic import BaseModel +from openai.types.completion import Completion +from pydantic import BaseModel, TypeAdapter import litellm from litellm import verbose_logger @@ -26,11 +27,37 @@ from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject +from litellm.litellm_core_utils.streaming_handler_types import ( + AI21StreamChunk, + AlephAlphaStreamChunk, + AzureStreamChunk, + BasetenModelOutput, + BasetenStreamChunk, + HasCompletionUsage, + HasFunctionsAttr, + HasModelAttr, + HasStatusCode, + HasTextChoices, + MaritalkStreamChunk, + NlpCloudStreamChunk, + OpenAIChatParsedChunk, + ParsedProviderChunk, + PredibaseStreamChunk, + SimpleParsedChunk, + StreamingCompletionObj, + TextCompletionChoiceLike, + TextCompletionParsedChunk, + TritonStreamChunk, + VertexProtoChunkLike, + VllmRequestOutputLike, +) from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.types.llms.openai import OpenAIChatCompletionChunk from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( Delta, + Function, + FunctionCall, LlmProviders, ModelResponse, ModelResponseStream, @@ -55,8 +82,13 @@ _SYNC_ITER_EXHAUSTED = object() _GCHUNK_FIELDS: frozenset = frozenset(GChunk.__annotations__) +_DICT_STR_OBJECT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) +_LIST_OBJECT_ADAPTER: TypeAdapter[list[object]] = TypeAdapter(list[object]) +_OBJECT_ADAPTER: TypeAdapter[object] = TypeAdapter(object) +_STR_ADAPTER: TypeAdapter[str] = TypeAdapter(str) -def _next_sync_or_exhausted(it: Any) -> Any: + +def _next_sync_or_exhausted(it: Any) -> object: # any-ok: sync provider streams are untyped third-party iterators """ Call next(it) from a thread and return _SYNC_ITER_EXHAUSTED on StopIteration. @@ -65,12 +97,20 @@ def _next_sync_or_exhausted(it: Any) -> Any: Returning a sentinel instead keeps StopIteration out of the coroutine boundary. """ try: - return next(it) + return next(it) # any-ok: third-party iterator protocol is only known at runtime except StopIteration: return _SYNC_ITER_EXHAUSTED -def is_async_iterable(obj: Any) -> bool: +def _delta_or_none(choice: StreamingChoices) -> Delta | None: + return choice.delta + + +def _arguments_or_none(function_obj: "FunctionCall | Function") -> str | None: + return function_obj.arguments + + +def is_async_iterable(obj: object) -> bool: """ Check if an object is an async iterable (can be used with 'async for'). @@ -93,12 +133,12 @@ def print_verbose(print_statement): @dataclass(frozen=True, slots=True) class _ProviderChunkParsed: - response_obj: dict[str, Any] + response_obj: ParsedProviderChunk @dataclass(frozen=True, slots=True) class _ProviderChunkEarlyReturn: - value: Any + value: ModelResponseStream | None _ProviderChunkResult = Union[_ProviderChunkParsed, _ProviderChunkEarlyReturn] @@ -109,7 +149,7 @@ class CustomStreamWrapper: self, completion_stream, model, - logging_obj: Any, + logging_obj: LiteLLMLoggingObject, custom_llm_provider: str | None = None, stream_options=None, make_call: Callable | None = None, @@ -147,20 +187,21 @@ class CustomStreamWrapper: self.holding_chunk = "" self.complete_response = "" self.response_uptil_now = "" - _model_info: dict = litellm_params.model_info or {} + _model_info = _DICT_STR_OBJECT_ADAPTER.validate_python(litellm_params.model_info or {}) + _raw_litellm_params = self.logging_obj.model_call_details.get("litellm_params", None) _api_base = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=_raw_litellm_params if isinstance(_raw_litellm_params, dict) else {}, ) - self._hidden_params = { + self._hidden_params: dict[str, object] = { "model_id": (_model_info.get("id", None)), "api_base": _api_base, } # returned as x-litellm-model-id response header in proxy - self._hidden_params["additional_headers"] = process_response_headers( - _response_headers or {} + self._hidden_params["additional_headers"] = _DICT_STR_OBJECT_ADAPTER.validate_python( + process_response_headers(_response_headers or {}) ) # GUARANTEE OPENAI HEADERS IN RESPONSE self._response_headers = _response_headers @@ -188,7 +229,7 @@ class CustomStreamWrapper: # Snapshot assumes self._hidden_params is populated from litellm_params # at init and never mutated during the stream. If that ever changes, # this cache must be removed. - self._base_hidden_params: dict[str, Any] = { + self._base_hidden_params: dict[str, object] = { **self._hidden_params, "response_cost": None, } @@ -307,7 +348,7 @@ class CustomStreamWrapper: llm_provider="", ) - def check_special_tokens(self, chunk: str, finish_reason: str | None): + def check_special_tokens(self, chunk: str | None, finish_reason: str | None) -> tuple[bool, str | None]: """ Output parse / special tokens for sagemaker + hf streaming. """ @@ -315,6 +356,9 @@ class CustomStreamWrapper: if self.custom_llm_provider != "sagemaker": return hold, chunk + if chunk is None: + return hold, chunk + if finish_reason: for token in self.special_tokens: if token in chunk: @@ -342,28 +386,28 @@ class CustomStreamWrapper: self.holding_chunk = "" return hold, curr_chunk - def handle_predibase_chunk(self, chunk): + def handle_predibase_chunk(self, chunk) -> SimpleParsedChunk: try: if not isinstance(chunk, str): chunk = chunk.decode("utf-8") # DO NOT REMOVE this: This is required for HF inference API + Streaming - text = "" + text: str | None = "" is_finished = False finish_reason = "" print_verbose(f"chunk: {chunk}") if chunk.startswith("data:"): - data_json = json.loads(chunk[5:]) + data_json = PredibaseStreamChunk.model_validate(json.loads(chunk[5:])) print_verbose(f"data json: {data_json}") - if "token" in data_json and "text" in data_json["token"]: - text = data_json["token"]["text"] - if data_json.get("details", False) and data_json["details"].get("finish_reason", False): + if data_json.token is not None and data_json.token.text is not None: + text = data_json.token.text + if data_json.details is not None and data_json.details.finish_reason: is_finished = True - finish_reason = data_json["details"]["finish_reason"] - elif data_json.get("generated_text", False): # if full generated text exists, then stream is complete + finish_reason = data_json.details.finish_reason + elif data_json.generated_text: # if full generated text exists, then stream is complete text = "" # don't return the final bos token is_finished = True finish_reason = "stop" - elif data_json.get("error", False): - raise Exception(data_json.get("error")) + elif data_json.error: + raise Exception(data_json.error) return { "text": text, "is_finished": is_finished, @@ -379,11 +423,12 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_ai21_chunk(self, chunk): # fake streaming + def handle_ai21_chunk(self, chunk) -> SimpleParsedChunk: # fake streaming chunk = chunk.decode("utf-8") - data_json = json.loads(chunk) + raw_json = _OBJECT_ADAPTER.validate_python(json.loads(chunk)) try: - text = data_json["completions"][0]["data"]["text"] + data_json = AI21StreamChunk.model_validate(raw_json) + text = data_json.completions[0].data.text is_finished = True finish_reason = "stop" return { @@ -394,11 +439,12 @@ class CustomStreamWrapper: except Exception: raise ValueError(f"Unable to parse response. Original response: {chunk}") - def handle_maritalk_chunk(self, chunk): # fake streaming + def handle_maritalk_chunk(self, chunk) -> SimpleParsedChunk: # fake streaming chunk = chunk.decode("utf-8") - data_json = json.loads(chunk) + raw_json = _OBJECT_ADAPTER.validate_python(json.loads(chunk)) try: - text = data_json["answer"] + data_json = MaritalkStreamChunk.model_validate(raw_json) + text = data_json.answer is_finished = True finish_reason = "stop" return { @@ -409,7 +455,7 @@ class CustomStreamWrapper: except Exception: raise ValueError(f"Unable to parse response. Original response: {chunk}") - def handle_nlp_cloud_chunk(self, chunk): + def handle_nlp_cloud_chunk(self, chunk) -> SimpleParsedChunk: text = "" is_finished = False finish_reason = "" @@ -417,8 +463,8 @@ class CustomStreamWrapper: if self.model and "dolphin" in self.model: chunk = self.process_chunk(chunk=chunk) else: - data_json = json.loads(chunk) - chunk = data_json["generated_text"] + data_json = NlpCloudStreamChunk.model_validate(json.loads(chunk)) + chunk = data_json.generated_text text = chunk if "[DONE]" in text: text = text.replace("[DONE]", "") @@ -432,11 +478,12 @@ class CustomStreamWrapper: except Exception: raise ValueError(f"Unable to parse response. Original response: {chunk}") - def handle_aleph_alpha_chunk(self, chunk): + def handle_aleph_alpha_chunk(self, chunk) -> SimpleParsedChunk: chunk = chunk.decode("utf-8") - data_json = json.loads(chunk) + raw_json = _OBJECT_ADAPTER.validate_python(json.loads(chunk)) try: - text = data_json["completions"][0]["completion"] + data_json = AlephAlphaStreamChunk.model_validate(raw_json) + text = data_json.completions[0].completion is_finished = True finish_reason = "stop" return { @@ -447,10 +494,10 @@ class CustomStreamWrapper: except Exception: raise ValueError(f"Unable to parse response. Original response: {chunk}") - def handle_azure_chunk(self, chunk): + def handle_azure_chunk(self, chunk) -> SimpleParsedChunk: is_finished = False finish_reason = "" - text = "" + text: str | None = "" print_verbose(f"chunk: {chunk}") if "data: [DONE]" in chunk: text = "" @@ -462,14 +509,15 @@ class CustomStreamWrapper: "finish_reason": finish_reason, } elif chunk.startswith("data:"): - data_json = json.loads(chunk[5:]) # chunk.startswith("data:"): + raw_json = _OBJECT_ADAPTER.validate_python(json.loads(chunk[5:])) # chunk.startswith("data:"): try: - if len(data_json["choices"]) > 0: - delta = data_json["choices"][0]["delta"] - text = "" if delta is None else delta.get("content", "") - if data_json["choices"][0].get("finish_reason", None): + data_json = AzureStreamChunk.model_validate(raw_json) + if len(data_json.choices) > 0: + delta = data_json.choices[0].delta + text = "" if delta is None else delta.content + if data_json.choices[0].finish_reason: is_finished = True - finish_reason = data_json["choices"][0]["finish_reason"] + finish_reason = data_json.choices[0].finish_reason print_verbose(f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}") return { "text": text, @@ -487,7 +535,7 @@ class CustomStreamWrapper: "finish_reason": finish_reason, } - def handle_replicate_chunk(self, chunk): + def handle_replicate_chunk(self, chunk) -> SimpleParsedChunk: try: text = "" is_finished = False @@ -508,7 +556,7 @@ class CustomStreamWrapper: except Exception: raise ValueError(f"Unable to parse response. Original response: {chunk}") - def handle_openai_chat_completion_chunk(self, chunk): + def handle_openai_chat_completion_chunk(self, chunk) -> OpenAIChatParsedChunk | None: try: str_line = chunk text = "" @@ -546,12 +594,12 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_azure_text_completion_chunk(self, chunk): + def handle_azure_text_completion_chunk(self, chunk: object) -> SimpleParsedChunk: try: text = "" is_finished = False finish_reason = None - choices = getattr(chunk, "choices", []) + choices: Sequence[TextCompletionChoiceLike] = chunk.choices if isinstance(chunk, HasTextChoices) else () if len(choices) > 0: text = choices[0].text if choices[0].finish_reason is not None: @@ -566,19 +614,18 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_openai_text_completion_chunk(self, chunk): + def handle_openai_text_completion_chunk(self, chunk: object) -> TextCompletionParsedChunk: try: text = "" is_finished = False finish_reason = None - usage = None - choices = getattr(chunk, "choices", []) + choices: Sequence[TextCompletionChoiceLike] = chunk.choices if isinstance(chunk, HasTextChoices) else () if len(choices) > 0: text = choices[0].text if choices[0].finish_reason is not None: is_finished = True finish_reason = choices[0].finish_reason - usage = getattr(chunk, "usage", None) + usage = chunk.usage if isinstance(chunk, HasCompletionUsage) else None return { "text": text, "is_finished": is_finished, @@ -589,28 +636,25 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_baseten_chunk(self, chunk): + def handle_baseten_chunk(self, chunk) -> str | None: try: chunk = chunk.decode("utf-8") if len(chunk) > 0: if chunk.startswith("data:"): - data_json = json.loads(chunk[5:]) - if "token" in data_json and "text" in data_json["token"]: - return data_json["token"]["text"] + data_json = BasetenStreamChunk.model_validate(json.loads(chunk[5:])) + if data_json.token is not None and "text" in data_json.token.model_fields_set: + return data_json.token.text else: return "" - data_json = json.loads(chunk) - if "model_output" in data_json: - if ( - isinstance(data_json["model_output"], dict) - and "data" in data_json["model_output"] - and isinstance(data_json["model_output"]["data"], list) - ): - return data_json["model_output"]["data"][0] - elif isinstance(data_json["model_output"], str): - return data_json["model_output"] - elif "completion" in data_json and isinstance(data_json["completion"], str): - return data_json["completion"] + data_json = BasetenStreamChunk.model_validate(json.loads(chunk)) + if "model_output" in data_json.model_fields_set: + model_output = data_json.model_output + if isinstance(model_output, BasetenModelOutput) and model_output.data is not None: + return model_output.data[0] + elif isinstance(model_output, str): + return model_output + elif "completion" in data_json.model_fields_set and data_json.completion is not None: + return data_json.completion else: raise ValueError(f"Unable to parse response. Original response: {chunk}") else: @@ -621,54 +665,53 @@ class CustomStreamWrapper: verbose_logger.exception(f"litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - {e!s}") return "" - def handle_triton_stream(self, chunk): + def handle_triton_stream(self, chunk) -> SimpleParsedChunk: try: if isinstance(chunk, dict): - parsed_response = chunk + parsed_response = TritonStreamChunk.model_validate(chunk) elif isinstance(chunk, (str, bytes)): if isinstance(chunk, bytes): chunk = chunk.decode("utf-8") if "text_output" in chunk: response = CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" response = response.strip() - parsed_response = json.loads(response) + parsed_response = TritonStreamChunk.model_validate(json.loads(response)) else: return { "text": "", "is_finished": False, + "finish_reason": None, "prompt_tokens": 0, "completion_tokens": 0, } else: print_verbose(f"chunk: {chunk} (Type: {type(chunk)})") raise ValueError(f"Unable to parse response. Original response: {chunk}") - text = parsed_response.get("text_output", "") - finish_reason = parsed_response.get("stop_reason") - is_finished = parsed_response.get("is_finished", False) return { - "text": text, - "is_finished": is_finished, - "finish_reason": finish_reason, - "prompt_tokens": parsed_response.get("input_token_count", 0), - "completion_tokens": parsed_response.get("generated_token_count", 0), + "text": parsed_response.text_output, + "is_finished": parsed_response.is_finished, + "finish_reason": parsed_response.stop_reason, + "prompt_tokens": parsed_response.input_token_count, + "completion_tokens": parsed_response.generated_token_count, } - return {"text": "", "is_finished": False} except Exception as e: raise e - def model_response_creator(self, chunk: dict | None = None, hidden_params: dict | None = None): + def model_response_creator( + self, chunk: dict | None = None, hidden_params: dict | None = None + ) -> ModelResponseStream: _model = self._cached_model_name _logging_obj_llm_provider = self._cached_logging_llm_provider if chunk is None: - args: dict[str, Any] = {"model": _model} + args: dict[str, object] = {"model": _model} else: chunk.pop("model", None) args = {"model": _model} if chunk: args.update({k: v for k, v in chunk.items() if k != "stream"}) - model_response = ModelResponseStream(**args) + model_response = ModelResponseStream.model_validate(args) if self.response_id is not None: model_response.id = self.response_id if self.system_fingerprint is not None: @@ -740,15 +783,15 @@ class CustomStreamWrapper: provider_specific_fields = getattr(original_chunk, "provider_specific_fields", None) if provider_specific_fields is not None: model_response.provider_specific_fields = provider_specific_fields - for k, v in provider_specific_fields.items(): + for k, v in _DICT_STR_OBJECT_ADAPTER.validate_python(provider_specific_fields).items(): setattr(model_response, k, v) return model_response def is_chunk_non_empty( self, - completion_obj: dict[str, Any], + completion_obj: StreamingCompletionObj, model_response: ModelResponseStream, - response_obj: dict[str, Any], + response_obj: ParsedProviderChunk, ) -> bool: if ( "content" in completion_obj @@ -761,8 +804,8 @@ class CustomStreamWrapper: or ("function_call" in completion_obj and completion_obj["function_call"] is not None) or ( "tool_calls" in model_response.choices[0].delta - and model_response.choices[0].delta["tool_calls"] is not None - and len(model_response.choices[0].delta["tool_calls"]) > 0 + and model_response.choices[0].delta.tool_calls is not None + and len(model_response.choices[0].delta.tool_calls) > 0 ) or ( "function_call" in model_response.choices[0].delta @@ -804,7 +847,7 @@ class CustomStreamWrapper: _initial_delta = model_response.choices[0].delta.model_dump() _initial_delta.pop("role", None) - model_response.choices[0].delta = Delta(**_initial_delta) + model_response.choices[0].delta = Delta.model_validate(_initial_delta) return model_response def _has_special_delta_content(self, model_response: ModelResponseStream) -> bool: @@ -872,10 +915,10 @@ class CustomStreamWrapper: def return_processed_chunk_logic( # noqa: C901 self, - completion_obj: dict[str, Any], + completion_obj: StreamingCompletionObj, model_response: ModelResponseStream, - response_obj: dict[str, Any], - ): + response_obj: ParsedProviderChunk, + ) -> ModelResponseStream | None: from litellm.litellm_core_utils.core_helpers import ( preserve_upstream_non_openai_attributes, ) @@ -897,15 +940,14 @@ class CustomStreamWrapper: choices = [] for choice in original_chunk.choices: try: - if isinstance(choice, BaseModel): - choice_json = choice.model_dump() # type: ignore - choice_json.pop( - "finish_reason", None - ) # for mistral etc. which return a value in their last chunk (not-openai compatible). - choices.append(StreamingChoices(**choice_json)) + choice_json = choice.model_dump() + choice_json.pop( + "finish_reason", None + ) # for mistral etc. which return a value in their last chunk (not-openai compatible). + choices.append(StreamingChoices.model_validate(choice_json)) except Exception: choices.append(StreamingChoices()) - setattr(model_response, "choices", choices) + model_response.choices = choices else: return model_response.system_fingerprint = original_chunk.system_fingerprint @@ -931,8 +973,9 @@ class CustomStreamWrapper: if self.sent_first_chunk is False: completion_obj["role"] = "assistant" self.sent_first_chunk = True - if response_obj.get("provider_specific_fields") is not None: - completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] + _response_psf = response_obj.get("provider_specific_fields") + if _response_psf is not None: + completion_obj["provider_specific_fields"] = _response_psf model_response.choices[0].delta = Delta(**completion_obj) _index: int | None = completion_obj.get("index") if _index is not None: @@ -1030,11 +1073,12 @@ class CustomStreamWrapper: def _dispatch_provider_chunk( self, - chunk: Any, + chunk: object, model_response: ModelResponseStream, - completion_obj: dict[str, Any], + completion_obj: StreamingCompletionObj, ) -> _ProviderChunkResult: - response_obj: dict[str, Any] = {} + empty_response_obj: ParsedProviderChunk = {} + response_obj = empty_response_obj if ( isinstance(chunk, ModelResponseStream) and self.custom_llm_provider is not None @@ -1081,11 +1125,7 @@ class CustomStreamWrapper: self.intermittent_finish_reason = anthropic_response_obj["finish_reason"] if anthropic_response_obj["usage"] is not None: - setattr( - model_response, - "usage", - litellm.Usage(**anthropic_response_obj["usage"]), - ) + model_response.usage = litellm.Usage(**anthropic_response_obj["usage"]) if "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None: completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] @@ -1094,10 +1134,20 @@ class CustomStreamWrapper: "provider_specific_fields" in anthropic_response_obj and anthropic_response_obj["provider_specific_fields"] is not None ): - for key, value in anthropic_response_obj["provider_specific_fields"].items(): + for key, value in _DICT_STR_OBJECT_ADAPTER.validate_python( + anthropic_response_obj["provider_specific_fields"] + ).items(): setattr(model_response, key, value) - response_obj = cast(dict[str, Any], anthropic_response_obj) + anthropic_parsed: ParsedProviderChunk = { + "text": anthropic_response_obj["text"], + "is_finished": anthropic_response_obj["is_finished"], + "finish_reason": anthropic_response_obj["finish_reason"], + "usage": anthropic_response_obj["usage"], + "tool_use": anthropic_response_obj.get("tool_use"), + "provider_specific_fields": anthropic_response_obj.get("provider_specific_fields"), + } + response_obj = anthropic_parsed elif self.model == "replicate" or self.custom_llm_provider == "replicate": response_obj = self.handle_replicate_chunk(chunk) completion_obj["content"] = response_obj["text"] @@ -1121,7 +1171,12 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider and self.custom_llm_provider == "vllm": - completion_obj["content"] = chunk[0].outputs[0].text + if not isinstance(chunk, list): + raise ValueError(f"Unable to parse response. Original response: {chunk}") + first_output = _OBJECT_ADAPTER.validate_python(chunk[0]) + if not isinstance(first_output, VllmRequestOutputLike): + raise ValueError(f"Unable to parse response. Original response: {chunk}") + completion_obj["content"] = first_output.outputs[0].text elif ( self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha" ): # aleph alpha doesn't provide streaming @@ -1143,31 +1198,26 @@ class CustomStreamWrapper: raise Exception("An unknown error occurred with the stream") self.received_finish_reason = "stop" elif self.custom_llm_provider == "vertex_ai" and not isinstance(chunk, ModelResponseStream): - chunk = cast(Any, chunk) - import proto # type: ignore - - if hasattr(chunk, "candidates") is True: + if isinstance(chunk, VertexProtoChunkLike): + vertex_chunk = chunk try: try: - completion_obj["content"] = chunk.text # type: ignore + vertex_text = _OBJECT_ADAPTER.validate_python(getattr(vertex_chunk, "text")) + if isinstance(vertex_text, str): + completion_obj["content"] = vertex_text except Exception as e: original_exception = e if "Part has no text." in str(e): ## check for function calling - function_call = ( - chunk.candidates[0].content.parts[0].function_call # type: ignore - ) + function_call = vertex_chunk.candidates[0].content.parts[0].function_call - args_dict = {} + args_dict: dict[str, object] = {} # Check if it's a RepeatedComposite instance for key, val in function_call.args.items(): - if isinstance( - val, - proto.marshal.collections.repeated.RepeatedComposite, # type: ignore - ): + if isinstance(val, Sequence) and not isinstance(val, str): # If so, convert to list - args_dict[key] = [v for v in val] + args_dict[key] = list(val) else: args_dict[key] = val @@ -1191,20 +1241,18 @@ class CustomStreamWrapper: _streaming_response = StreamingChoices(delta=_delta_obj) _model_response = ModelResponseStream() _model_response.choices = [_streaming_response] - response_obj = {"original_chunk": _model_response} + vertex_parsed: ParsedProviderChunk = {"original_chunk": _model_response} + response_obj = vertex_parsed else: raise original_exception if ( - hasattr(chunk.candidates[0], "finish_reason") # type: ignore - and chunk.candidates[0].finish_reason.name # type: ignore - != "FINISH_REASON_UNSPECIFIED" + hasattr(vertex_chunk.candidates[0], "finish_reason") + and vertex_chunk.candidates[0].finish_reason.name != "FINISH_REASON_UNSPECIFIED" ): # every non-final chunk in vertex ai has this - self.received_finish_reason = map_finish_reason( # type: ignore - chunk.candidates[0].finish_reason.name - ) + self.received_finish_reason = map_finish_reason(vertex_chunk.candidates[0].finish_reason.name) except Exception: - if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore - raise Exception(f"The response was blocked by VertexAI. {chunk!s}") + if vertex_chunk.candidates[0].finish_reason.name == "SAFETY": + raise Exception(f"The response was blocked by VertexAI. {vertex_chunk!s}") else: completion_obj["content"] = str(chunk) elif self.custom_llm_provider == "petals": @@ -1214,23 +1262,22 @@ class CustomStreamWrapper: else: self.received_finish_reason = "stop" chunk_size = 30 - stream = cast(Any, self.completion_stream) + stream = _STR_ADAPTER.validate_python(self.completion_stream) new_chunk = stream[:chunk_size] completion_obj["content"] = new_chunk - self.completion_stream = stream[chunk_size:] + setattr(self, "completion_stream", stream[chunk_size:]) elif self.custom_llm_provider == "palm": # fake streaming - response_obj = {} if self.completion_stream is None or len(self.completion_stream) == 0: if self.received_finish_reason is not None: raise StopIteration else: self.received_finish_reason = "stop" chunk_size = 30 - stream = cast(Any, self.completion_stream) + stream = _STR_ADAPTER.validate_python(self.completion_stream) new_chunk = stream[:chunk_size] completion_obj["content"] = new_chunk - self.completion_stream = stream[chunk_size:] + setattr(self, "completion_stream", stream[chunk_size:]) elif self.custom_llm_provider == "triton": response_obj = self.handle_triton_stream(chunk) completion_obj["content"] = response_obj["text"] @@ -1238,42 +1285,41 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider == "text-completion-openai": - response_obj = self.handle_openai_text_completion_chunk(chunk) - completion_obj["content"] = response_obj["text"] + text_completion_response_obj = self.handle_openai_text_completion_chunk(chunk) + response_obj = text_completion_response_obj + completion_obj["content"] = text_completion_response_obj["text"] print_verbose(f"completion obj content: {completion_obj['content']}") - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - if response_obj["usage"] is not None: - setattr( - model_response, - "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].prompt_tokens, - completion_tokens=response_obj["usage"].completion_tokens, - total_tokens=response_obj["usage"].total_tokens, - ), + if text_completion_response_obj["is_finished"]: + self.received_finish_reason = text_completion_response_obj["finish_reason"] + _text_completion_usage = text_completion_response_obj["usage"] + if _text_completion_usage is not None: + model_response.usage = litellm.Usage( + prompt_tokens=_text_completion_usage.prompt_tokens, + completion_tokens=_text_completion_usage.completion_tokens, + total_tokens=_text_completion_usage.total_tokens, ) elif self.custom_llm_provider == "text-completion-codestral": if not isinstance(chunk, str): raise ValueError(f"chunk is not a string: {chunk}") - response_obj = cast( - dict[str, Any], - litellm.CodestralTextCompletionConfig()._chunk_parser(chunk), - ) - completion_obj["content"] = response_obj["text"] + codestral_response_obj = litellm.CodestralTextCompletionConfig()._chunk_parser(chunk) + _codestral_original_chunk = codestral_response_obj.get("original_chunk") + codestral_parsed: ParsedProviderChunk = { + "text": codestral_response_obj["text"], + "is_finished": codestral_response_obj["is_finished"], + "finish_reason": codestral_response_obj["finish_reason"], + "logprobs": codestral_response_obj.get("logprobs"), + "original_chunk": ( + _codestral_original_chunk if isinstance(_codestral_original_chunk, ModelResponseStream) else None + ), + } + response_obj = codestral_parsed + completion_obj["content"] = codestral_response_obj["text"] print_verbose(f"completion obj content: {completion_obj['content']}") - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - if "usage" in response_obj is not None: - setattr( - model_response, - "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].prompt_tokens, - completion_tokens=response_obj["usage"].completion_tokens, - total_tokens=response_obj["usage"].total_tokens, - ), - ) + if codestral_response_obj["is_finished"]: + self.received_finish_reason = codestral_response_obj["finish_reason"] + _codestral_usage = codestral_response_obj.get("usage") + if _codestral_usage is not None: + model_response.usage = litellm.Usage.model_validate(_codestral_usage.model_dump()) elif self.custom_llm_provider == "azure_text": response_obj = self.handle_azure_text_completion_chunk(chunk) completion_obj["content"] = response_obj["text"] @@ -1281,92 +1327,94 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] elif self.custom_llm_provider == "cached_response": - chunk = cast(ModelResponseStream, chunk) - chunk_finish_reason = chunk.choices[0].finish_reason - response_obj = { - "text": chunk.choices[0].delta.content, + cached_chunk = ModelResponseStream.model_validate(chunk) + chunk_finish_reason = cached_chunk.choices[0].finish_reason + cached_tool_calls = ( + cached_chunk.choices[0].delta.tool_calls + if hasattr(cached_chunk.choices[0].delta, "tool_calls") + else None + ) + cached_parsed: ParsedProviderChunk = { + "text": cached_chunk.choices[0].delta.content, "is_finished": chunk_finish_reason is not None, "finish_reason": chunk_finish_reason, - "original_chunk": chunk, - "tool_calls": ( - chunk.choices[0].delta.tool_calls if hasattr(chunk.choices[0].delta, "tool_calls") else None - ), + "original_chunk": cached_chunk, + "tool_calls": cached_tool_calls, } + response_obj = cached_parsed - completion_obj["content"] = response_obj["text"] - if response_obj["tool_calls"] is not None: - completion_obj["tool_calls"] = response_obj["tool_calls"] + completion_obj["content"] = cached_chunk.choices[0].delta.content + if cached_tool_calls is not None: + completion_obj["tool_calls"] = cached_tool_calls print_verbose(f"completion obj content: {completion_obj['content']}") - if hasattr(chunk, "id"): - model_response.id = chunk.id - self.response_id = chunk.id - if hasattr(chunk, "system_fingerprint"): - self.system_fingerprint = chunk.system_fingerprint - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] + if hasattr(cached_chunk, "id"): + model_response.id = cached_chunk.id + self.response_id = cached_chunk.id + if hasattr(cached_chunk, "system_fingerprint"): + self.system_fingerprint = cached_chunk.system_fingerprint + if chunk_finish_reason is not None: + self.received_finish_reason = chunk_finish_reason else: # openai / azure chat model if self.custom_llm_provider in [ LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value, ]: - if isinstance(chunk, BaseModel) and hasattr(chunk, "model"): + if isinstance(chunk, BaseModel) and isinstance(chunk, HasModelAttr): # for azure, we need to pass the model from the original chunk - self.model = getattr(chunk, "model", self.model) - response_obj = self.handle_openai_chat_completion_chunk(chunk) - if response_obj is None: + setattr(self, "model", chunk.model) + openai_response_obj = self.handle_openai_chat_completion_chunk(chunk) + if openai_response_obj is None: return _ProviderChunkEarlyReturn(None) - completion_obj["content"] = response_obj["text"] - self.intermittent_finish_reason = response_obj.get("finish_reason", None) - if response_obj["is_finished"]: - if response_obj["finish_reason"] == "error": + response_obj = openai_response_obj + completion_obj["content"] = openai_response_obj["text"] + self.intermittent_finish_reason = openai_response_obj.get("finish_reason", None) + if openai_response_obj["is_finished"]: + if openai_response_obj["finish_reason"] == "error": raise Exception( - f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}" + f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={openai_response_obj}" ) - self.received_finish_reason = response_obj["finish_reason"] - if response_obj.get("original_chunk", None) is not None: - if hasattr(response_obj["original_chunk"], "id"): - model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) - if hasattr(response_obj["original_chunk"], "system_fingerprint"): - model_response.system_fingerprint = response_obj["original_chunk"].system_fingerprint - self.system_fingerprint = response_obj["original_chunk"].system_fingerprint - if response_obj["logprobs"] is not None: - model_response.choices[0].logprobs = response_obj["logprobs"] + self.received_finish_reason = openai_response_obj["finish_reason"] + _openai_original_chunk = openai_response_obj.get("original_chunk", None) + if _openai_original_chunk is not None: + if hasattr(_openai_original_chunk, "id"): + model_response = self.set_model_id(_openai_original_chunk.id, model_response) + if hasattr(_openai_original_chunk, "system_fingerprint"): + model_response.system_fingerprint = _openai_original_chunk.system_fingerprint + self.system_fingerprint = _openai_original_chunk.system_fingerprint + if openai_response_obj["logprobs"] is not None: + model_response.choices[0].logprobs = openai_response_obj["logprobs"] - if response_obj["usage"] is not None: - if isinstance(response_obj["usage"], dict): - setattr( - model_response, - "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None, - completion_tokens=response_obj["usage"].get("completion_tokens", None) or None, - total_tokens=response_obj["usage"].get("total_tokens", None) or None, - ), - ) - elif isinstance(response_obj["usage"], Usage): - setattr( - model_response, - "usage", - response_obj["usage"], - ) - elif isinstance(response_obj["usage"], BaseModel): - setattr( - model_response, - "usage", - litellm.Usage(**response_obj["usage"].model_dump()), + _openai_usage = openai_response_obj["usage"] + if _openai_usage is not None: + if isinstance(_openai_usage, dict): + _usage_dict = _DICT_STR_OBJECT_ADAPTER.validate_python(_openai_usage) + _pt = _usage_dict.get("prompt_tokens", None) + _ct = _usage_dict.get("completion_tokens", None) + _tt = _usage_dict.get("total_tokens", None) + model_response.usage = litellm.Usage( + prompt_tokens=(_pt if isinstance(_pt, int) else None) or None, + completion_tokens=(_ct if isinstance(_ct, int) else None) or None, + total_tokens=(_tt if isinstance(_tt, int) else None) or None, ) + elif isinstance(_openai_usage, Usage): + model_response.usage = _openai_usage + elif isinstance(_openai_usage, BaseModel): + model_response.usage = litellm.Usage.model_validate(_openai_usage.model_dump()) return _ProviderChunkParsed(response_obj) - def chunk_creator(self, chunk: Any): # type: ignore - if hasattr(chunk, "id"): - self.response_id = chunk.id + def chunk_creator( + self, + chunk: Any, # any-ok: provider streams yield heterogeneous third-party SDK objects typed only at runtime + ) -> ModelResponseStream | None: + chunk_obj = _OBJECT_ADAPTER.validate_python(chunk) + if isinstance(chunk_obj, (ModelResponse, ModelResponseStream, OpenAIChatCompletionChunk, Completion)): + self.response_id = chunk_obj.id model_response = self.model_response_creator() - response_obj: dict[str, Any] = {} try: # return this for all models - completion_obj: dict[str, Any] = {"content": ""} + completion_obj: StreamingCompletionObj = {"content": ""} dispatch_result = self._dispatch_provider_chunk( - chunk=chunk, + chunk=chunk_obj, model_response=model_response, completion_obj=completion_obj, ) @@ -1376,7 +1424,7 @@ class CustomStreamWrapper: model_response.model = self.model ## FUNCTION CALL PARSING - original_chunk = response_obj.get("original_chunk") if response_obj is not None else None + original_chunk = response_obj.get("original_chunk") if ( original_chunk is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints @@ -1388,37 +1436,25 @@ class CustomStreamWrapper: original_chunk, model_response ) if original_chunk.choices and len(original_chunk.choices) > 0: - delta = original_chunk.choices[0].delta + delta = _delta_or_none(original_chunk.choices[0]) if delta is not None and (delta.function_call is not None or delta.tool_calls is not None): try: model_response.system_fingerprint = original_chunk.system_fingerprint ## AZURE - check if arguments is not None - if original_chunk.choices[0].delta.function_call is not None: - if ( - getattr( - original_chunk.choices[0].delta.function_call, - "arguments", - ) - is None - ): - original_chunk.choices[0].delta.function_call.arguments = "" - elif original_chunk.choices[0].delta.tool_calls is not None: - if isinstance(original_chunk.choices[0].delta.tool_calls, list): - for t in original_chunk.choices[0].delta.tool_calls: - if hasattr(t, "functions") and hasattr(t.functions, "arguments"): - if ( - getattr( - t.function, - "arguments", - ) - is None - ): - t.function.arguments = "" + if delta.function_call is not None: + if _arguments_or_none(delta.function_call) is None: + delta.function_call.arguments = "" + elif delta.tool_calls is not None: + for t in delta.tool_calls: + if isinstance(t, HasFunctionsAttr) and hasattr(t.functions, "arguments"): + if _arguments_or_none(t.function) is None: + t.function.arguments = "" _json_delta = delta.model_dump() if "role" not in _json_delta or _json_delta["role"] is None: _json_delta["role"] = "assistant" # mistral's api returns role as None - if "tool_calls" in _json_delta and isinstance(_json_delta["tool_calls"], list): - for tool in _json_delta["tool_calls"]: + _json_tool_calls = _json_delta.get("tool_calls") + if isinstance(_json_tool_calls, list): + for tool in _LIST_OBJECT_ADAPTER.validate_python(_json_tool_calls): if ( isinstance(tool, dict) and "function" in tool @@ -1427,7 +1463,7 @@ class CustomStreamWrapper: ): # if function returned but type set to None - mistral's api returns type: None tool["type"] = "function" - model_response.choices[0].delta = Delta(**_json_delta) + model_response.choices[0].delta = Delta.model_validate(_json_delta) except Exception as e: verbose_logger.exception( f"litellm.CustomStreamWrapper.chunk_creator(): Exception occured - {e!s}" @@ -1437,12 +1473,8 @@ class CustomStreamWrapper: self._handle_special_delta_attributes(delta, model_response) else: try: - delta = ( - dict() - if original_chunk.choices[0].delta is None - else dict(original_chunk.choices[0].delta) - ) - model_response.choices[0].delta = Delta(**delta) + delta_dict: dict[str, object] = {} if delta is None else dict(delta) + model_response.choices[0].delta = Delta.model_validate(delta_dict) except Exception: model_response.choices[0].delta = Delta() else: @@ -1452,20 +1484,26 @@ class CustomStreamWrapper: return ## CHECK FOR TOOL USE - if "tool_calls" in completion_obj and len(completion_obj["tool_calls"]) > 0: + _pending_tool_calls = completion_obj.get("tool_calls") + if _pending_tool_calls is not None and len(_pending_tool_calls) > 0: if self.is_function_call is True: # user passed in 'functions' param - completion_obj["function_call"] = completion_obj["tool_calls"][0]["function"] + _first_tool = _pending_tool_calls[0] + completion_obj["function_call"] = ( + _first_tool["function"] if isinstance(_first_tool, dict) else _first_tool.function + ) completion_obj["tool_calls"] = None self.tool_call = True - if hasattr(chunk, "usage") and chunk.usage is not None: - model_response.usage = chunk.usage + if isinstance(chunk_obj, (ModelResponse, ModelResponseStream, OpenAIChatCompletionChunk, Completion)): + _chunk_usage = getattr(chunk_obj, "usage", None) + if _chunk_usage is not None: + model_response.usage = _chunk_usage ## RETURN ARG result = self.return_processed_chunk_logic( completion_obj=completion_obj, - model_response=model_response, # type: ignore + model_response=model_response, response_obj=response_obj, ) return result @@ -1656,8 +1694,8 @@ class CustomStreamWrapper: else: asyncio.run(self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit)) ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler - litellm_params = self.logging_obj.model_call_details.get("litellm_params", {}) - if self.logging_obj._is_sync_litellm_request(litellm_params): + litellm_params = self.logging_obj.model_call_details.get("litellm_params", None) + if self.logging_obj._is_sync_litellm_request(litellm_params if isinstance(litellm_params, dict) else {}): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) def finish_reason_handler(self): @@ -1685,10 +1723,11 @@ class CustomStreamWrapper: calculator uses it instead of a token-based estimate. """ _usage = getattr(response, "usage", None) - if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None: + _usage_cost = _usage.cost if isinstance(_usage, Usage) and hasattr(_usage, "cost") else None + if _usage_cost is not None: if "additional_headers" not in response._hidden_params: response._hidden_params["additional_headers"] = {} - response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost) + response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage_cost) def __next__(self) -> "ModelResponseStream": cache_hit = False @@ -1966,7 +2005,7 @@ class CustomStreamWrapper: if isinstance(self.completion_stream, str) or isinstance(self.completion_stream, bytes): chunk = self.completion_stream else: - chunk = await asyncio.to_thread(_next_sync_or_exhausted, self.completion_stream) # type: ignore[arg-type] + chunk = await asyncio.to_thread(_next_sync_or_exhausted, self.completion_stream) if chunk is _SYNC_ITER_EXHAUSTED: raise StopAsyncIteration if chunk is not None and chunk != b"": @@ -2166,19 +2205,17 @@ class CustomStreamWrapper: """Best-effort status_code extraction.""" try: code = getattr(exc, "status_code", None) - if code is not None: + if isinstance(code, (int, str, float)): return int(code) except Exception: pass response = getattr(exc, "response", None) - if response is not None: - try: - status_code = getattr(response, "status_code", None) - if status_code is not None: - return int(status_code) - except Exception: - pass + try: + if isinstance(response, HasStatusCode): + return int(response.status_code) + except Exception: + pass return None mapped_status_code = _normalize_status_code(mapped_exception) @@ -2242,16 +2279,31 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" prompt_tokens: int = 0 completion_tokens: int = 0 - latest_usage_chunk = None + latest_usage_chunk: Usage | dict[str, object] | None = None for chunk in chunks: - if "usage" in chunk and chunk["usage"] is not None: - usage = chunk["usage"] - latest_usage_chunk = usage - if "prompt_tokens" in usage: - prompt_tokens = usage.get("prompt_tokens", 0) or 0 - if "completion_tokens" in usage: - completion_tokens = usage.get("completion_tokens", 0) or 0 + usage_obj = ( + _DICT_STR_OBJECT_ADAPTER.validate_python(chunk).get("usage") + if isinstance(chunk, dict) + else getattr(chunk, "usage", None) + ) + if usage_obj is None: + continue + if isinstance(usage_obj, Usage): + latest_usage_chunk = usage_obj + if hasattr(usage_obj, "prompt_tokens"): + prompt_tokens = getattr(usage_obj, "prompt_tokens", 0) or 0 + if hasattr(usage_obj, "completion_tokens"): + completion_tokens = getattr(usage_obj, "completion_tokens", 0) or 0 + elif isinstance(usage_obj, dict): + usage_dict = _DICT_STR_OBJECT_ADAPTER.validate_python(usage_obj) + latest_usage_chunk = usage_dict + if "prompt_tokens" in usage_dict: + _pt = usage_dict.get("prompt_tokens", 0) + prompt_tokens = _pt if isinstance(_pt, int) else 0 + if "completion_tokens" in usage_dict: + _ct = usage_dict.get("completion_tokens", 0) + completion_tokens = _ct if isinstance(_ct, int) else 0 returned_usage_chunk = Usage( prompt_tokens=prompt_tokens, @@ -2265,7 +2317,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: if isinstance(latest_usage_chunk, dict) else getattr(latest_usage_chunk, "cost", None) ) - if latest_cost is not None: + if isinstance(latest_cost, (int, float)): returned_usage_chunk.cost = latest_cost return returned_usage_chunk diff --git a/litellm/litellm_core_utils/streaming_handler_types.py b/litellm/litellm_core_utils/streaming_handler_types.py new file mode 100644 index 00000000000..7371413613c --- /dev/null +++ b/litellm/litellm_core_utils/streaming_handler_types.py @@ -0,0 +1,233 @@ +from collections.abc import Mapping, Sequence +from typing import Protocol, Union, runtime_checkable + +from pydantic import BaseModel +from typing_extensions import ReadOnly, Required, TypedDict + +from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, +) +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Function, + ModelResponseStream, +) + + +class CompletionUsageLike(Protocol): + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ParsedProviderChunk(TypedDict, total=False): + text: ReadOnly[str | None] + is_finished: ReadOnly[bool] + finish_reason: ReadOnly[str | None] + logprobs: ReadOnly[object] + original_chunk: ReadOnly[ModelResponseStream | None] + usage: ReadOnly[object] + tool_use: ReadOnly[ChatCompletionToolCallChunk | None] + tool_calls: ReadOnly[list[ChatCompletionDeltaToolCall] | None] + provider_specific_fields: ReadOnly[dict[str, object] | None] + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + index: ReadOnly[int] + + +class SimpleParsedChunk(ParsedProviderChunk): + text: Required[ReadOnly[str | None]] + is_finished: Required[ReadOnly[bool]] + finish_reason: Required[ReadOnly[str | None]] + + +class TextCompletionParsedChunk(SimpleParsedChunk): + usage: Required[ReadOnly[CompletionUsageLike | None]] + + +class OpenAIChatParsedChunk(SimpleParsedChunk): + logprobs: Required[ReadOnly[object]] + original_chunk: Required[ReadOnly[ModelResponseStream | None]] + usage: Required[ReadOnly[object]] + + +class StreamingCompletionObj(TypedDict, total=False): + content: Required[str | None] + role: str + tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionDeltaToolCall] | None + function_call: ChatCompletionToolCallFunctionChunk | Function | None + provider_specific_fields: dict[str, object] | None + index: int + + +class TextCompletionChoiceLike(Protocol): + @property + def text(self) -> str: ... + + @property + def finish_reason(self) -> str | None: ... + + +@runtime_checkable +class HasTextChoices(Protocol): + @property + def choices(self) -> Sequence[TextCompletionChoiceLike]: ... + + +@runtime_checkable +class HasCompletionUsage(Protocol): + usage: CompletionUsageLike | None + + +@runtime_checkable +class HasFunctionsAttr(Protocol): + functions: object + + +@runtime_checkable +class HasStatusCode(Protocol): + @property + def status_code(self) -> int: ... + + +VertexArgValue = Union[ + str, + int, + float, + bool, + None, + Sequence["VertexArgValue"], + Mapping[str, "VertexArgValue"], +] + + +class VertexFunctionCallLike(Protocol): + @property + def args(self) -> Mapping[str, VertexArgValue]: ... + + @property + def name(self) -> str: ... + + +class VertexPartLike(Protocol): + @property + def function_call(self) -> VertexFunctionCallLike: ... + + +class VertexContentLike(Protocol): + @property + def parts(self) -> Sequence[VertexPartLike]: ... + + +class VertexEnumNameLike(Protocol): + @property + def name(self) -> str: ... + + +class VertexCandidateLike(Protocol): + @property + def content(self) -> VertexContentLike: ... + + @property + def finish_reason(self) -> VertexEnumNameLike: ... + + +@runtime_checkable +class VertexProtoChunkLike(Protocol): + @property + def candidates(self) -> Sequence[VertexCandidateLike]: ... + + +class VllmCompletionOutputLike(Protocol): + @property + def text(self) -> str: ... + + +@runtime_checkable +class VllmRequestOutputLike(Protocol): + @property + def outputs(self) -> Sequence[VllmCompletionOutputLike]: ... + + +@runtime_checkable +class HasModelAttr(Protocol): + model: str | None + + +class PredibaseStreamToken(BaseModel): + text: str | None = None + + +class PredibaseStreamDetails(BaseModel): + finish_reason: str | None = None + + +class PredibaseStreamChunk(BaseModel): + token: PredibaseStreamToken | None = None + details: PredibaseStreamDetails | None = None + generated_text: str | None = None + error: str | None = None + + +class AI21StreamData(BaseModel): + text: str + + +class AI21StreamCompletion(BaseModel): + data: AI21StreamData + + +class AI21StreamChunk(BaseModel): + completions: list[AI21StreamCompletion] + + +class MaritalkStreamChunk(BaseModel): + answer: str + + +class NlpCloudStreamChunk(BaseModel): + generated_text: str + + +class AlephAlphaStreamCompletion(BaseModel): + completion: str + + +class AlephAlphaStreamChunk(BaseModel): + completions: list[AlephAlphaStreamCompletion] + + +class AzureStreamDelta(BaseModel): + content: str | None = "" + + +class AzureStreamChoice(BaseModel): + delta: AzureStreamDelta | None = None + finish_reason: str | None = None + + +class AzureStreamChunk(BaseModel): + choices: list[AzureStreamChoice] + + +class BasetenStreamToken(BaseModel): + text: str | None = None + + +class BasetenModelOutput(BaseModel): + data: list[str] | None = None + + +class BasetenStreamChunk(BaseModel): + token: BasetenStreamToken | None = None + model_output: BasetenModelOutput | str | None = None + completion: str | None = None + + +class TritonStreamChunk(BaseModel): + text_output: str | None = "" + stop_reason: str | None = None + is_finished: bool = False + input_token_count: int = 0 + generated_token_count: int = 0 diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a5aa1509969..80918e10eb6 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,9 +1,11 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from typing import ( Any, cast, ) +from pydantic import TypeAdapter, ValidationError + import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function @@ -27,8 +29,27 @@ 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 = dict[str, object] | list[dict[str, object]] | None -def _messages_have_compaction_block(messages: list[dict]) -> bool: +_CONTEXT_MANAGEMENT_SPEC_ADAPTER: TypeAdapter[ContextManagementSpec] = TypeAdapter(ContextManagementSpec) + + +def _validate_context_management_spec(raw: object) -> ContextManagementSpec: + """Validate the raw ``context_management`` request value into a typed spec. + + Shapes the dispatcher cannot use (non-dict scalars, lists with non-dict + entries, non-string keys) normalize to ``None``, matching the dispatcher's + own isinstance-based rejection of malformed specs. + """ + if raw is None: + return None + try: + return _CONTEXT_MANAGEMENT_SPEC_ADAPTER.validate_python(raw, strict=True) + except ValidationError: + return None + + +def _messages_have_compaction_block(messages: Sequence[Mapping[str, object]]) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: content = msg.get("content") @@ -40,7 +61,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: Mapping[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 +82,14 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | async def _prepare_context_managed_request( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, - additional_drop_params: list[str] | None, - llm_router: Any, - user_api_key_auth: Any = None, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: str | list[dict[str, object]] | None, + context_management_spec: ContextManagementSpec, + litellm_metadata: dict[str, object] | None, + additional_drop_params: Sequence[str] | 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 +109,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: str | list[dict[str, object]] | None = 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 +143,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,8 +151,8 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, - context_management_spec: Any, - additional_drop_params: list[str] | None, + context_management_spec: ContextManagementSpec, + additional_drop_params: Sequence[str] | None, ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. @@ -157,8 +178,8 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, - context_management_spec: Any, - additional_drop_params: list[str] | None, + context_management_spec: ContextManagementSpec, + additional_drop_params: Sequence[str] | None, ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -184,7 +205,7 @@ def _spec_has_non_compact_edits( ) -def _context_management_explicitly_dropped(additional_drop_params: list[str] | None) -> bool: +def _context_management_explicitly_dropped(additional_drop_params: Sequence[str] | None) -> bool: """True when the caller opted out of context_management via ``additional_drop_params``. ``drop_params`` deliberately does NOT gate the polyfill: ``context_management`` @@ -198,9 +219,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N def _normalize_spec_edits( *, - context_management_spec: Any, - additional_drop_params: list[str] | None, -) -> list[dict[str, Any]] | None: + context_management_spec: ContextManagementSpec, + additional_drop_params: Sequence[str] | None, +) -> Sequence[Mapping[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 +246,14 @@ def _normalize_spec_edits( async def _run_polyfill_if_enabled( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, - additional_drop_params: list[str] | None, - llm_router: Any, - user_api_key_auth: Any = None, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: str | list[dict[str, object]] | None, + context_management_spec: ContextManagementSpec, + litellm_metadata: dict[str, object] | None, + additional_drop_params: Sequence[str] | 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 +314,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: Mapping[str, object] | None, ) -> None: """ When users call `litellm.anthropic.messages.*` with a non-Anthropic model and @@ -308,8 +329,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: If the user provides a `summary` field in the thinking dict, it is passed through to the OpenAI reasoning params (opt-in per OpenAI spec). """ - custom_llm_provider = completion_kwargs.get("custom_llm_provider") - if custom_llm_provider is None: + raw_provider = completion_kwargs.get("custom_llm_provider") + custom_llm_provider: str | None + if raw_provider is None: try: _, inferred_provider, _, _ = litellm.utils.get_llm_provider( model=cast(str, completion_kwargs.get("model")) @@ -317,6 +339,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: custom_llm_provider = inferred_provider except Exception: custom_llm_provider = None + else: + custom_llm_provider = raw_provider if isinstance(raw_provider, str) else None if custom_llm_provider != "openai": return @@ -342,7 +366,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: @@ -358,7 +382,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def _normalize_reasoning_effort( - completion_kwargs: dict[str, Any], + completion_kwargs: dict[str, object], ) -> None: """ Normalize reasoning_effort values based on target model capabilities. @@ -375,7 +399,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: return model = cast(str, completion_kwargs.get("model", "")) - custom_llm_provider = completion_kwargs.get("custom_llm_provider") + raw_provider = completion_kwargs.get("custom_llm_provider") + custom_llm_provider = raw_provider if isinstance(raw_provider, str) else None if isinstance(reasoning_effort, str): normalized = normalize_reasoning_effort_value( @@ -396,21 +421,21 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _prepare_completion_kwargs( *, max_tokens: int, - messages: list[dict], + messages: Sequence[Mapping[str, object]], model: str, - metadata: dict | None = None, - stop_sequences: list[str] | None = None, + metadata: Mapping[str, object] | None = None, + stop_sequences: Sequence[str] | None = None, stream: bool | None = False, - system: str | list[dict[str, Any]] | None = None, + system: str | Sequence[Mapping[str, object]] | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: Mapping[str, object] | None = None, + tool_choice: Mapping[str, object] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, - ) -> tuple[dict[str, Any], dict[str, str]]: + output_format: Mapping[str, object] | None = None, + extra_kwargs: Mapping[str, object] | None = None, + ) -> tuple[Mapping[str, Any], dict[str, str]]: # any-ok: kwargs spread into litellm.completion's open params """Prepare kwargs for litellm.completion/acompletion. Returns: @@ -422,7 +447,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data = { + request_data: dict[str, object] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -467,7 +492,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, Any] = dict(openai_request) # any-ok: heterogeneous per-provider kwargs if stream: completion_kwargs["stream"] = stream @@ -517,24 +542,25 @@ 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, - stop_sequences: list[str] | None = None, + metadata: Mapping[str, object] | None = None, + stop_sequences: Sequence[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: Mapping[str, object] | None = None, + tool_choice: Mapping[str, object] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, - **kwargs, - ) -> AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]: + output_format: Mapping[str, object] | None = None, + **kwargs: object, + ) -> 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) + context_management = _validate_context_management_spec(kwargs.pop("context_management", None)) + raw_drop_params = kwargs.get("additional_drop_params", None) + additional_drop_params: list[str] | None = raw_drop_params if isinstance(raw_drop_params, list) else None litellm_router = kwargs.pop("litellm_router", None) if litellm_router is None: try: @@ -585,7 +611,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) - completion_response = await litellm.acompletion(**completion_kwargs) + completion_response = await litellm.acompletion(**completion_kwargs) # any-ok: open per-provider kwargs if stream: transformed_stream = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( @@ -611,26 +637,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, - stop_sequences: list[str] | None = None, + metadata: Mapping[str, object] | None = None, + stop_sequences: Sequence[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: Mapping[str, object] | None = None, + tool_choice: Mapping[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: Mapping[str, object] | None = None, _is_async: bool = False, - **kwargs, + **kwargs: object, ) -> ( 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: @@ -657,8 +683,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: # ``clear_tool_uses_20250919``. The dispatcher is async (so the # ``compact_20260112`` editor can ``await`` the summarization model); # bridge to it via ``run_async_function``. - context_management = kwargs.pop("context_management", None) - additional_drop_params: list[str] | None = kwargs.get("additional_drop_params", None) + context_management = _validate_context_management_spec(kwargs.pop("context_management", None)) + raw_drop_params = kwargs.get("additional_drop_params", None) + additional_drop_params: list[str] | None = raw_drop_params if isinstance(raw_drop_params, list) else None # Deliberately do NOT auto-attach the proxy ``llm_router`` here: # ``run_async_function`` spawns a new event loop in a worker thread # to bridge to the async dispatcher, but the proxy router's httpx @@ -722,7 +749,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: extra_kwargs=kwargs, ) - completion_response = litellm.completion(**completion_kwargs) + completion_response = litellm.completion(**completion_kwargs) # any-ok: open per-provider kwargs if stream: transformed_stream = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4d4f459a080..3dd8c0bcdf4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,13 +4,17 @@ import logging import math import time import traceback -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping from datetime import datetime from functools import lru_cache from typing import ( TYPE_CHECKING, Any, Literal, + Protocol, + TypedDict, + TypeVar, + runtime_checkable, ) import anyio @@ -18,6 +22,7 @@ import httpx import orjson from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import TypeAdapter from starlette.types import Receive, Scope, Send import litellm @@ -58,10 +63,76 @@ from litellm.types.router import RouterRateLimitError from litellm.types.utils import ServerToolUse # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) -StreamChunkSerializer = Callable[[Any], str] +StreamChunkSerializer = Callable[[Any], str] # any-ok: chunks are heterogeneous provider frames forwarded verbatim # Type alias for streaming error serializer (ProxyException -> wire format) StreamErrorSerializer = Callable[[ProxyException], str] +_T = TypeVar("_T") + +_OBJECT_ADAPTER: TypeAdapter[object] = TypeAdapter(object) + + +@runtime_checkable +class _AsyncHttpResponseLike(Protocol): + status_code: int + headers: httpx.Headers + + async def aread(self) -> bytes: ... + + +class _ModelInfoIdPayload(TypedDict, total=False): + id: str + + +class _ModelInfoMetadataPayload(TypedDict, total=False): + model_info: _ModelInfoIdPayload + + +class _ModelInfoLitellmParamsPayload(TypedDict, total=False): + model_info: _ModelInfoIdPayload + metadata: _ModelInfoMetadataPayload + + +class _LoggingKwargsPayload(TypedDict, total=False): + litellm_params: _ModelInfoLitellmParamsPayload + + +class _StreamUsageKwargsPayload(TypedDict, total=False): + prompt_tokens: int + completion_tokens: int + total_tokens: int + completion_tokens_details: dict[str, object] + prompt_tokens_details: dict[str, object] + server_tool_use: ServerToolUse + cache_creation_input_tokens: int + cache_read_input_tokens: int + + +async def _await_str_stream_coroutine( + coro: "Coroutine[None, None, AsyncGenerator[str, None]]", +) -> AsyncGenerator[str, None]: + return await coro + + +async def _await_bytes_stream_coroutine( + coro: "Coroutine[None, None, AsyncIterator[bytes]]", +) -> AsyncIterator[bytes]: + return await coro + + +@runtime_checkable +class _BytesAsyncIterator(Protocol): + def __aiter__(self) -> AsyncIterator[bytes]: ... + + def __anext__(self) -> Awaitable[bytes]: ... + + +def _as_bytes_async_iterator(response: object) -> AsyncIterator[bytes]: + if isinstance(response, _BytesAsyncIterator): + return response + raise TypeError(f"expected an async iterator of bytes, got {type(response).__name__}") + + if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig @@ -232,7 +303,7 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return True -async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None: +async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[object]"]) -> None: pending_tasks = [task for task in tasks if not task.done()] for task in pending_tasks: task.cancel() @@ -309,8 +380,8 @@ def _stream_usage_tracking_updates( def _serialize_http_exception_detail( - detail: Any, -) -> tuple[str, dict | None]: + detail: object, +) -> tuple[str, dict[str, object] | None]: """ Convert an HTTPException.detail value into (message, structured_fields) for ProxyException / SSE error frames. @@ -339,7 +410,7 @@ def _serialize_http_exception_detail( return str(detail), None -def _collect_response_file_search_vector_store_ids(data: dict[str, Any]) -> set[str]: +def _collect_response_file_search_vector_store_ids(data: dict[str, object]) -> set[str]: vector_store_ids: set[str] = set() tools = data.get("tools") if not isinstance(tools, list): @@ -366,7 +437,7 @@ def _collect_response_file_search_vector_store_ids(data: dict[str, Any]) -> set[ async def _authorize_response_file_search_vector_stores( - data: dict[str, Any], + data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, ) -> None: vector_store_ids = _collect_response_file_search_vector_store_ids(data) @@ -392,7 +463,7 @@ async def _parse_event_data_for_error(event_line: str | bytes) -> int | None: if not json_str or json_str == "[DONE]": # handle empty data or [DONE] message return None try: - data = orjson.loads(json_str) + data = _OBJECT_ADAPTER.validate_python(orjson.loads(json_str)) if isinstance(data, dict) and "error" in data and isinstance(data["error"], dict): error_code_raw = data["error"].get("code") error_code: int | None = None @@ -445,7 +516,7 @@ def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict: return default_error try: - data = orjson.loads(json_str) + data = _OBJECT_ADAPTER.validate_python(orjson.loads(json_str)) if isinstance(data, dict) and "error" in data: error_obj = data["error"] if isinstance(error_obj, dict): @@ -596,7 +667,7 @@ async def create_response( try: # Handle coroutine that returns a generator if asyncio.iscoroutine(generator): - generator = await generator + generator = await _await_str_stream_coroutine(generator) # Now get the first chunk from the actual generator first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) @@ -668,13 +739,13 @@ async def create_response( existing_fields = getattr(e, "provider_specific_fields", None) or {} if structured_fields: - merged_fields: dict | None = {**existing_fields, **structured_fields} + merged_fields: dict[str, object] | None = {**existing_fields, **structured_fields} else: merged_fields = existing_fields or None # Match ProxyException.to_dict() shape so streaming and non-streaming # error frames are byte-identical. - error_obj: dict[str, Any] = { + error_obj: dict[str, object] = { "message": message, "type": getattr(e, "type", "None"), "param": getattr(e, "param", "None"), @@ -740,7 +811,7 @@ def _is_azure_model_router_request(model: str) -> bool: def _override_openai_response_model( *, - response_obj: Any, + response_obj: object, requested_model: str, log_context: str, return_raw_model_name: bool = False, @@ -913,7 +984,7 @@ def _log_llm_api_exception(e: Exception) -> None: async def _cancel_llm_call_on_client_disconnect( request: Request, - llm_api_call: "asyncio.Future[Any]", + llm_api_call: "asyncio.Future[_T]", disconnect_event: asyncio.Event, ) -> None: try: @@ -932,8 +1003,8 @@ async def _cancel_llm_call_on_client_disconnect( async def _await_llm_call_cancelling_on_disconnect( request: Request, - llm_api_call: "asyncio.Future[Any]", -) -> Any: + llm_api_call: "asyncio.Future[_T]", +) -> _T: disconnect_event = asyncio.Event() monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event)) try: @@ -1054,7 +1125,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod async def build_litellm_proxy_success_headers_from_llm_response( *, - response: Any, + response: object, request_data: dict, request: Request, user_api_key_dict: UserAPIKeyAuth, @@ -1309,7 +1380,7 @@ class ProxyBaseLLMRequestProcessing: ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call ## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse. - logging_obj, self.data = litellm.utils.function_setup( + logging_obj, self.data = litellm.utils.function_setup( # any-ok: function_setup (litellm/utils.py) is untyped original_function=route_type, rules_obj=litellm.utils.Rules(), start_time=start_time, @@ -1357,7 +1428,7 @@ class ProxyBaseLLMRequestProcessing: self.data["router_settings_override"] = router_settings if "messages" in self.data and self.data["messages"]: - logging_obj.update_messages(self.data["messages"]) + logging_obj.update_messages(self.data["messages"]) # any-ok: logging_obj/messages come from untyped seams return self.data, logging_obj @@ -1509,7 +1580,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _response_cost_from_logging_obj( *, - response: Any, + response: Any, # any-ok: forwarded verbatim to the legacy calculator union in litellm_logging.py logging_obj: LiteLLMLoggingObj, ) -> float | str: """ @@ -1523,7 +1594,7 @@ class ProxyBaseLLMRequestProcessing: stored_cost = logging_obj.model_call_details.get("response_cost") if isinstance(stored_cost, (int, float)): return float(stored_cost) - recomputed_cost = logging_obj._response_cost_calculator(result=response) + recomputed_cost = logging_obj._response_cost_calculator(result=response) # any-ok: untyped response seam return recomputed_cost if isinstance(recomputed_cost, (int, float)) else "" def _debug_log_request_payload(self) -> None: @@ -1654,7 +1725,7 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request: bool | None = False, contents: list | None = None, # Add contents parameter skip_pre_call_logic: bool = False, - ) -> Any: + ) -> object: """ Common request processing logic for both chat completions and responses API endpoints """ @@ -1664,7 +1735,7 @@ class ProxyBaseLLMRequestProcessing: self._debug_log_request_payload() if skip_pre_call_logic: - logging_obj = self.data.get("litellm_logging_obj") + logging_obj: LiteLLMLoggingObj | None = self.data.get("litellm_logging_obj") if logging_obj is None: raise ValueError( "skip_pre_call_logic=True requires litellm_logging_obj to be set in data. " @@ -1737,7 +1808,9 @@ class ProxyBaseLLMRequestProcessing: llm_call_task = asyncio.create_task(llm_call) tasks.append(llm_call_task) - llm_responses = asyncio.gather(*tasks) # run the moderation check in parallel to the actual llm api call + llm_responses: asyncio.Future[list[object]] = asyncio.gather( + *tasks + ) # run the moderation check in parallel to the actual llm api call try: if general_settings.get("cancel_on_disconnect", False): @@ -1747,7 +1820,7 @@ class ProxyBaseLLMRequestProcessing: finally: await _cancel_pending_gather_tasks(tasks) - response = responses[1] + response: object = responses[1] _exception_raised = False try: @@ -1830,7 +1903,9 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict = user_api_key_dict _captured_logging_obj = logging_obj - async def _on_deferred_stream_complete(assembled_response, cache_hit): + async def _on_deferred_stream_complete( + assembled_response: ModelResponse, cache_hit: bool | None + ) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( captured_data=_captured_data, captured_user_api_key_dict=_captured_user_api_key_dict, @@ -1845,17 +1920,15 @@ class ProxyBaseLLMRequestProcessing: # Check if response is an async generator if self._is_streaming_response(response): if asyncio.iscoroutine(response): - generator = await response + generator = await _await_bytes_stream_coroutine(response) else: - generator = response + generator = _as_bytes_async_iterator(response) if ( self._has_post_call_guardrails_for_passthrough() and self._passthrough_endpoint_has_stream_guardrail_handler() ): - body_bytes = b"".join( - [chunk async for chunk in generator] # type: ignore[union-attr] - ) + body_bytes = b"".join([chunk async for chunk in generator]) modified_bytes = await self._handle_event_stream_allm_passthrough_route( body_bytes=body_bytes, proxy_logging_obj=proxy_logging_obj, @@ -1874,7 +1947,7 @@ class ProxyBaseLLMRequestProcessing: # For passthrough routes, stream directly without error parsing # since we're dealing with raw binary data (e.g., AWS event streams) return StreamingResponse( - content=generator, # type: ignore[arg-type] + content=generator, status_code=status.HTTP_200_OK, headers=custom_headers, ) @@ -1888,11 +1961,12 @@ class ProxyBaseLLMRequestProcessing: ) if _early is not None: return _early - return StreamingResponse( - content=response.aiter_bytes(), # type: ignore[union-attr] - status_code=response.status_code, # type: ignore[union-attr] - headers=custom_headers, - ) + if isinstance(response, httpx.Response): + return StreamingResponse( + content=response.aiter_bytes(), + status_code=response.status_code, + headers=custom_headers, + ) elif route_type == "anthropic_messages": # Check if response is actually a streaming response (async generator) # Non-streaming responses (dict) should be returned directly @@ -1914,7 +1988,7 @@ class ProxyBaseLLMRequestProcessing: ) # Non-streaming response - fall through to normal response handling elif select_data_generator: - selected_data_generator = select_data_generator( + selected_data_generator: AsyncGenerator[str, None] = select_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=self.data, @@ -1978,10 +2052,12 @@ class ProxyBaseLLMRequestProcessing: if _early is not None: return _early - response = await proxy_logging_obj.post_call_success_hook( - data=self.data, - user_api_key_dict=user_api_key_dict, - response=response, # type: ignore[arg-type] + response = _OBJECT_ADAPTER.validate_python( + await proxy_logging_obj.post_call_success_hook( + data=self.data, + user_api_key_dict=user_api_key_dict, + response=response, # type: ignore[arg-type] + ) ) except Exception: _exception_raised = True @@ -2080,7 +2156,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod async def _record_container_owners_from_responses_if_needed( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, ) -> None: """Register code-interpreter containers so follow-up file APIs pass ownership checks.""" @@ -2103,7 +2179,7 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _extract_completed_responses_response(stream_response: Any) -> Any: + def _extract_completed_responses_response(stream_response: object) -> object: """Pull the assembled ``ResponsesAPIResponse`` off a streaming iterator. ``ResponsesAPIStreamingIterator`` stores the terminal stream event @@ -2113,20 +2189,20 @@ class ProxyBaseLLMRequestProcessing: ``ResponsesAPIResponse`` directly. Handle both shapes so the container-ownership recording path can walk ``.output`` either way. """ - completed = getattr(stream_response, "completed_response", None) + completed: object = getattr(stream_response, "completed_response", None) if completed is None: return None - response_obj = getattr(completed, "response", None) + response_obj: object = getattr(completed, "response", None) if response_obj is not None: return response_obj return completed @staticmethod async def _wrap_responses_stream_for_container_ownership( - original_stream_response: Any, - wrapped_generator: Any, + original_stream_response: object, + wrapped_generator: AsyncGenerator[str, None], user_api_key_dict: UserAPIKeyAuth, - ): + ) -> AsyncGenerator[str, None]: """Forward SSE chunks, then record container ownership at stream end. Streaming ``/v1/responses`` short-circuits out of @@ -2222,6 +2298,9 @@ class ProxyBaseLLMRequestProcessing: if isinstance(result, Response): return result + if not isinstance(result, httpx.Response): + return result + content = await result.aread() return Response( content=content, @@ -2232,7 +2311,7 @@ class ProxyBaseLLMRequestProcessing: ), ) - def _is_streaming_response(self, response: Any) -> bool: + def _is_streaming_response(self, response: object) -> bool: """ Check if the response object is actually a streaming response by inspecting its type. @@ -2340,7 +2419,7 @@ class ProxyBaseLLMRequestProcessing: async def _handle_non_streaming_allm_passthrough_route( self, - response: Any, + response: object, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", custom_headers: dict, @@ -2358,12 +2437,12 @@ class ProxyBaseLLMRequestProcessing: HttpPassThroughEndpointHelpers, ) - try: - response_status: int = response.status_code # type: ignore[union-attr] - content_type: str = response.headers.get("content-type", "") # type: ignore[union-attr] - except AttributeError: + if not isinstance(response, _AsyncHttpResponseLike): return None + response_status: int = response.status_code + content_type: str = dict(response.headers).get("content-type", "") + if response_status >= 300: return None @@ -2374,7 +2453,7 @@ class ProxyBaseLLMRequestProcessing: return None response_headers = HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, # type: ignore[union-attr] + headers=response.headers, custom_headers=custom_headers, ) callback_headers = await proxy_logging_obj.post_call_response_headers_hook( @@ -2387,7 +2466,7 @@ class ProxyBaseLLMRequestProcessing: response_headers.update(callback_headers) if is_event_stream: - body_bytes = await response.aread() # type: ignore[union-attr] + body_bytes = await response.aread() modified_bytes = await self._handle_event_stream_allm_passthrough_route( body_bytes=body_bytes, proxy_logging_obj=proxy_logging_obj, @@ -2400,9 +2479,9 @@ class ProxyBaseLLMRequestProcessing: headers=response_headers, ) - body_bytes = await response.aread() # type: ignore[union-attr] + body_bytes = await response.aread() try: - parsed = _json.loads(body_bytes) + parsed = _json.loads(body_bytes) # any-ok: provider-shaped JSON consumed by duck-typed guardrail hooks except (_json.JSONDecodeError, UnicodeDecodeError): return Response( content=body_bytes, @@ -2413,7 +2492,7 @@ class ProxyBaseLLMRequestProcessing: processed = await proxy_logging_obj.post_call_success_hook( data=self.data, user_api_key_dict=user_api_key_dict, - response=parsed, + response=parsed, # any-ok: hook contract declares the OpenAI model union; passthrough bodies are dicts ) if isinstance(processed, dict): content = _json.dumps(processed).encode() @@ -2449,7 +2528,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _flush_deferred_async_logging( - logging_obj: Any, + logging_obj: object, exception_raised: bool, ) -> None: """ @@ -2474,10 +2553,10 @@ class ProxyBaseLLMRequestProcessing: Extracted as a static method so tests can exercise the production gating logic directly rather than reimplementing the finally block. """ - _enqueue_fn = getattr(logging_obj, "_enqueue_deferred_logging", None) + _enqueue_fn: Callable[[], None] | None = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is None: return - logging_obj._enqueue_deferred_logging = None # type: ignore[union-attr] + setattr(logging_obj, "_enqueue_deferred_logging", None) if exception_raised: return try: @@ -2489,9 +2568,9 @@ class ProxyBaseLLMRequestProcessing: async def _run_deferred_stream_guardrails( captured_data: dict, captured_user_api_key_dict: "UserAPIKeyAuth", - captured_logging_obj: Any, - assembled_response: Any, - cache_hit: Any, + captured_logging_obj: LiteLLMLoggingObj, + assembled_response: ModelResponse, + cache_hit: bool | None, ) -> None: """ Run non-streaming post-call guardrail hooks on an assembled streaming @@ -2526,7 +2605,7 @@ class ProxyBaseLLMRequestProcessing: ): continue try: - guardrail_result = None + guardrail_result: ModelResponse | None = None if "apply_guardrail" in type(cb).__dict__: # Skip — apply_guardrail guardrails already ran via # unified_guardrail's end-of-stream block in the @@ -2541,7 +2620,7 @@ class ProxyBaseLLMRequestProcessing: # here would duplicate the scan and can spuriously block the guardrail that already passed / failed. continue else: - guardrail_result = await cb.async_post_call_success_hook( + guardrail_result = await cb.async_post_call_success_hook( # any-ok: hook declared -> Any user_api_key_dict=captured_user_api_key_dict, data=guardrail_data, response=_response, @@ -2640,9 +2719,9 @@ class ProxyBaseLLMRequestProcessing: headers = getattr(e, "headers", None) or {} if not headers: # Try to get headers from e.response.headers (httpx.Response) - _response = getattr(e, "response", None) + _response: httpx.Response | None = getattr(e, "response", None) if _response is not None: - _response_headers = getattr(_response, "headers", None) + _response_headers: httpx.Headers | None = getattr(_response, "headers", None) if _response_headers: headers = get_response_headers(dict(_response_headers)) headers.update(custom_headers) @@ -2736,7 +2815,7 @@ class ProxyBaseLLMRequestProcessing: ######################################################### @staticmethod - def return_sse_chunk(chunk: Any) -> str: + def return_sse_chunk(chunk: Any) -> str: # any-ok: chunks are provider frames (str/bytes/dict) forwarded verbatim """ Helper function to format streaming chunks for Anthropic API format @@ -2751,13 +2830,13 @@ class ProxyBaseLLMRequestProcessing: chunk_str = safe_dumps(chunk) return f"{STREAM_SSE_DATA_PREFIX}{chunk_str}\n\n" else: - return chunk + return chunk # any-ok: non-dict frames (str/bytes) pass through unchanged @staticmethod async def _finalize_streaming_generator_cleanup( request: Request | None, request_data: dict, - response: Any, + response: object, stream_completed: bool = False, client_disconnected: bool = False, user_api_key_dict: UserAPIKeyAuth | None = None, @@ -2794,9 +2873,10 @@ class ProxyBaseLLMRequestProcessing: ): await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect(user_api_key_dict) - if hasattr(response, "aclose"): + aclose_fn: Callable[[], Awaitable[object]] | None = getattr(response, "aclose", None) + if aclose_fn is not None: try: - await response.aclose() + await aclose_fn() except BaseException as e: # noqa: BLE001 verbose_proxy_logger.debug( "async_streaming_data_generator: error closing response stream: %s", @@ -2805,7 +2885,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod async def async_streaming_data_generator( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, request_data: dict, proxy_logging_obj: ProxyLogging, @@ -2859,7 +2939,7 @@ class ProxyBaseLLMRequestProcessing: str_so_far += response_str elif hasattr(chunk, "model_dump"): try: - d = chunk.model_dump(mode="json", exclude_none=True) + d: dict[str, object] = chunk.model_dump(mode="json", exclude_none=True) if isinstance(d, dict): str_so_far += str(d.get("content", "")) except Exception: @@ -2936,7 +3016,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def async_sse_data_generator( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, request_data: dict, proxy_logging_obj: ProxyLogging, @@ -2963,7 +3043,7 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any: + def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: """ Process a streaming chunk and inject cost information if enabled. @@ -3023,12 +3103,13 @@ class ProxyBaseLLMRequestProcessing: if stripped_ln.startswith("data:"): json_part = stripped_ln.split("data:", 1)[1].strip() if json_part and json_part != "[DONE]": - obj = json.loads(json_part) - maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) - if maybe_modified is not None: - # Replace just this line with updated JSON using safe_dumps - lines[idx] = f"data: {safe_dumps(maybe_modified)}" - return "\n".join(lines) + obj = _OBJECT_ADAPTER.validate_python(json.loads(json_part)) + if isinstance(obj, dict): + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + if maybe_modified is not None: + # Replace just this line with updated JSON using safe_dumps + lines[idx] = f"data: {safe_dumps(maybe_modified)}" + return "\n".join(lines) return None except Exception: return None @@ -3054,13 +3135,13 @@ class ProxyBaseLLMRequestProcessing: ) # Extract additional usage fields - cache_creation_input_tokens = _usage.get("cache_creation_input_tokens") - cache_read_input_tokens = _usage.get("cache_read_input_tokens") - web_search_requests = _usage.get("web_search_requests") - completion_tokens_details = _usage.get("completion_tokens_details") - prompt_tokens_details = _usage.get("prompt_tokens_details") + cache_creation_input_tokens: int | None = _usage.get("cache_creation_input_tokens") + cache_read_input_tokens: int | None = _usage.get("cache_read_input_tokens") + web_search_requests: int | None = _usage.get("web_search_requests") + completion_tokens_details: dict[str, object] | None = _usage.get("completion_tokens_details") + prompt_tokens_details: dict[str, object] | None = _usage.get("prompt_tokens_details") - usage_kwargs: dict[str, Any] = { + usage_kwargs: _StreamUsageKwargsPayload = { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens, @@ -3120,18 +3201,18 @@ class ProxyBaseLLMRequestProcessing: # 2. Fallback to kwargs (initial) if not model_id: - _kwargs = getattr(_logging_obj, "kwargs", None) + _kwargs: _LoggingKwargsPayload | None = getattr(_logging_obj, "kwargs", None) if _kwargs: - litellm_params = _kwargs.get("litellm_params", {}) + kwargs_litellm_params = _kwargs.get("litellm_params", {}) # First check direct model_info path - model_info = litellm_params.get("model_info") or {} - model_id = model_info.get("id", None) + kwargs_model_info: _ModelInfoIdPayload = kwargs_litellm_params.get("model_info") or {} + model_id = kwargs_model_info.get("id", None) # Fallback to nested metadata path if not model_id: - metadata = litellm_params.get("metadata") or {} - model_info = metadata.get("model_info") or {} - model_id = model_info.get("id", None) + kwargs_metadata: _ModelInfoMetadataPayload = kwargs_litellm_params.get("metadata") or {} + kwargs_model_info = kwargs_metadata.get("model_info") or {} + model_id = kwargs_model_info.get("id", None) # 3. Final fallback to self.data["litellm_metadata"] (for routes like /v1/responses that populate data before error) if not model_id: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7253a684b3c..2cbfbe1b955 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, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime @@ -16,11 +16,12 @@ from typing import ( TYPE_CHECKING, Any, Literal, + Protocol, TypedDict, - Union, - cast, ) +from pydantic import TypeAdapter + from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE @@ -54,9 +55,10 @@ 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] + Span = _Span InternalUsageCache = _InternalUsageCache else: Span = Any @@ -342,6 +344,71 @@ class RateLimitResponseWithDescriptors(TypedDict): response: RateLimitResponse +class _WindowLimitMetadata(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 + + +class _TokenUsageDict(TypedDict, total=False): + prompt_tokens: int | None + completion_tokens: int | None + total_tokens: int | None + prompt_tokens_details: Mapping[str, int | None] | None + + +class _RedisLuaScript(Protocol): + def __call__(self, keys: Sequence[str], args: Sequence[int | float | str]) -> Awaitable[Sequence[object]]: ... + + +class _CallTypeRateLimiter(Protocol): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, + ) -> Exception | str | dict[str, object] | None: ... + + +_CACHED_VALUE_ADAPTER: TypeAdapter[object] = TypeAdapter(object) + + +def _as_int(value: object) -> int: + if isinstance(value, (int, float, str, bytes)): + return int(value) + raise TypeError(f"Expected an int-convertible rate limiter value, got {type(value).__name__}") + + +def _str_or_none(value: object) -> str | None: + return value if isinstance(value, str) else None + + +def _object_dict_or_none(value: object) -> dict[object, object] | None: + return value if isinstance(value, dict) else None + + +def _str_object_dict_or_none(value: object) -> dict[str, object] | None: + return value if isinstance(value, dict) else None + + @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -391,7 +458,7 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: +def claim_request_stash_for_data(data: Mapping) -> RequestRateLimiterStash: stash = get_or_create_request_stash() owner_call_id = data.get("litellm_call_id") if isinstance(owner_call_id, str): @@ -424,23 +491,23 @@ 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( - BATCH_RATE_LIMITER_SCRIPT + self.batch_rate_limiter_script: _RedisLuaScript | None = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script(BATCH_RATE_LIMITER_SCRIPT) ) - self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - TOKEN_INCREMENT_SCRIPT + self.token_increment_script: _RedisLuaScript | None = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script(TOKEN_INCREMENT_SCRIPT) ) - self.check_and_increment_by_n_script = ( + self.check_and_increment_by_n_script: _RedisLuaScript | None = ( self.internal_usage_cache.dual_cache.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( - PARALLEL_ACQUIRE_SCRIPT + self.parallel_acquire_script: _RedisLuaScript | None = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script(PARALLEL_ACQUIRE_SCRIPT) ) - self.parallel_release_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - PARALLEL_RELEASE_SCRIPT + self.parallel_release_script: _RedisLuaScript | None = ( + self.internal_usage_cache.dual_cache.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: _RedisLuaScript | None = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script(PARALLEL_COUNT_SCRIPT) ) else: self.batch_rate_limiter_script = None @@ -459,7 +526,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: _CallTypeRateLimiter | None = None # Serializes multi-phase check+increment sequences (batch + dynamic # limiters) within this process to close the TOCTOU window between @@ -477,7 +544,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) -> _CallTypeRateLimiter | None: """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: try: @@ -497,6 +564,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """Return the current time for rate limiting calculations.""" return self._time_provider() + async def _get_local_cache_value(self, key: str, parent_otel_span: Span | None) -> object: + return _CACHED_VALUE_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=key, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + ) + @staticmethod def _no_max_tokens_output_floor( min_configured_tpm_limit: int | None, @@ -514,7 +590,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _estimate_tokens_for_request( self, - data: dict, + data: Mapping, model: str | None = None, min_configured_tpm_limit: int | None = None, ) -> int: @@ -601,15 +677,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def in_memory_cache_sliding_window( self, - keys: list[str], + keys: Sequence[str], now_int: int, window_size: int, - ) -> list[Any]: + ) -> Sequence[object]: """ 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[object] = [] # Process each window/counter pair for i in range(0, len(keys), 2): @@ -618,14 +694,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): increment_value = 1 # Get the window start time - window_start = await self.internal_usage_cache.async_get_cache( - key=window_key, - litellm_parent_otel_span=None, - local_only=True, - ) + window_start = await self._get_local_cache_value(key=window_key, parent_otel_span=None) # Check if window exists and is valid - if window_start is None or (now_int - int(window_start)) >= window_size: + if window_start is None or (now_int - _as_int(window_start)) >= window_size: # Reset window and counter await self.internal_usage_cache.async_set_cache( key=window_key, @@ -645,12 +717,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): results.append(increment_value) # counter else: # Increment the counter - current_counter = await self.internal_usage_cache.async_get_cache( - key=counter_key, - litellm_parent_otel_span=None, - local_only=True, - ) - new_counter_value = (int(current_counter) if current_counter is not None else 0) + increment_value + current_counter = await self._get_local_cache_value(key=counter_key, parent_otel_span=None) + new_counter_value = (_as_int(current_counter) if current_counter is not None else 0) + increment_value await self.internal_usage_cache.async_set_cache( key=counter_key, value=new_counter_value, @@ -678,9 +746,9 @@ 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], + keys_to_fetch: Sequence[str], + cache_values: Sequence[object], + key_metadata: Mapping[str, _WindowLimitMetadata], ) -> RateLimitResponse: """ Check if the cache values are over the limit. @@ -709,12 +777,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if current_limit is None or rate_limit_type is None: continue - if counter_value is not None and int(counter_value) > current_limit: + if counter_value is not None and _as_int(counter_value) > current_limit: overall_code = "OVER_LIMIT" item_code = "OVER_LIMIT" # Only compute limit_remaining if current_limit is not None - limit_remaining = current_limit - int(counter_value) if counter_value is not None else current_limit + limit_remaining = current_limit - _as_int(counter_value) if counter_value is not None else current_limit statuses.append( { @@ -783,7 +851,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, keys_to_fetch: list[str], now_int: int, - ) -> list[Any]: + ) -> Sequence[object]: """ Execute Redis operations grouped by hash tag for cluster compatibility. @@ -792,13 +860,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): now_int: int - Current timestamp Returns: - List[Any] - List of cache values + 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[object] = [] for hash_tag, group_keys in key_groups.items(): try: @@ -821,7 +889,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def should_rate_limit( self, - descriptors: list[RateLimitDescriptor], + descriptors: Sequence[RateLimitDescriptor], parent_otel_span: Span | None = None, read_only: bool = False, skip_tpm_check: bool = False, @@ -866,7 +934,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[object] | 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 +956,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 @@ -947,16 +1013,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _collect_windowed_keys_and_gauges( self, - descriptors: list[RateLimitDescriptor], + descriptors: Sequence[RateLimitDescriptor], skip_tpm_check: bool, - ) -> tuple[list[str], dict[str, dict[str, Any]], list[ParallelRequestGauge]]: + ) -> tuple[list[str], dict[str, _WindowLimitMetadata], 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, _WindowLimitMetadata] = {} gauges: list[ParallelRequestGauge] = [] for descriptor in descriptors: descriptor_key = descriptor["key"] @@ -1012,7 +1078,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: object) -> int: """ In-flight count from a cached gauge value: a dict of slot_id -> acquire timestamp when the in-memory registry is authoritative, or @@ -1020,14 +1086,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if raw_value is None: return 0 - if isinstance(raw_value, dict): + slot_acquire_times = _object_dict_or_none(raw_value) + if slot_acquire_times is not None: cutoff = self._get_current_time().timestamp() - PARALLEL_REQUEST_SLOT_TTL_SECONDS - return sum(1 for ts in raw_value.values() if isinstance(ts, (int, float)) and ts >= cutoff) - return max(0, int(raw_value)) + return sum(1 for ts in slot_acquire_times.values() if isinstance(ts, (int, float)) and ts >= cutoff) + return max(0, _as_int(raw_value)) async def _check_parallel_request_gauges( self, - gauges: list[ParallelRequestGauge], + gauges: Sequence[ParallelRequestGauge], slot_id: str, parent_otel_span: Span | None = None, read_only: bool = False, @@ -1053,7 +1120,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): keys=gauge_keys, args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges], ) - counts = [max(0, int(value)) for value in raw_counts] + counts = [max(0, _as_int(value)) for value in raw_counts] except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500 verbose_proxy_logger.warning(f"parallel_count_script failed, using local mirror: {e!s}") counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) @@ -1088,22 +1155,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.warning(f"parallel_acquire_script failed, falling back to in-memory gauge: {e!s}") async with self._check_and_increment_lock: return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) - if int(raw[0]) == 1: - gauge = gauges[int(raw[1]) - 1] + if _as_int(raw[0]) == 1: + gauge = gauges[_as_int(raw[1]) - 1] return RateLimitResponse( overall_code="OVER_LIMIT", - statuses=[self._gauge_status(gauge, int(raw[2]), "OVER_LIMIT")], + statuses=[self._gauge_status(gauge, _as_int(raw[2]), "OVER_LIMIT")], ) statuses = [] for gauge, in_flight in zip(gauges, raw[1:]): await self.internal_usage_cache.async_set_cache( key=gauge["counter_key"], - value=int(in_flight), + value=_as_int(in_flight), ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, litellm_parent_otel_span=parent_otel_span, local_only=True, ) - statuses.append(self._gauge_status(gauge, int(in_flight), "OK")) + statuses.append(self._gauge_status(gauge, _as_int(in_flight), "OK")) return RateLimitResponse(overall_code="OK", statuses=statuses) async with self._check_and_increment_lock: @@ -1113,8 +1180,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, gauge_keys: list[str], parent_otel_span: Span | None = None, - ) -> list[int]: - values = await self.internal_usage_cache.async_batch_get_cache( + ) -> Sequence[int]: + values: Sequence[object] | None = await self.internal_usage_cache.async_batch_get_cache( keys=gauge_keys, parent_otel_span=parent_otel_span, local_only=True, @@ -1125,7 +1192,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _acquire_parallel_slots_in_memory( self, - gauges: list[ParallelRequestGauge], + gauges: Sequence[ParallelRequestGauge], slot_id: str, parent_otel_span: Span | None = None, ) -> RateLimitResponse: @@ -1143,14 +1210,16 @@ 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 = await self._get_local_cache_value( key=gauge["counter_key"], - litellm_parent_otel_span=parent_otel_span, - local_only=True, + parent_otel_span=parent_otel_span, ) - if isinstance(raw_value, dict): + slot_acquire_times = _object_dict_or_none(raw_value) + if slot_acquire_times is not None: registry: dict[str, float] | None = { - key: float(ts) for key, ts in raw_value.items() if isinstance(ts, (int, float)) and ts >= cutoff + key: float(ts) + for key, ts in slot_acquire_times.items() + if isinstance(key, str) and isinstance(ts, (int, float)) and ts >= cutoff } in_flight = len(registry or {}) elif raw_value is None: @@ -1158,7 +1227,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): in_flight = 0 else: registry = None - in_flight = max(0, int(raw_value)) + in_flight = max(0, _as_int(raw_value)) if in_flight + 1 > gauge["limit"]: return RateLimitResponse( overall_code="OVER_LIMIT", @@ -1205,7 +1274,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for counter_key, remaining in zip(counter_keys, raw): await self.internal_usage_cache.async_set_cache( key=counter_key, - value=max(0, int(remaining)), + value=max(0, _as_int(remaining)), ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, litellm_parent_otel_span=parent_otel_span, local_only=True, @@ -1218,19 +1287,21 @@ 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 = await self._get_local_cache_value( key=counter_key, - litellm_parent_otel_span=parent_otel_span, - local_only=True, + parent_otel_span=parent_otel_span, ) - if isinstance(raw_value, dict): - if slot_id not in raw_value: + slot_acquire_times = _object_dict_or_none(raw_value) + if slot_acquire_times is not None: + if slot_id not in slot_acquire_times: continue - new_value: dict[str, float] | int = {key: ts for key, ts in raw_value.items() if key != slot_id} + new_value: dict[object, object] | int = { + key: ts for key, ts in slot_acquire_times.items() if key != slot_id + } elif raw_value is None: continue else: - new_value = max(0, int(raw_value) - 1) + new_value = max(0, _as_int(raw_value) - 1) await self.internal_usage_cache.async_set_cache( key=counter_key, value=new_value, @@ -1241,8 +1312,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def atomic_check_and_increment_by_n( self, - descriptors: list[RateLimitDescriptor], - increments: list[dict[Literal["requests", "tokens"], int]], + descriptors: Sequence[RateLimitDescriptor], + increments: Sequence[Mapping[Literal["requests", "tokens"], int]], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """ @@ -1277,7 +1348,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[Sequence[str], Sequence[int], Sequence[_AtomicCounterMeta]]] = [] for descriptor, increment_amounts in zip(descriptors, increments): keys, args, meta = self._build_descriptor_atomic_payload( descriptor=descriptor, @@ -1300,7 +1371,9 @@ 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: Sequence[_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, @@ -1310,8 +1383,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_descriptor_atomic_payload( self, descriptor: RateLimitDescriptor, - increment_amounts: dict[Literal["requests", "tokens"], int], - ) -> tuple[list[str], list[Any], list[dict[str, Any]]]: + increment_amounts: Mapping[Literal["requests", "tokens"], int], + ) -> 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 +1398,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) + rlt: Literal["requests", "tokens"] = rate_limit_type if rlt == "requests": limit_value = rate_limit.get("requests_per_unit") inc_amount = int(increment_amounts.get("requests", 0) or 0) @@ -1365,7 +1438,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _atomic_lua_per_descriptor( self, - descriptor_groups: list[tuple[list[str], list[Any], list[dict[str, Any]]]], + descriptor_groups: Sequence[tuple[Sequence[str], Sequence[int], Sequence[_AtomicCounterMeta]]], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """ @@ -1374,7 +1447,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[Sequence[_AtomicCounterMeta]] = [] statuses: list[RateLimitStatus] = [] for _idx, (keys, args, meta) in enumerate(descriptor_groups): @@ -1397,7 +1470,9 @@ 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: Sequence[_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 +1490,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _refund_applied_descriptor_groups( self, - applied: list[list[dict[str, Any]]], + applied: Sequence[Sequence[_AtomicCounterMeta]], ) -> None: """ Decrement counters for descriptor groups already applied via Lua. @@ -1441,8 +1516,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_atomic_response( self, - raw: list[Any], - per_counter_meta: list[dict[str, Any]], + raw: Sequence[object], + per_counter_meta: Sequence[_AtomicCounterMeta], ) -> RateLimitResponse: """Convert Lua script return value to RateLimitResponse. @@ -1458,12 +1533,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not raw: return RateLimitResponse(overall_code="OK", statuses=[]) - status_code = int(raw[0]) + status_code = _as_int(raw[0]) if status_code == 1: # Over limit: { 1, counter_index (1-based), current_counter, limit } - descriptor_index = int(raw[1]) - 1 - current_counter = int(raw[2]) - limit = int(raw[3]) + descriptor_index = _as_int(raw[1]) - 1 + current_counter = _as_int(raw[2]) + limit = _as_int(raw[3]) meta = per_counter_meta[descriptor_index] return RateLimitResponse( overall_code="OVER_LIMIT", @@ -1484,7 +1559,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): RateLimitStatus( code="OK", current_limit=meta["current_limit"], - limit_remaining=max(0, meta["current_limit"] - int(new_counter)), + limit_remaining=max(0, meta["current_limit"] - _as_int(new_counter)), rate_limit_type=meta["rate_limit_type"], descriptor_key=meta["descriptor_key"], ) @@ -1493,7 +1568,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: Sequence[_AtomicCounterMeta], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """In-memory all-or-nothing check-and-increment. Caller holds lock. @@ -1508,27 +1583,23 @@ 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 = await self._get_local_cache_value( key=meta["window_key"], - litellm_parent_otel_span=parent_otel_span, - local_only=True, + parent_otel_span=parent_otel_span, ) - window_expired = window_start is None or (now_int - int(window_start)) >= window_size - current_counter = ( - 0 + window_expired = window_start is None or (now_int - _as_int(window_start)) >= window_size + counter_raw = ( + None if window_expired - else int( - await self.internal_usage_cache.async_get_cache( - key=meta["counter_key"], - litellm_parent_otel_span=parent_otel_span, - local_only=True, - ) - or 0 + else await self._get_local_cache_value( + key=meta["counter_key"], + parent_otel_span=parent_otel_span, ) ) + current_counter = 0 if window_expired else _as_int(counter_raw or 0) over_limit = ( current_counter + meta["increment"] > meta["current_limit"] if meta["increment"] > 0 @@ -1581,7 +1652,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def reserve_tpm_tokens( self, - descriptors: list[RateLimitDescriptor], + descriptors: Sequence[RateLimitDescriptor], estimated_tokens: int, parent_otel_span: Span | None = None, ) -> RateLimitResponse: @@ -1595,7 +1666,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): atomicity (Lua on Redis, asyncio-locked DualCache otherwise) to the shared primitive. """ - tpm_descriptors: list[RateLimitDescriptor] = [ + tpm_descriptors: Sequence[RateLimitDescriptor] = [ RateLimitDescriptor( key=d["key"], value=d["value"], @@ -1610,7 +1681,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not tpm_descriptors: return RateLimitResponse(overall_code="OK", statuses=[]) - increments: list[dict[Literal["requests", "tokens"], int]] = [ + increments: Sequence[Mapping[Literal["requests", "tokens"], int]] = [ {"tokens": estimated_tokens} for _ in tpm_descriptors ] return await self.atomic_check_and_increment_by_n( @@ -1621,7 +1692,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def create_organization_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None = None - ) -> list[RateLimitDescriptor]: + ) -> Sequence[RateLimitDescriptor]: descriptors: list[RateLimitDescriptor] = [] # Global org rate limits @@ -1920,7 +1991,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 @@ -1931,7 +2002,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Resolve the agent_id from either the API key or request metadata. Key-level agent_id takes precedence over metadata/header-supplied agent_id. """ - key_agent_id = getattr(user_api_key_dict, "agent_id", None) + key_agent_id_raw: object = getattr(user_api_key_dict, "agent_id", None) + key_agent_id = _str_or_none(key_agent_id_raw) if key_agent_id: return key_agent_id metadata = data.get("metadata") or {} @@ -1969,8 +2041,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 +2056,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 +2316,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) -> _CallTypeRateLimiter | None: """Get the rate limiter for the call type.""" if call_type == "acreate_batch": batch_limiter = self._get_batch_rate_limiter() @@ -2371,7 +2443,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, - data: dict, + data: dict[str, object], call_type: str, ): """ @@ -2402,7 +2474,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # For dynamic mode, check if the model has recent failures model_has_failures = False - requested_model = data.get("model", None) + requested_model = _str_or_none(data.get("model", None)) if ( self._is_dynamic_rate_limiting_enabled( @@ -2604,7 +2676,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return pipeline_operations def _get_total_tokens_from_usage( - self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"] + self, usage: "Usage | _TokenUsageDict | None", rate_limit_type: Literal["output", "input", "total"] ) -> int: """ Get total tokens from response usage for rate limiting. @@ -2616,34 +2688,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): total_tokens = 0 cached_tokens = 0 - if usage: - if isinstance(usage, Usage): - if rate_limit_type == "output": - total_tokens = usage.completion_tokens or 0 - elif rate_limit_type == "input": - total_tokens = usage.prompt_tokens or 0 - elif rate_limit_type == "total": - total_tokens = usage.total_tokens or 0 + if isinstance(usage, Usage): + if rate_limit_type == "output": + total_tokens = usage.completion_tokens or 0 + elif rate_limit_type == "input": + total_tokens = usage.prompt_tokens or 0 + elif rate_limit_type == "total": + total_tokens = usage.total_tokens or 0 - # Get cached tokens to exclude from input/total - if rate_limit_type in ("input", "total"): - if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: - cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + # Get cached tokens to exclude from input/total + if rate_limit_type in ("input", "total"): + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: + cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 - elif isinstance(usage, dict): - # Responses API usage comes as a dict - if rate_limit_type == "output": - total_tokens = usage.get("completion_tokens", 0) or 0 - elif rate_limit_type == "input": - total_tokens = usage.get("prompt_tokens", 0) or 0 - elif rate_limit_type == "total": - total_tokens = usage.get("total_tokens", 0) or 0 + elif isinstance(usage, dict) and usage: + # Responses API usage comes as a dict + if rate_limit_type == "output": + total_tokens = usage.get("completion_tokens", 0) or 0 + elif rate_limit_type == "input": + total_tokens = usage.get("prompt_tokens", 0) or 0 + elif rate_limit_type == "total": + total_tokens = usage.get("total_tokens", 0) or 0 - # Get cached tokens from dict - if rate_limit_type in ("input", "total"): - prompt_details = usage.get("prompt_tokens_details") or {} - if isinstance(prompt_details, dict): - cached_tokens = prompt_details.get("cached_tokens", 0) or 0 + # Get cached tokens from dict + if rate_limit_type in ("input", "total"): + prompt_details = usage.get("prompt_tokens_details") or {} + if isinstance(prompt_details, dict): + cached_tokens = prompt_details.get("cached_tokens", 0) or 0 # Subtract cached tokens for input/total (providers don't count them) if cached_tokens > 0: @@ -2652,7 +2723,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return total_tokens @staticmethod - def _aggregate_only_total_tokens(usage: Usage | dict | None) -> int: + def _aggregate_only_total_tokens(usage: "Usage | _TokenUsageDict | None") -> int: """Total for usage that carries no input/output split, else 0. A source that can only report one number for the whole request (a @@ -2767,15 +2838,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @staticmethod def _merge_ratelimit_statuses_into_additional_headers( - additional_headers: dict[str, Any], + additional_headers: Mapping[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"] @@ -2784,8 +2855,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _collect_tpm_scope_targets( self, - standard_logging_metadata: dict[str, Any], - kwargs: Any, + standard_logging_metadata: Mapping[str, object], + kwargs: Mapping[str, object], model_group: str | None, ) -> list[tuple[str, str]]: """ @@ -2795,16 +2866,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): the emitter; this helper just lists the candidate scopes so callers can split reserved-vs-unreserved. """ - user_api_key = standard_logging_metadata.get("user_api_key_hash") - user_api_key_user_id = standard_logging_metadata.get("user_api_key_user_id") - user_api_key_team_id = standard_logging_metadata.get("user_api_key_team_id") - user_api_key_organization_id = standard_logging_metadata.get("user_api_key_org_id") - user_api_key_project_id = standard_logging_metadata.get("user_api_key_project_id") - user_api_key_end_user_id = ( + user_api_key = _str_or_none(standard_logging_metadata.get("user_api_key_hash")) + user_api_key_user_id = _str_or_none(standard_logging_metadata.get("user_api_key_user_id")) + user_api_key_team_id = _str_or_none(standard_logging_metadata.get("user_api_key_team_id")) + user_api_key_organization_id = _str_or_none(standard_logging_metadata.get("user_api_key_org_id")) + user_api_key_project_id = _str_or_none(standard_logging_metadata.get("user_api_key_project_id")) + user_api_key_end_user_id = _str_or_none( kwargs.get("user") if isinstance(kwargs, dict) else None - ) or standard_logging_metadata.get("user_api_key_end_user_id") - agent_id = standard_logging_metadata.get("agent_id") - session_id = standard_logging_metadata.get("session_id") or standard_logging_metadata.get("trace_id") + ) or _str_or_none(standard_logging_metadata.get("user_api_key_end_user_id")) + agent_id = _str_or_none(standard_logging_metadata.get("agent_id")) + session_id = _str_or_none(standard_logging_metadata.get("session_id")) or _str_or_none( + standard_logging_metadata.get("trace_id") + ) targets: list[tuple[str, str]] = [] if user_api_key: @@ -2882,8 +2955,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.""" @@ -2893,8 +2966,8 @@ 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_object = _str_object_dict_or_none(kwargs.get("standard_logging_object")) or {} + standard_logging_metadata = _str_object_dict_or_none(standard_logging_object.get("metadata")) or {} model_group = get_model_group_from_litellm_kwargs(kwargs) @@ -2903,7 +2976,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # than parsed out of the body) carry their usage in # ``combined_usage_object`` instead, and would otherwise never charge # the TPM window. - _usage: Usage | dict | None = None + _usage: Usage | _TokenUsageDict | None = None if isinstance( response_obj, ( @@ -2967,7 +3040,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return pipeline_operations - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ): """ Update TPM usage on successful API calls by incrementing counters using pipeline """ @@ -3007,10 +3082,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def async_logging_hook( self, - kwargs: dict, - result: Any, + kwargs: dict[str, object], + result: object, call_type: str, - ) -> tuple[dict, Any]: + ) -> tuple[dict[str, object], object]: """ Mirror the pre-call rate-limit snapshot into the SLP so streaming success callbacks see the same ``x-ratelimit-*`` headers the @@ -3027,8 +3102,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 @@ -3056,7 +3131,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) standard_logging_object["hidden_params"] = hidden_params - response_hidden = getattr(response_obj, "_hidden_params", None) + response_hidden: object = getattr(response_obj, "_hidden_params", None) if isinstance(response_hidden, dict): existing = response_hidden.get("additional_headers") response_hidden["additional_headers"] = self._merge_ratelimit_statuses_into_additional_headers( @@ -3064,7 +3139,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): statuses=statuses, ) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ): """ On failure: decrement max_parallel_requests and refund the upfront TPM reservation only against the scopes the reservation actually @@ -3162,19 +3239,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if litellm_proxy_rate_limit_response is not None: # Update response headers - if hasattr(response, "_hidden_params"): - _hidden_params = getattr(response, "_hidden_params") + _hidden_params_raw: object = getattr(response, "_hidden_params", None) + + if isinstance(_hidden_params_raw, BaseModel): + _hidden_params: dict[str, object] | None = _hidden_params_raw.model_dump() else: - _hidden_params = None - - if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) - ): - if isinstance(_hidden_params, BaseModel): - _hidden_params = _hidden_params.model_dump() + _hidden_params = _str_object_dict_or_none(_hidden_params_raw) + if _hidden_params is not None: + existing_headers = _str_object_dict_or_none(_hidden_params.get("additional_headers")) _additional_headers = self._merge_ratelimit_statuses_into_additional_headers( - additional_headers=_hidden_params.get("additional_headers", {}) or {}, + additional_headers=existing_headers or {}, statuses=litellm_proxy_rate_limit_response["statuses"], ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 35b64963d4d..2248985279c 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 Callable, Coroutine, Mapping, Sequence +from typing import TYPE_CHECKING, Literal, Protocol, TypeVar, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field @@ -76,8 +76,99 @@ from litellm.types.router import ( ) from litellm.utils import get_utc_datetime +if TYPE_CHECKING: + from prisma.models import LiteLLM_ProxyModelTable as PrismaProxyModelTable + from prisma.types import LiteLLM_ProxyModelTableWhereInput + router = APIRouter() +_TableT_co = TypeVar("_TableT_co", covariant=True) + + +class _RepositoryTable(Protocol[_TableT_co]): + @property + def table(self) -> _TableT_co: ... + + +class _ProxyModelRow(Protocol): + @property + def model_id(self) -> str: ... + + @property + def litellm_params(self) -> Mapping[str, object]: ... + + def model_dump(self, *, exclude_none: bool = ...) -> dict[str, object]: ... + + def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... + + +class _ProxyModelQueryClient(Protocol): + async def find_many(self, *, where: "LiteLLM_ProxyModelTableWhereInput") -> "Sequence[PrismaProxyModelTable]": ... + + +class _ProxyModelTableClient(_ProxyModelQueryClient, Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> "_ProxyModelRow | None": ... + + async def create(self, *, data: Mapping[str, object]) -> LiteLLM_ProxyModelTable: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_ProxyModelRow": ... + + async def delete(self, *, where: Mapping[str, object]) -> "_ProxyModelRow | None": ... + + +class _TeamRow(Protocol): + @property + def team_id(self) -> str: ... + + @property + def models(self) -> Sequence[str]: ... + + def model_dump(self) -> dict[str, object]: ... + + +class _TeamTableClient(Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> "_TeamRow | None": ... + + async def update( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = ..., + ) -> LiteLLM_TeamTable: ... + + +class _ModelAliasRow(Protocol): + @property + def id(self) -> int: ... + + @property + def model_aliases(self) -> dict[str, str]: ... + + @property + def team(self) -> "_TeamRow | None": ... + + +class _ModelAliasTableClient(Protocol): + async def find_many(self, *, include: Mapping[str, object]) -> "Sequence[_ModelAliasRow]": ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_ModelAliasRow": ... + + +def _proxy_model_table(repository: "_RepositoryTable[_ProxyModelTableClient]") -> "_ProxyModelTableClient": + return repository.table + + +def _team_table(repository: "_RepositoryTable[_TeamTableClient]") -> "_TeamTableClient": + return repository.table + + +def _model_alias_table(repository: "_RepositoryTable[_ModelAliasTableClient]") -> "_ModelAliasTableClient": + return repository.table + + +_loads_json: Callable[[str], object] = json.loads + async def update_team(*args, **kwargs): """ @@ -96,15 +187,12 @@ 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 _proxy_model_table(ModelRepository(prisma_client)).find_unique(where={"model_id": model_id}) if not db_model: return None - deployment_pydantic_obj = Deployment(**db_model.model_dump(exclude_none=True)) + deployment_pydantic_obj = Deployment.model_validate(db_model.model_dump(exclude_none=True)) return deployment_pydantic_obj @@ -151,7 +239,9 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr merged_deployment_dict = DeploymentTypedDict( model_name=db_model.model_name, litellm_params=LiteLLMParamsTypedDict( - **db_model.litellm_params.model_dump(exclude_none=True) # type: ignore + **db_model.litellm_params.model_dump( + exclude_none=True + ) # any-ok: raw pydantic dump is Any-valued; TypedDict validation would drop extra="allow" params and reject encrypted values ), model_info=db_model.model_info.model_dump(exclude_none=True), ) @@ -163,7 +253,12 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.litellm_params: # Encrypt any sensitive values encrypted_params = { - k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() + k: encrypt_value_helper( + v + ) # any-ok: v is Any from pydantic model_dump; encrypt_value_helper handles non-str at runtime + for k, v in updated_patch.litellm_params.model_dump( + exclude_none=True + ).items() # any-ok: same pydantic model_dump seam } merged_deployment_dict["litellm_params"].update(encrypted_params) # type: ignore @@ -322,7 +417,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 _proxy_model_table(ModelRepository(prisma_client)).update( where={"model_id": model_id}, data=update_data, ) @@ -425,7 +520,7 @@ async def _set_model_blocked_status( param=None, ) - updated_model = await ModelRepository(prisma_client).table.update( + updated_model = await _proxy_model_table(ModelRepository(prisma_client)).update( where={"model_id": data.model_id}, data={ "blocked": blocked, @@ -459,7 +554,7 @@ async def _set_model_blocked_status( still_desired=still_desired_ids, ) - return updated_model + return LiteLLM_ProxyModelTable.model_validate(updated_model.model_dump()) except Exception as e: verbose_proxy_logger.exception(f"Error in model {action}: {e!s}") @@ -549,27 +644,30 @@ async def _add_model_to_db( # encrypt litellm params # _litellm_params_dict = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name = model_params.litellm_params.model - for k, v in _litellm_params_dict.items(): - encrypted_value = encrypt_value_helper(value=v, new_encryption_key=new_encryption_key) + for ( + k, + v, + ) in ( + _litellm_params_dict.items() + ): # any-ok: v is Any from pydantic dict(); encrypt_value_helper handles non-str at runtime + encrypted_value = encrypt_value_helper( + value=v, new_encryption_key=new_encryption_key + ) # any-ok: same pydantic dict() seam model_params.litellm_params[k] = encrypted_value - _data: dict = { + _data: dict[str, str | None] = { "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 _proxy_model_table(ModelRepository(prisma_client)).create(data=_data) else: - model_response = LiteLLM_ProxyModelTable(**_data) + model_response = LiteLLM_ProxyModelTable.model_validate(_data) return model_response @@ -757,8 +855,8 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient, table: Any | None = None -) -> list[LiteLLM_ProxyModelTable]: + team_id: str, prisma_client: PrismaClient, table: "_ProxyModelQueryClient | None" = None +) -> "list[PrismaProxyModelTable]": """ Fetch all deployments for a given team_id from the database. @@ -773,7 +871,7 @@ async def _get_team_deployments( existing transaction. """ prefix = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table + table = table or _proxy_model_table(ModelRepository(prisma_client)) response = await table.find_many( where={ "model_name": {"startswith": prefix}, @@ -783,7 +881,7 @@ async def _get_team_deployments( return [] # Confirm team_id in model_info (defensive check) - result = [] + result: list[PrismaProxyModelTable] = [] 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 +892,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. @@ -808,7 +906,7 @@ async def delete_team_models( Returns the model_ids that were deleted. """ deleted_model_ids: list[str] = [] - async with prisma_client.db.tx() as tx: + async with prisma_client.tx() as tx: for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] @@ -908,14 +1006,14 @@ 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}) + existing_team_row = await _team_table(TeamRepository(prisma_client)).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(TeamRepository(prisma_client)).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, @@ -941,7 +1039,7 @@ async def _update_existing_team_model_assignment( """ def _get_team_public_model_name( - model_info: dict | str | None, + model_info: object, ) -> str | None: parsed = model_info_as_mapping(model_info) if parsed is None: @@ -1050,7 +1148,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique( + _existing_team_row = await _team_table(TeamRepository(prisma_client)).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1079,7 +1177,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 _team_table(TeamRepository(prisma_client)).find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -1157,14 +1255,16 @@ async def delete_model( }, ) - model_in_db = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_info.id}) + model_in_db = await _proxy_model_table(ModelRepository(prisma_client)).find_unique( + where={"model_id": model_info.id} + ) if model_in_db is None: raise HTTPException( status_code=400, detail={"error": f"Model with id={model_info.id} not found in db"}, ) - model_params = Deployment(**model_in_db.model_dump()) + model_params = Deployment.model_validate(model_in_db.model_dump()) await ModelManagementAuthChecks.can_user_make_model_call( model_params=model_params, user_api_key_dict=user_api_key_dict, @@ -1180,7 +1280,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 _proxy_model_table(ModelRepository(prisma_client)).delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1253,9 +1353,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(ModelTableRepository(prisma_client)).find_many(include={"team": True}) + tasks: list[Coroutine[None, None, _ModelAliasRow]] = [] + 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 +1366,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(ModelTableRepository(prisma_client)).update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1481,7 +1581,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 +1589,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 _proxy_model_table(ModelRepository(prisma_client)).find_unique( + where={"model_id": _model_id} + ) if _existing_litellm_params is None: if llm_router is not None and llm_router.get_deployment(model_id=_model_id) is not None: @@ -1499,7 +1601,7 @@ async def update_model( ) else: raise Exception("model not found") - deployment = Deployment(**_existing_litellm_params.model_dump()) + deployment = Deployment.model_validate(_existing_litellm_params.model_dump()) await ModelManagementAuthChecks.can_user_make_model_call( model_params=deployment, @@ -1523,15 +1625,23 @@ async def update_model( _new_litellm_params_dict = model_params.litellm_params.dict(exclude_none=True) ### ENCRYPT PARAMS ### - for k, v in _new_litellm_params_dict.items(): - encrypted_value = encrypt_value_helper(value=v) + for ( + k, + v, + ) in ( + _new_litellm_params_dict.items() + ): # any-ok: v is Any from pydantic dict(); encrypt_value_helper handles non-str at runtime + encrypted_value = encrypt_value_helper(value=v) # any-ok: same pydantic dict() seam model_params.litellm_params[k] = encrypted_value ### MERGE WITH EXISTING DATA ### merged_dictionary = {} _mp = model_params.litellm_params.dict() - for key, value in _mp.items(): + for ( + key, + value, + ) in _mp.items(): # any-ok: value is Any from pydantic dict(); merged verbatim into the stored params if value is not None: merged_dictionary[key] = value elif key in _existing_litellm_params_dict and _existing_litellm_params_dict[key] is not None: @@ -1539,13 +1649,13 @@ async def update_model( else: pass - _data: dict = { - "litellm_params": json.dumps(merged_dictionary), # type: ignore + _data: dict[str, str] = { + "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response = await ModelRepository(prisma_client).table.update( + model_response = await _proxy_model_table(ModelRepository(prisma_client)).update( where={"model_id": _model_id}, - data=_data, # type: ignore + data=_data, ) # Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates) @@ -1787,7 +1897,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 = _loads_json(model_info) except (TypeError, ValueError): return None return parsed if isinstance(parsed, Mapping) else None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d274879f82a..284a2517c65 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -2383,7 +2383,7 @@ async def insert_sso_user( ) if result_openid and hasattr(result_openid, "provider"): - new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")} + new_user_request.metadata = {"auth_provider": result_openid.provider} response = await new_user( data=new_user_request, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1395fc9d32f..30c7234060b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -725,7 +725,7 @@ async def handle_bedrock_passthrough_router_model( user_max_tokens: int | None, user_api_base: str | None, version: str | None, -) -> Response | StreamingResponse: +) -> object: """ Handle Bedrock passthrough for router models (models defined in config.yaml). diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 78fa732a67b..7cb81be24d5 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -32,10 +32,12 @@ from __future__ import annotations import json import re -from typing import Any +from collections.abc import Mapping +from typing import TYPE_CHECKING, Literal, TypeVar from urllib.parse import quote, unquote from fastapi import HTTPException +from pydantic import JsonValue, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.managed_resources.isolation import ( @@ -43,13 +45,26 @@ from litellm.llms.base_llm.managed_resources.isolation import ( can_access_resource, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.repositories.table_repositories import ( - ManagedFileRepository, - ManagedObjectRepository, -) from litellm.types.llms.openai import OpenAIFileObject from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id +from .managed_id_rewriter_types import ( + ManagedFileRowLike, + ManagedObjectRowLike, + PassthroughListResponse, + as_file_row, + as_file_rows, + as_managed_file_reader, + as_managed_file_writer, + as_object_row, + as_object_rows, + file_table_actions, + object_table_actions, +) + +if TYPE_CHECKING: + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.utils import PrismaClient # --------------------------------------------------------------------------- # 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 @@ -196,8 +211,10 @@ class _RawIdGuardBudget: # to the upstream provider, so each caller only sees IDs they own. # --------------------------------------------------------------------------- +_ListResourceKind = Literal["files", "batches"] + # Maps (provider, canonical_path) -> "files" | "batches" -_LIST_ROUTE_TABLE: dict[tuple[str, str], str] = { +_LIST_ROUTE_TABLE: dict[tuple[str, str], _ListResourceKind] = { ("openai", "/v1/files"): "files", ("openai", "/v1/batches"): "batches", ("azure", "/v1/files"): "files", @@ -263,8 +280,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: CustomLogger | None, ) -> str: """ Resolve a single value that may be a passthrough managed ID. @@ -307,7 +324,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( + file_row = await as_managed_file_reader(managed_files_hook).get_unified_file_id( managed_id, litellm_parent_otel_span=None, ) @@ -322,8 +339,8 @@ 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 = as_file_row( + await file_table_actions(prisma_client).find_first(where={"unified_file_id": managed_id}) ) if db_row is not None: row_created_by = db_row.created_by @@ -338,8 +355,8 @@ 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 = as_object_row( + await object_table_actions(prisma_client).find_first(where={"unified_object_id": managed_id}) ) if obj_row is not None: row_created_by = obj_row.created_by @@ -372,7 +389,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,15 +415,15 @@ 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( - where={"flat_model_file_ids": {"has": raw_id}}, + candidates = as_file_rows( + await file_table_actions(prisma_client).find_many( + where={"flat_model_file_ids": {"has": raw_id}}, + ) ) except Exception: verbose_proxy_logger.debug("managed_id_rewriter: raw file-id guard lookup failed", exc_info=True) return - provider_rows = [ - row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider) - ] + provider_rows = [row for row in candidates if _managed_id_matches_provider(row.unified_file_id, provider)] if provider_rows and not any( can_access_resource(user_api_key_dict, row.created_by, row.team_id) for row in provider_rows ): @@ -419,8 +436,10 @@ 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( - where={"model_object_id": f"passthrough:{provider}:{raw_id}"} + existing = as_object_row( + await object_table_actions(prisma_client).find_first( + where={"model_object_id": f"passthrough:{provider}:{raw_id}"} + ) ) except Exception: verbose_proxy_logger.debug("managed_id_rewriter: raw object-id guard lookup failed", exc_info=True) @@ -434,7 +453,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: Mapping[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 +461,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 +474,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: CustomLogger | None, + file_object_snapshot: Mapping[str, object] | None = None, is_create_route: bool = True, ) -> str: """Return an existing managed file ID or mint + store a new one.""" @@ -479,16 +498,16 @@ 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( - where={"flat_model_file_ids": {"has": raw_id}}, - order={"created_at": "asc"}, + candidates = as_file_rows( + await file_table_actions(prisma_client).find_many( + where={"flat_model_file_ids": {"has": raw_id}}, + order={"created_at": "asc"}, + ) ) except Exception: candidates = [] verbose_proxy_logger.debug("managed_id_rewriter: file dedup lookup failed", exc_info=True) - provider_rows = [ - row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider) - ] + provider_rows = [row for row in candidates if _managed_id_matches_provider(row.unified_file_id, provider)] owned_row = next( (row for row in provider_rows if can_access_resource(user_api_key_dict, row.created_by, row.team_id)), None, @@ -525,7 +544,7 @@ async def _mint_or_reuse_file( ) if managed_files_hook is not None: try: - await managed_files_hook.store_unified_file_id( + await as_managed_file_writer(managed_files_hook).store_unified_file_id( file_id=managed_id, file_object=_build_managed_file_object(file_object_snapshot, managed_id), litellm_parent_otel_span=None, @@ -551,14 +570,15 @@ 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.""" if prisma_client is None: return raw_id + client = prisma_client # Namespace raw_id with provider so two providers that happen to issue # the same raw batch/response ID get distinct rows. The @unique constraint @@ -569,7 +589,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: ManagedObjectRowLike, 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 +618,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 object_table_actions(client).update( where={"unified_object_id": existing.unified_object_id}, data={ "file_object": json.dumps(body_snapshot), @@ -618,8 +638,8 @@ 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( - where={"model_object_id": namespaced_model_object_id} + existing = as_object_row( + await object_table_actions(client).find_first(where={"model_object_id": namespaced_model_object_id}) ) except Exception: verbose_proxy_logger.debug("managed_id_rewriter: object dedup lookup failed", exc_info=True) @@ -635,7 +655,7 @@ async def _mint_or_reuse_object( raw_id.split("_", 1)[0], ) try: - await ManagedObjectRepository(prisma_client).table.upsert( + await object_table_actions(client).upsert( where={"unified_object_id": managed_id}, data={ "create": { @@ -659,8 +679,8 @@ 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( - where={"model_object_id": namespaced_model_object_id} + raced = as_object_row( + await object_table_actions(client).find_first(where={"model_object_id": namespaced_model_object_id}) ) except Exception: raced = None @@ -677,14 +697,17 @@ async def _mint_or_reuse_object( return managed_id +_BODY_DICT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + async def rewrite_response_ids( provider: str, method: str, route: str, body: dict, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, + prisma_client: PrismaClient | None, + managed_files_hook: CustomLogger | None, ) -> dict: """ Mint managed IDs for raw provider values listed in @@ -714,7 +737,8 @@ async def rewrite_response_ids( # creates may degrade to a raw id on a cross-owner collision. is_create_route = "{" not in canonical - mutated = dict(body) # shallow copy; only return if something changed + typed_body = _BODY_DICT_ADAPTER.validate_python(body) + mutated = dict(typed_body) # shallow copy; only return if something changed changed = False def _record(field_name: str, raw_value: str, managed_id: str) -> None: @@ -746,7 +770,7 @@ async def rewrite_response_ids( managed_files_hook, # The file's own ``id`` carries the full upstream metadata; nested # references do not, so only the former is persisted as a snapshot. - file_object_snapshot=body if field_name == "id" else None, + file_object_snapshot=typed_body if field_name == "id" else None, is_create_route=is_create_route, ) _record(field_name, raw_value, managed_id) @@ -795,21 +819,24 @@ 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: +_FILE_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) + + +def _parse_file_object(file_object: object) -> dict[str, JsonValue] | None: """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 treat the result uniformly. """ - if isinstance(file_object, str): - try: - return json.loads(file_object) - except (TypeError, ValueError): - return None - return file_object + try: + if isinstance(file_object, str): + return _FILE_OBJECT_DICT_ADAPTER.validate_json(file_object) + return _FILE_OBJECT_DICT_ADAPTER.validate_python(file_object) + except (TypeError, ValueError): + return None -def _empty_list_response() -> dict[str, Any]: +def _empty_list_response() -> PassthroughListResponse: return { "object": "list", "data": [], @@ -819,7 +846,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: Mapping[str, str] | None) -> tuple[int, int]: params = query_params or {} try: raw_limit = int(params.get("limit", 20)) @@ -830,17 +857,17 @@ def _parse_list_limit(query_params: dict[str, Any] | None) -> tuple[int, int]: async def _build_list_where_with_cursor( - prisma_client: Any, - resource_kind: str, + prisma_client: PrismaClient, + resource_kind: _ListResourceKind, provider: str, - owner_filter: dict[str, Any], - query_params: dict[str, Any] | None, -) -> tuple[dict[str, Any], str]: + owner_filter: dict[str, object], + query_params: Mapping[str, str] | None, +) -> tuple[dict[str, object], 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: dict[str, object] = dict(owner_filter) fetch_order = "desc" cursor_id = after_id or before_id @@ -850,14 +877,14 @@ async def _build_list_where_with_cursor( if not cursor_id or not _managed_id_matches_provider(cursor_id, provider): return where, fetch_order - cursor_table = ( - ManagedFileRepository(prisma_client).table - if resource_kind == "files" - else ManagedObjectRepository(prisma_client).table - ) cursor_field = "unified_file_id" if resource_kind == "files" else "unified_object_id" try: - cursor_row = await cursor_table.find_first(where={**owner_filter, cursor_field: cursor_id}) + cursor_where: dict[str, object] = {**owner_filter, cursor_field: cursor_id} + cursor_row: ManagedFileRowLike | ManagedObjectRowLike | None + if resource_kind == "files": + cursor_row = as_file_row(await file_table_actions(prisma_client).find_first(where=cursor_where)) + else: + cursor_row = as_object_row(await object_table_actions(prisma_client).find_first(where=cursor_where)) if cursor_row is not None: if after_id: op = "lt" @@ -867,7 +894,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: dict[str, object] = { "OR": [ {"created_at": {op: cursor_row.created_at}}, { @@ -884,110 +911,117 @@ async def _build_list_where_with_cursor( return where, fetch_order -async def _fetch_list_rows( - prisma_client: Any, - resource_kind: str, - where: dict[str, Any], +async def _fetch_file_list_rows( + prisma_client: PrismaClient, + provider: str, + where: dict[str, object], fetch_order: str, fetch_limit: int, -) -> list[Any] | 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. +) -> list[ManagedFileRowLike] | None: + """Fetch list rows scoped to *provider* at the DB level: file rows carry + ``_passthrough_provider:{provider}`` in ``flat_model_file_ids`` (see + ``_mint_or_reuse_file``), since the file table has no provider column. + Pushing the scope into the query means a single DB round-trip serves the + page, with no application-layer scanning that could truncate large pools. + + 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. + + A DB failure returns ``None`` (fail closed) so the caller serves an empty + page and never falls through to the upstream provider. + """ + scoped_where: dict[str, object] = { + **where, + "flat_model_file_ids": {"has": _passthrough_provider_marker(provider)}, + } try: - if resource_kind == "files": - return await ManagedFileRepository(prisma_client).table.find_many( - where=where, + return as_file_rows( + await file_table_actions(prisma_client).find_many( + where=scoped_where, order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}], take=fetch_limit, ) - return await ManagedObjectRepository(prisma_client).table.find_many( - where={**where, "file_purpose": "batch"}, - order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}], - take=fetch_limit, ) except Exception: verbose_proxy_logger.warning("managed_id_rewriter: list DB query failed", exc_info=True) return None -async def _fetch_provider_scoped_list_rows( - prisma_client: Any, - resource_kind: str, +async def _fetch_batch_list_rows( + prisma_client: PrismaClient, provider: str, - where: dict[str, Any], + where: dict[str, object], fetch_order: str, - raw_limit: int, fetch_limit: int, -) -> tuple[list[Any], 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 - filters on directly: object rows namespace ``model_object_id`` as - ``passthrough:{provider}:{raw}`` (see ``_mint_or_reuse_object``) and file - rows carry ``_passthrough_provider:{provider}`` in ``flat_model_file_ids`` - (see ``_mint_or_reuse_file``), since the file table has no provider column. - Pushing the scope into the query means a single DB round-trip serves the - page, with no application-layer scanning that could truncate large pools. - - A DB failure returns an empty page (fail closed) so the caller never falls - through to the upstream provider. +) -> list[ManagedObjectRowLike] | None: + """Fetch list rows scoped to *provider* at the DB level: object rows + namespace ``model_object_id`` as ``passthrough:{provider}:{raw}`` (see + ``_mint_or_reuse_object``), so the query filters on it directly. Ordering + and failure semantics mirror ``_fetch_file_list_rows``. """ - scoped_where = dict(where) - if resource_kind == "files": - scoped_where["flat_model_file_ids"] = {"has": _passthrough_provider_marker(provider)} - else: - scoped_where["model_object_id"] = {"startswith": f"passthrough:{provider}:"} + scoped_where: dict[str, object] = { + **where, + "model_object_id": {"startswith": f"passthrough:{provider}:"}, + "file_purpose": "batch", + } + try: + return as_object_rows( + await object_table_actions(prisma_client).find_many( + where=scoped_where, + order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}], + take=fetch_limit, + ) + ) + except Exception: + verbose_proxy_logger.warning("managed_id_rewriter: list DB query failed", exc_info=True) + return None - rows = await _fetch_list_rows(prisma_client, resource_kind, scoped_where, fetch_order, fetch_limit) + +_RowT = TypeVar("_RowT") + + +def _page_from_rows(rows: list[_RowT] | None, raw_limit: int, fetch_order: str) -> tuple[list[_RowT], bool]: if rows is None: return [], False - effective_limit = min(raw_limit, 100) has_more = len(rows) > effective_limit page = rows[:effective_limit] if fetch_order == "asc": - page = list(reversed(page)) + return list(reversed(page)), has_more return page, has_more -def _serialize_file_list_item(row: Any) -> dict[str, Any]: - item: dict[str, Any] = { +def _serialize_file_list_item(row: ManagedFileRowLike) -> dict[str, JsonValue]: + item: dict[str, JsonValue] = { "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): + 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] = {} +def _serialize_batch_list_item(row: ManagedObjectRowLike) -> dict[str, JsonValue]: + item: dict[str, JsonValue] = {} file_object = _parse_file_object(row.file_object) - if isinstance(file_object, dict): + 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]: - 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) - - 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: Mapping[str, str] | None = None, +) -> PassthroughListResponse | None: """Query the DB for managed IDs the caller owns and return an OpenAI-style paginated list response. @@ -1020,21 +1054,25 @@ async def list_passthrough_ids_from_db( where, fetch_order = await _build_list_where_with_cursor( prisma_client, resource_kind, provider, owner_filter, query_params ) - page, has_more = await _fetch_provider_scoped_list_rows( - prisma_client, - resource_kind, - provider, - where, - fetch_order, - raw_limit, - fetch_limit, - ) if resource_kind == "files": - data = [_serialize_file_list_item(row) for row in page] + file_page, has_more = _page_from_rows( + await _fetch_file_list_rows(prisma_client, provider, where, fetch_order, fetch_limit), + raw_limit, + fetch_order, + ) + data = [_serialize_file_list_item(row) for row in file_page] + first_id = file_page[0].unified_file_id if file_page else None + last_id = file_page[-1].unified_file_id if file_page else None else: - data = [_serialize_batch_list_item(row) for row in page] + batch_page, has_more = _page_from_rows( + await _fetch_batch_list_rows(prisma_client, provider, where, fetch_order, fetch_limit), + raw_limit, + fetch_order, + ) + data = [_serialize_batch_list_item(row) for row in batch_page] + first_id = batch_page[0].unified_object_id if batch_page else None + last_id = batch_page[-1].unified_object_id if batch_page else None - first_id, last_id = _list_boundary_ids(page, resource_kind) verbose_proxy_logger.debug( "managed_id_rewriter: list served from DB provider=%s kind=%s count=%d admin=%s", provider, @@ -1060,8 +1098,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: CustomLogger | None, ) -> str: """ Walk URL path segments and resolve any passthrough managed IDs to raw @@ -1092,12 +1130,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: CustomLogger | None, +) -> dict[str, object] | None: """ Walk query param values and resolve any passthrough managed IDs. Returns *params* unchanged (same object) when nothing is resolved. @@ -1123,13 +1161,19 @@ async def rewrite_query_ids( return mutated if rewritten_keys else params +def _is_litellm_internal_key(key: object) -> bool: + if not isinstance(key, str): + return False + return key.startswith("litellm_") + + async def rewrite_body_ids( - body: dict[str, Any] | None, + body: dict[str, JsonValue] | 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: CustomLogger | None, +) -> dict[str, JsonValue] | None: """ Recursively walk a request body dict/list and resolve any passthrough managed IDs. Skips litellm internal keys (``litellm_*``). @@ -1140,15 +1184,15 @@ async def rewrite_body_ids( budget = _RawIdGuardBudget() - async def _walk(node: Any, depth: int) -> Any: + async def _walk(node: JsonValue, depth: int) -> JsonValue: if depth >= _MAX_BODY_REWRITE_DEPTH: return node if isinstance(node, dict): - result: dict[str, Any] = {} + result: dict[str, JsonValue] = {} 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_"): + if _is_litellm_internal_key(k): result[k] = v continue new_v = await _walk(v, depth + 1) @@ -1156,12 +1200,12 @@ async def rewrite_body_ids( if new_v is not v: changed_inner = True return result if changed_inner else node - elif isinstance(node, list): + if 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 - elif isinstance(node, str): + if isinstance(node, str): if is_managed(node): return await _resolve_one(node, provider, user_api_key_dict, prisma_client, managed_files_hook) await _guard_raw_provider_id(node, provider, user_api_key_dict, prisma_client, budget) @@ -1169,6 +1213,7 @@ async def rewrite_body_ids( return node rewritten = await _walk(body, 0) - if rewritten is not body: + if isinstance(rewritten, dict) and rewritten is not body: verbose_proxy_logger.debug("managed_id_rewriter: body ids rewritten provider=%s", provider) - return rewritten + return rewritten + return body diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter_types.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter_types.py new file mode 100644 index 00000000000..b260f88d37d --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter_types.py @@ -0,0 +1,148 @@ +""" +Typed boundary for ``managed_id_rewriter``'s untyped dependencies. + +The managed file/object Prisma tables are reached through +``PrismaTableRepository.table`` (untyped) and the enterprise managed-files +hook arrives as a bare ``CustomLogger``. This module narrows both to +structural protocols once, at the boundary, so the rewriter itself works +with fully typed values. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Mapping, Sequence +from datetime import datetime +from typing import TYPE_CHECKING, Literal, Protocol, TypedDict, runtime_checkable + +from pydantic import JsonValue + +from litellm.repositories.table_repositories import ( + ManagedFileRepository, + ManagedObjectRepository, +) + +if TYPE_CHECKING: + from litellm.integrations.custom_logger import CustomLogger + from litellm.models.managed_files import LiteLLM_ManagedFileTable + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import PrismaClient + from litellm.types.llms.openai import OpenAIFileObject + + +@runtime_checkable +class ManagedFileRowLike(Protocol): + unified_file_id: str + created_by: str | None + team_id: str | None + created_at: datetime | None + file_object: object + + +@runtime_checkable +class ManagedObjectRowLike(Protocol): + unified_object_id: str + created_by: str | None + team_id: str | None + created_at: datetime | None + file_object: object + + +@runtime_checkable +class ManagedTableActions(Protocol): + def find_first(self, where: Mapping[str, object]) -> Awaitable[object]: ... + + def find_many( + self, + where: Mapping[str, object], + order: Sequence[Mapping[str, str]] | Mapping[str, str] | None = None, + take: int | None = None, + ) -> Awaitable[Sequence[object]]: ... + + def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ... + + def upsert(self, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]) -> Awaitable[object]: ... + + +@runtime_checkable +class ManagedFileReaderLike(Protocol): + def get_unified_file_id( + self, + file_id: str, + litellm_parent_otel_span: object = None, + ) -> Awaitable[LiteLLM_ManagedFileTable | None]: ... + + +@runtime_checkable +class ManagedFileWriterLike(Protocol): + def store_unified_file_id( + self, + file_id: str, + file_object: OpenAIFileObject | None, + litellm_parent_otel_span: object, + model_mappings: dict[str, str], + user_api_key_dict: UserAPIKeyAuth, + ) -> Awaitable[None]: ... + + +class PassthroughListResponse(TypedDict): + object: Literal["list"] + data: list[dict[str, JsonValue]] + first_id: str | None + last_id: str | None + has_more: bool + + +class _TableHolder(Protocol): + @property + def table(self) -> object: ... + + +def _table_actions(holder: _TableHolder) -> ManagedTableActions: + table = holder.table + if isinstance(table, ManagedTableActions): + return table + raise TypeError("Prisma table object does not expose find/update/upsert actions") + + +def file_table_actions(prisma_client: PrismaClient) -> ManagedTableActions: + return _table_actions(ManagedFileRepository(prisma_client)) + + +def object_table_actions(prisma_client: PrismaClient) -> ManagedTableActions: + return _table_actions(ManagedObjectRepository(prisma_client)) + + +def as_managed_file_reader(hook: CustomLogger) -> ManagedFileReaderLike: + if isinstance(hook, ManagedFileReaderLike): + return hook + raise TypeError("managed_files hook does not expose get_unified_file_id") + + +def as_managed_file_writer(hook: CustomLogger) -> ManagedFileWriterLike: + if isinstance(hook, ManagedFileWriterLike): + return hook + raise TypeError("managed_files hook does not expose store_unified_file_id") + + +def as_file_row(record: object) -> ManagedFileRowLike | None: + if isinstance(record, ManagedFileRowLike): + return record + return None + + +def as_object_row(record: object) -> ManagedObjectRowLike | None: + if isinstance(record, ManagedObjectRowLike): + return record + return None + + +def as_file_rows(records: Sequence[object] | None) -> list[ManagedFileRowLike]: + if records is None: + return [] + return [record for record in records if isinstance(record, ManagedFileRowLike)] + + +def as_object_rows(records: Sequence[object] | None) -> list[ManagedObjectRowLike]: + if records is None: + return [] + return [record for record in records if isinstance(record, ManagedObjectRowLike)] diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index b957618d776..c36819d8b42 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -5,10 +5,11 @@ import json import posixpath import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from datetime import datetime +from functools import partial from itertools import groupby -from typing import Any, cast +from typing import Literal, Protocol, TypedDict, cast, runtime_checkable from urllib.parse import urlencode, urlparse import httpx @@ -23,7 +24,9 @@ from fastapi import ( WebSocket, status, ) +from fastapi import params as fastapi_params from fastapi.responses import StreamingResponse +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.websockets import WebSocketState from websockets.asyncio.client import connect @@ -91,13 +94,94 @@ router = APIRouter() pass_through_endpoint_logging = PassThroughEndpointLogging() + +@runtime_checkable +class RouteRegistrableApp(Protocol): + @property + def routes(self) -> Sequence[object]: ... + + def add_api_route( + self, + path: str, + endpoint: Callable[..., object], + *, + methods: list[str] | None = None, + dependencies: Sequence[fastapi_params.Depends] | None = None, + ) -> None: ... + + +@runtime_checkable +class TeamTableClient(Protocol): + def find_unique(self, where: dict[str, str]) -> Awaitable[object]: ... + + +class RegisteredPassthroughParams(TypedDict): + target: str + custom_headers: dict[str, object] | None + forward_headers: bool | None + merge_query_params: bool | None + default_query_params: dict[str, object] | None + dependencies: Sequence[object] | None + cost_per_request: float | None + guardrails: dict[str, object] | None + timeout: float | None + + +class RegisteredPassthroughRoute(TypedDict): + endpoint_id: str + path: str + type: Literal["exact", "subpath"] + methods: list[str] + auth: bool + passthrough_params: RegisteredPassthroughParams + + # Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: dict[str, dict[str, str | bool | list[str] | dict[str, Any]]] = {} +_registered_pass_through_routes: dict[str, RegisteredPassthroughRoute] = {} + +_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) +_OPTIONAL_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, object] | None] = TypeAdapter(dict[str, object] | None) + +_EMPTY_ROUTE_PARAMS: Mapping[str, object] = {} +_JSON_VALUE_ADAPTER: TypeAdapter[JsonValue] = TypeAdapter(JsonValue) +_JSON_DICT_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) +_CUSTOM_LOGGER_ADAPTER: TypeAdapter[CustomLogger] = TypeAdapter( + CustomLogger, config=ConfigDict(arbitrary_types_allowed=True) +) +_APP_ADAPTER: TypeAdapter[RouteRegistrableApp] = TypeAdapter( + RouteRegistrableApp, config=ConfigDict(arbitrary_types_allowed=True) +) +_TEAM_TABLE_ADAPTER: TypeAdapter[TeamTableClient] = TypeAdapter( + TeamTableClient, config=ConfigDict(arbitrary_types_allowed=True) +) +_TEAM_METADATA_ADAPTER: TypeAdapter[dict[str, JsonValue] | None] = TypeAdapter(dict[str, JsonValue] | None) +_ALLOWED_ROUTES_ADAPTER: TypeAdapter[list[JsonValue] | dict[str, JsonValue] | str | None] = TypeAdapter( + list[JsonValue] | dict[str, JsonValue] | str | None +) +_ENDPOINT_LIST_ADAPTER: TypeAdapter[list[dict[str, object] | PassThroughGenericEndpoint] | None] = TypeAdapter( + list[dict[str, object] | PassThroughGenericEndpoint] | None +) -def get_response_body(response: httpx.Response) -> dict | None: +def _parse_json_value(raw: str | bytes) -> JsonValue: + return _JSON_VALUE_ADAPTER.validate_python(json.loads(raw)) + + +def _json_dict_or_none(value: JsonValue) -> dict[str, JsonValue] | None: + return value if isinstance(value, dict) else None + + +def _truthiness_or_none(value: JsonValue) -> bool | None: + return None if value is None else bool(value) + + +def _read_header(headers: Mapping[str, str], key: str) -> str | None: + return headers.get(key) + + +def get_response_body(response: httpx.Response) -> dict[str, object] | None: try: - return response.json() + return _OBJECT_DICT_ADAPTER.validate_python(response.json()) except Exception: return None @@ -174,9 +258,9 @@ async def chat_completion_pass_through_endpoint( body = await request.body() body_str = body.decode() try: - data = ast.literal_eval(body_str) + data = _OBJECT_DICT_ADAPTER.validate_python(ast.literal_eval(body_str)) except Exception: - data = json.loads(body_str) + data = _OBJECT_DICT_ADAPTER.validate_python(json.loads(body_str)) data["adapter_id"] = adapter_id @@ -848,7 +932,7 @@ async def pass_through_request( # parsed dict (hooks mutate it, breaking the signature / Content-Length). # Tolerate request objects without `state` (test fixtures) and only honor # values httpx accepts for `content=`. - _request_state = getattr(request, "state", None) + _request_state = request.state if hasattr(request, "state") else None state_raw_body: str | bytes | None = ( getattr(_request_state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, None) if _request_state is not None @@ -1125,15 +1209,17 @@ async def pass_through_request( else: # SigV4-signed callers (Bedrock) supply the exact pre-signed bytes; # otherwise httpx encodes the parsed JSON dict as before. - body_kwargs: dict[str, Any] = ( - {"content": state_raw_body} if state_raw_body is not None else {"json": _parsed_body} - ) - req = async_client.build_request( + build_streaming_request = partial( + async_client.build_request, request.method, url, params=requested_query_params, headers=headers, - **body_kwargs, + ) + req = ( + build_streaming_request(content=state_raw_body) + if state_raw_body is not None + else build_streaming_request(json=_OPTIONAL_OBJECT_DICT_ADAPTER.validate_python(_parsed_body)) ) response = await async_client.send(req, stream=stream) @@ -1583,7 +1669,12 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di async def _parse_request_data_by_content_type( request: Request, -) -> tuple[Any | None, Any | None, Any | None, Any | None]: +) -> tuple[ + dict[str, JsonValue] | dict[str, str] | str | StarletteUploadFile | None, + dict[str, JsonValue] | str | StarletteUploadFile | None, + None, + bool | None, +]: """ Parse request data based on content type. @@ -1594,18 +1685,18 @@ async def _parse_request_data_by_content_type( """ content_type = request.headers.get("content-type", "") - query_params_data = None - custom_body_data = None + query_params_data: dict[str, JsonValue] | dict[str, str] | str | StarletteUploadFile | None = None + custom_body_data: dict[str, JsonValue] | str | StarletteUploadFile | None = None file_data = None - stream = None + stream: bool | None = None if "application/json" in content_type: # ✅ Handle JSON try: - body = await request.json() - query_params_data = body.get("query_params") - custom_body_data = body.get("custom_body") - stream = body.get("stream") + body = _JSON_DICT_ADAPTER.validate_python(await request.json()) + query_params_data = _json_dict_or_none(body.get("query_params")) + custom_body_data = _json_dict_or_none(body.get("custom_body")) + stream = _truthiness_or_none(body.get("stream")) except json.JSONDecodeError: # Handle requests with no body (e.g., DELETE requests) pass @@ -1613,13 +1704,14 @@ async def _parse_request_data_by_content_type( # ✅ Try to parse as JSON first (handles misconfigured clients sending JSON with multipart content-type) # If that fails, skip parsing - pass_through_request will handle actual multipart try: - body = await request.json() + body = _JSON_DICT_ADAPTER.validate_python(await request.json()) # Successfully parsed as JSON - treat as JSON body - query_params_data = body.get("query_params") - custom_body_data = body.get("custom_body") - stream = body.get("stream") + raw_custom_body = body.get("custom_body") + query_params_data = _json_dict_or_none(body.get("query_params")) + custom_body_data = _json_dict_or_none(raw_custom_body) + stream = _truthiness_or_none(body.get("stream")) # If custom_body is not set, use the entire body - if custom_body_data is None and body: + if raw_custom_body is None and body: custom_body_data = body except (json.JSONDecodeError, Exception): # Not JSON - this is actual multipart data @@ -1643,17 +1735,17 @@ async def _parse_request_data_by_content_type( def create_pass_through_route( endpoint, target: str, - custom_headers: Mapping[str, Any] | None = None, + custom_headers: Mapping[str, object] | None = None, _forward_headers: bool | None = False, _merge_query_params: bool | None = False, - dependencies: list | None = None, + dependencies: Sequence[fastapi_params.Depends] | None = None, include_subpath: bool | None = False, cost_per_request: float | None = None, custom_llm_provider: str | None = None, is_streaming_request: bool | None = False, query_params: dict | None = None, default_query_params: dict | None = None, - guardrails: dict[str, Any] | None = None, + guardrails: dict[str, object] | None = None, config_file_path: str | None = None, timeout: float | None = None, ): @@ -1665,7 +1757,9 @@ def create_pass_through_route( if isinstance(target, CustomLogger): adapter = target else: - adapter = get_instance_fn(value=target, config_file_path=config_file_path) + adapter = _CUSTOM_LOGGER_ADAPTER.validate_python( + get_instance_fn(value=target, config_file_path=config_file_path) + ) adapter_id = str(uuid.uuid4()) litellm.adapters = [{"id": adapter_id, "adapter": adapter}] @@ -1725,7 +1819,7 @@ def create_pass_through_route( status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail=f"Method {request.method} is not allowed for pass-through endpoint {path}.", ) - target_params = { + target_params: dict[str, object] = { "target": target, "custom_headers": custom_headers, "forward_headers": _forward_headers, @@ -1736,7 +1830,7 @@ def create_pass_through_route( } if passthrough_params is not None: - target_params.update(passthrough_params.get("passthrough_params", {})) + target_params.update(passthrough_params.get("passthrough_params", _EMPTY_ROUTE_PARAMS)) # Extract and cast parameters with proper types param_target = target_params.get("target") or target @@ -1757,12 +1851,18 @@ def create_pass_through_route( # Ensure custom_headers is a dict. Botocore returns a HeadersDict # for SigV4-prepared requests, which is a Mapping but not a dict. - headers_dict = dict(param_custom_headers) if isinstance(param_custom_headers, Mapping) else {} + headers_dict = ( + _OBJECT_DICT_ADAPTER.validate_python(param_custom_headers) + if isinstance(param_custom_headers, Mapping) + else {} + ) # Ensure query_params and custom_body are dicts or None - final_query_params = query_params_data if isinstance(query_params_data, dict) else {} + final_query_params: dict[str, object] = ( + {key: value for key, value in query_params_data.items()} if isinstance(query_params_data, dict) else {} + ) if query_params: - final_query_params.update(query_params) + final_query_params.update(_OBJECT_DICT_ADAPTER.validate_python(query_params)) # Programmatic callers set LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY on # request.state (see Bedrock proxy). Parsed JSON envelope otherwise. state_custom_body: dict | None = getattr( @@ -1884,7 +1984,7 @@ async def websocket_passthrough_request( # Initialize tracking variables start_time = datetime.now() - websocket_messages: list[dict[str, Any]] = [] + websocket_messages: list[JsonValue] = [] litellm_call_id = str(uuid.uuid4()) verbose_proxy_logger.info(f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}") @@ -1977,10 +2077,9 @@ async def websocket_passthrough_request( ) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_data: dict[str, Any] = {} - websocket_data = await proxy_logging_obj.pre_call_hook( + await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, - data=websocket_data, + data={}, call_type="pass_through_endpoint", ) @@ -2007,14 +2106,14 @@ async def websocket_passthrough_request( text_data = message.get("text") bytes_data = message.get("bytes") - if text_data is not None: + if isinstance(text_data, str): # Try to extract model from client setup message for Vertex AI Live if endpoint and "/vertex_ai/live" in endpoint: verbose_proxy_logger.debug( f"WebSocket passthrough ({endpoint}): Processing client message for model extraction" ) try: - client_message = json.loads(text_data) + client_message = _parse_json_value(text_data) if isinstance(client_message, dict) and "setup" in client_message: setup_data = client_message["setup"] verbose_proxy_logger.debug( @@ -2044,14 +2143,14 @@ async def websocket_passthrough_request( verbose_proxy_logger.debug( f"WebSocket passthrough ({endpoint}): Client message does not contain setup data" ) - except (json.JSONDecodeError, KeyError, TypeError) as e: + except (json.JSONDecodeError, ValidationError, KeyError, TypeError) as e: verbose_proxy_logger.debug( f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}" ) # Not a JSON message or doesn't contain setup data await upstream_ws.send(text_data) - elif bytes_data is not None: + elif isinstance(bytes_data, bytes): await upstream_ws.send(bytes_data) except asyncio.CancelledError: raise @@ -2069,7 +2168,7 @@ async def websocket_passthrough_request( # Ensure raw_response is bytes before decoding if isinstance(raw_response, str): raw_response = raw_response.encode("ascii") - setup_response = json.loads(raw_response.decode("ascii")) + setup_response = _parse_json_value(raw_response.decode("ascii")) verbose_proxy_logger.debug(f"Setup response: {setup_response}") # Extract model and provider from setup response for Vertex AI Live @@ -2106,17 +2205,17 @@ async def websocket_passthrough_request( await websocket.send_bytes(upstream_message) # Parse and collect for cost tracking try: - message_data = json.loads(upstream_message.decode()) + message_data = _parse_json_value(upstream_message.decode()) websocket_messages.append(message_data) - except (json.JSONDecodeError, UnicodeDecodeError): + except (json.JSONDecodeError, ValidationError, UnicodeDecodeError): pass else: await websocket.send_text(upstream_message) # Parse and collect for cost tracking try: - message_data = json.loads(upstream_message) + message_data = _parse_json_value(upstream_message) websocket_messages.append(message_data) - except json.JSONDecodeError: + except (json.JSONDecodeError, ValidationError): pass except (ConnectionClosedOK, ConnectionClosedError) as e: @@ -2272,7 +2371,7 @@ async def websocket_passthrough_request( def _is_streaming_response(response: httpx.Response) -> bool: - _content_type = response.headers.get("content-type") + _content_type = _read_header(response.headers, "content-type") if _content_type is not None and "text/event-stream" in _content_type: return True return False @@ -2290,7 +2389,7 @@ def _should_buffer_passthrough_response(response: httpx.Response) -> bool: """ if response.status_code >= 400: return True - media_type = response.headers.get("content-type", "").split(";")[0].strip().lower() + media_type = (_read_header(response.headers, "content-type") or "").split(";")[0].strip().lower() return media_type in ("", "application/json") or media_type.endswith("+json") @@ -2342,7 +2441,7 @@ async def _relay_passthrough_response_bytes( ) -def _extract_model_from_vertex_ai_setup(setup_response: dict) -> str | None: +def _extract_model_from_vertex_ai_setup(setup_response: JsonValue) -> str | None: """ Extract the model name from Vertex AI Live setup response. @@ -2381,7 +2480,7 @@ class SafeRouteAdder: """ @staticmethod - def _is_path_registered(app: FastAPI, path: str, methods: list[str]) -> bool: + def _is_path_registered(app: RouteRegistrableApp, path: str, methods: list[str]) -> bool: """ Check if a path with any of the specified methods is already registered on the app. @@ -2406,11 +2505,11 @@ class SafeRouteAdder: @staticmethod def add_api_route_if_not_exists( - app: FastAPI, + app: RouteRegistrableApp, path: str, - endpoint: Any, + endpoint: Callable[..., object], methods: list[str], - dependencies: list | None = None, + dependencies: Sequence[fastapi_params.Depends] | None = None, ) -> bool: """ Add an API route to the app only if it doesn't already exist. @@ -2450,13 +2549,13 @@ class SafeRouteAdder: class InitPassThroughEndpointHelpers: @staticmethod def add_exact_path_route( - app: FastAPI, + app: RouteRegistrableApp, path: str, target: str, custom_headers: dict | None, forward_headers: bool | None, merge_query_params: bool | None, - dependencies: list | None, + dependencies: Sequence[fastapi_params.Depends] | None, cost_per_request: float | None, endpoint_id: str, guardrails: dict | None = None, @@ -2533,13 +2632,13 @@ class InitPassThroughEndpointHelpers: @staticmethod def add_subpath_route( - app: FastAPI, + app: RouteRegistrableApp, path: str, target: str, custom_headers: dict | None, forward_headers: bool | None, merge_query_params: bool | None, - dependencies: list | None, + dependencies: Sequence[fastapi_params.Depends] | None, cost_per_request: float | None, endpoint_id: str, guardrails: dict | None = None, @@ -2624,15 +2723,14 @@ class InitPassThroughEndpointHelpers: ] for key in keys_to_remove: route_info = _registered_pass_through_routes[key] - path = route_info.get("path") - if isinstance(path, str): - openai_routes = LiteLLMRoutes.openai_routes.value - if path in openai_routes: - openai_routes.remove(path) - if route_info.get("type") == "subpath": - wildcard_path = path.rstrip("/") + "/*" - if wildcard_path in openai_routes: - openai_routes.remove(wildcard_path) + path = route_info["path"] + openai_routes = LiteLLMRoutes.openai_routes.value + if path in openai_routes: + openai_routes.remove(path) + if route_info.get("type") == "subpath": + wildcard_path = path.rstrip("/") + "/*" + if wildcard_path in openai_routes: + openai_routes.remove(wildcard_path) del _registered_pass_through_routes[key] verbose_proxy_logger.debug("Removed pass-through route from registry: %s", key) @@ -2699,7 +2797,7 @@ class InitPassThroughEndpointHelpers: return False @staticmethod - def get_registered_pass_through_route(route: str, method: str | None = None) -> dict[str, Any] | None: + def get_registered_pass_through_route(route: str, method: str | None = None) -> RegisteredPassthroughRoute | None: """Get passthrough params for a given route and optionally filter by HTTP method""" comparison_route = InitPassThroughEndpointHelpers._route_for_registry_lookup(route) for key in _registered_pass_through_routes: @@ -2712,7 +2810,7 @@ class InitPassThroughEndpointHelpers: # but keep supporting test fixtures / older registry entries that # only encoded methods in the route key. methods_entry = _registered_pass_through_routes[key].get("methods", []) - route_methods: list[str] = methods_entry if isinstance(methods_entry, list) else [] + route_methods: list[str] = methods_entry if methods_entry else [] if not route_methods and len(parts) == 4: route_methods = parts[3].split(",") @@ -2740,14 +2838,27 @@ def _get_combined_pass_through_endpoints( return pass_through_endpoints + config_pass_through_endpoints +class _RegisteredEndpointFields(BaseModel): + path: str | None = None + target: str | None = None + headers: dict[str, object] | None = None + forward_headers: bool | None = None + merge_query_params: bool | None = None + default_query_params: dict[str, object] | None = None + guardrails: dict[str, object] | None = None + methods: list[str] | None = None + cost_per_request: float | None = None + timeout: float | None = None + + async def _register_pass_through_endpoint( - endpoint: dict[str, Any] | PassThroughGenericEndpoint, + endpoint: dict[str, object] | PassThroughGenericEndpoint, app: FastAPI, premium_user: bool, visited_endpoints: set[str], config_file_path: str | None = None, ) -> None: - endpoint_data: dict[str, Any] + endpoint_data: dict[str, object] if isinstance(endpoint, PassThroughGenericEndpoint): endpoint_data = endpoint.model_dump() else: @@ -2757,17 +2868,18 @@ async def _register_pass_through_endpoint( endpoint_data["id"] = str(uuid.uuid4()) endpoint_id = cast(str, endpoint_data["id"]) - target = endpoint_data.get("target") - path = endpoint_data.get("path") + endpoint_fields = _RegisteredEndpointFields.model_validate(endpoint_data) + target = endpoint_fields.target + path = endpoint_fields.path if path is None: raise ValueError("Path is required for pass-through endpoint") - custom_headers = await set_env_variables_in_header(custom_headers=endpoint_data.get("headers")) - forward_headers = endpoint_data.get("forward_headers") - merge_query_params = endpoint_data.get("merge_query_params") - default_query_params = endpoint_data.get("default_query_params") + custom_headers = await set_env_variables_in_header(custom_headers=endpoint_fields.headers) + forward_headers = endpoint_fields.forward_headers + merge_query_params = endpoint_fields.merge_query_params + default_query_params = endpoint_fields.default_query_params auth = endpoint_data.get("auth") - dependencies = None + dependencies: list[fastapi_params.Depends] | None = None auth_enforced = auth is not None and str(auth).lower() == "true" if auth_enforced: @@ -2782,10 +2894,10 @@ async def _register_pass_through_endpoint( if target is None: return - guardrails = endpoint_data.get("guardrails") - methods = endpoint_data.get("methods") - cost_per_request = endpoint_data.get("cost_per_request") - timeout = endpoint_data.get("timeout") + guardrails = endpoint_fields.guardrails + methods = endpoint_fields.methods + cost_per_request = endpoint_fields.cost_per_request + timeout = endpoint_fields.timeout verbose_proxy_logger.debug("Initializing pass through endpoint: %s (ID: %s)", path, endpoint_id) InitPassThroughEndpointHelpers.add_exact_path_route( @@ -2925,12 +3037,12 @@ def _get_pass_through_endpoints_from_config() -> list[PassThroughGenericEndpoint if isinstance(endpoint, dict): endpoint_dict = dict(endpoint) endpoint_dict["is_from_config"] = True - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) elif isinstance(endpoint, PassThroughGenericEndpoint): # Create a copy with is_from_config=True endpoint_dict = endpoint.model_dump() endpoint_dict["is_from_config"] = True - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) except ValidationError as e: verbose_proxy_logger.warning( "Skipping malformed pass-through endpoint from config: %s", @@ -2954,10 +3066,10 @@ async def _get_pass_through_endpoints_from_db( response: ConfigFieldInfo = await get_config_general_settings( field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict ) + pass_through_endpoint_data = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value")) except Exception: return [] - pass_through_endpoint_data: list | None = response.field_value if pass_through_endpoint_data is None: return [] @@ -2965,14 +3077,11 @@ async def _get_pass_through_endpoints_from_db( if endpoint_id is None: # Return all endpoints from DB, mark as not from config for endpoint in pass_through_endpoint_data: - if isinstance(endpoint, dict): - endpoint_dict = dict(endpoint) - endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) - elif isinstance(endpoint, PassThroughGenericEndpoint): - endpoint_dict = endpoint.model_dump() - endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + endpoint_dict = ( + endpoint.model_dump() if isinstance(endpoint, PassThroughGenericEndpoint) else dict(endpoint) + ) + endpoint_dict["is_from_config"] = False + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) else: # Find specific endpoint by ID found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) @@ -2983,7 +3092,7 @@ async def _get_pass_through_endpoints_from_db( else dict(found_endpoint) ) endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) return returned_endpoints @@ -3008,9 +3117,8 @@ async def _filter_endpoints_by_team_allowed_routes( HTTPException: If team is not found """ # retrieve team from db - team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id}, - ) + team_table = _TEAM_TABLE_ADAPTER.validate_python(getattr(TeamRepository(prisma_client), "table")) + team = await team_table.find_unique(where={"team_id": team_id}) if team is None: raise HTTPException( status_code=404, @@ -3018,14 +3126,15 @@ async def _filter_endpoints_by_team_allowed_routes( ) # retrieve team metadata - team_metadata = team.metadata - if team_metadata is not None and team_metadata.get("allowed_passthrough_routes") is not None: + team_metadata = _TEAM_METADATA_ADAPTER.validate_python(getattr(team, "metadata")) + allowed_routes = ( + _ALLOWED_ROUTES_ADAPTER.validate_python(team_metadata.get("allowed_passthrough_routes")) + if team_metadata is not None + else None + ) + if allowed_routes is not None: ## FILTER pass_through_endpoints by allowed_passthrough_routes - pass_through_endpoints = [ - endpoint - for endpoint in pass_through_endpoints - if endpoint.path in team_metadata.get("allowed_passthrough_routes") - ] + pass_through_endpoints = [endpoint for endpoint in pass_through_endpoints if endpoint.path in allowed_routes] return pass_through_endpoints @@ -3115,7 +3224,7 @@ async def update_pass_through_endpoints( detail={"error": "No pass-through endpoints found"}, ) - pass_through_endpoint_data: list | None = response.field_value + pass_through_endpoint_data = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value")) if pass_through_endpoint_data is None: raise HTTPException( status_code=404, @@ -3134,7 +3243,7 @@ async def update_pass_through_endpoints( # Find the index for updating the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint + _endpoint = PassThroughGenericEndpoint.model_validate(endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -3165,7 +3274,7 @@ async def update_pass_through_endpoints( endpoint_dict.pop("is_from_config", None) # Create updated endpoint object - updated_endpoint = PassThroughGenericEndpoint(**endpoint_dict) + updated_endpoint = PassThroughGenericEndpoint.model_validate(endpoint_dict) # Update the list pass_through_endpoint_data[endpoint_index] = endpoint_dict @@ -3186,9 +3295,10 @@ async def update_pass_through_endpoints( _custom_headers: dict | None = updated_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) + request_app = _APP_ADAPTER.validate_python(getattr(request, "app")) if updated_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( - app=request.app, + app=request_app, path=updated_endpoint.path, target=updated_endpoint.target, custom_headers=_custom_headers, @@ -3205,7 +3315,7 @@ async def update_pass_through_endpoints( ) else: InitPassThroughEndpointHelpers.add_exact_path_route( - app=request.app, + app=request_app, path=updated_endpoint.path, target=updated_endpoint.target, custom_headers=_custom_headers, @@ -3257,29 +3367,28 @@ async def create_pass_through_endpoints( if data_dict.get("id") is None: data_dict["id"] = str(uuid.uuid4()) - if response.field_value is None: - response.field_value = [data_dict] - elif isinstance(response.field_value, list): - response.field_value.append(data_dict) + existing_field_value = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value")) + updated_field_value = [data_dict] if existing_field_value is None else [*existing_field_value, data_dict] ## Update db updated_data = ConfigFieldUpdate( field_name="pass_through_endpoints", - field_value=response.field_value, + field_value=updated_field_value, config_type="general_settings", ) await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) # Return the created endpoint with the generated ID - created_endpoint = PassThroughGenericEndpoint(**data_dict) + created_endpoint = PassThroughGenericEndpoint.model_validate(data_dict) # Register the new route _custom_headers: dict | None = created_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) + request_app = _APP_ADAPTER.validate_python(getattr(request, "app")) if created_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( - app=request.app, + app=request_app, path=created_endpoint.path, target=created_endpoint.target, custom_headers=_custom_headers, @@ -3296,7 +3405,7 @@ async def create_pass_through_endpoints( ) else: InitPassThroughEndpointHelpers.add_exact_path_route( - app=request.app, + app=request_app, path=created_endpoint.path, target=created_endpoint.target, custom_headers=_custom_headers, @@ -3344,8 +3453,8 @@ async def delete_pass_through_endpoints( response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Update field by removing endpoint - pass_through_endpoint_data: list | None = response.field_value - if response.field_value is None or pass_through_endpoint_data is None: + pass_through_endpoint_data = _ENDPOINT_LIST_ADAPTER.validate_python(getattr(response, "field_value")) + if pass_through_endpoint_data is None: raise HTTPException( status_code=400, detail={"error": "There are no pass-through endpoints setup."}, @@ -3363,7 +3472,7 @@ async def delete_pass_through_endpoints( # Find the index for deleting from the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint + _endpoint = PassThroughGenericEndpoint.model_validate(endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -3393,7 +3502,7 @@ async def delete_pass_through_endpoints( def _find_endpoint_by_id( - endpoints_data: list, + endpoints_data: list[dict[str, object] | PassThroughGenericEndpoint], endpoint_id: str, ) -> PassThroughGenericEndpoint | None: """ @@ -3407,14 +3516,14 @@ def _find_endpoint_by_id( Found endpoint or None if not found """ for endpoint in endpoints_data: - _endpoint: PassThroughGenericEndpoint | None = None - if isinstance(endpoint, dict): - _endpoint = PassThroughGenericEndpoint(**endpoint) - elif isinstance(endpoint, PassThroughGenericEndpoint): - _endpoint = endpoint + _endpoint = ( + endpoint + if isinstance(endpoint, PassThroughGenericEndpoint) + else PassThroughGenericEndpoint.model_validate(endpoint) + ) # Only compare IDs to IDs - if _endpoint is not None and _endpoint.id == endpoint_id: + if _endpoint.id == endpoint_id: return _endpoint return None diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 4faf75b3951..52e264e33d1 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -137,7 +137,7 @@ class ResponsesSessionHandler: model_response = ModelResponse(**_response_output) for choice in model_response.choices: if hasattr(choice, "message"): - chat_completion_message_history.append(getattr(choice, "message")) + chat_completion_message_history.append(choice.message) return chat_completion_message_history @staticmethod diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 6b1ca3564e3..25079e21754 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,12 +4,13 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence -from typing import Any, Literal, cast +from collections.abc import Iterable, Mapping, Sequence +from typing import Literal, Protocol, cast from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_param import FunctionToolParam +from pydantic import BaseModel, InstanceOf, TypeAdapter from typing_extensions import TypedDict from litellm._logging import verbose_logger @@ -27,6 +28,7 @@ from litellm.types.llms.openai import ( ChatCompletionImageUrlObject, ChatCompletionResponseMessage, ChatCompletionSystemMessage, + ChatCompletionTextObject, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolMessage, @@ -72,6 +74,52 @@ from .custom_tools import ( unwrap_custom_tool_arguments, ) +_STR_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) +_STR_OBJECT_DICT_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[dict[str, object]]) +_OBJECT_LIST_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[list[object]]) +_OBJECT_SEQUENCE_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[Sequence[object]]) +_OBJECT_ITERABLE_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[Iterable[object]]) +_TEXT_OBJECT_LIST_INSTANCE_ADAPTER = TypeAdapter(InstanceOf[list[ChatCompletionTextObject]]) + + +def _as_str_object_dict(value: object) -> dict[str, object] | None: + return _STR_OBJECT_DICT_INSTANCE_ADAPTER.validate_python(value) if isinstance(value, dict) else None + + +def _as_object_list(value: object) -> list[object] | None: + return _OBJECT_LIST_INSTANCE_ADAPTER.validate_python(value) if isinstance(value, list) else None + + +def _as_object_sequence(value: object) -> Sequence[object] | None: + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + return None + return _OBJECT_SEQUENCE_INSTANCE_ADAPTER.validate_python(value) + + +def _as_object_iterable(value: object) -> Iterable[object] | None: + if isinstance(value, (str, bytes)) or not isinstance(value, Iterable): + return None + return _OBJECT_ITERABLE_INSTANCE_ADAPTER.validate_python(value) + + +def _as_tool_message_content(value: object) -> str | list[ChatCompletionTextObject] | None: + if isinstance(value, str): + return value + return _TEXT_OBJECT_LIST_INSTANCE_ADAPTER.validate_python(value) if isinstance(value, list) else None + + +def _model_as_dict(model: BaseModel) -> dict[str, object]: + return _STR_OBJECT_DICT_ADAPTER.validate_python(model.model_dump()) + + +class SupportsApplyPatchOperation(Protocol): + @property + def call_id(self) -> str: ... + + @property + def operation(self) -> BaseModel: ... + + ########### Initialize Classes used for Responses API ########### TOOL_CALLS_CACHE = InMemoryCache() @@ -116,8 +164,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_tool_choice( - tool_choice: Any, - ) -> str | dict[str, Any] | None: + tool_choice: str | Mapping[str, object] | None, + ) -> str | Mapping[str, object] | None: """ Transform tool_choice from various formats to OpenAI Chat Completion format. @@ -145,7 +193,8 @@ class LiteLLMCompletionResponsesConfig: tool_choice_type = tool_choice.get("type") # If it has a function with name, it's standard OpenAI format - pass through - if tool_choice.get("function") and tool_choice.get("function", {}).get("name"): + function_dict = _as_str_object_dict(tool_choice.get("function")) + if function_dict is not None and function_dict.get("name"): return tool_choice # Handle Cursor IDE dict formats without function name @@ -189,7 +238,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, str] | None = None, **kwargs, ) -> dict: """ @@ -446,7 +495,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,19 +507,17 @@ 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 = "" - if ( - isinstance(tool_calls, Sequence) - and not isinstance(tool_calls, (str, bytes)) - and len(tool_calls) > 0 - ): - first_call = tool_calls[0] - call_id_raw = ( - first_call.get("id") - if isinstance(first_call, dict) + tool_calls_seq = _as_object_sequence(tool_calls) + if tool_calls_seq is not None and len(tool_calls_seq) > 0: + first_call = tool_calls_seq[0] + first_call_dict = _as_str_object_dict(first_call) + call_id_raw: object = ( + first_call_dict.get("id") + if first_call_dict is not None else getattr(first_call, "id", None) ) if call_id_raw: @@ -518,23 +567,19 @@ 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: tool_calls = getattr(tool_call_message, "tool_calls", None) - if ( - isinstance(tool_calls, Sequence) - and not isinstance(tool_calls, (str, bytes)) - and len(tool_calls) > 0 - ): - first_call = tool_calls[0] - call_id_raw = None - if isinstance(first_call, dict): - call_id_raw = first_call.get("id") - else: - call_id_raw = getattr(first_call, "id", None) + tool_calls_seq = _as_object_sequence(tool_calls) + if tool_calls_seq is not None and len(tool_calls_seq) > 0: + first_call = tool_calls_seq[0] + first_call_dict = _as_str_object_dict(first_call) + call_id_raw: object = ( + first_call_dict.get("id") if first_call_dict is not None else getattr(first_call, "id", None) + ) if call_id_raw: call_id = str(call_id_raw) @@ -562,25 +607,40 @@ 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": + candidate = messages[j] + candidate_dict = _as_str_object_dict(candidate) + role = candidate_dict.get("role") if candidate_dict is not None else getattr(candidate, "role", None) + if role == "assistant": return j 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) + assistant_dict = _as_str_object_dict(assistant_message) + tool_calls_raw: object = ( + assistant_dict.get("tool_calls") + if assistant_dict is not None 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_list = _as_object_list(tool_calls_raw) + if tool_calls_list is not None and len(tool_calls_list) > 0: + first_tool_call = tool_calls_list[0] + first_tool_call_dict = _as_str_object_dict(first_tool_call) + if first_tool_call_dict is not None: + tool_call_id_raw: object = first_tool_call_dict.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,28 +648,32 @@ 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") - if isinstance(assistant_message, dict) + assistant_dict = _as_str_object_dict(assistant_message) + tool_calls_raw: object = ( + assistant_dict.get("tool_calls") + if assistant_dict is not None else getattr(assistant_message, "tool_calls", None) ) if tool_calls_raw is None: return [] - if isinstance(tool_calls_raw, list): - return tool_calls_raw - if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)): - return list(tool_calls_raw) + tool_calls_list = _as_object_list(tool_calls_raw) + if tool_calls_list is not None: + return tool_calls_list + tool_calls_iterable = _as_object_iterable(tool_calls_raw) + if tool_calls_iterable is not None: + return list(tool_calls_iterable) 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 - if isinstance(tool_call, dict): - tool_call_id_to_check = tool_call.get("id") + tool_call_id_to_check: object = None + tool_call_dict = _as_str_object_dict(tool_call) + if tool_call_dict is not None: + tool_call_id_to_check = tool_call_dict.get("id") elif hasattr(tool_call, "id"): tool_call_id_to_check = getattr(tool_call, "id", None) if tool_call_id_to_check == tool_call_id: @@ -617,52 +681,51 @@ 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): - tool_function = tool.get("function") or {} - tool_name = tool_function.get("name") or tool.get("name") or "" - if tool_name: - return { - "id": tool_call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": "{}", # We don't know the arguments, use empty - }, - } + tool_dict = _as_str_object_dict(tool) + if tool_dict is None: + continue + tool_function_dict: Mapping[str, object] = _as_str_object_dict(tool_dict.get("function")) or {} + tool_name = tool_function_dict.get("name") or tool_dict.get("name") or "" + if tool_name: + return { + "id": tool_call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": "{}", # We don't know the arguments, use empty + }, + } 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. """ if obj is None: return default - if isinstance(obj, dict): - return obj.get(key, default) + obj_dict = _as_str_object_dict(obj) + if obj_dict is not None: + return obj_dict.get(key, default) - getter = getattr(obj, "get", None) - if callable(getter): - try: - return getter(key, default) - except (TypeError, AttributeError): - pass + if isinstance(obj, BaseModel): + return _model_as_dict(obj).get(key, default) return getattr(obj, key, default) @staticmethod def _create_tool_call_chunk( - tool_use_definition: dict[str, Any], tool_call_id: str, index: int + tool_use_definition: dict[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,15 +744,16 @@ 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. """ if not tool_use_definition: return None - if isinstance(tool_use_definition, dict): - normalized_definition: dict[str, Any] = dict(tool_use_definition) + tool_use_definition_dict = _as_str_object_dict(tool_use_definition) + if tool_use_definition_dict is not None: + normalized_definition: dict[str, object] = dict(tool_use_definition_dict) 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 +786,22 @@ 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) - 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): + assistant_dict = _as_str_object_dict(assistant_message) + if assistant_dict is not None: + if "tool_calls" not in assistant_dict: + assistant_dict["tool_calls"] = [] + tool_calls_list = _as_object_list(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) + if getattr(assistant_message, "tool_calls", None) is None: + setattr(assistant_message, "tool_calls", []) + existing_tool_calls: object = getattr(assistant_message, "tool_calls", None) + existing_tool_calls_list = _as_object_list(existing_tool_calls) + if existing_tool_calls_list is not None: + existing_tool_calls_list.append(tool_call_chunk) @staticmethod def _ensure_tool_results_have_corresponding_tool_calls( @@ -746,7 +812,7 @@ class LiteLLMCompletionResponsesConfig: | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: Sequence[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -808,9 +874,8 @@ class LiteLLMCompletionResponsesConfig: ) if tool_call_id: # 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 = _as_str_object_dict(message) + if message_dict is not None: message_dict["tool_call_id"] = tool_call_id elif hasattr(message, "tool_call_id"): setattr(message, "tool_call_id", tool_call_id) @@ -862,7 +927,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 +962,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=role_value if isinstance(role_value, str) and role_value else "user", content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( content ), @@ -907,7 +973,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 +986,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 +996,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,8 +1008,8 @@ class LiteLLMCompletionResponsesConfig: return [] def _normalize_function_call_output_to_tool_content( - output: Any, - ) -> Any: + output: object, + ) -> str | list[dict[str, object]]: """ Normalize Responses API function_call_output.output into a shape that downstream chat adapters (esp. Gemini) can reliably consume. @@ -964,22 +1030,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_list = _as_object_list(output) + if output_list 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_list: + part_dict = _as_str_object_dict(part) + if part_dict is None: continue - part_type = part.get("type") + part_type = part_dict.get("type") if part_type in ("input_text", "output_text", "text"): - txt = part.get("text") + txt = part_dict.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_dict.get("image_url") or part_dict.get("url") + image_url_dict = _as_str_object_dict(image_url_val) + if image_url_dict is not None: + url = image_url_dict.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: @@ -1012,9 +1081,11 @@ class LiteLLMCompletionResponsesConfig: except Exception: return str(output) + normalized_output = _normalize_function_call_output_to_tool_content(tool_call_output.get("output")) + tool_message_content = _as_tool_message_content(normalized_output) tool_output_message = ChatCompletionToolMessage( role="tool", - content=_normalize_function_call_output_to_tool_content(tool_call_output.get("output")), + content=tool_message_content if tool_message_content is not None else "", tool_call_id=str(call_id), ) @@ -1066,7 +1137,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 @@ -1087,15 +1158,17 @@ class LiteLLMCompletionResponsesConfig: # Create a tool call for the function call. Custom tool calls # store their payload in "input" (raw string) rather than # "arguments" (JSON string), so normalize to arguments here. - raw_arguments = function_call.get("arguments") + raw_arguments: object = function_call.get("arguments") 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,16 +1184,17 @@ 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]) -> str | None: """ Return the effective file_id for a Responses API input_file item. Explicit file_id takes precedence; file_url is used as fallback so downstream providers (Anthropic, Gemini) can handle the URL natively. """ - return item.get("file_id") or item.get("file_url") or None + resolved = item.get("file_id") or item.get("file_url") or None + return resolved if isinstance(resolved, str) else 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,35 +1204,38 @@ 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 """ + 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=url_value if isinstance(url_value, str) and url_value else "", + detail=detail_value if isinstance(detail_value, str) and 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 @@ -1171,37 +1248,41 @@ class LiteLLMCompletionResponsesConfig: 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 := _as_object_list(content)) is not None: + content_list: list[str | dict[str, object]] = [] + for item in content_items: if isinstance(item, str): content_list.append(item) - elif isinstance(item, dict): - if item.get("type") == "input_file": - content_list.append( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) - ) - elif item.get("type") == "input_image": - image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) - ) - if "cache_control" in item: - image_block["cache_control"] = item["cache_control"] - content_list.append(image_block) - else: - # Skip text blocks with None text to avoid downstream errors - text_value = item.get("text") - if text_value is None: - continue - content_block: dict[str, Any] = { - "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( - item.get("type") or "text" - ), - "text": text_value, - } - if "cache_control" in item: - content_block["cache_control"] = item["cache_control"] - content_list.append(content_block) + continue + item_dict = _as_str_object_dict(item) + if item_dict is None: + continue + if item_dict.get("type") == "input_file": + content_list.append( + LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item_dict) + ) + elif item_dict.get("type") == "input_image": + image_block: dict[str, object] = dict( + LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item_dict) + ) + if "cache_control" in item_dict: + image_block["cache_control"] = item_dict["cache_control"] + content_list.append(image_block) + else: + # Skip text blocks with None text to avoid downstream errors + text_value = item_dict.get("text") + if text_value is None: + continue + block_type = item_dict.get("type") + content_block: dict[str, object] = { + "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( + block_type if isinstance(block_type, str) and block_type else "text" + ), + "text": text_value, + } + if "cache_control" in item_dict: + content_block["cache_control"] = item_dict["cache_control"] + content_list.append(content_block) return content_list else: raise ValueError(f"Invalid content type: {type(content)}") @@ -1283,7 +1364,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 +1405,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 +1413,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: Mapping[str, object] = _as_str_object_dict(tool.get("function")) or {} + parameters_dict = _as_str_object_dict(fn.get("parameters")) + parameters: dict[str, object] = dict(parameters_dict) if parameters_dict is not None 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 "", @@ -1360,9 +1442,17 @@ class LiteLLMCompletionResponsesConfig: result.append(responses_tool) else: # mcp or other: pass through unchanged - result.append(dict(tool)) + passthrough_tool: dict[str, object] = dict(tool) + result.append(passthrough_tool) return result + @staticmethod + def _read_provider_specific_fields(source: BaseModel) -> dict[str, object] | None: + raw_dict = _as_str_object_dict(_model_as_dict(source).get("provider_specific_fields")) + if raw_dict: + return raw_dict + return None + @staticmethod def transform_chat_completion_tools_to_responses_tools( chat_completion_response: ModelResponse, @@ -1414,25 +1504,9 @@ class LiteLLMCompletionResponsesConfig: responses_tools.append(custom_item) else: # Build regular function_call output item - provider_specific_fields: dict | None = None - if hasattr(tool, "provider_specific_fields") and getattr(tool, "provider_specific_fields", None): - provider_specific_fields = getattr(tool, "provider_specific_fields") - if not isinstance(provider_specific_fields, dict): - provider_specific_fields = ( - dict(provider_specific_fields) # type: ignore - if hasattr(provider_specific_fields, "__dict__") - else {} - ) - elif hasattr(function_definition, "provider_specific_fields") and getattr( - function_definition, "provider_specific_fields", None - ): - provider_specific_fields = getattr(function_definition, "provider_specific_fields") - if not isinstance(provider_specific_fields, dict): - provider_specific_fields = ( - dict(provider_specific_fields) # type: ignore - if hasattr(provider_specific_fields, "__dict__") - else {} - ) + provider_specific_fields = LiteLLMCompletionResponsesConfig._read_provider_specific_fields( + tool + ) or LiteLLMCompletionResponsesConfig._read_provider_specific_fields(function_definition) output_tool_call: ResponseFunctionToolCall = ResponseFunctionToolCall( name=tool_name, @@ -1496,9 +1570,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. @@ -1509,37 +1583,36 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary in ChatCompletionToolCallChunk format """ - # Extract provider_specific_fields if present - provider_specific_fields = getattr(tool_call_item, "provider_specific_fields", 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 - 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 {} - ) - ) + if isinstance(tool_call_item, BaseModel): + dumped = _model_as_dict(tool_call_item) + name_value: object = dumped.get("name") + arguments_value: object = dumped.get("arguments") + item_id_value: object = dumped.get("id") + call_id_value: object = dumped.get("call_id") + psf_value: object = dumped.get("provider_specific_fields") + else: + name_value = getattr(tool_call_item, "name", None) + arguments_value = getattr(tool_call_item, "arguments", None) + item_id_value = getattr(tool_call_item, "id", None) + call_id_value = getattr(tool_call_item, "call_id", None) + psf_value = getattr(tool_call_item, "provider_specific_fields", None) - function_dict: dict[str, Any] = { - "name": tool_call_item.name, - "arguments": tool_call_item.arguments, + # Extract provider_specific_fields if present + psf_dict = _as_str_object_dict(psf_value) + provider_specific_fields = psf_dict if psf_dict else None + + function_dict: dict[str, object] = { + "name": name_value, + "arguments": arguments_value, } 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), + item_id_value if isinstance(item_id_value, str) else None, + call_id_value if isinstance(call_id_value, str) else None, ), "function": function_dict, "type": "function", @@ -1553,9 +1626,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_apply_patch_tool_call_to_chat_completion_tool_call( - tool_call_item: Any, + tool_call_item: SupportsApplyPatchOperation, index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. @@ -1570,10 +1643,8 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary in ChatCompletionToolCallChunk format """ - import json - - operation_dict = tool_call_item.operation.model_dump() - tool_call_dict: dict[str, Any] = { + operation_dict = _model_as_dict(tool_call_item.operation) + tool_call_dict: dict[str, object] = { "id": tool_call_item.call_id, "function": { "name": "apply_patch", @@ -1711,11 +1782,11 @@ class LiteLLMCompletionResponsesConfig: """ output_items: list = [] for choice in chat_completion_response.choices or []: - message = getattr(choice, "message", None) + message: Message | None = getattr(choice, "message", None) if not message: continue - psf = getattr(message, "provider_specific_fields", None) - if not psf or not isinstance(psf, dict): + psf = message.provider_specific_fields + if not psf: continue results = psf.get("code_interpreter_results") if results and isinstance(results, list): @@ -1783,13 +1854,16 @@ class LiteLLMCompletionResponsesConfig: """ image_generation_items: list[OutputImageGenerationCall] = [] - images = getattr(choice.message, "images", []) + images = choice.message.images or [] 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_value: object = image_item.get("image_url") + image_url_dict = _as_str_object_dict(image_url_value) + url_value = image_url_dict.get("url", "") if image_url_dict is not None else "" + image_url = url_value if isinstance(url_value, str) else "" base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: @@ -2034,8 +2108,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,18 +2136,19 @@ 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): - format_type = format_param.get("type") + text_param_dict = _as_str_object_dict(text_param) + if text_param_dict is not None: + format_param_dict = _as_str_object_dict(text_param_dict.get("format")) + if format_param_dict: + format_type = format_param_dict.get("type") if format_type == "json_schema": return { "type": "json_schema", "json_schema": { - "name": format_param.get("name", "response_schema"), - "schema": format_param.get("schema", {}), - "strict": format_param.get("strict", False), + "name": format_param_dict.get("name", "response_schema"), + "schema": format_param_dict.get("schema", {}), + "strict": format_param_dict.get("strict", False), }, } elif format_type == "json_object": diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index c68628429da..468ad8279cb 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -64,7 +64,7 @@ async def create_mcp_list_tools_events( # Convert tools to dict format for the event mcp_tools_dict = [] for tool in filtered_mcp_tools: - if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")): + if hasattr(tool, "model_dump") and callable(tool.model_dump): # Type cast to help mypy understand this is safe after hasattr check mcp_tools_dict.append(cast(Any, tool).model_dump()) elif hasattr(tool, "__dict__"): @@ -103,9 +103,9 @@ async def create_mcp_list_tools_events( # Add input_schema if available if hasattr(tool, "inputSchema"): - tool_dict["input_schema"] = getattr(tool, "inputSchema") + tool_dict["input_schema"] = tool.inputSchema elif hasattr(tool, "input_schema"): - tool_dict["input_schema"] = getattr(tool, "input_schema") + tool_dict["input_schema"] = tool.input_schema formatted_tools.append(tool_dict) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 3bcc19822a6..28e56ed6f96 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,14 +5,15 @@ import json import time import traceback import uuid -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import Any, Literal +from typing import Any, Literal, Protocol import httpx from openai._streaming import SSEDecoder +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError import litellm from litellm.constants import ( @@ -30,7 +31,20 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils -from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, + ContentPartDoneEvent, + ContentPartDonePartOutputText, + ContentPartDonePartReasoningText, + ContentPartDonePartRefusal, + OpenAIChatCompletionLogprobsContent, + ResponseAPIUsage, + ResponseCreatedEvent, + ResponseInProgressEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ResponsesAPIStreamingResponse, +) from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook @@ -42,7 +56,70 @@ def _get_openai_response_types(): return openai_types -def _log_background_task_failure(task: asyncio.Task[Any], *, task_name: str) -> None: +_JSON_DICT_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) +_DICT_OBJECT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) +_OPTIONAL_DICT_OBJECT_ADAPTER: TypeAdapter[Mapping[str, object] | None] = TypeAdapter(Mapping[str, object] | None) +_OBJECT_LIST_ADAPTER: TypeAdapter[Sequence[object]] = TypeAdapter(Sequence[object]) +_STR_ADAPTER: TypeAdapter[str] = TypeAdapter(str) +_BOOL_ADAPTER: TypeAdapter[bool] = TypeAdapter(bool) +_OPTIONAL_FLOAT_ADAPTER: TypeAdapter[float | None] = TypeAdapter(float | None) +_USAGE_ADAPTER: TypeAdapter[ResponseAPIUsage] = TypeAdapter(ResponseAPIUsage) +_LOGPROBS_ADAPTER: TypeAdapter[list[OpenAIChatCompletionLogprobsContent] | None] = TypeAdapter( + list[OpenAIChatCompletionLogprobsContent] | None +) +_LOOSE_KWARGS_ADAPTER: TypeAdapter[Mapping[str, Any]] = TypeAdapter( + Mapping[str, Any] +) # any-ok: boundary adapter for litellm.aresponses(**kwargs); its ~30 typed named params reject JsonValue values and typing them would require deep pydantic validation of the client wire payload + + +def _parse_json_dict(raw: str | bytes) -> dict[str, JsonValue] | None: + try: + return _JSON_DICT_ADAPTER.validate_json(raw) + except ValidationError: + return None + + +def _json_list(value: JsonValue | None) -> Sequence[JsonValue]: + return value if isinstance(value, list) else [] + + +def _model_id_from_metadata(litellm_metadata: Mapping[str, JsonValue] | None) -> str | None: + model_info = litellm_metadata.get("model_info") if litellm_metadata else None + model_id = model_info.get("id") if isinstance(model_info, dict) else None + return model_id if isinstance(model_id, str) else None + + +class ClientWebSocketLike(Protocol): + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + +class BackendWebSocketLike(Protocol): + async def recv(self, decode: bool = ...) -> str | bytes: ... + + async def send(self, message: str) -> None: ... + + async def close(self) -> None: ... + + +class PresidioGuardrailLike(Protocol): + def get_presidio_settings_from_request_data(self, request_data: dict[str, JsonValue]) -> object: ... + + async def check_pii( + self, + text: str, + output_parse_pii: bool, + presidio_config: object, + request_data: dict[str, JsonValue], + ) -> str: ... + + +def _unmask_text(cb: PresidioGuardrailLike, text: str, pii_tokens: Mapping[str, str]) -> str: + return _STR_ADAPTER.validate_python(getattr(cb, "_unmask_pii_text")(text, pii_tokens)) + + +def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str) -> None: if task.cancelled(): return exception = task.exception() @@ -121,9 +198,9 @@ class BaseResponsesAPIStreamingIterator: model: str, responses_api_provider_config: BaseResponsesAPIConfig | None, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: dict[str, JsonValue] | None = None, custom_llm_provider: str | None = None, - request_data: dict[str, Any] | None = None, + request_data: dict[str, object] | None = None, call_type: str | None = None, ): self.response = response @@ -131,7 +208,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 @@ -145,23 +222,26 @@ class BaseResponsesAPIStreamingIterator: # track request context for hooks self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - self.request_data: dict[str, Any] = request_data or {} + self.request_data: dict[str, object] = request_data or {} self.call_type: str | None = call_type # set hidden params for response headers (e.g., x-litellm-model-id) # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py _api_base = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=_DICT_OBJECT_ADAPTER.validate_python( + self.logging_obj.model_call_details.get("litellm_params") or {} + ), ) - _model_info: dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} - self._hidden_params = { + _raw_model_info = litellm_metadata.get("model_info") if litellm_metadata else None + _model_info: Mapping[str, JsonValue] = _raw_model_info if isinstance(_raw_model_info, dict) else {} + self._hidden_params: dict[str, object] = { "model_id": _model_info.get("id", None), "api_base": _api_base, "custom_llm_provider": custom_llm_provider, } - self._hidden_params["additional_headers"] = process_response_headers( - self.response.headers or {} + self._hidden_params["additional_headers"] = _DICT_OBJECT_ADAPTER.validate_python( + process_response_headers(self.response.headers or {}) ) # GUARANTEE OPENAI HEADERS IN RESPONSE def _check_max_streaming_duration(self) -> None: @@ -176,7 +256,7 @@ class BaseResponsesAPIStreamingIterator: llm_provider=self.custom_llm_provider or "", ) - def _process_chunk(self, chunk) -> Any | None: + def _process_chunk(self, chunk: str | None) -> ResponsesAPIStreamingResponse | None: """Process a single chunk of data from the stream""" if not chunk: return None @@ -194,134 +274,121 @@ class BaseResponsesAPIStreamingIterator: if self.logging_obj.completion_start_time is None: self.logging_obj._update_completion_start_time(completion_start_time=datetime.now()) + # Parse the JSON chunk; skip chunks that are not JSON objects + parsed_chunk = _parse_json_dict(chunk) + if parsed_chunk is None: + return None + try: - # Parse the JSON chunk - parsed_chunk = json.loads(chunk) + if self.responses_api_provider_config is None: + raise ValueError("responses_api_provider_config is required to process live streaming chunks") + openai_responses_api_chunk = self.responses_api_provider_config.transform_streaming_response( + model=self.model, + parsed_chunk=parsed_chunk, + logging_obj=self.logging_obj, + ) - # Format as ResponsesAPIStreamingResponse - if isinstance(parsed_chunk, dict): - if self.responses_api_provider_config is None: - raise ValueError("responses_api_provider_config is required to process live streaming chunks") - openai_responses_api_chunk = self.responses_api_provider_config.transform_streaming_response( - model=self.model, - parsed_chunk=parsed_chunk, - logging_obj=self.logging_obj, - ) + # Only when the SSE JSON carries a response body (delta events do not). + # 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) + if isinstance(response_object, ResponsesAPIResponse): + response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=response_object, + litellm_metadata=self.litellm_metadata, + custom_llm_provider=self.custom_llm_provider, + ) + setattr(openai_responses_api_chunk, "response", response) - # Only when the SSE JSON carries a response body (delta events do not). - # 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) - if response_object is not None: - response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=response_object, - litellm_metadata=self.litellm_metadata, - custom_llm_provider=self.custom_llm_provider, + # Encode container_id on streaming events so proxy/UI follow-ups route correctly + _event_type = getattr(openai_responses_api_chunk, "type", None) + if _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: + _delta = getattr(openai_responses_api_chunk, "delta", None) + if isinstance(_delta, str): + self._generated_content += _delta + _stream_model_id = _model_id_from_metadata(self.litellm_metadata) + if _event_type in ( + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + ): + _item: object = getattr(openai_responses_api_chunk, "item", None) + if _item is not None: + ResponsesAPIRequestUtils._encode_container_id_on_output_item( + item=_item, + custom_llm_provider=self.custom_llm_provider, + model_id=_stream_model_id, + ) + elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED: + _annotation: object = getattr(openai_responses_api_chunk, "annotation", None) + if _annotation is not None: + ResponsesAPIRequestUtils._encode_container_id_on_output_item( + item=_annotation, + custom_llm_provider=self.custom_llm_provider, + model_id=_stream_model_id, + ) + elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: + _part: object = getattr(openai_responses_api_chunk, "part", None) + if _part is not None: + if isinstance(_part, dict): + ResponsesAPIRequestUtils._encode_container_ids_in_annotations( + _part.get("annotations"), + self.custom_llm_provider, + _stream_model_id, + ) + else: + ResponsesAPIRequestUtils._encode_container_ids_in_annotations( + getattr(_part, "annotations", None), + self.custom_llm_provider, + _stream_model_id, ) - setattr(openai_responses_api_chunk, "response", response) - # Encode container_id on streaming events so proxy/UI follow-ups route correctly - _event_type = getattr(openai_responses_api_chunk, "type", None) - if _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: - _delta = getattr(openai_responses_api_chunk, "delta", None) - if isinstance(_delta, str): - self._generated_content += _delta - _stream_model_id = ( - self.litellm_metadata.get("model_info", {}).get("id") if self.litellm_metadata else None - ) - if _event_type in ( + # Wrap encrypted_content in streaming events (output_item.added, output_item.done) + if self.litellm_metadata and self.litellm_metadata.get("encrypted_content_affinity_enabled"): + event_type = getattr(openai_responses_api_chunk, "type", None) + if event_type in ( ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - _item = getattr(openai_responses_api_chunk, "item", None) - if _item is not None: - ResponsesAPIRequestUtils._encode_container_id_on_output_item( - item=_item, - custom_llm_provider=self.custom_llm_provider, - model_id=_stream_model_id, - ) - elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED: - _annotation = getattr(openai_responses_api_chunk, "annotation", None) - if _annotation is not None: - ResponsesAPIRequestUtils._encode_container_id_on_output_item( - item=_annotation, - custom_llm_provider=self.custom_llm_provider, - model_id=_stream_model_id, - ) - elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: - _part = getattr(openai_responses_api_chunk, "part", None) - if _part is not None: - if isinstance(_part, dict): - ResponsesAPIRequestUtils._encode_container_ids_in_annotations( - _part.get("annotations"), - self.custom_llm_provider, - _stream_model_id, - ) - else: - ResponsesAPIRequestUtils._encode_container_ids_in_annotations( - getattr(_part, "annotations", None), - self.custom_llm_provider, - _stream_model_id, - ) - - # Wrap encrypted_content in streaming events (output_item.added, output_item.done) - if self.litellm_metadata and self.litellm_metadata.get("encrypted_content_affinity_enabled"): - openai_types = _get_openai_response_types() - event_type = getattr(openai_responses_api_chunk, "type", None) - if event_type in ( - openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - ): - item = getattr(openai_responses_api_chunk, "item", None) - if item: - encrypted_content = getattr(item, "encrypted_content", None) - if encrypted_content and isinstance(encrypted_content, str): - model_id = ( - self.litellm_metadata.get("model_info", {}).get("id") - if self.litellm_metadata - else None + item: object = getattr(openai_responses_api_chunk, "item", None) + if item: + encrypted_content: object = getattr(item, "encrypted_content", None) + if encrypted_content and isinstance(encrypted_content, str): + model_id = _model_id_from_metadata(self.litellm_metadata) + if model_id: + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id ) - if model_id: - wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - encrypted_content, model_id - ) - setattr(item, "encrypted_content", wrapped_content) + setattr(item, "encrypted_content", wrapped_content) - # Store the completed response (also for incomplete/failed so logging still fires) - _chunk_type = getattr(openai_responses_api_chunk, "type", None) - openai_types = _get_openai_response_types() - if openai_responses_api_chunk and _chunk_type in ( - openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, - openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, - ): - 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) - if response_obj: - usage_obj: Any | 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) - if cost is not None: - setattr(usage_obj, "cost", cost) - except Exception: - # Best-effort usage cost annotation should not break stream replay. - pass + # Store the completed response (also for incomplete/failed so logging still fires) + _chunk_type = getattr(openai_responses_api_chunk, "type", None) + if openai_responses_api_chunk and _chunk_type in ( + ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + ResponsesAPIStreamEvents.RESPONSE_FAILED, + ): + 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 = getattr(openai_responses_api_chunk, "response", None) + if isinstance(response_obj, ResponsesAPIResponse): + usage_obj: object = getattr(response_obj, "usage", None) + if usage_obj is not None: + try: + cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) + if cost is not None: + setattr(usage_obj, "cost", cost) + except Exception: + # Best-effort usage cost annotation should not break stream replay. + pass - if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: - self._handle_logging_failed_response() - else: - self._handle_logging_completed_response() + if _chunk_type == ResponsesAPIStreamEvents.RESPONSE_FAILED: + self._handle_logging_failed_response() + else: + self._handle_logging_completed_response() - return openai_responses_api_chunk - - return None - except json.JSONDecodeError: - # If we can't parse the chunk, continue - return None + return openai_responses_api_chunk except Exception as e: # Trigger failure hooks before re-raising # This ensures failures are logged even when _process_chunk is called directly @@ -389,8 +456,8 @@ 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: object = getattr(self.completed_response, "response", None) if self.completed_response else None + error_info: object = 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,15 +468,17 @@ 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: object | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj = getattr(response_obj, "usage", None) + usage_obj: object = getattr(response_obj, "usage", None) if usage_obj is None: return try: self.logging_obj.model_call_details["combined_usage_object"] = ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage_obj) + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + _USAGE_ADAPTER.validate_python(usage_obj) + ) ) except (TypeError, ValueError) as usage_error: verbose_logger.debug( @@ -418,7 +487,10 @@ class BaseResponsesAPIStreamingIterator: ) return self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj._response_cost_calculator(result=response_obj) or 0.0 + _OPTIONAL_FLOAT_ADAPTER.validate_python( + getattr(self.logging_obj, "_response_cost_calculator")(result=response_obj) + ) + or 0.0 ) def _maybe_raise_for_error_event(self, result: object) -> None: @@ -451,14 +523,13 @@ class BaseResponsesAPIStreamingIterator: is_pre_first_chunk=not self._yielded_first_chunk, ) - def _get_completed_response_object(self) -> Any | None: - openai_types = _get_openai_response_types() + def _get_completed_response_object(self) -> ResponsesAPIResponse | None: completed_response = self.completed_response - if isinstance(completed_response, openai_types.ResponsesAPIResponse): + if isinstance(completed_response, ResponsesAPIResponse): return completed_response response_obj = getattr(completed_response, "response", None) - if isinstance(response_obj, openai_types.ResponsesAPIResponse): + if isinstance(response_obj, ResponsesAPIResponse): return response_obj return None @@ -468,23 +539,27 @@ class BaseResponsesAPIStreamingIterator: return completed_response = self.completed_response - openai_types = _get_openai_response_types() - if getattr(completed_response, "type", None) != openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + if getattr(completed_response, "type", None) != ResponsesAPIStreamEvents.RESPONSE_COMPLETED: return response_obj = self._get_completed_response_object() if response_obj is None: return - caching_handler = getattr(self.logging_obj, "_llm_caching_handler", None) + caching_handler: object = getattr(self.logging_obj, "_llm_caching_handler", None) if caching_handler is None: return - request_kwargs = getattr(caching_handler, "request_kwargs", None) - if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: + try: + raw_request_kwargs = _OPTIONAL_DICT_OBJECT_ADAPTER.validate_python( + getattr(caching_handler, "request_kwargs", None) + ) + except ValidationError: return - request_kwargs = request_kwargs.copy() - preset_cache_key = getattr(caching_handler, "preset_cache_key", None) + if raw_request_kwargs is None or raw_request_kwargs.get("stream") is not True: + return + request_kwargs = dict(raw_request_kwargs) + 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,9 +569,11 @@ class BaseResponsesAPIStreamingIterator: if preset_cache_key is not None: request_kwargs["cache_key"] = preset_cache_key - if not caching_handler._should_store_result_in_cache( - original_function=caching_handler.original_function, - kwargs=request_kwargs, + if not _BOOL_ADAPTER.validate_python( + getattr(caching_handler, "_should_store_result_in_cache")( + original_function=getattr(caching_handler, "original_function", None), + kwargs=request_kwargs, + ) ): return @@ -527,7 +604,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). """ @@ -556,7 +635,7 @@ class BaseResponsesAPIStreamingIterator: response_chunk=chunk, call_type=typed_call_type, ) - if result is not None: + if isinstance(result, BaseLiteLLMOpenAIResponseObject): chunk = result if hooks_ran: setattr(chunk, "_post_streaming_hooks_ran", True) @@ -564,7 +643,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). """ @@ -577,9 +658,8 @@ class BaseResponsesAPIStreamingIterator: if self.completed_response is None: return - request_payload: dict[str, Any] = {} - if isinstance(self.request_data, dict): - request_payload.update(self.request_data) + request_payload: dict[str, object] = {} + request_payload.update(self.request_data) try: if hasattr(self.logging_obj, "model_call_details"): request_payload.update(self.logging_obj.model_call_details) @@ -587,8 +667,8 @@ class BaseResponsesAPIStreamingIterator: pass if "litellm_params" not in request_payload: try: - request_payload["litellm_params"] = getattr(self.logging_obj, "model_call_details", {}).get( - "litellm_params", {} + request_payload["litellm_params"] = _DICT_OBJECT_ADAPTER.validate_python( + self.logging_obj.model_call_details.get("litellm_params") or {} ) except Exception: request_payload["litellm_params"] = {} @@ -666,14 +746,15 @@ class BaseResponsesAPIStreamingIterator: pass -async def call_post_streaming_hooks_for_testing(iterator, chunk): +async def call_post_streaming_hooks_for_testing( + iterator: object, chunk: ResponsesAPIStreamingResponse +) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ - hook_fn = getattr(iterator, "_call_post_streaming_deployment_hook", None) - if hook_fn is None: + if not isinstance(iterator, BaseResponsesAPIStreamingIterator): return chunk - return await hook_fn(chunk) + return await iterator.call_post_streaming_hooks_for_testing(chunk) class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @@ -687,9 +768,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: dict[str, JsonValue] | None = None, custom_llm_provider: str | None = None, - request_data: dict[str, Any] | None = None, + request_data: dict[str, object] | None = None, call_type: str | None = None, ): super().__init__( @@ -707,7 +788,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: @@ -769,9 +850,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: dict[str, JsonValue] | None = None, custom_llm_provider: str | None = None, - request_data: dict[str, Any] | None = None, + request_data: dict[str, object] | None = None, call_type: str | None = None, ): super().__init__( @@ -789,7 +870,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __iter__(self): return self - def __next__(self): + def __next__(self) -> ResponsesAPIStreamingResponse: try: self._check_max_streaming_duration() while True: @@ -856,9 +937,9 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: LiteLLMLoggingObj, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: dict[str, JsonValue] | None = None, custom_llm_provider: str | None = None, - request_data: dict[str, Any] | None = None, + request_data: dict[str, object] | None = None, call_type: str | None = None, ): transformed = responses_api_provider_config.transform_response_api_response( @@ -880,7 +961,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,13 +975,12 @@ 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] self._idx += 1 - openai_types = _get_openai_response_types() - if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + if getattr(evt, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=True) return evt @@ -908,13 +988,12 @@ 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] self._idx += 1 - openai_types = _get_openai_response_types() - if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + if getattr(evt, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=False) return evt @@ -923,15 +1002,15 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __init__( self, - response: Any, + response: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, - request_data: dict[str, Any] | None = None, + request_data: dict[str, object] | None = None, call_type: str | None = None, ): BaseResponsesAPIStreamingIterator.__init__( self, response=httpx.Response(200), - model=getattr(response, "model", ""), + model=response.model or "", responses_api_provider_config=None, logging_obj=logging_obj, litellm_metadata=None, @@ -941,13 +1020,13 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): ) self._completed_response_cache_hit = True self._persist_completed_response_before_logging = False - self._events: list[Any] = [] + self._events: Sequence[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,13 +1040,12 @@ 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] self._idx += 1 - openai_types = _get_openai_response_types() - if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + if getattr(evt, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=True) return evt @@ -975,41 +1053,45 @@ 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] self._idx += 1 - openai_types = _get_openai_response_types() - if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + if getattr(evt, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=False) return evt -def _dump_response_object(obj: Any) -> dict[str, Any]: - if hasattr(obj, "model_dump"): - return obj.model_dump() +def _dump_response_object(obj: object) -> dict[str, JsonValue]: + if isinstance(obj, BaseModel): + try: + return _JSON_DICT_ADAPTER.validate_python(obj.model_dump()) + except ValidationError: + return _JSON_DICT_ADAPTER.validate_json(obj.model_dump_json()) if isinstance(obj, dict): - return obj + try: + return _JSON_DICT_ADAPTER.validate_python(obj) + except ValidationError: + return {} return {} def _build_response_status_event( event_type: Literal[ - "response.created", - "response.in_progress", + ResponsesAPIStreamEvents.RESPONSE_CREATED, + ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, ], - transformed: Any, -) -> Any: - openai_types = _get_openai_response_types() + transformed: ResponsesAPIResponse, +) -> ResponseCreatedEvent | ResponseInProgressEvent: in_progress_response = transformed.model_copy( deep=True, update={"status": "in_progress", "output": []}, ) - if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED: - return openai_types.ResponseCreatedEvent(type=event_type, response=in_progress_response) - return openai_types.ResponseInProgressEvent(type=event_type, response=in_progress_response) + if event_type == ResponsesAPIStreamEvents.RESPONSE_CREATED: + return ResponseCreatedEvent(type=event_type, response=in_progress_response) + return ResponseInProgressEvent(type=event_type, response=in_progress_response) def _build_content_part_done_event( @@ -1017,37 +1099,37 @@ def _build_content_part_done_event( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], -) -> Any | None: - openai_types = _get_openai_response_types() + part_payload: dict[str, JsonValue], +) -> ContentPartDoneEvent | None: part_type = part_payload.get("type") - part: Any + part: ContentPartDonePartOutputText | ContentPartDonePartRefusal | ContentPartDonePartReasoningText if part_type == "output_text": annotations = [ - openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) - for annotation in part_payload.get("annotations", []) or [] + BaseLiteLLMOpenAIResponseObject(**annotation) + for annotation in _json_list(part_payload.get("annotations")) + if isinstance(annotation, dict) ] - part = openai_types.ContentPartDonePartOutputText( + part = ContentPartDonePartOutputText( type="output_text", text=str(part_payload.get("text") or ""), annotations=annotations, - logprobs=part_payload.get("logprobs"), + logprobs=_LOGPROBS_ADAPTER.validate_python(part_payload.get("logprobs")), ) elif part_type == "refusal": - part = openai_types.ContentPartDonePartRefusal( + part = ContentPartDonePartRefusal( type="refusal", refusal=str(part_payload.get("refusal") or ""), ) elif part_type == "reasoning_text": - part = openai_types.ContentPartDonePartReasoningText( + part = ContentPartDonePartReasoningText( type="reasoning_text", reasoning=str(part_payload.get("reasoning") or ""), ) else: return None - return openai_types.ContentPartDoneEvent( - type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_DONE, + return ContentPartDoneEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, item_id=item_id, output_index=output_index, content_index=content_index, @@ -1057,11 +1139,11 @@ 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, - part_payload: dict[str, Any], + part_payload: dict[str, JsonValue], chunk_size: int, ) -> None: openai_types = _get_openai_response_types() @@ -1078,7 +1160,8 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []): + annotation_dicts = [a for a in _json_list(part_payload.get("annotations")) if isinstance(a, dict)] + for annotation_index, annotation in enumerate(annotation_dicts): events.append( openai_types.OutputTextAnnotationAddedEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, @@ -1123,13 +1206,13 @@ def _add_text_like_part_events( def _build_synthetic_response_events( *, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, chunk_size: int, -) -> list[Any]: +) -> Sequence[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 = getattr(transformed, "usage", None) if usage_obj is not None: try: cost: float | None = logging_obj._response_cost_calculator(result=transformed) @@ -1138,13 +1221,14 @@ def _build_synthetic_response_events( except Exception: pass - events: list[Any] = [ - _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), - _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), + events: list[ResponsesAPIStreamingResponse] = [ + _build_response_status_event(ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), + _build_response_status_event(ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), ] sequence_number = 0 - for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): + output_items = _OBJECT_LIST_ADAPTER.validate_python(getattr(transformed, "output", None) 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 +1242,7 @@ def _build_synthetic_response_events( ) if item_type == "message": - for content_index, part in enumerate(output_item_payload.get("content", []) or []): + for content_index, part in enumerate(_json_list(output_item_payload.get("content"))): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( @@ -1205,7 +1289,7 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): + for summary_index, summary in enumerate(_json_list(output_item_payload.get("summary"))): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1292,56 +1376,57 @@ class ResponsesWebSocketStreaming: def __init__( self, - websocket: Any, - backend_ws: Any, + websocket: ClientWebSocketLike, + backend_ws: BackendWebSocketLike, 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, JsonValue] | None = None, first_message: str | None = None, - guardrail_callbacks: list[Any] | None = None, - output_guardrail_callbacks: list[Any] | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailLike] | None = None, + output_guardrail_callbacks: Sequence[PresidioGuardrailLike] | 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, JsonValue] = request_data or {} + self.messages: list[Mapping[str, JsonValue]] = [] 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: Sequence[PresidioGuardrailLike] = guardrail_callbacks or [] + self.output_guardrail_callbacks: Sequence[PresidioGuardrailLike] = 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: Mapping[str, JsonValue]) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES - def _store_event(self, event: Any) -> None: + def _store_event(self, event: str | bytes | Mapping[str, JsonValue]) -> None: if isinstance(event, bytes): event = event.decode("utf-8") if isinstance(event, str): - try: - event_obj = json.loads(event) - except (json.JSONDecodeError, TypeError): + parsed_event = _parse_json_dict(event) + if parsed_event is None: return + event_obj = parsed_event else: event_obj = event 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: str | Mapping[str, JsonValue]) -> None: """Extract user input content from response.create for logging.""" try: if isinstance(message, str): - msg_obj = json.loads(message) - elif isinstance(message, dict): - msg_obj = message + parsed_message = _parse_json_dict(message) else: + parsed_message = message + if parsed_message is None: return + msg_obj = parsed_message if msg_obj.get("type") != "response.create": return @@ -1363,12 +1448,12 @@ class ResponsesWebSocketStreaming: for c in content: if isinstance(c, dict) and c.get("type") == "input_text": text = c.get("text", "") - if text: + if isinstance(text, str) and text: self.input_messages.append({"role": "user", "content": text}) except (json.JSONDecodeError, AttributeError, TypeError): pass - def _store_input(self, message: Any) -> None: + def _store_input(self, message: str) -> None: self._collect_input_from_client_event(message) if self.logging_obj: self.logging_obj.pre_call(input=message, api_key="") @@ -1388,9 +1473,9 @@ class ResponsesWebSocketStreaming: try: while True: try: - raw_response = await self.backend_ws.recv(decode=False) # type: ignore[union-attr] + raw_response = await self.backend_ws.recv(decode=False) except TypeError: - raw_response = await self.backend_ws.recv() # type: ignore[union-attr, assignment] + raw_response = await self.backend_ws.recv() if isinstance(raw_response, bytes): response_str = raw_response.decode("utf-8") @@ -1406,10 +1491,8 @@ class ResponsesWebSocketStreaming: # before response.completed arrives. The client receives only the # masked response.completed. if self.output_guardrail_callbacks: - try: - _evt_type = json.loads(response_str).get("type") - except (json.JSONDecodeError, TypeError): - _evt_type = None + parsed_event = _parse_json_dict(response_str) + _evt_type = parsed_event.get("type") if parsed_event is not None else None if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES: continue @@ -1429,7 +1512,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, JsonValue]) -> bool: """ Overwrite any ``model`` field in a ``response.create`` frame with the connection-authorized model to prevent deployment-substitution attacks. @@ -1470,9 +1553,8 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ - try: - msg_obj = json.loads(message) - except (json.JSONDecodeError, TypeError): + msg_obj = _parse_json_dict(message) + if msg_obj is None: return message if msg_obj.get("type") != "response.create": @@ -1495,8 +1577,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]] = [] + nested_candidate = msg_obj.get("response") + nested_response = nested_candidate if isinstance(nested_candidate, dict) else None + text_containers: list[tuple[dict[str, JsonValue], str]] = [] for container in (msg_obj, nested_response): if container is None: continue @@ -1533,13 +1616,14 @@ class ResponsesWebSocketStreaming: modified = True elif isinstance(value, list): for block in value: - if ( - isinstance(block, dict) - and block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES - and isinstance(block.get("text"), str) + if not isinstance(block, dict): + continue + block_text = block.get("text") + if block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES and isinstance( + block_text, str ): block["text"] = await cb.check_pii( - text=block["text"], + text=block_text, output_parse_pii=True, presidio_config=presidio_config, request_data=self.request_data, @@ -1590,13 +1674,18 @@ 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") + raw_pii_tokens = metadata.get("pii_tokens") if isinstance(metadata, dict) else None + pii_tokens: Mapping[str, str] = ( + {token: value for token, value in raw_pii_tokens.items() if isinstance(value, str)} + if isinstance(raw_pii_tokens, dict) + else {} + ) if not pii_tokens: return response_str - try: - evt_obj = json.loads(response_str) - except (json.JSONDecodeError, TypeError): + evt_obj = _parse_json_dict(response_str) + if evt_obj is None: return response_str cb = self.guardrail_callbacks[0] @@ -1604,21 +1693,18 @@ class ResponsesWebSocketStreaming: if event_type == "response.completed": modified = False - response_obj = evt_obj.get("response") or {} + response_obj = evt_obj.get("response") if not isinstance(response_obj, dict): return response_str - for output_item in response_obj.get("output") or []: + for output_item in _json_list(response_obj.get("output")): if not isinstance(output_item, dict): continue - content = output_item.get("content") or [] - if not isinstance(content, list): - continue - for content_block in content: + for content_block in _json_list(output_item.get("content")): if not isinstance(content_block, dict): continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text(text, pii_tokens) + unmasked = _unmask_text(cb, text, pii_tokens) if unmasked != text: content_block["text"] = unmasked modified = True @@ -1627,7 +1713,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 = _unmask_text(cb, delta, pii_tokens) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) @@ -1649,9 +1735,8 @@ class ResponsesWebSocketStreaming: if not self.output_guardrail_callbacks: return response_str - try: - evt_obj = json.loads(response_str) - except (json.JSONDecodeError, TypeError): + evt_obj = _parse_json_dict(response_str) + if evt_obj is None: return response_str if evt_obj.get("type") != "response.completed": @@ -1660,10 +1745,10 @@ class ResponsesWebSocketStreaming: modified = False for cb in self.output_guardrail_callbacks: presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) - response_obj = evt_obj.get("response") or {} + response_obj = evt_obj.get("response") if not isinstance(response_obj, dict): continue - for output_item in response_obj.get("output") or []: + for output_item in _json_list(response_obj.get("output")): if not isinstance(output_item, dict): continue arguments = output_item.get("arguments") @@ -1677,26 +1762,21 @@ class ResponsesWebSocketStreaming: if masked_args != arguments: output_item["arguments"] = masked_args modified = True - summary = output_item.get("summary") or [] - if isinstance(summary, list): - for summary_block in summary: - if not isinstance(summary_block, dict): - continue - summary_text = summary_block.get("text") - if isinstance(summary_text, str): - masked_summary = await cb.check_pii( - text=summary_text, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=self.request_data, - ) - if masked_summary != summary_text: - summary_block["text"] = masked_summary - modified = True - content = output_item.get("content") or [] - if not isinstance(content, list): - continue - for content_block in content: + for summary_block in _json_list(output_item.get("summary")): + if not isinstance(summary_block, dict): + continue + summary_text = summary_block.get("text") + if isinstance(summary_text, str): + masked_summary = await cb.check_pii( + text=summary_text, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=self.request_data, + ) + if masked_summary != summary_text: + summary_block["text"] = masked_summary + modified = True + for content_block in _json_list(output_item.get("content")): if not isinstance(content_block, dict): continue text = content_block.get("text") @@ -1720,14 +1800,14 @@ class ResponsesWebSocketStreaming: masked_first = await self._mask_response_create(self.first_message) self._store_input(masked_first) self._store_event(masked_first) - await self.backend_ws.send(masked_first) # type: ignore[union-attr] + await self.backend_ws.send(masked_first) while True: message = await self.websocket.receive_text() masked = await self._mask_response_create(message) self._store_input(masked) self._store_event(masked) - await self.backend_ws.send(masked) # type: ignore[union-attr] + await self.backend_ws.send(masked) except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) @@ -1756,12 +1836,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,26 +1873,25 @@ class ManagedResponsesWebSocketHandler: def __init__( self, - websocket: Any, + websocket: ClientWebSocketLike, model: str, logging_obj: LiteLLMLoggingObj, - user_api_key_dict: Any | None = None, - litellm_metadata: dict[str, Any] | None = None, + user_api_key_dict: object | None = None, + litellm_metadata: Mapping[str, JsonValue] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, - **kwargs: Any, + **kwargs: JsonValue, ) -> None: self.websocket = websocket self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.litellm_metadata: dict[str, Any] = litellm_metadata or {} - self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( - "deployment_model_name" - ) + self.litellm_metadata: Mapping[str, JsonValue] = litellm_metadata or {} + raw_model_group = self.litellm_metadata.get("model_group") or self.litellm_metadata.get("deployment_model_name") + self.model_group: str | None = raw_model_group if isinstance(raw_model_group, str) else None self.api_key = api_key self.api_base = api_base self.timeout = timeout @@ -1820,25 +1899,25 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message # Carry through safe pass-through kwargs (e.g. extra_headers) - self.extra_kwargs: dict[str, Any] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} + self.extra_kwargs: Mapping[str, JsonValue] = { + k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS + } # In-memory session history: response_id → full accumulated message list. # Keyed by the DECODED (pre-encoding) response ID from response.completed. # This avoids the async DB-write race condition where spend logs haven't # been committed yet when the next response.create arrives. - self._session_history: dict[str, list[dict[str, Any]]] = {} + self._session_history: dict[str, Sequence[JsonValue]] = {} # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ @staticmethod - def _serialize_chunk(chunk: Any) -> str | None: + def _serialize_chunk(chunk: object) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: - if hasattr(chunk, "model_dump_json"): + if isinstance(chunk, BaseModel): return chunk.model_dump_json(exclude_none=True) - if hasattr(chunk, "model_dump"): - return json.dumps(chunk.model_dump(exclude_none=True), default=str) if isinstance(chunk, dict): return json.dumps(chunk, default=str) return json.dumps(str(chunk)) @@ -1854,7 +1933,7 @@ class ManagedResponsesWebSocketHandler: except Exception: pass - def _get_history_messages(self, previous_response_id: str) -> list[dict[str, Any]]: + def _get_history_messages(self, previous_response_id: str) -> list[JsonValue]: """ Return accumulated message history for *previous_response_id*. @@ -1865,7 +1944,7 @@ class ManagedResponsesWebSocketHandler: raw_id = decoded.get("response_id", previous_response_id) return list(self._session_history.get(raw_id, [])) - def _store_history(self, response_id: str, messages: list[dict[str, Any]]) -> None: + def _store_history(self, response_id: str, messages: Sequence[JsonValue]) -> None: """ Store the complete accumulated message history for *response_id*. @@ -1875,13 +1954,14 @@ class ManagedResponsesWebSocketHandler: self._session_history[response_id] = messages @staticmethod - def _extract_response_id(completed_event: dict[str, Any]) -> str | None: + def _extract_response_id(completed_event: Mapping[str, JsonValue]) -> str | None: """ Pull the raw (decoded) response ID out of a ``response.completed`` event. Returns *None* if the event doesn't contain a usable ID. """ - resp_obj = completed_event.get("response", {}) - encoded_id: str | None = resp_obj.get("id") if isinstance(resp_obj, dict) else None + resp_obj = completed_event.get("response") + raw_encoded_id = resp_obj.get("id") if isinstance(resp_obj, dict) else None + encoded_id = raw_encoded_id if isinstance(raw_encoded_id, str) else None if not encoded_id: return None decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id) @@ -1889,27 +1969,29 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, Any], - ) -> list[dict[str, Any]]: + completed_event: Mapping[str, JsonValue], + ) -> list[JsonValue]: """ Convert the output items in a ``response.completed`` event into Responses API message dicts suitable for the next turn's ``input``. """ - resp_obj = completed_event.get("response", {}) + resp_obj = completed_event.get("response") if not isinstance(resp_obj, dict): return [] - messages: list[dict[str, Any]] = [] - for item in resp_obj.get("output", []) or []: + messages: list[JsonValue] = [] + for item in _json_list(resp_obj.get("output")): if not isinstance(item, dict): continue item_type = item.get("type") role = item.get("role", "assistant") if item_type == "message": - content_parts = item.get("content") or [] + content_parts = _json_list(item.get("content")) text_parts = [ - p.get("text", "") + part_text for p in content_parts - if isinstance(p, dict) and p.get("type") in ("output_text", "text") + if isinstance(p, dict) + and p.get("type") in ("output_text", "text") + and isinstance(part_text := p.get("text", ""), str) ] text = "".join(text_parts) if text: @@ -1925,7 +2007,7 @@ class ManagedResponsesWebSocketHandler: return messages @staticmethod - def _input_to_messages(input_val: Any) -> list[dict[str, Any]]: + def _input_to_messages(input_val: JsonValue | None) -> list[JsonValue]: """ Normalise the ``input`` field of a ``response.create`` event to a list of Responses API message dicts. @@ -1946,11 +2028,10 @@ class ManagedResponsesWebSocketHandler: # _process_response_create sub-methods # ------------------------------------------------------------------ - async def _parse_message(self, raw_message: str) -> dict[str, Any] | None: + async def _parse_message(self, raw_message: str) -> Mapping[str, JsonValue] | None: """Parse raw WS text; return the message dict or None (JSON error / ignored type).""" - try: - msg_obj = json.loads(raw_message) - except json.JSONDecodeError: + msg_obj = _parse_json_dict(raw_message) + if msg_obj is None: await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None if msg_obj.get("type") != "response.create": @@ -1959,7 +2040,7 @@ class ManagedResponsesWebSocketHandler: return msg_obj @staticmethod - def _is_warmup_frame(msg_obj: dict[str, Any]) -> bool: + def _is_warmup_frame(msg_obj: Mapping[str, JsonValue]) -> bool: """Return True for a response.create whose generate flag is false.""" nested = msg_obj.get("response") source = nested if isinstance(nested, dict) and nested else msg_obj @@ -1975,13 +2056,13 @@ class ManagedResponsesWebSocketHandler: return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX) @staticmethod - def _warmup_source_params(msg_obj: dict[str, Any]) -> dict[str, Any]: + def _warmup_source_params(msg_obj: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: nested = msg_obj.get("response") if isinstance(nested, dict) and nested: return nested return {k: v for k, v in msg_obj.items() if k != "type"} - def _build_warmup_response(self, msg_obj: dict[str, Any]) -> dict[str, Any]: + def _build_warmup_response(self, msg_obj: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: """Build a minimal completed Responses API object for a warmup ack.""" source = self._warmup_source_params(msg_obj) wire_model = source.get("model") or self.model_group or self.model @@ -1999,7 +2080,7 @@ class ManagedResponsesWebSocketHandler: }, } - async def _send_warmup_ack(self, msg_obj: dict[str, Any]) -> None: + async def _send_warmup_ack(self, msg_obj: Mapping[str, JsonValue]) -> None: """ Acknowledge a generate=false prewarm without calling the provider. @@ -2022,14 +2103,14 @@ class ManagedResponsesWebSocketHandler: await self.websocket.send_text(serialized) @staticmethod - def _build_base_call_kwargs(msg_obj: dict[str, Any]) -> dict[str, Any]: + def _build_base_call_kwargs(msg_obj: Mapping[str, JsonValue]) -> dict[str, JsonValue]: """ Extract Responses API params from the event, handling both wire formats: Nested: {"type": "response.create", "response": {"input": [...], ...}} Flat: {"type": "response.create", "input": [...], "model": "...", ...} """ nested = msg_obj.get("response") - response_params: dict[str, Any] = ( + response_params: Mapping[str, JsonValue] = ( nested if isinstance(nested, dict) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { @@ -2040,10 +2121,10 @@ class ManagedResponsesWebSocketHandler: def _apply_history( self, - call_kwargs: dict[str, Any], + call_kwargs: dict[str, JsonValue], previous_response_id: str | None, - current_messages: list[dict[str, Any]], - prior_history: list[dict[str, Any]], + current_messages: list[JsonValue], + prior_history: list[JsonValue], ) -> None: """Prepend in-memory turn history, or fall back to DB-based reconstruction.""" if not previous_response_id: @@ -2093,7 +2174,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None: + def _inject_credentials(self, call_kwargs: dict[str, JsonValue], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key @@ -2112,26 +2193,35 @@ class ManagedResponsesWebSocketHandler: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) @staticmethod - def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: + def _update_proxy_request(call_kwargs: dict[str, JsonValue], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" - proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get("proxy_server_request") or {} - if not isinstance(proxy_server_request, dict): + existing_metadata = call_kwargs.get("litellm_metadata") + raw_proxy_server_request = ( + existing_metadata.get("proxy_server_request") if isinstance(existing_metadata, dict) else None + ) + if raw_proxy_server_request and not isinstance(raw_proxy_server_request, dict): return - body = dict(proxy_server_request.get("body") or {}) + proxy_server_request = raw_proxy_server_request if isinstance(raw_proxy_server_request, dict) else {} + raw_body = proxy_server_request.get("body") + body = dict(raw_body) if isinstance(raw_body, dict) else {} body["input"] = call_kwargs.get("input") body["store"] = call_kwargs.get("store") body["model"] = model for k in ("tools", "tool_choice", "instructions", "metadata"): if k in call_kwargs and call_kwargs[k] is not None: body[k] = call_kwargs[k] - proxy_server_request = {**proxy_server_request, "body": body} - if "litellm_metadata" not in call_kwargs: - call_kwargs["litellm_metadata"] = {} - call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request - call_kwargs.setdefault("litellm_params", {}) - call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request + updated_proxy_server_request: dict[str, JsonValue] = {**proxy_server_request, "body": body} + if isinstance(existing_metadata, dict): + existing_metadata["proxy_server_request"] = updated_proxy_server_request + else: + call_kwargs["litellm_metadata"] = {"proxy_server_request": updated_proxy_server_request} + litellm_params = call_kwargs.setdefault("litellm_params", {}) + if isinstance(litellm_params, dict): + litellm_params["proxy_server_request"] = updated_proxy_server_request - async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, Any] | None: + async def _stream_and_forward( + self, model: str, call_kwargs: Mapping[str, JsonValue] + ) -> Mapping[str, JsonValue] | None: """ Stream ``litellm.aresponses`` and forward every chunk over the WebSocket. @@ -2139,21 +2229,21 @@ class ManagedResponsesWebSocketHandler: directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the completed event dict, or ``None``. """ - completed_event: dict[str, Any] | None = None - stream_response = await litellm.aresponses(model=model, **call_kwargs) - async for chunk in stream_response: # type: ignore[union-attr] - if chunk is None: + completed_event: Mapping[str, JsonValue] | None = None + loose_kwargs = _LOOSE_KWARGS_ADAPTER.validate_python(call_kwargs) + stream_response = await litellm.aresponses( + model=model, **loose_kwargs + ) # any-ok: aresponses fans dynamic client JSON into ~30 typed named params; per-param typing would require deep pydantic validation of the wire payload + async for chunk in stream_response: + if not isinstance(chunk, BaseLiteLLMOpenAIResponseObject): continue # Read type from the object before serializing to avoid double JSON parse - chunk_type = getattr(chunk, "type", None) or (chunk.get("type") if isinstance(chunk, dict) else None) + chunk_type = getattr(chunk, "type", None) serialized = self._serialize_chunk(chunk) if serialized is None: continue if chunk_type == "response.completed" and completed_event is None: - try: - completed_event = json.loads(serialized) - except Exception: - pass + completed_event = _parse_json_dict(serialized) try: await self.websocket.send_text(serialized) except Exception as send_exc: @@ -2163,9 +2253,9 @@ class ManagedResponsesWebSocketHandler: def _save_turn_history( self, - completed_event: dict[str, Any] | None, - prior_history: list[dict[str, Any]], - current_messages: list[dict[str, Any]], + completed_event: Mapping[str, JsonValue] | None, + prior_history: list[JsonValue], + current_messages: list[JsonValue], ) -> None: """Store this turn in in-memory history for future previous_response_id lookups.""" if completed_event is None: @@ -2228,13 +2318,17 @@ class ManagedResponsesWebSocketHandler: # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. - requested_model = call_kwargs.pop("model", None) + popped_model = call_kwargs.pop("model", None) + requested_model = popped_model if isinstance(popped_model, str) else None if requested_model is None or requested_model == self.model_group: model = self.model else: model = requested_model - previous_response_id: str | None = call_kwargs.pop("previous_response_id", None) + popped_previous_response_id = call_kwargs.pop("previous_response_id", None) + previous_response_id: str | None = ( + popped_previous_response_id if isinstance(popped_previous_response_id, str) else None + ) current_messages = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index d3ef01940bb..3c8761c6898 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,30 +1,30 @@ { "ANN001": { - "limit": 3104 + "limit": 3080 }, "ANN002": { "limit": 69 }, "ANN003": { - "limit": 831 + "limit": 829 }, "ANN201": { - "limit": 2137 + "limit": 2119 }, "ANN202": { - "limit": 941 + "limit": 939 }, "ANN204": { - "limit": 724 + "limit": 723 }, "ANN205": { - "limit": 127 + "limit": 126 }, "ANN206": { "limit": 130 }, "ANN401": { - "limit": 1851 + "limit": 1726 }, "ASYNC230": { "limit": 14 @@ -42,7 +42,7 @@ "limit": 84 }, "B010": { - "limit": 194 + "limit": 192 }, "B018": { "limit": 5 @@ -72,7 +72,7 @@ "limit": 22 }, "C408": { - "limit": 14 + "limit": 13 }, "C414": { "limit": 7 @@ -81,7 +81,7 @@ "limit": 4 }, "C901": { - "limit": 312 + "limit": 311 }, "D419": { "limit": 9 @@ -255,16 +255,16 @@ "limit": 480 }, "S110": { - "limit": 236 + "limit": 235 }, "S112": { "limit": 24 }, "SIM101": { - "limit": 61 + "limit": 59 }, "SIM102": { - "limit": 324 + "limit": 323 }, "SIM103": { "limit": 129 @@ -321,7 +321,7 @@ "limit": 121 }, "TRY300": { - "limit": 879 + "limit": 878 }, "UP006": { "limit": 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 8dbf38555b3..ab9efa9c7cf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -2649,7 +2649,7 @@ class _TxProxyModelTable: class _TxPrismaClient: - """Minimal prisma stub whose ``db.tx()`` yields a transaction and records commit.""" + """Minimal prisma stub whose ``tx()`` yields a transaction and records commit.""" def __init__(self, rows): self.events: list = [] @@ -2668,6 +2668,7 @@ class _TxPrismaClient: self.db = MagicMock() self.db.tx = MagicMock(return_value=_TxCM()) + self.tx = self.db.tx class _RecordingRouter: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f071c381916..bcd9109289a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23350 + "limit": 23263 }, "LIT002": { - "limit": 27256 + "limit": 27216 }, "LIT003": { "limit": 292 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1105 + "limit": 1085 }, "LIT007": { "limit": 0 @@ -24,6 +24,6 @@ "limit": 1004 }, "LIT009": { - "limit": 2465 + "limit": 2426 } }