diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d3259f88dce..af80cb7616b 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,18 +1,18 @@ { "reportAny": { - "limit": 29809 + "limit": 26731 }, "reportArgumentType": { - "limit": 2645 + "limit": 2627 }, "reportAssignmentType": { "limit": 329 }, "reportAttributeAccessIssue": { - "limit": 516 + "limit": 501 }, "reportCallIssue": { - "limit": 123 + "limit": 115 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 24 }, "reportExplicitAny": { - "limit": 9473 + "limit": 8505 }, "reportFunctionMemberAccess": { "limit": 7 @@ -57,7 +57,7 @@ "limit": 5855 }, "reportMissingTypeArgument": { - "limit": 15849 + "limit": 15793 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1079 + "limit": 1042 }, "reportOptionalOperand": { "limit": 0 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 219 + "limit": 213 }, "reportTypedDictNotRequiredAccess": { "limit": 27 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45262 + "limit": 45217 }, "reportUnknownLambdaType": { - "limit": 113 + "limit": 106 }, "reportUnknownMemberType": { - "limit": 40452 + "limit": 39883 }, "reportUnknownParameterType": { - "limit": 20309 + "limit": 20210 }, "reportUnknownVariableType": { - "limit": 31978 + "limit": 31466 }, "reportUnnecessaryCast": { "limit": 124 @@ -123,10 +123,10 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 866 + "limit": 863 }, "reportUntypedBaseClass": { - "limit": 165 + "limit": 164 }, "reportUntypedFunctionDecorator": { "limit": 33 diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index cb9d36f2dfd..04120c5118b 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -10,6 +10,7 @@ from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason if TYPE_CHECKING: + from openai.types.chat import ChatCompletionChunk from opentelemetry.trace import Span as _Span from litellm.types.utils import ModelResponseStream @@ -326,7 +327,8 @@ def process_response_headers( def preserve_upstream_non_openai_attributes( - model_response: "ModelResponseStream", original_chunk: "ModelResponseStream" + model_response: "ModelResponseStream", + original_chunk: "ModelResponseStream | ChatCompletionChunk", ): """ Preserve non-OpenAI attributes from the original chunk. diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..4c36425aa00 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,13 +3,12 @@ import time from collections.abc import Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Final, TypedDict, Union, cast + +from pydantic import TypeAdapter from litellm._logging import verbose_logger -from litellm.types.llms.openai import ( - ChatCompletionAssistantContentValue, - ChatCompletionAudioDelta, -) +from litellm.types.llms.openai import ChatCompletionAssistantContentValue from litellm.types.utils import ( CacheCreationTokenDetails, ChatCompletionAudioResponse, @@ -30,6 +29,7 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, ) @@ -39,6 +39,54 @@ if TYPE_CHECKING: ) +_MAPPING_ADAPTER: Final = TypeAdapter(dict[str, object]) +_SEQUENCE_ADAPTER: Final = TypeAdapter(list[object]) + + +def _as_mapping(value: object) -> dict[str, object] | None: + if not isinstance(value, dict): + return None + return _MAPPING_ADAPTER.validate_python(value) + + +def _as_list(value: object) -> list[object] | None: + if not isinstance(value, list): + return None + return _SEQUENCE_ADAPTER.validate_python(value) + + +def _str_attr(obj: object, name: str) -> str | None: + value: Final = getattr(obj, name, None) + return value if isinstance(value, str) else None + + +def _int_attr(obj: object, name: str) -> int | None: + value: Final = getattr(obj, name, None) + return value if isinstance(value, int) else None + + +def _dict_style_get(obj: object, key: str, default: object = None) -> object: + getter: Final = getattr(obj, "get", None) + if callable(getter): + return getter(key, default) + return getattr(obj, key, default) + + +def _hidden_params_of(chunk: object) -> dict[str, object] | None: + mapping: Final = _as_mapping(chunk) + if mapping is not None: + return _as_mapping(mapping.get("_hidden_params")) + return _as_mapping(getattr(chunk, "_hidden_params", None)) + + +class _ToolCallAccumulator(TypedDict): + id: str | None + name: str | None + type: str | None + custom_name: str | None + provider_specific_fields: dict[str, object] | None + + def capture_cache_creation_token_details( prompt_tokens_details: PromptTokensDetailsWrapper | None, current: CacheCreationTokenDetails | None, @@ -71,63 +119,54 @@ class ChunkProcessor: def __init__(self, chunks: list, messages: list | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages - self.first_chunk = chunks[0] + self.first_chunk: Mapping[str, object] = chunks[0] def _sort_chunks(self, chunks: list) -> list: if not chunks: return [] first_chunk: Final = chunks[0] - first_hidden_params: dict[str, Any] = {} - if isinstance(first_chunk, dict): - candidate = first_chunk.get("_hidden_params", {}) - if isinstance(candidate, dict): - first_hidden_params = candidate - else: - candidate = getattr(first_chunk, "_hidden_params", {}) - if isinstance(candidate, dict): - first_hidden_params = candidate + first_hidden_params: Final = _hidden_params_of(first_chunk) or {} if first_hidden_params.get("created_at"): - def _created_at(chunk: Any) -> int | float: - if isinstance(chunk, dict): - params = chunk.get("_hidden_params", {}) - else: - params = getattr(chunk, "_hidden_params", {}) - if isinstance(params, dict): - return cast(int | float, params.get("created_at", float("inf"))) - return float("inf") + def _created_at(chunk: object) -> int | float: + params: Final = _hidden_params_of(chunk) + if params is None: + return float("inf") + created_at: Final = params.get("created_at", float("inf")) + return created_at if isinstance(created_at, (int, float)) else float("inf") return sorted(chunks, key=_created_at) return chunks def update_model_response_with_hidden_params( - self, model_response: ModelResponse, chunk: dict[str, Any] | None = None + self, model_response: ModelResponse, chunk: Mapping[str, object] | None = None ) -> ModelResponse: if chunk is None: return model_response # set hidden params from chunk to model_response if model_response is not None and hasattr(model_response, "_hidden_params"): - model_response._hidden_params = chunk.get("_hidden_params", {}) + model_response._hidden_params = _as_mapping(chunk.get("_hidden_params")) or {} return model_response @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, - chunks: list[Any], - logging_obj: Any | None = None, + chunks: list[object], + logging_obj: "LiteLLMLoggingObject | None" = None, ) -> None: if not chunks: return - model: Final = getattr(response, "model", None) + model: Final = response.model if not model: return - custom_llm_provider = None - if logging_obj is not None: - custom_llm_provider = logging_obj.model_call_details.get("custom_llm_provider") + provider_value: Final = ( + logging_obj.model_call_details.get("custom_llm_provider") if logging_obj is not None else None + ) + custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else None try: from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -159,18 +198,19 @@ class ChunkProcessor: ) @staticmethod - def _get_chunk_id(chunks: list[dict[str, Any]]) -> str: + def _get_chunk_id(chunks: list[Mapping[str, object]]) -> str: """ Chunks: [{"id": ""}, {"id": "1"}, {"id": "1"}] """ for chunk in chunks: - if chunk.get("id"): - return chunk["id"] + chunk_id = chunk.get("id") + if isinstance(chunk_id, str) and chunk_id: + return chunk_id return "" @staticmethod - def _get_model_from_chunks(chunks: list[dict[str, Any]], first_chunk_model: str) -> str: + def _get_model_from_chunks(chunks: list[Mapping[str, object]], first_chunk_model: str) -> str: """ Get the actual model from chunks, preferring a model that differs from the first chunk. @@ -181,33 +221,43 @@ class ChunkProcessor: # Look for a model in chunks that differs from the first chunk's model for chunk in chunks: chunk_model = chunk.get("model") - if chunk_model and chunk_model != first_chunk_model: + if isinstance(chunk_model, str) and chunk_model and chunk_model != first_chunk_model: return chunk_model # Fall back to first chunk's model if no different model found return first_chunk_model - def build_base_response(self, chunks: list[dict[str, Any]]) -> ModelResponse: + def build_base_response(self, chunks: list[Mapping[str, object]]) -> ModelResponse: chunk = self.first_chunk id: Final = ChunkProcessor._get_chunk_id(chunks) object: Final = chunk["object"] created: Final = chunk["created"] - first_chunk_model: Final = chunk["model"] + first_chunk_model_value: Final = chunk["model"] + first_chunk_model: Final = first_chunk_model_value if isinstance(first_chunk_model_value, str) else "" # Get the actual model - for Azure Model Router, this finds the real model from later chunks model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model) system_fingerprint: Final = chunk.get("system_fingerprint", None) first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), chunk) - role: Final = first_chunk_with_choices["choices"][0]["delta"]["role"] + first_role_choices: Final = _as_list(first_chunk_with_choices.get("choices")) or [] + first_role_choice: Final = first_role_choices[0] if first_role_choices else None + first_role_delta: Final = _dict_style_get(first_role_choice, "delta") if first_role_choice is not None else None + role: Final = _dict_style_get(first_role_delta, "role") if first_role_delta is not None else None finish_reason = "stop" for chunk in chunks: - if "choices" in chunk and len(chunk["choices"]) > 0: + choices_value = _as_list(chunk.get("choices")) + if choices_value is None or len(choices_value) == 0: + continue + first_choice = choices_value[0] + if hasattr(first_choice, "finish_reason"): + chunk_finish_reason = _str_attr(first_choice, "finish_reason") + else: 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"] - if chunk_finish_reason is not None: - finish_reason = chunk_finish_reason + first_choice_mapping = _as_mapping(first_choice) + if first_choice_mapping is not None and "finish_reason" in first_choice_mapping: + candidate = first_choice_mapping.get("finish_reason") + chunk_finish_reason = candidate if isinstance(candidate, str) else None + if chunk_finish_reason is not None: + finish_reason = chunk_finish_reason # Initialize the response dictionary response = ModelResponse( @@ -237,35 +287,47 @@ class ChunkProcessor: @staticmethod def _iter_tool_call_fragments( - tool_call_chunks: Sequence[Mapping[str, Any]], + tool_call_chunks: Sequence[Mapping[str, object]], ) -> Iterator[tuple[int, str, str]]: for chunk in tool_call_chunks: - for choice in chunk["choices"]: - delta = choice.get("delta") + for choice in _as_list(chunk.get("choices")) or (): + delta = _dict_style_get(choice, "delta") if not delta: continue - for tool_call in delta.get("tool_calls", ()): + tool_calls = _as_list(_dict_style_get(delta, "tool_calls", ())) + for tool_call in tool_calls or (): if not tool_call: continue - if isinstance(tool_call, dict): - index = tool_call.get("index", 0) - function = tool_call.get("function") - if isinstance(function, dict): - if function.get("arguments"): - yield index, "arguments", function["arguments"] - elif getattr(function, "arguments", None): - yield index, "arguments", function.arguments - custom = tool_call.get("custom") - if isinstance(custom, dict) and custom.get("input"): - yield index, "custom_input", custom["input"] + tool_call_mapping = _as_mapping(tool_call) + if tool_call_mapping is not None: + index_value = tool_call_mapping.get("index", 0) + index = index_value if isinstance(index_value, int) else 0 + function = tool_call_mapping.get("function") + function_mapping = _as_mapping(function) + if function_mapping is not None: + arguments = function_mapping.get("arguments") + if isinstance(arguments, str) and arguments: + yield index, "arguments", arguments + else: + arguments = _str_attr(function, "arguments") + if arguments: + yield index, "arguments", arguments + custom = tool_call_mapping.get("custom") + custom_mapping = _as_mapping(custom) + if custom_mapping is not None: + custom_input = custom_mapping.get("input") + if isinstance(custom_input, str) and custom_input: + yield index, "custom_input", custom_input else: - index = getattr(tool_call, "index", 0) + index = _int_attr(tool_call, "index") or 0 function = getattr(tool_call, "function", None) - if getattr(function, "arguments", None): - yield index, "arguments", function.arguments + arguments = _str_attr(function, "arguments") + if arguments: + yield index, "arguments", arguments custom = getattr(tool_call, "custom", None) - if getattr(custom, "input", None): - yield index, "custom_input", custom.input + custom_input = _str_attr(custom, "input") + if custom_input: + yield index, "custom_input", custom_input @staticmethod def _join_fragments_by_index_and_field( @@ -282,44 +344,49 @@ class ChunkProcessor: ) def get_combined_tool_content( - self, tool_call_chunks: Sequence[Mapping[str, Any]] + self, tool_call_chunks: Sequence[Mapping[str, object]] ) -> list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ]: # mutable-ok: assigned verbatim to Message.tool_calls, a list field tool_calls_list: list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ] = [] # mutable-ok: see return type - tool_call_map: Final[dict[int, dict[str, Any]]] = {} # Map to store tool calls by index + tool_call_map: Final[dict[int, _ToolCallAccumulator]] = {} # Map to store tool calls by index for chunk in tool_call_chunks: - choices = chunk["choices"] + choices = _as_list(chunk.get("choices")) or () for choice in choices: - delta = choice.get("delta", {}) - tool_calls = delta.get("tool_calls", []) + delta = _dict_style_get(choice, "delta", {}) + tool_calls = _as_list(_dict_style_get(delta, "tool_calls", [])) or () for tool_call in tool_calls: # Handle both dict and object formats if not tool_call: continue + tool_call_mapping = _as_mapping(tool_call) + # Check if tool_call has function (either as attribute or dict key) has_function = False has_custom = False - if isinstance(tool_call, dict): - has_function = "function" in tool_call and tool_call["function"] is not None - has_custom = "custom" in tool_call and tool_call["custom"] is not None + if tool_call_mapping is not None: + has_function = "function" in tool_call_mapping and tool_call_mapping["function"] is not None + has_custom = "custom" in tool_call_mapping and tool_call_mapping["custom"] is not None else: - has_function = hasattr(tool_call, "function") and tool_call.function is not None + has_function = ( + hasattr(tool_call, "function") and getattr(tool_call, "function", None) is not None + ) has_custom = getattr(tool_call, "custom", None) is not None if not has_function and not has_custom: continue # Get index (handle both dict and object) - if isinstance(tool_call, dict): - index = tool_call.get("index", 0) + if tool_call_mapping is not None: + index_value = tool_call_mapping.get("index", 0) else: - index = getattr(tool_call, "index", 0) + index_value = getattr(tool_call, "index", 0) + index = index_value if isinstance(index_value, int) else 0 if index not in tool_call_map: tool_call_map[index] = { @@ -331,62 +398,81 @@ 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"] + if tool_call_mapping is not None: + tool_call_id = tool_call_mapping.get("id") + if isinstance(tool_call_id, str) and tool_call_id: + tool_call_map[index]["id"] = tool_call_id + tool_call_type = tool_call_mapping.get("type") + if isinstance(tool_call_type, str) and 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"] + function = tool_call_mapping.get("function", {}) + function_mapping = _as_mapping(function) + if function_mapping is not None: + function_name = function_mapping.get("name") + if isinstance(function_name, str) and function_name: + tool_call_map[index]["name"] = function_name else: # function is an object - if hasattr(function, "name") and function.name: - tool_call_map[index]["name"] = function.name + function_name = _str_attr(function, "name") + if function_name: + tool_call_map[index]["name"] = function_name - custom = tool_call.get("custom") - if isinstance(custom, dict): - if custom.get("name"): - tool_call_map[index]["custom_name"] = custom["name"] + custom = tool_call_mapping.get("custom") + custom_mapping = _as_mapping(custom) + if custom_mapping is not None: + custom_name = custom_mapping.get("name") + if isinstance(custom_name, str) and custom_name: + tool_call_map[index]["custom_name"] = custom_name else: # tool_call is an object - if hasattr(tool_call, "id") and tool_call.id: - tool_call_map[index]["id"] = tool_call.id - if hasattr(tool_call, "type") and tool_call.type: - tool_call_map[index]["type"] = tool_call.type + tool_call_id = _str_attr(tool_call, "id") + if tool_call_id: + tool_call_map[index]["id"] = tool_call_id + tool_call_type = _str_attr(tool_call, "type") + if tool_call_type: + tool_call_map[index]["type"] = tool_call_type if hasattr(tool_call, "function"): - if hasattr(tool_call.function, "name") and tool_call.function.name: - tool_call_map[index]["name"] = tool_call.function.name + tool_call_function = getattr(tool_call, "function", None) + function_name = _str_attr(tool_call_function, "name") + if function_name: + tool_call_map[index]["name"] = function_name - custom = getattr(tool_call, "custom", None) - if custom is not None: - if getattr(custom, "name", None): - tool_call_map[index]["custom_name"] = custom.name + custom_name = _str_attr(getattr(tool_call, "custom", None), "name") + if custom_name: + tool_call_map[index]["custom_name"] = custom_name # Preserve provider_specific_fields from streaming chunks - provider_fields = None - if isinstance(tool_call, dict): - provider_fields = tool_call.get("provider_specific_fields") - if not provider_fields and isinstance(tool_call.get("function"), dict): - provider_fields = tool_call["function"].get("provider_specific_fields") + provider_fields: object = None + if tool_call_mapping is not None: + provider_fields = tool_call_mapping.get("provider_specific_fields") + if not provider_fields: + nested_function_mapping = _as_mapping(tool_call_mapping.get("function")) + if nested_function_mapping is not None: + provider_fields = nested_function_mapping.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 + if hasattr(tool_call, "provider_specific_fields") and getattr( + tool_call, "provider_specific_fields", None ): - provider_fields = tool_call.function.provider_specific_fields + provider_fields = getattr(tool_call, "provider_specific_fields", None) + else: + tool_call_function = getattr(tool_call, "function", None) + if ( + hasattr(tool_call, "function") + and hasattr(tool_call_function, "provider_specific_fields") + and getattr(tool_call_function, "provider_specific_fields", None) + ): + provider_fields = getattr(tool_call_function, "provider_specific_fields", None) if provider_fields: # Merge provider_specific_fields if multiple chunks have them - if tool_call_map[index]["provider_specific_fields"] is None: - tool_call_map[index]["provider_specific_fields"] = {} - if isinstance(provider_fields, dict): - tool_call_map[index]["provider_specific_fields"].update(provider_fields) + provider_fields_mapping = _as_mapping(provider_fields) + if provider_fields_mapping is not None: + existing_fields = tool_call_map[index]["provider_specific_fields"] or {} + tool_call_map[index]["provider_specific_fields"] = { + **existing_fields, + **provider_fields_mapping, + } joined_fragments: Final = self._join_fragments_by_index_and_field( self._iter_tool_call_fragments(tool_call_chunks) @@ -414,39 +500,45 @@ 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) + provider_specific_fields = tool_call_data.get("provider_specific_fields") + if provider_specific_fields: + tool_call = ChatCompletionMessageToolCall( + id=tool_call_data["id"], + function=function, + type=tool_call_data["type"] or "function", + provider_specific_fields=provider_specific_fields, + ) + else: + tool_call = ChatCompletionMessageToolCall( + id=tool_call_data["id"], + function=function, + type=tool_call_data["type"] or "function", + ) tool_calls_list.append(tool_call) return tool_calls_list - def get_combined_function_call_content(self, function_call_chunks: list[dict[str, Any]]) -> FunctionCall: - argument_list: Final = [] - delta = function_call_chunks[0]["choices"][0]["delta"] - function_call = delta.get("function_call", "") - function_call_name: Final = function_call.name + def get_combined_function_call_content(self, function_call_chunks: list[Mapping[str, object]]) -> FunctionCall: + argument_list: Final[list[str]] = [] + first_choices: Final = _as_list(function_call_chunks[0].get("choices")) or [] + first_choice: Final = first_choices[0] if first_choices else None + first_delta: Final = _dict_style_get(first_choice, "delta") if first_choice is not None else None + first_function_call: Final = _dict_style_get(first_delta, "function_call") if first_delta is not None else None + function_call_name: Final = _str_attr(first_function_call, "name") for chunk in function_call_chunks: - choices = chunk["choices"] + choices = _as_list(chunk.get("choices")) or () for choice in choices: - delta = choice.get("delta", {}) - function_call = delta.get("function_call", "") + delta = _dict_style_get(choice, "delta", {}) + function_call = _dict_style_get(delta, "function_call", "") # Check if a function call is present if function_call: # Now, function_call is expected to be a dictionary - arguments = function_call.arguments - argument_list.append(arguments) + arguments = _str_attr(function_call, "arguments") + if arguments is not None: + argument_list.append(arguments) combined_arguments: Final = "".join(argument_list) @@ -456,17 +548,18 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: list[Mapping[str, object]], delta_key: str = "content" ) -> ChatCompletionAssistantContentValue: content_list: Final[list[str]] = [] for chunk in chunks: - choices = chunk["choices"] + choices = _as_list(chunk.get("choices")) or () for choice in choices: - delta = choice.get("delta", {}) - content = delta.get(delta_key, "") + delta = _dict_style_get(choice, "delta", {}) + content = _dict_style_get(delta, delta_key, "") if content is None: continue # openai v1.0.0 sets content = None for chunks - content_list.append(content) + if isinstance(content, str): + content_list.append(content) # Combine the "content" strings into a single string || combine the 'function' strings into a single string combined_content: Final = "".join(content_list) @@ -475,7 +568,7 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] + self, chunks: list[Mapping[str, object]] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -500,17 +593,20 @@ class ChunkProcessor: current_signature = None for chunk in chunks: - choices = chunk["choices"] + choices = _as_list(chunk.get("choices")) or () for choice in choices: - delta = choice.get("delta", {}) - thinking = delta.get("thinking_blocks", None) - if thinking and isinstance(thinking, list): - for thinking_block in thinking: + delta = _dict_style_get(choice, "delta", {}) + thinking = _as_list(_dict_style_get(delta, "thinking_blocks", None)) + if thinking: + for thinking_block_value in thinking: + thinking_block = _as_mapping(thinking_block_value) + if thinking_block is None: + continue thinking_type = thinking_block.get("type", None) if thinking_type and thinking_type == "redacted_thinking": _flush_thinking_block() redacted_data = thinking_block.get("data", None) - if redacted_data: + if isinstance(redacted_data, str) and redacted_data: thinking_blocks.append( ChatCompletionRedactedThinkingBlock( type="redacted_thinking", @@ -519,10 +615,10 @@ class ChunkProcessor: ) else: thinking_text = thinking_block.get("thinking", None) - if thinking_text: + if isinstance(thinking_text, str) and thinking_text: current_thinking_text_parts.append(thinking_text) signature = thinking_block.get("signature", None) - if signature: + if isinstance(signature, str) and signature: current_signature = signature _flush_thinking_block() @@ -532,20 +628,20 @@ class ChunkProcessor: return thinking_blocks return None - def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue: + def get_combined_reasoning_content(self, chunks: list[Mapping[str, object]]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: list[Mapping[str, object]]) -> ChatCompletionAudioResponse: base64_data_list: Final[list[str]] = [] transcript_list: Final[list[str]] = [] expires_at: int | None = None id: str | None = None for chunk in chunks: - choices = chunk["choices"] + choices = _as_list(chunk.get("choices")) or () for choice in choices: - delta = choice.get("delta") or {} - audio: ChatCompletionAudioDelta | None = delta.get("audio") + delta = _dict_style_get(choice, "delta", {}) + audio = _as_mapping(_dict_style_get(delta, "audio")) if audio is not None: for k, v in audio.items(): if k == "data" and v is not None and isinstance(v, str): @@ -623,24 +719,28 @@ class ChunkProcessor: 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: Mapping[str, object] | ModelResponse | ModelResponseStream) -> Usage | None: + def _usage_value() -> object: + if hasattr(chunk, "usage") and getattr(chunk, "usage", None) is not None: + return getattr(chunk, "usage", None) + if "usage" in chunk: + return chunk.get("usage") + if (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr( + chunk, "_hidden_params" + ): + return chunk._hidden_params.get("usage", None) + return None - if isinstance(usage_chunk, dict): - return Usage(**usage_chunk) - return usage_chunk + usage_value: Final = _usage_value() + if isinstance(usage_value, Usage): + return usage_value + if isinstance(usage_value, dict): + return Usage(**usage_value) + return None def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: list[Mapping[str, object] | ModelResponse], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -707,19 +807,12 @@ 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", + if usage_chunk_dict["prompt_tokens_details"] is not None: + candidate_web_search_requests = _int_attr( + usage_chunk_dict["prompt_tokens_details"], "web_search_requests" ) + if candidate_web_search_requests is not None: + web_search_requests = candidate_web_search_requests prompt_tokens_details = ( cast( @@ -758,7 +851,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: list[Mapping[str, object] | ModelResponse], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -781,15 +874,18 @@ class ChunkProcessor: if saw_non_cursor_completion: return completion_tokens - custom_llm_provider: str | None = None - if chunks: - first_chunk: Final = 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") + first_chunk: Final = chunks[0] if chunks else None + hp: Final = ( + ( + first_chunk.get("_hidden_params") + if isinstance(first_chunk, dict) + else getattr(first_chunk, "_hidden_params", None) + ) + if first_chunk is not None + else None + ) + provider_value: Final = hp.get("custom_llm_provider") if isinstance(hp, dict) else None + custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else None if custom_llm_provider == "anthropic" and completion_tokens == 1: return 0 @@ -797,7 +893,7 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: list[Mapping[str, object] | ModelResponse], model: str, completion_output: str, messages: list | None = None, @@ -852,7 +948,10 @@ 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() + accepted_prediction_tokens=completion_tokens_details.accepted_prediction_tokens, + audio_tokens=completion_tokens_details.audio_tokens, + reasoning_tokens=completion_tokens_details.reasoning_tokens, + rejected_prediction_tokens=completion_tokens_details.rejected_prediction_tokens, ) else: returned_usage.completion_tokens_details = completion_tokens_details diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index be5694b7554..39ae3905e83 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -8,11 +8,12 @@ import time import traceback from collections.abc import AsyncIterator, Callable, Iterator from dataclasses import dataclass -from typing import Any, Final, NoReturn, TypeVar, Union, cast +from typing import Any, Final, NoReturn, TypedDict, TypeVar, Union, cast import anyio import httpx -from pydantic import BaseModel +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, TypeAdapter import litellm from litellm import verbose_logger @@ -22,10 +23,12 @@ from litellm.litellm_core_utils.model_response_utils import ( ) from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.types.llms.openai import OpenAIChatCompletionChunk +from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CacheCreationTokenDetails, + ChatCompletionDeltaCustomToolCall, + ChatCompletionDeltaToolCall, CompletionTokensDetailsWrapper, Delta, LlmProviders, @@ -68,7 +71,7 @@ def _next_sync_or_exhausted(it: Any) -> Any: return _SYNC_ITER_EXHAUSTED -def is_async_iterable(obj: Any) -> bool: +def is_async_iterable(obj: object) -> bool: """ Check if an object is an async iterable (can be used with 'async for'). @@ -89,14 +92,107 @@ def print_verbose(print_statement): pass +class _ProviderResponseObj(TypedDict, total=False): + original_chunk: "ModelResponseStream | ChatCompletionChunk" + provider_specific_fields: dict[str, object] + + +def _extract_provider_response_obj(original_chunk: object, provider_specific_fields: object) -> _ProviderResponseObj: + if isinstance(original_chunk, (ModelResponseStream, ChatCompletionChunk)) and isinstance( + provider_specific_fields, dict + ): + return {"original_chunk": original_chunk, "provider_specific_fields": provider_specific_fields} + if isinstance(original_chunk, (ModelResponseStream, ChatCompletionChunk)): + return {"original_chunk": original_chunk} + if isinstance(provider_specific_fields, dict): + return {"provider_specific_fields": provider_specific_fields} + return {} + + +class _PredibaseTokenPayload(BaseModel): + text: str = "" + + +class _PredibaseDetailsPayload(BaseModel): + finish_reason: str | None = None + + +class _PredibaseChunkPayload(BaseModel): + token: _PredibaseTokenPayload | None = None + details: _PredibaseDetailsPayload | None = None + generated_text: str | None = None + error: str | None = None + + +class _Ai21CompletionDataPayload(BaseModel): + text: str + + +class _Ai21CompletionPayload(BaseModel): + data: _Ai21CompletionDataPayload + + +class _Ai21ChunkPayload(BaseModel): + completions: list[_Ai21CompletionPayload] + + +class _MaritalkChunkPayload(BaseModel): + answer: str + + +class _NlpCloudChunkPayload(BaseModel): + generated_text: str + + +class _AlephAlphaCompletionPayload(BaseModel): + completion: str + + +class _AlephAlphaChunkPayload(BaseModel): + completions: list[_AlephAlphaCompletionPayload] + + +class _AzureChunkDeltaPayload(BaseModel): + content: str | None = None + + +class _AzureChunkChoicePayload(BaseModel): + delta: _AzureChunkDeltaPayload | None = None + finish_reason: str | None = None + + +class _AzureChunkPayload(BaseModel): + choices: list[_AzureChunkChoicePayload] + + +class _TritonChunkPayload(BaseModel): + text_output: str = "" + stop_reason: str | None = None + is_finished: bool = False + input_token_count: int = 0 + generated_token_count: int = 0 + + +class _CompletionObjRequiredFields(TypedDict): + content: str + + +class _CompletionObjFields(_CompletionObjRequiredFields, total=False): + role: str + tool_calls: "list[ChatCompletionToolCallChunk] | list[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall] | None" + function_call: object + provider_specific_fields: dict[str, object] + index: int + + @dataclass(frozen=True, slots=True) class _ProviderChunkParsed: - response_obj: dict[str, Any] + response_obj: _ProviderResponseObj @dataclass(frozen=True, slots=True) class _ProviderChunkEarlyReturn: - value: Any + value: "ModelResponseStream | None" _ProviderChunkResult = Union[_ProviderChunkParsed, _ProviderChunkEarlyReturn] @@ -107,7 +203,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, @@ -349,19 +445,21 @@ class CustomStreamWrapper: finish_reason = "" print_verbose(f"chunk: {chunk}") if chunk.startswith("data:"): - data_json: Final = json.loads(chunk[5:]) + data_json: Final = TypeAdapter(_PredibaseChunkPayload).validate_json(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: + text = data_json.token.text # rebind-ok: conditionally overrides the default set above + 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 - text = "" # don't return the final bos token + finish_reason = ( + data_json.details.finish_reason + ) # rebind-ok: conditionally overrides the default set above + elif data_json.generated_text: + text = "" # rebind-ok: conditionally overrides the default set above 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, @@ -381,7 +479,8 @@ class CustomStreamWrapper: chunk = chunk.decode("utf-8") data_json: Final = json.loads(chunk) try: - text: Final = data_json["completions"][0]["data"]["text"] + parsed: Final = _Ai21ChunkPayload.model_validate(data_json) + text: Final = parsed.completions[0].data.text is_finished: Final = True finish_reason: Final = "stop" return { @@ -396,7 +495,8 @@ class CustomStreamWrapper: chunk = chunk.decode("utf-8") data_json: Final = json.loads(chunk) try: - text: Final = data_json["answer"] + parsed: Final = _MaritalkChunkPayload.model_validate(data_json) + text: Final = parsed.answer is_finished: Final = True finish_reason: Final = "stop" return { @@ -415,8 +515,8 @@ class CustomStreamWrapper: if self.model and "dolphin" in self.model: chunk = self.process_chunk(chunk=chunk) else: - data_json: Final = json.loads(chunk) - chunk = data_json["generated_text"] + data_json: Final = TypeAdapter(_NlpCloudChunkPayload).validate_json(chunk) + chunk = data_json.generated_text text = chunk if "[DONE]" in text: text = text.replace("[DONE]", "") @@ -434,7 +534,8 @@ class CustomStreamWrapper: chunk = chunk.decode("utf-8") data_json: Final = json.loads(chunk) try: - text: Final = data_json["completions"][0]["completion"] + parsed: Final = _AlephAlphaChunkPayload.model_validate(data_json) + text: Final = parsed.completions[0].completion is_finished: Final = True finish_reason: Final = "stop" return { @@ -462,12 +563,17 @@ class CustomStreamWrapper: elif chunk.startswith("data:"): data_json: Final = json.loads(chunk[5:]) # chunk.startswith("data:"): try: - if len(data_json["choices"]) > 0: - delta: Final = data_json["choices"][0]["delta"] - text = "" if delta is None else delta.get("content", "") - if data_json["choices"][0].get("finish_reason", None): + parsed: Final = _AzureChunkPayload.model_validate(data_json) + if len(parsed.choices) > 0: + delta: Final = parsed.choices[0].delta + text = ( + "" if delta is None else (delta.content or "") + ) # rebind-ok: conditionally overrides the default set above + if parsed.choices[0].finish_reason: is_finished = True - finish_reason = data_json["choices"][0]["finish_reason"] + finish_reason = parsed.choices[ + 0 + ].finish_reason # rebind-ok: conditionally overrides the default set above print_verbose(f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}") return { "text": text, @@ -621,34 +727,37 @@ class CustomStreamWrapper: def handle_triton_stream(self, chunk): try: - if isinstance(chunk, dict): - parsed_response = 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) - else: - return { - "text": "", - "is_finished": False, - "prompt_tokens": 0, - "completion_tokens": 0, - } - else: + + def _parsed_response() -> _TritonChunkPayload | None: + if isinstance(chunk, dict): + return _TritonChunkPayload.model_validate(chunk) + if isinstance(chunk, (str, bytes)): + working_chunk: Final = chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk + if "text_output" not in working_chunk: + return None + response: Final = (CustomStreamWrapper._strip_sse_data_from_chunk(working_chunk) or "").strip() + return TypeAdapter(_TritonChunkPayload).validate_json(response) print_verbose(f"chunk: {chunk} (Type: {type(chunk)})") raise ValueError(f"Unable to parse response. Original response: {chunk}") - text: Final = parsed_response.get("text_output", "") - finish_reason: Final = parsed_response.get("stop_reason") - is_finished: Final = parsed_response.get("is_finished", False) + + parsed_response: Final = _parsed_response() + if parsed_response is None: + return { + "text": "", + "is_finished": False, + "finish_reason": None, + "prompt_tokens": 0, + "completion_tokens": 0, + } + text: Final = parsed_response.text_output + finish_reason: Final = parsed_response.stop_reason + is_finished: Final = parsed_response.is_finished 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), + "prompt_tokens": parsed_response.input_token_count, + "completion_tokens": parsed_response.generated_token_count, } return {"text": "", "is_finished": False} except Exception as e: @@ -729,7 +838,7 @@ class CustomStreamWrapper: def copy_model_response_level_provider_specific_fields( self, - original_chunk: ModelResponseStream | OpenAIChatCompletionChunk, + original_chunk: ModelResponseStream | ChatCompletionChunk, model_response: ModelResponseStream, ) -> ModelResponseStream: """ @@ -744,13 +853,13 @@ class CustomStreamWrapper: def is_chunk_non_empty( self, - completion_obj: dict[str, Any], + completion_obj: _CompletionObjFields, model_response: ModelResponseStream, - response_obj: dict[str, Any], + response_obj: _ProviderResponseObj, ) -> bool: if ( "content" in completion_obj - and (isinstance(completion_obj["content"], str) and len(completion_obj["content"]) > 0) + and len(completion_obj["content"]) > 0 or ( "tool_calls" in completion_obj and completion_obj["tool_calls"] is not None @@ -775,7 +884,7 @@ class CustomStreamWrapper: "provider_specific_fields" in model_response and model_response.choices[0].delta.provider_specific_fields is not None ) - or ("provider_specific_fields" in response_obj and response_obj["provider_specific_fields"] is not None) + or ("provider_specific_fields" in response_obj) or ( "annotations" in model_response.choices[0].delta and model_response.choices[0].delta.annotations is not None @@ -870,9 +979,9 @@ class CustomStreamWrapper: def return_processed_chunk_logic( # noqa: C901 self, - completion_obj: dict[str, Any], + completion_obj: _CompletionObjFields, model_response: ModelResponseStream, - response_obj: dict[str, Any], + response_obj: _ProviderResponseObj, ): from litellm.litellm_core_utils.core_helpers import ( preserve_upstream_non_openai_attributes, @@ -895,12 +1004,11 @@ class CustomStreamWrapper: choices: Final = [] for choice in original_chunk.choices: try: - if isinstance(choice, BaseModel): - 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(**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(**choice_json)) except Exception: choices.append(StreamingChoices()) setattr(model_response, "choices", choices) @@ -929,8 +1037,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"] + _provider_specific_fields: Final = response_obj.get("provider_specific_fields") + if _provider_specific_fields is not None: + completion_obj["provider_specific_fields"] = _provider_specific_fields model_response.choices[0].delta = Delta(**completion_obj) _index: Final[int | None] = completion_obj.get("index") if _index is not None: @@ -1030,7 +1139,7 @@ class CustomStreamWrapper: self, chunk: Any, model_response: ModelResponseStream, - completion_obj: dict[str, Any], + completion_obj: _CompletionObjFields, ) -> _ProviderChunkResult: response_obj: dict[str, Any] = {} if ( @@ -1348,16 +1457,21 @@ class CustomStreamWrapper: "usage", litellm.Usage(**response_obj["usage"].model_dump()), ) - return _ProviderChunkParsed(response_obj) + return _ProviderChunkParsed( + _extract_provider_response_obj( + original_chunk=response_obj.get("original_chunk"), + provider_specific_fields=response_obj.get("provider_specific_fields"), + ) + ) def chunk_creator(self, chunk: Any): if hasattr(chunk, "id"): self.response_id = chunk.id model_response = self.model_response_creator() - response_obj: dict[str, Any] = {} + response_obj: _ProviderResponseObj = {} try: # return this for all models - completion_obj: Final[dict[str, Any]] = {"content": ""} + completion_obj: Final[_CompletionObjFields] = {"content": ""} dispatch_result: Final = self._dispatch_provider_chunk( chunk=chunk, model_response=model_response, @@ -1369,7 +1483,7 @@ class CustomStreamWrapper: model_response.model = self.model ## FUNCTION CALL PARSING - original_chunk: Final = response_obj.get("original_chunk") if response_obj is not None else None + original_chunk: Final = response_obj.get("original_chunk") if ( original_chunk is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints @@ -1382,7 +1496,7 @@ class CustomStreamWrapper: ) if original_chunk.choices and len(original_chunk.choices) > 0: delta = original_chunk.choices[0].delta - if delta is not None and (delta.function_call is not None or delta.tool_calls is not None): + if 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 @@ -1396,17 +1510,12 @@ class CustomStreamWrapper: ): 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 = "" + for t in original_chunk.choices[0].delta.tool_calls: + _t_functions = getattr(t, "functions", None) + if _t_functions is not None and hasattr(_t_functions, "arguments"): + _t_function = getattr(t, "function", None) + if getattr(_t_function, "arguments", None) is None: + setattr(_t_function, "arguments", "") # noqa: B010 # _t_function's type is unknown (defensive getattr chain above); direct assignment would break None-narrowing _json_delta: Final = 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 @@ -1430,11 +1539,7 @@ 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) - ) + delta = dict(original_chunk.choices[0].delta) model_response.choices[0].delta = Delta(**delta) except Exception: model_response.choices[0].delta = Delta() @@ -1445,7 +1550,11 @@ class CustomStreamWrapper: return ## CHECK FOR TOOL USE - if "tool_calls" in completion_obj and len(completion_obj["tool_calls"]) > 0: + if ( + "tool_calls" in completion_obj + and completion_obj["tool_calls"] is not None + and len(completion_obj["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"] completion_obj["tool_calls"] = None @@ -2067,9 +2176,10 @@ class CustomStreamWrapper: # end-of-stream blocks complete. Scheduling here via # create_task would race with unified_guardrail's # end-of-stream block for short-stream providers. - self.logging_obj._deferred_stream_complete_args = ( - complete_streaming_response, - cache_hit, + setattr( # noqa: B010 # _deferred_stream_complete_args is a dynamic cross-module signal, not a declared attribute + self.logging_obj, + "_deferred_stream_complete_args", + (complete_streaming_response, cache_hit), ) else: # prefer_async_handlers routes CustomLogger to async_success_handler diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 27754012ea5..aaf606e0227 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,5 +1,5 @@ from collections.abc import AsyncIterator, Coroutine, Iterator -from typing import Any, Final, cast +from typing import Any, Final, TypeAlias, cast import litellm from litellm._logging import verbose_logger @@ -24,8 +24,10 @@ 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: Final[frozenset[str]] = frozenset({"output_config"}) +AnthropicSystemContent: TypeAlias = str | list[dict[str, object]] | None -def _messages_have_compaction_block(messages: list[dict]) -> bool: + +def _messages_have_compaction_block(messages: list[dict[str, object]]) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: content = msg.get("content") @@ -37,7 +39,7 @@ def _messages_have_compaction_block(messages: list[dict]) -> bool: return False -def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: +def _extract_proxy_litellm_metadata(kwargs: dict[str, object]) -> dict[str, object] | None: """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise. The proxy attaches its auth/spend-attribution fields (``user_api_key``, @@ -58,9 +60,9 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | async def _prepare_context_managed_request( *, model: str, - messages: list[dict], + messages: list[dict[str, object]], tools: list[dict] | None, - system: Any | None, + system: AnthropicSystemContent, context_management_spec: Any, litellm_metadata: dict | None, additional_drop_params: list[str] | None, @@ -85,11 +87,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 # rebind-ok: else branch below assigns the same names once + working_system: AnthropicSystemContent = system # rebind-ok: else branch below assigns the same names once 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 @@ -119,7 +121,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 @@ -222,9 +224,9 @@ def _normalize_spec_edits( async def _run_polyfill_if_enabled( *, model: str, - messages: list[dict], + messages: list[dict[str, object]], tools: list[dict] | None, - system: Any | None, + system: AnthropicSystemContent, context_management_spec: Any, litellm_metadata: dict | None, additional_drop_params: list[str] | None, @@ -292,7 +294,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _route_openai_thinking_to_responses_api_if_needed( completion_kwargs: dict[str, Any], *, - thinking: dict[str, Any] | None, + thinking: dict[str, object] | None, ) -> None: """ When users call `litellm.anthropic.messages.*` with a non-Anthropic model and @@ -339,7 +341,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: reasoning_effort: Final = completion_kwargs.get("reasoning_effort") summary: Final = thinking.get("summary") if isinstance(reasoning_effort, str) and reasoning_effort: - reasoning_dict: Final[dict[str, Any]] = {"effort": reasoning_effort} + reasoning_dict: Final[dict[str, object]] = {"effort": reasoning_effort} if summary: reasoning_dict["summary"] = summary elif auto_summary: @@ -398,7 +400,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: metadata: dict | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, - system: str | list[dict[str, Any]] | None = None, + system: AnthropicSystemContent = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, @@ -406,7 +408,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, + extra_kwargs: dict[str, object] | None = None, ) -> tuple[dict[str, Any], dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. @@ -419,7 +421,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -514,7 +516,7 @@ 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, @@ -528,7 +530,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_p: float | None = None, output_format: dict | None = None, **kwargs, - ) -> AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]: + ) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]: """Handle non-Anthropic models asynchronously using the adapter""" context_management: Final = kwargs.pop("context_management", None) additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) @@ -608,7 +610,7 @@ 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, @@ -626,8 +628,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: ) -> ( AnthropicMessagesResponse | Iterator[bytes] - | AsyncIterator[Any] - | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]] + | AsyncIterator[bytes] + | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]] ): """Handle non-Anthropic models using the adapter.""" if _is_async is True: diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 79c4362055e..9035f93196e 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -8,11 +8,13 @@ import json import os from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, Union, cast +from typing import TYPE_CHECKING, Final, Literal, TypeVar, Union, cast, get_args, get_origin from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue, TypeAdapter +from pydantic.fields import FieldInfo +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -612,11 +614,11 @@ class RegisterGuardrailRequest(BaseModel): """Request body for POST /guardrails/register. Follows Generic Guardrail API config.""" guardrail_name: str - litellm_params: dict[str, Any] # guardrail, mode, api_base required; api_key, headers, etc. optional + litellm_params: dict[str, object] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: dict[str, object] | None = None team_id: str | None = None - def get_litellm_params_dict(self) -> dict[str, Any]: + def get_litellm_params_dict(self) -> dict[str, object]: return dict(self.litellm_params) @@ -708,6 +710,11 @@ async def register_guardrail( status_code=400, detail="litellm_params.api_base is required for generic_guardrail_api", ) + if not isinstance(api_base, str): + raise HTTPException( + status_code=400, + detail="litellm_params.api_base must be a string", + ) parsed: Final = urlparse(api_base) if parsed.scheme not in ("http", "https"): raise HTTPException( @@ -772,16 +779,17 @@ async def register_guardrail( raise HTTPException(status_code=500, detail=str(e)) -def _parse_json_field(value: object) -> dict[str, Any] | None: +def _parse_json_field(value: object) -> dict[str, object] | None: if value is None: return None if isinstance(value, dict): return value if isinstance(value, str): try: - return json.loads(value) + parsed: Final = json.loads(value) except Exception: return None + return parsed if isinstance(parsed, dict) else None return None @@ -809,6 +817,10 @@ async def _get_user_team_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str]: return [t for t in user_obj.teams if t] +def _str_or_none(value: object) -> str | None: + return value if isinstance(value, str) else None + + def _row_to_submission_item(row: "LiteLLM_GuardrailsTable") -> GuardrailSubmissionItem: from litellm.litellm_core_utils.litellm_logging import _get_masked_values @@ -824,8 +836,8 @@ def _row_to_submission_item(row: "LiteLLM_GuardrailsTable") -> GuardrailSubmissi team_guardrail=team_guardrail, litellm_params=masked_params, guardrail_info=guardrail_info, - submitted_by_user_id=guardrail_info.get("submitted_by_user_id"), - submitted_by_email=guardrail_info.get("submitted_by_email"), + submitted_by_user_id=_str_or_none(guardrail_info.get("submitted_by_user_id")), + submitted_by_email=_str_or_none(guardrail_info.get("submitted_by_email")), submitted_at=getattr(row, "submitted_at", None), reviewed_at=getattr(row, "reviewed_at", None), created_at=row.created_at, @@ -1171,12 +1183,18 @@ async def patch_guardrail( ) # Update litellm_params if default_on is provided or pii_entities_config is provided - litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {}))) - if request.litellm_params is not None: - requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True) - litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) - litellm_params_dict.update(requested_litellm_params) - litellm_params = LitellmParams(**litellm_params_dict) + existing_litellm_params: Final = existing_guardrail.get("litellm_params", {}) + + def _merged_litellm_params() -> LitellmParams: + base: Final = LitellmParams.model_validate(existing_litellm_params) + if request.litellm_params is None: + return base + requested: Final = request.litellm_params.model_dump(exclude_unset=True) + merged: Final = base.model_dump(exclude_unset=True) + merged.update(requested) + return LitellmParams.model_validate(merged) + + litellm_params: Final = _merged_litellm_params() # Update guardrail_info if provided guardrail_info: Final = ( @@ -1428,6 +1446,15 @@ async def get_category_yaml(category_name: str): raise HTTPException(status_code=500, detail=f"Error reading category file: {e}") +class _AirlineEntry(TypedDict): + id: str + match: str + tags: list[str] + + +_AIRLINES_ADAPTER: Final = TypeAdapter(list[_AirlineEntry]) + + @router.get( "/guardrails/ui/major_airlines", tags=["Guardrails"], @@ -1452,9 +1479,7 @@ async def get_major_airlines(): ) try: with open(airlines_path, "r", encoding="utf-8") as f: - import json - - airlines: Final = json.load(f) + airlines: Final = _AIRLINES_ADAPTER.validate_json(f.read()) return {"airlines": airlines} except Exception as e: raise HTTPException(status_code=500, detail=f"Error reading major_airlines.json: {e}") from e @@ -1551,36 +1576,38 @@ async def validate_blocked_words_file(request: dict[str, str]): return {"valid": False, "error": f"Validation error: {e}"} -def _get_field_type_from_annotation(field_annotation: Any) -> str: +def _annotation_origin(annotation: object) -> object | None: + return get_origin(annotation) + + +def _annotation_args(annotation: object) -> tuple[object, ...]: + return get_args(annotation) + + +def _get_field_type_from_annotation(field_annotation: object) -> str: """ Convert a Python type annotation to a UI-friendly type string """ # Handle Union types (like Optional[T]) - if ( - hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is Union - and hasattr(field_annotation, "__args__") - ): + if _annotation_origin(field_annotation) is Union: # For Optional[T], get the non-None type - args: Final = field_annotation.__args__ - non_none_args: Final = [arg for arg in args if arg is not type(None)] + non_none_args: Final = tuple(arg for arg in _annotation_args(field_annotation) if arg is not type(None)) if non_none_args: field_annotation = non_none_args[0] + origin: Final = _annotation_origin(field_annotation) + # Handle List types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is list: + if origin is list: return "array" # Handle Dict types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is dict: + if origin is dict: return "dict" - # Handle Literal types - if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"): - # Check for Literal types (Python 3.8+) - origin: Final = field_annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": - return "select" # For dropdown/select inputs + # Handle Literal types (Python 3.8+) + if origin is Literal: + return "select" # For dropdown/select inputs # Handle basic types if field_annotation is str: @@ -1598,66 +1625,49 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str: return "string" -def _extract_literal_values(annotation: Any) -> list[str]: +def _extract_literal_values(annotation: object) -> list[str]: """ Extract literal values from a Literal type annotation """ - if hasattr(annotation, "__origin__") and hasattr(annotation, "__args__"): - origin: Final = annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": - return list(annotation.__args__) + if _annotation_origin(annotation) is Literal: + return [arg for arg in _annotation_args(annotation) if isinstance(arg, str)] return [] -def _get_dict_key_options(field_annotation: Any) -> list[str] | None: +def _get_dict_key_options(field_annotation: object) -> list[str] | None: """ Extract key options from Dict[Literal[...], T] types """ - if ( - hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict - and hasattr(field_annotation, "__args__") - ): - args: Final = field_annotation.__args__ + if _annotation_origin(field_annotation) is dict: + args: Final = _annotation_args(field_annotation) if len(args) >= 2: - key_type: Final = args[0] - return _extract_literal_values(key_type) + return _extract_literal_values(args[0]) return None -def _get_dict_value_type(field_annotation: Any) -> str: +def _get_dict_value_type(field_annotation: object) -> str: """ Get the value type from Dict[K, V] types """ - if ( - hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict - and hasattr(field_annotation, "__args__") - ): - args: Final = field_annotation.__args__ + if _annotation_origin(field_annotation) is dict: + args: Final = _annotation_args(field_annotation) if len(args) >= 2: - value_type: Final = args[1] - return _get_field_type_from_annotation(value_type) + return _get_field_type_from_annotation(args[1]) return "string" -def _get_list_element_options(field_annotation: Any) -> list[str] | None: +def _get_list_element_options(field_annotation: object) -> list[str] | None: """ Extract element options from List[Literal[...]] types """ - if ( - hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is list - and hasattr(field_annotation, "__args__") - ): - args: Final = field_annotation.__args__ + if _annotation_origin(field_annotation) is list: + args: Final = _annotation_args(field_annotation) if len(args) >= 1: - element_type: Final = args[0] - return _extract_literal_values(element_type) + return _extract_literal_values(args[0]) return None -def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool: +def _should_skip_optional_params(field_name: str, field_annotation: object) -> bool: """Check if optional_params field should be skipped (not meaningfully overridden).""" if field_name != "optional_params": return False @@ -1666,62 +1676,55 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool return True # Check if the annotation is still a generic TypeVar (not specialized) - if isinstance(field_annotation, TypeVar) or ( - hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar - ): + if isinstance(field_annotation, TypeVar) or _annotation_origin(field_annotation) is TypeVar: return True # Also skip if it's a generic type that wasn't specialized - if hasattr(field_annotation, "__name__") and field_annotation.__name__ in ( - "T", - "TypeVar", - ): + if getattr(field_annotation, "__name__", None) in ("T", "TypeVar"): return True # Handle Optional[T] where T is still a TypeVar - if hasattr(field_annotation, "__args__"): - non_none_args: Final = [arg for arg in field_annotation.__args__ if arg is not type(None)] - if non_none_args and isinstance(non_none_args[0], TypeVar): - return True - - return False + non_none_args: Final = tuple(arg for arg in _annotation_args(field_annotation) if arg is not type(None)) + return bool(non_none_args and isinstance(non_none_args[0], TypeVar)) -def _unwrap_optional_type(field_annotation: Any) -> Any: +def _unwrap_optional_type(field_annotation: object) -> object: """Unwrap Optional types to get the actual type.""" - if ( - hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is Union - and hasattr(field_annotation, "__args__") - ): + if _annotation_origin(field_annotation) is Union: # For Optional[BaseModel], get the non-None type - args: Final = field_annotation.__args__ - non_none_args: Final = [arg for arg in args if arg is not type(None)] + non_none_args: Final = tuple(arg for arg in _annotation_args(field_annotation) if arg is not type(None)) if non_none_args: return non_none_args[0] return field_annotation +def _field_json_schema_extra(field: FieldInfo) -> dict[str, JsonValue] | None: + extra: Final = field.json_schema_extra + return extra if isinstance(extra, dict) else None + + def _build_field_dict( - field: Any, - field_annotation: Any, + field: FieldInfo, + field_annotation: object, description: str, required: bool, -) -> dict[str, Any]: +) -> dict[str, object]: """Build field dictionary for non-nested fields.""" - # Determine the field type from annotation - field_type = _get_field_type_from_annotation(field_annotation) - # Check for custom UI type override - field_json_schema_extra: Final = getattr(field, "json_schema_extra", {}) - if field_json_schema_extra and "ui_type" in field_json_schema_extra: - ui_type: Final = field_json_schema_extra["ui_type"] - field_type = ui_type.value if hasattr(ui_type, "value") else ui_type - elif field_json_schema_extra and "type" in field_json_schema_extra: - field_type = field_json_schema_extra["type"] + field_json_schema_extra: Final = _field_json_schema_extra(field) + + def _field_type() -> object: + if field_json_schema_extra and "ui_type" in field_json_schema_extra: + ui_type: Final = field_json_schema_extra["ui_type"] + return getattr(ui_type, "value", ui_type) + if field_json_schema_extra and "type" in field_json_schema_extra: + return field_json_schema_extra["type"] + return _get_field_type_from_annotation(field_annotation) + + field_type: Final = _field_type() # Add the field to the dictionary - field_dict: Final = { + field_dict: Final[dict[str, object]] = { "description": description, "required": required, "type": field_type, @@ -1770,7 +1773,7 @@ def _build_field_dict( def _extract_fields_recursive( model: type[BaseModel], depth: int = 0, -) -> dict[str, Any]: +) -> dict[str, object]: # Check if we've exceeded the maximum recursion depth if depth > DEFAULT_MAX_RECURSE_DEPTH: raise HTTPException( @@ -1778,7 +1781,7 @@ def _extract_fields_recursive( detail=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing model fields. Please check the model structure for excessive nesting.", ) - fields: Final = {} + fields: Final[dict[str, object]] = {} for field_name, field in model.model_fields.items(): field_annotation = field.annotation @@ -1798,15 +1801,13 @@ def _extract_fields_recursive( required = field.is_required() # Check if this is a BaseModel subclass - is_basemodel_subclass = ( + if ( inspect.isclass(field_annotation) and issubclass(field_annotation, BaseModel) and field_annotation is not BaseModel - ) - - if is_basemodel_subclass: + ): # Recursively get fields from the nested model - nested_fields = _extract_fields_recursive(cast(type[BaseModel], field_annotation), depth + 1) + nested_fields = _extract_fields_recursive(field_annotation, depth + 1) fields[field_name] = { "description": description, "required": required, @@ -1824,7 +1825,7 @@ def _extract_fields_recursive( return fields -def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, Any]: +def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, object]: """ Get the fields from a Pydantic model as a nested dictionary structure """ @@ -1951,6 +1952,10 @@ class TestCustomCodeGuardrailResponse(BaseModel): """Type of error: 'compilation' or 'execution'.""" +def _widen_str_object_dict(mapping: dict[str, object]) -> dict[str, object]: + return mapping + + @router.post( "/guardrails/test_custom_code", tags=["Guardrails"], @@ -2042,7 +2047,7 @@ async def test_custom_code_guardrail( EXECUTION_TIMEOUT_SECONDS: Final = 5 try: - exec_globals: Final = build_sandbox_globals() + exec_globals: Final = _widen_str_object_dict(build_sandbox_globals()) try: compiled: Final[CodeType] = compile_sandboxed(request.custom_code) @@ -2069,7 +2074,7 @@ async def test_custom_code_guardrail( error_type="compilation", ) - apply_fn: Final[object] = exec_globals["apply_guardrail"] + apply_fn: Final = exec_globals["apply_guardrail"] if not callable(apply_fn): return TestCustomCodeGuardrailResponse( success=False, @@ -2148,18 +2153,20 @@ def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: return "response" if input_type == "response" else "request" -def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None: +def _patch_logging_obj_for_guardrail(litellm_logging_obj: object, request: ApplyGuardrailRequest) -> None: """Configure the logging object so Langfuse/OTEL extract input and output correctly.""" - litellm_logging_obj.call_type = "pass_through_endpoint" - litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint" - litellm_logging_obj.update_messages( - request.messages if request.messages else [{"role": "user", "content": request.text}] - ) + setattr(litellm_logging_obj, "call_type", "pass_through_endpoint") # noqa: B010 # litellm_logging_obj is typed `object` to avoid a logging-module import cycle + model_call_details: Final = getattr(litellm_logging_obj, "model_call_details", None) + if isinstance(model_call_details, dict): + model_call_details["call_type"] = "pass_through_endpoint" + update_messages: Final = getattr(litellm_logging_obj, "update_messages", None) + if callable(update_messages): + update_messages(request.messages if request.messages else [{"role": "user", "content": request.text}]) async def _emit_guardrail_success_logs( - proxy_logging_obj: Any, - litellm_logging_obj: Any, + proxy_logging_obj: object, + litellm_logging_obj: object | None, data: dict, user_api_key_dict: UserAPIKeyAuth, response: ApplyGuardrailResponse, @@ -2176,7 +2183,7 @@ async def _emit_guardrail_success_logs( ) try: - modified: Final = await proxy_logging_obj.post_call_success_hook( + modified: Final = await getattr(proxy_logging_obj, "post_call_success_hook")( # noqa: B009 # proxy_logging_obj is typed `object` to avoid an import cycle data=data, user_api_key_dict=user_api_key_dict, response=response, @@ -2194,7 +2201,7 @@ async def _emit_guardrail_success_logs( if litellm_logging_obj is not None: end_time: Final = datetime.now(timezone.utc) try: - await litellm_logging_obj.async_success_handler( + await getattr(litellm_logging_obj, "async_success_handler")( # noqa: B009 # litellm_logging_obj is typed `object | None` to avoid an import cycle result=response_for_logging, start_time=start_time, end_time=end_time, @@ -2204,7 +2211,7 @@ async def _emit_guardrail_success_logs( verbose_proxy_logger.exception("apply_guardrail: async_success_handler failed") try: thread_pool_executor.submit( - litellm_logging_obj.success_handler, + getattr(litellm_logging_obj, "success_handler"), # noqa: B009 # litellm_logging_obj is typed `object | None` to avoid an import cycle response_for_logging, start_time, end_time, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 2725da1ee12..4d8e3a5af5a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -12,7 +12,9 @@ from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypedDict, Union, cast + +from pydantic import TypeAdapter from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -47,6 +49,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + from litellm.types.agents import AgentResponse from litellm.types.caching import RedisPipelineIncrementOperation Span = Union[_Span, Any] @@ -55,6 +58,17 @@ else: Span = Any InternalUsageCache = Any + +class BatchRateLimiterProtocol(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: ... + + BATCH_RATE_LIMITER_SCRIPT: Final = """ local results = {} local now = tonumber(ARGV[1]) @@ -335,6 +349,46 @@ class RateLimitResponseWithDescriptors(TypedDict): response: RateLimitResponse +WindowedCacheValue: TypeAlias = int | str | None +GaugeCacheValue: TypeAlias = dict[str, float] | int | None + + +class WindowKeyMeta(TypedDict): + requests_limit: int | None + tokens_limit: int | None + window_size: int + descriptor_key: str + + +class DescriptorCounterMeta(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 DescriptorRuntimeState(TypedDict): + window_expired: bool + current: int + + +_WINDOWED_CACHE_VALUE_ADAPTER: Final = TypeAdapter(WindowedCacheValue) +_WINDOWED_CACHE_VALUE_LIST_ADAPTER: Final = TypeAdapter(list[WindowedCacheValue] | None) +_LUA_BATCH_RESULT_ADAPTER: Final = TypeAdapter(list[WindowedCacheValue]) +_GAUGE_CACHE_VALUE_ADAPTER: Final = TypeAdapter(GaugeCacheValue) +_GAUGE_CACHE_VALUE_LIST_ADAPTER: Final = TypeAdapter(list[GaugeCacheValue] | None) +_INT_LIST_ADAPTER: Final = TypeAdapter(list[int]) +_OPTIONAL_STR_ADAPTER: Final = TypeAdapter(str | None) + + +def _as_int(value: object, default: int = 0) -> int: + return int(value) if isinstance(value, (int, float)) else default + + @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -452,7 +506,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: BatchRateLimiterProtocol | None = None # Serializes multi-phase check+increment sequences (batch + dynamic # limiters) within this process to close the TOCTOU window between @@ -470,7 +524,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) -> BatchRateLimiterProtocol | None: """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: try: @@ -599,12 +653,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): keys: list[str], now_int: int, window_size: int, - ) -> list[Any]: + ) -> list[WindowedCacheValue]: """ 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: Final[list[Any]] = [] + results: Final[list[WindowedCacheValue]] = [] # Process each window/counter pair for i in range(0, len(keys), 2): @@ -613,10 +667,12 @@ 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 = _WINDOWED_CACHE_VALUE_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=window_key, + litellm_parent_otel_span=None, + local_only=True, + ) ) # Check if window exists and is valid @@ -640,10 +696,12 @@ 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, + current_counter = _WINDOWED_CACHE_VALUE_ADAPTER.validate_python( + 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 await self.internal_usage_cache.async_set_cache( @@ -674,8 +732,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def is_cache_list_over_limit( self, keys_to_fetch: list[str], - cache_values: list[Any], - key_metadata: dict[str, Any], + cache_values: list[WindowedCacheValue], + key_metadata: dict[str, WindowKeyMeta], ) -> RateLimitResponse: """ Check if the cache values are over the limit. @@ -778,7 +836,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, keys_to_fetch: list[str], now_int: int, - ) -> list[Any]: + ) -> list[WindowedCacheValue]: """ Execute Redis operations grouped by hash tag for cluster compatibility. @@ -787,19 +845,21 @@ 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: Final = self._group_keys_by_hash_tag(keys_to_fetch) - all_cache_values: Final = [] + all_cache_values: Final[list[WindowedCacheValue]] = [] for hash_tag, group_keys in key_groups.items(): try: - group_cache_values = await self.batch_rate_limiter_script( - keys=group_keys, - args=[now_int, self.window_size], # Use integer timestamp + group_cache_values = _LUA_BATCH_RESULT_ADAPTER.validate_python( + await self.batch_rate_limiter_script( + keys=group_keys, + args=[now_int, self.window_size], # Use integer timestamp + ) ) all_cache_values.extend(group_cache_values) except Exception as e: @@ -861,10 +921,14 @@ 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( - keys=keys_to_fetch, - parent_otel_span=parent_otel_span, - local_only=True, + cache_values: list[WindowedCacheValue] | None = ( + _WINDOWED_CACHE_VALUE_LIST_ADAPTER.validate_python( # rebind-ok: multi-stage cache-then-redis accumulator reassigned further below + await self.internal_usage_cache.async_batch_get_cache( + keys=keys_to_fetch, + parent_otel_span=parent_otel_span, + local_only=True, + ) + ) ) if cache_values is not None: @@ -875,10 +939,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ## IF under limit in-memory, check Redis if read_only: # READ-ONLY MODE: Just read current values without incrementing - cache_values = await self.internal_usage_cache.async_batch_get_cache( - keys=keys_to_fetch, - parent_otel_span=parent_otel_span, - local_only=False, # Check Redis too + cache_values = _WINDOWED_CACHE_VALUE_LIST_ADAPTER.validate_python( # rebind-ok: multi-stage cache-then-redis accumulator + await self.internal_usage_cache.async_batch_get_cache( + keys=keys_to_fetch, + parent_otel_span=parent_otel_span, + local_only=False, # Check Redis too + ) ) # For keys that don't exist yet, set them to 0 @@ -944,14 +1010,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, descriptors: list[RateLimitDescriptor], skip_tpm_check: bool, - ) -> tuple[list[str], dict[str, dict[str, Any]], list[ParallelRequestGauge]]: + ) -> tuple[list[str], dict[str, WindowKeyMeta], 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: Final[list[str]] = [] - key_metadata: Final[dict[str, dict[str, Any]]] = {} + key_metadata: Final[dict[str, WindowKeyMeta]] = {} gauges: Final[list[ParallelRequestGauge]] = [] for descriptor in descriptors: descriptor_key = descriptor["key"] @@ -1007,7 +1073,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: GaugeCacheValue) -> int: """ In-flight count from a cached gauge value: a dict of slot_id -> acquire timestamp when the in-memory registry is authoritative, or @@ -1017,7 +1083,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return 0 if isinstance(raw_value, dict): cutoff: Final = 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 sum(1 for ts in raw_value.values() if ts >= cutoff) return max(0, int(raw_value)) async def _check_parallel_request_gauges( @@ -1044,11 +1110,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if read_only: if self.parallel_count_script is not None: try: - raw_counts: Final = await self.parallel_count_script( - keys=gauge_keys, - args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges], + raw_counts: Final = _INT_LIST_ADAPTER.validate_python( + await self.parallel_count_script( + 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, value) for value in raw_counts + ] # rebind-ok: except/else branches below assign the same name once except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500 verbose_proxy_logger.warning("parallel_count_script failed, using local mirror: %s", e) counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) @@ -1073,11 +1143,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if self.parallel_acquire_script is not None: try: - raw: Final = await self.parallel_acquire_script( - keys=gauge_keys, - args=[ - arg for gauge in gauges for arg in (gauge["limit"], PARALLEL_REQUEST_SLOT_TTL_SECONDS, slot_id) - ], + raw: Final = _INT_LIST_ADAPTER.validate_python( + await self.parallel_acquire_script( + keys=gauge_keys, + args=[ + arg + for gauge in gauges + for arg in (gauge["limit"], PARALLEL_REQUEST_SLOT_TTL_SECONDS, slot_id) + ], + ) ) except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500 verbose_proxy_logger.warning("parallel_acquire_script failed, falling back to in-memory gauge: %s", e) @@ -1109,10 +1183,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): gauge_keys: list[str], parent_otel_span: Span | None = None, ) -> list[int]: - values: Final = await self.internal_usage_cache.async_batch_get_cache( - keys=gauge_keys, - parent_otel_span=parent_otel_span, - local_only=True, + values: Final = _GAUGE_CACHE_VALUE_LIST_ADAPTER.validate_python( + await self.internal_usage_cache.async_batch_get_cache( + keys=gauge_keys, + parent_otel_span=parent_otel_span, + local_only=True, + ) ) if values is None: return [0 for _ in gauge_keys] @@ -1138,10 +1214,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): cutoff: Final = now - PARALLEL_REQUEST_SLOT_TTL_SECONDS states: Final[list[tuple[dict[str, float] | None, int]]] = [] for gauge in gauges: - raw_value = await self.internal_usage_cache.async_get_cache( - key=gauge["counter_key"], - litellm_parent_otel_span=parent_otel_span, - local_only=True, + raw_value = _GAUGE_CACHE_VALUE_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=gauge["counter_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) ) if isinstance(raw_value, dict): registry: dict[str, float] | None = { @@ -1193,9 +1271,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return if self.parallel_release_script is not None: try: - raw: Final = await self.parallel_release_script( - keys=counter_keys, - args=[slot_id for _ in counter_keys], + raw: Final = _INT_LIST_ADAPTER.validate_python( + await self.parallel_release_script( + keys=counter_keys, + args=[slot_id for _ in counter_keys], + ) ) for counter_key, remaining in zip(counter_keys, raw): await self.internal_usage_cache.async_set_cache( @@ -1211,10 +1291,12 @@ 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( - key=counter_key, - litellm_parent_otel_span=parent_otel_span, - local_only=True, + raw_value = _GAUGE_CACHE_VALUE_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=counter_key, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) ) if isinstance(raw_value, dict): if slot_id not in raw_value: @@ -1270,7 +1352,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: Final[list[tuple[list[str], list[Any], list[dict[str, Any]]]]] = [] + descriptor_groups: Final[list[tuple[list[str], list[int], list[DescriptorCounterMeta]]]] = [] for descriptor, increment_amounts in zip(descriptors, increments): keys, args, meta = self._build_descriptor_atomic_payload( descriptor=descriptor, @@ -1293,7 +1375,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: Final[list[DescriptorCounterMeta]] = [ + 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, @@ -1304,7 +1388,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, descriptor: RateLimitDescriptor, increment_amounts: dict[Literal["requests", "tokens"], int], - ) -> tuple[list[str], list[Any], list[dict[str, Any]]]: + ) -> tuple[list[str], list[int], list[DescriptorCounterMeta]]: """ Build (KEYS, ARGV, per-counter meta) for a single descriptor's Lua call. All keys returned share the descriptor's {key:value} hash tag. @@ -1318,8 +1402,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): window_key: Final = f"{{{descriptor_key}:{descriptor_value}}}:window" keys: Final[list[str]] = [] - args: Final[list[Any]] = [] - meta: Final[list[dict[str, Any]]] = [] + args: Final[list[int]] = [] + meta: Final[list[DescriptorCounterMeta]] = [] for rate_limit_type in ("requests", "tokens"): rlt: Literal["requests", "tokens"] = cast(Literal["requests", "tokens"], rate_limit_type) @@ -1358,7 +1442,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: list[tuple[list[str], list[int], list[DescriptorCounterMeta]]], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """ @@ -1367,14 +1451,16 @@ 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: Final[list[list[dict[str, Any]]]] = [] + applied: Final[list[list[DescriptorCounterMeta]]] = [] statuses: Final[list[RateLimitStatus]] = [] for _idx, (keys, args, meta) in enumerate(descriptor_groups): try: - raw = await self.check_and_increment_by_n_script( # pyright: ignore[reportOptionalCall] # sole caller guards it is not None - keys=keys, - args=args, + raw = _INT_LIST_ADAPTER.validate_python( + await self.check_and_increment_by_n_script( # pyright: ignore[reportOptionalCall] # sole caller guards it is not None + keys=keys, + args=args, + ) ) except Exception as e: # Lua failure (timeout, OOM, network partition) leaves Redis @@ -1389,7 +1475,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.window_size, ) await self._refund_applied_descriptor_groups(applied) - flat_meta: list[dict[str, Any]] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta] + flat_meta: list[DescriptorCounterMeta] = [ + 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, @@ -1407,7 +1495,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _refund_applied_descriptor_groups( self, - applied: list[list[dict[str, Any]]], + applied: list[list[DescriptorCounterMeta]], ) -> None: """ Decrement counters for descriptor groups already applied via Lua. @@ -1433,8 +1521,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_atomic_response( self, - raw: list[Any], - per_counter_meta: list[dict[str, Any]], + raw: list[int], + per_counter_meta: list[DescriptorCounterMeta], ) -> RateLimitResponse: """Convert Lua script return value to RateLimitResponse. @@ -1485,7 +1573,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _atomic_check_and_increment_in_memory( self, - per_counter_meta: list[dict[str, Any]], + per_counter_meta: list[DescriptorCounterMeta], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """In-memory all-or-nothing check-and-increment. Caller holds lock. @@ -1500,23 +1588,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): now_int: Final = int(self._get_current_time().timestamp()) # Pass 1: read state, validate. - descriptor_state: Final[list[dict[str, Any]]] = [] + descriptor_state: Final[list[DescriptorRuntimeState]] = [] for meta in per_counter_meta: window_size = meta["window_size"] - window_start = await self.internal_usage_cache.async_get_cache( - key=meta["window_key"], - litellm_parent_otel_span=parent_otel_span, - local_only=True, + window_start = _WINDOWED_CACHE_VALUE_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=meta["window_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) ) window_expired = window_start is None or (now_int - int(window_start)) >= window_size current_counter = ( 0 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, + _WINDOWED_CACHE_VALUE_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=meta["counter_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) ) or 0 ) @@ -1912,7 +2004,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 @@ -1923,11 +2015,11 @@ 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: Final = getattr(user_api_key_dict, "agent_id", None) + key_agent_id: Final = _OPTIONAL_STR_ADAPTER.validate_python(getattr(user_api_key_dict, "agent_id", None)) if key_agent_id: return key_agent_id metadata: Final = data.get("metadata") or {} - return metadata.get("agent_id") + return _OPTIONAL_STR_ADAPTER.validate_python(metadata.get("agent_id")) def _get_session_id_from_data(self, data: dict) -> str | None: """Extract session_id from request metadata or litellm_session_id.""" @@ -1961,8 +2053,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if agent is None: return descriptors - agent_rpm: Final = getattr(agent, "rpm_limit", None) - agent_tpm: Final = getattr(agent, "tpm_limit", None) + agent_rpm: Final = agent.rpm_limit + agent_tpm: Final = agent.tpm_limit if agent_rpm is not None or agent_tpm is not None: descriptors.append( RateLimitDescriptor( @@ -1976,8 +2068,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - session_rpm: Final = getattr(agent, "session_rpm_limit", None) - session_tpm: Final = getattr(agent, "session_tpm_limit", None) + session_rpm: Final = agent.session_rpm_limit + session_tpm: Final = agent.session_tpm_limit if session_rpm is not None or session_tpm is not None: session_id: Final = self._get_session_id_from_data(data) if session_id is not None: @@ -2238,7 +2330,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) -> BatchRateLimiterProtocol | None: """Get the rate limiter for the call type.""" if call_type == "acreate_batch": batch_limiter: Final = self._get_batch_rate_limiter() @@ -2600,7 +2692,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 | dict[str, object] | None, rate_limit_type: Literal["output", "input", "total"] ) -> int: """ Get total tokens from response usage for rate limiting. @@ -2622,24 +2714,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 + if rate_limit_type in ("input", "total") and usage.prompt_tokens_details is not None: + cached_tokens = ( + usage.prompt_tokens_details.cached_tokens or 0 + ) # rebind-ok: conditionally overrides the default set above - elif isinstance(usage, dict): + else: # Responses API usage comes as a dict if rate_limit_type == "output": - total_tokens = usage.get("completion_tokens", 0) or 0 + total_tokens = _as_int( + usage.get("completion_tokens") + ) # rebind-ok: conditionally overrides the default set above elif rate_limit_type == "input": - total_tokens = usage.get("prompt_tokens", 0) or 0 + total_tokens = _as_int( + usage.get("prompt_tokens") + ) # rebind-ok: conditionally overrides the default set above elif rate_limit_type == "total": - total_tokens = usage.get("total_tokens", 0) or 0 + total_tokens = _as_int( + usage.get("total_tokens") + ) # rebind-ok: conditionally overrides the default set above # Get cached tokens from dict if rate_limit_type in ("input", "total"): - prompt_details: Final = usage.get("prompt_tokens_details") or {} + prompt_details: Final = usage.get("prompt_tokens_details") if isinstance(prompt_details, dict): - cached_tokens = prompt_details.get("cached_tokens", 0) or 0 + cached_tokens = _as_int( + prompt_details.get("cached_tokens") + ) # rebind-ok: conditionally overrides the default set above # Subtract cached tokens for input/total (providers don't count them) if cached_tokens > 0: @@ -2765,15 +2866,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @staticmethod def _merge_ratelimit_statuses_into_additional_headers( - additional_headers: dict[str, Any], + additional_headers: dict[str, object], statuses: list[RateLimitStatus], - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Return ``additional_headers`` extended with ``x-ratelimit-{descriptor_key}-{remaining|limit}-{rate_limit_type}`` entries. Non-mutating so callers pick their own target dict. """ - merged: Final[dict[str, Any]] = dict(additional_headers) + merged: Final[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"] @@ -2782,8 +2883,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _collect_tpm_scope_targets( self, - standard_logging_metadata: dict[str, Any], - kwargs: Any, + standard_logging_metadata: dict[str, object], + kwargs: object, model_group: str | None, ) -> list[tuple[str, str]]: """ @@ -2793,16 +2894,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): the emitter; this helper just lists the candidate scopes so callers can split reserved-vs-unreserved. """ - user_api_key: Final = standard_logging_metadata.get("user_api_key_hash") - user_api_key_user_id: Final = standard_logging_metadata.get("user_api_key_user_id") - user_api_key_team_id: Final = standard_logging_metadata.get("user_api_key_team_id") - user_api_key_organization_id: Final = standard_logging_metadata.get("user_api_key_org_id") - user_api_key_project_id: Final = standard_logging_metadata.get("user_api_key_project_id") - user_api_key_end_user_id: Final = ( - kwargs.get("user") if isinstance(kwargs, dict) else None - ) or standard_logging_metadata.get("user_api_key_end_user_id") - agent_id: Final = standard_logging_metadata.get("agent_id") - session_id: Final = standard_logging_metadata.get("session_id") or standard_logging_metadata.get("trace_id") + user_api_key: Final = _OPTIONAL_STR_ADAPTER.validate_python(standard_logging_metadata.get("user_api_key_hash")) + user_api_key_user_id: Final = _OPTIONAL_STR_ADAPTER.validate_python( + standard_logging_metadata.get("user_api_key_user_id") + ) + user_api_key_team_id: Final = _OPTIONAL_STR_ADAPTER.validate_python( + standard_logging_metadata.get("user_api_key_team_id") + ) + user_api_key_organization_id: Final = _OPTIONAL_STR_ADAPTER.validate_python( + standard_logging_metadata.get("user_api_key_org_id") + ) + user_api_key_project_id: Final = _OPTIONAL_STR_ADAPTER.validate_python( + standard_logging_metadata.get("user_api_key_project_id") + ) + kwargs_user: Final = kwargs.get("user") if isinstance(kwargs, dict) else None + user_api_key_end_user_id: Final = _OPTIONAL_STR_ADAPTER.validate_python( + kwargs_user + ) or _OPTIONAL_STR_ADAPTER.validate_python(standard_logging_metadata.get("user_api_key_end_user_id")) + agent_id: Final = _OPTIONAL_STR_ADAPTER.validate_python(standard_logging_metadata.get("agent_id")) + session_id: Final = _OPTIONAL_STR_ADAPTER.validate_python( + standard_logging_metadata.get("session_id") + ) or _OPTIONAL_STR_ADAPTER.validate_python(standard_logging_metadata.get("trace_id")) targets: Final[list[tuple[str, str]]] = [] if user_api_key: @@ -2880,8 +2992,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_success_event_pipeline_operations( self, - kwargs: Any, - response_obj: Any, + kwargs: object, + response_obj: object, rate_limit_type: Literal["output", "input", "total"], ) -> list[RedisPipelineIncrementOperation]: """Build Redis pipeline increment ops for TPM / parallel-request counters.""" @@ -2889,12 +3001,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): get_model_group_from_litellm_kwargs, ) + kwargs_dict: Final[dict[str, object]] = kwargs if isinstance(kwargs, dict) else {} + # Get metadata from standard_logging_object - this correctly handles both # 'metadata' and 'litellm_metadata' fields from litellm_params - standard_logging_object: Final = kwargs.get("standard_logging_object") or {} - standard_logging_metadata: Final = standard_logging_object.get("metadata") or {} + _standard_logging_object: Final = kwargs_dict.get("standard_logging_object") + standard_logging_object: Final[dict[str, object]] = ( + _standard_logging_object if isinstance(_standard_logging_object, dict) else {} + ) + _standard_logging_metadata: Final = standard_logging_object.get("metadata") + standard_logging_metadata: Final[dict[str, object]] = ( + _standard_logging_metadata if isinstance(_standard_logging_metadata, dict) else {} + ) - model_group: Final = get_model_group_from_litellm_kwargs(kwargs) + model_group: Final = get_model_group_from_litellm_kwargs(kwargs_dict) # Get total tokens from response. Responses LiteLLM does not model # (e.g. pass-through, whose usage is reported by the upstream rather @@ -2911,9 +3031,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): BaseLiteLLMOpenAIResponseObject, ), ): - _usage = getattr(response_obj, "usage", None) + _usage_raw: Final = getattr(response_obj, "usage", None) + _usage = ( + _usage_raw if isinstance(_usage_raw, (Usage, dict)) else None + ) # rebind-ok: else branch below assigns the same name once else: - _combined_usage: Final = kwargs.get("combined_usage_object") + _combined_usage: Final = kwargs_dict.get("combined_usage_object") if isinstance(_combined_usage, Usage): _usage = _combined_usage total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type) @@ -3026,8 +3149,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 @@ -3045,9 +3168,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): standard_logging_object: Final = kwargs.get("standard_logging_object") if isinstance(standard_logging_object, dict): - hidden_params = standard_logging_object.get("hidden_params") - if not isinstance(hidden_params, dict): - hidden_params = {} + _hidden_params_raw: Final = standard_logging_object.get("hidden_params") + hidden_params: Final[dict[str, object]] = _hidden_params_raw if isinstance(_hidden_params_raw, dict) else {} existing = hidden_params.get("additional_headers") hidden_params["additional_headers"] = self._merge_ratelimit_statuses_into_additional_headers( additional_headers=existing if isinstance(existing, dict) else {}, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index a31687692d3..9b886f679a7 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -14,10 +14,10 @@ import asyncio import datetime import json from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -78,6 +78,7 @@ from litellm.types.router import ( from litellm.utils import get_utc_datetime router: Final = APIRouter() +_JSON_VALUE_ADAPTER: Final = TypeAdapter(object) async def update_team(*args, **kwargs): @@ -760,8 +761,17 @@ async def _setup_new_team_model_assignment( ) +class _ProxyModelTableRow(Protocol): + model_id: str + model_info: object + + +class _ProxyModelTableReader(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ProxyModelTableRow]: ... + + async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient, table: Any | None = None + team_id: str, prisma_client: PrismaClient, table: _ProxyModelTableReader | None = None ) -> list[LiteLLM_ProxyModelTable]: """ Fetch all deployments for a given team_id from the database. @@ -777,8 +787,8 @@ async def _get_team_deployments( existing transaction. """ prefix: Final = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table - response: Final = await table.find_many( + resolved_table: Final = table or ModelRepository(prisma_client).table + response: Final = await resolved_table.find_many( where={ "model_name": {"startswith": prefix}, } @@ -798,7 +808,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. @@ -1791,7 +1801,7 @@ def model_info_as_mapping(model_info: object) -> Mapping[str, object] | None: if not isinstance(model_info, str): return None try: - parsed: Final = json.loads(model_info) + parsed: Final = _JSON_VALUE_ADAPTER.validate_json(model_info) except (TypeError, ValueError): return None return parsed if isinstance(parsed, Mapping) else None diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fe5a0e06d2e..136af6dbe24 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -19,7 +19,7 @@ from typing import Annotated, Final, Protocol, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( UI_TEAM_ID, BlockTeamRequest, + BudgetNewRequest, CommonProxyErrors, DeleteTeamRequest, LiteLLM_AccessGroupTable, @@ -274,6 +275,23 @@ def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteL return tokens_table +class _RawTeamRow(Protocol): + """Shape of a team row as returned directly by Prisma, before the + litellm-side model re-validates JSON columns (e.g. ``members_with_roles`` + may still be a list of plain dicts rather than ``Member`` instances).""" + + team_id: str + organization_id: str | None + members_with_roles: Sequence[Mapping[str, object]] | None + + def model_dump(self) -> Mapping[str, object]: ... + + +def _team_db_raw(prisma_client: PrismaClient | None) -> "_PrismaTableActions[_RawTeamRow]": + raw_team_table: Final[_PrismaTableActions[_RawTeamRow]] = TeamRepository(prisma_client).table + return raw_team_table + + def _sanitize_for_log(value: object) -> str: """Strip CR/LF from user-controlled values to prevent log injection.""" try: @@ -339,6 +357,40 @@ async def _verify_team_access( ) +_STR_LIST_ADAPTER: Final = TypeAdapter(list[str]) + + +class _BudgetIdRow(Protocol): + """Shape actually read off the ``new_budget``/``update_budget`` endpoint + return value at these call sites: only ``budget_id`` is ever accessed.""" + + budget_id: str | None + + +async def _new_budget_typed( + budget_obj: BudgetNewRequest, + user_api_key_dict: UserAPIKeyAuth, +) -> _BudgetIdRow: + from litellm.proxy.management_endpoints.budget_management_endpoints import ( + new_budget, + ) + + budget_row: Final[_BudgetIdRow] = await new_budget(budget_obj=budget_obj, user_api_key_dict=user_api_key_dict) + return budget_row + + +async def _update_budget_typed( + budget_obj: BudgetNewRequest, + user_api_key_dict: UserAPIKeyAuth, +) -> _BudgetIdRow: + from litellm.proxy.management_endpoints.budget_management_endpoints import ( + update_budget, + ) + + budget_row: Final[_BudgetIdRow] = await update_budget(budget_obj=budget_obj, user_api_key_dict=user_api_key_dict) + return budget_row + + class TeamMemberBudgetHandler: """Helper class to handle team member budget, RPM, and TPM limit operations""" @@ -382,11 +434,6 @@ class TeamMemberBudgetHandler: team_member_budget_duration: str | None = None, ) -> dict: """Create team member budget table with provided limits""" - from litellm.proxy._types import BudgetNewRequest - from litellm.proxy.management_endpoints.budget_management_endpoints import ( - new_budget, - ) - if data.team_alias is not None: budget_id = f"team-{data.team_alias.replace(' ', '-')}-budget-{uuid.uuid4().hex}" else: @@ -407,7 +454,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - team_member_budget_table: Final = await new_budget( + team_member_budget_table: Final = await _new_budget_typed( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -433,11 +480,6 @@ class TeamMemberBudgetHandler: team_member_budget_duration: str | None = None, ) -> dict: """Upsert team member budget table with provided limits""" - from litellm.proxy._types import BudgetNewRequest - from litellm.proxy.management_endpoints.budget_management_endpoints import ( - update_budget, - ) - if team_table.metadata is None: team_table.metadata = {} @@ -455,7 +497,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - budget_row: Final = await update_budget( + budget_row: Final = await _update_budget_typed( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -501,7 +543,6 @@ class TeamMemberBudgetHandler: explicitly_set_fields: set, ) -> dict: """Clear explicitly-nulled fields on the team member budget row.""" - from litellm.proxy._types import BudgetNewRequest from litellm.proxy.management_endpoints.budget_management_endpoints import ( update_budget, ) @@ -557,20 +598,15 @@ class TeamMemberBudgetHandler: # Identify members with no existing membership row. # members_with_roles may contain Member instances or raw dicts depending # on how the team was fetched/deserialized. - missing: Final = [] - for m in members_with_roles: - user_id = m.get("user_id") if isinstance(m, dict) else m.user_id - if user_id is not None and user_id not in existing_user_ids: - missing.append( - { - "team_id": team_id, - "user_id": user_id, - "budget_id": team_member_budget_id, - } - ) + missing: Final[list[Mapping[str, object]]] = [ + {"team_id": team_id, "user_id": user_id, "budget_id": team_member_budget_id} + for m in members_with_roles + if (user_id := (m.get("user_id") if isinstance(m, dict) else m.user_id)) is not None + and user_id not in existing_user_ids + ] if missing: - await TeamMembershipRepository(prisma_client).table.create_many( + await _team_membership_db(prisma_client).create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -1393,7 +1429,11 @@ async def new_team( members_with_roles = complete_team_data.members_with_roles complete_team_data.members_with_roles = [] - complete_team_data_dict = complete_team_data.model_dump(exclude_none=True) + complete_team_data_dict: dict[str, object] = ( + complete_team_data.model_dump( # rebind-ok: reassigned via prisma_client.jsonify_team_object below + exclude_none=True + ) + ) # Serialize router_settings to JSON (matching key creation pattern) router_settings_value: Final = getattr(data, "router_settings", None) @@ -1407,7 +1447,7 @@ async def new_team( complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) - team_row: Final[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.create( + team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create( data=complete_team_data_dict, include={"litellm_model_table": True}, ) @@ -1855,7 +1895,7 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row: Final = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team_row is None: raise HTTPException( @@ -2431,7 +2471,7 @@ async def _process_team_members( # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models member_allowed_models = data.allowed_models - team_default_member_models: Final = getattr(complete_team_data, "default_team_member_models", None) + team_default_member_models: Final = complete_team_data.default_team_member_models if member_allowed_models is None and team_default_member_models: member_allowed_models = team_default_member_models @@ -2602,14 +2642,14 @@ async def _resolve_existing_member_user_ids( if not requested_user_ids: return frozenset() - found: Final = await UserRepository(prisma_client).table.find_many( + found: Final = await _user_db(prisma_client).find_many( where={ # mutable-ok: Prisma query filters are dict-shaped "user_id": { # mutable-ok: Prisma query filters are dict-shaped "in": sorted(requested_user_ids) } } ) - return frozenset(user.user_id for user in found or () if user.user_id is not None) + return frozenset(user.user_id for user in found or ()) def _pre_existing_user_ids( @@ -3092,20 +3132,22 @@ async def team_member_delete( ## DELETE TEAM ID from USER ROW, IF EXISTS ## # get user row - key_val: Final = {} - if data.user_id is not None: - key_val["user_id"] = data.user_id - elif data.user_email is not None: - key_val["user_email"] = data.user_email - existing_user_rows: Final = await UserRepository(prisma_client).table.find_many(where=key_val) + key_val: Final[dict[str, object]] = ( + {"user_id": data.user_id} + if data.user_id is not None + else {"user_email": data.user_email} + if data.user_email is not None + else {} + ) + existing_user_rows: Final = await _user_db(prisma_client).find_many(where=key_val) - if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): + if len(existing_user_rows) > 0: for existing_user in existing_user_rows: - team_list = [] + team_list: list[str] = [] if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await UserRepository(prisma_client).table.update( + await _user_db(prisma_client).update( where={ "user_id": existing_user.user_id, }, @@ -3113,18 +3155,12 @@ async def team_member_delete( ) # Also clean up any existing team membership rows for this user and team - user_ids_to_delete: Final = set() - if data.user_id is not None: - user_ids_to_delete.add(data.user_id) - if existing_user_rows is not None and isinstance(existing_user_rows, list): - for existing_user in existing_user_rows: - if getattr(existing_user, "user_id", None): - user_ids_to_delete.add(existing_user.user_id) + user_ids_to_delete: Final[frozenset[str]] = ( + frozenset([data.user_id]) if data.user_id is not None else frozenset[str]() + ) | frozenset(existing_user.user_id for existing_user in existing_user_rows if existing_user.user_id) for _uid in user_ids_to_delete: - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"team_id": data.team_id, "user_id": _uid} - ) + await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid}) ## DELETE KEYS CREATED BY USER FOR THIS TEAM if user_ids_to_delete: @@ -3133,9 +3169,7 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many( + keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3150,7 +3184,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await VerificationTokenRepository(prisma_client).table.delete_many( + await _tokens_db(prisma_client).delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3310,7 +3344,7 @@ async def team_member_update( ### upsert new budget budget_patch: Final = _build_member_budget_patch(data) - async with prisma_client.db.tx() as tx: + async with prisma_client.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, @@ -3653,7 +3687,7 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( where={"team_id": {"in": data.team_ids}} ) @@ -4452,6 +4486,16 @@ def _convert_teams_to_response_models( return team_list +class _TeamKeyCountGroup(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + team_id: str | None = None + count_by_field: dict[str, int] = Field(default_factory=dict, alias="_count") + + +_TEAM_KEY_COUNT_GROUP_ADAPTER: Final = TypeAdapter(list[_TeamKeyCountGroup]) + + async def _get_keys_count_by_team( prisma_client: PrismaClient, teams: Sequence[LiteLLM_TeamTable], @@ -4466,12 +4510,14 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped: Final = await VerificationTokenRepository(prisma_client).table.group_by( - by=["team_id"], - where={"team_id": {"in": page_team_ids}}, - count={"team_id": True}, + grouped: Final = _TEAM_KEY_COUNT_GROUP_ADAPTER.validate_python( + await VerificationTokenRepository(prisma_client).table.group_by( + by=["team_id"], + where={"team_id": {"in": page_team_ids}}, + count={"team_id": True}, + ) ) - return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")} + return {row.team_id: row.count_by_field.get("team_id", 0) for row in grouped if row.team_id} async def _enforce_list_team_v2_access( @@ -4735,7 +4781,7 @@ async def _authorize_and_filter_teams( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, -) -> list: +) -> list[_RawTeamRow]: """ Authorize the /team/list request and return filtered teams. @@ -4780,7 +4826,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams: Final = await TeamRepository(prisma_client).table.find_many( + org_teams: Final = await _team_db_raw(prisma_client).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4794,7 +4840,7 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response: Final = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}) + response: Final = await _team_db_raw(prisma_client).find_many(include={"litellm_model_table": True}) return [ team for team in response @@ -4802,7 +4848,7 @@ async def _authorize_and_filter_teams( ] else: # Proxy admin: all teams - return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) + return list(await _team_db_raw(prisma_client).find_many(include={"litellm_model_table": True})) @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -4854,14 +4900,16 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id}) + keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id}) try: returned_responses.append( - TeamListResponseObject( - **team.model_dump(), - team_memberships=_team_memberships, - keys=keys, + TeamListResponseObject.model_validate( + { + **team.model_dump(), + "team_memberships": _team_memberships, + "keys": keys, + } ) ) except Exception as e: @@ -4905,7 +4953,7 @@ async def get_paginated_teams( total_count: Final = await _team_db(prisma_client).count() # Get paginated teams - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4931,7 +4979,7 @@ async def ui_view_teams( page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=100), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> list[LiteLLM_TeamTable]: """ [PROXY-ADMIN ONLY] Filter teams based on partial match of team_id or team_alias with pagination. @@ -4955,7 +5003,7 @@ async def ui_view_teams( skip: Final = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Final = {} + where_conditions: Final[dict[str, object]] = {} if team_id: where_conditions["team_id"] = { @@ -4970,7 +5018,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -5160,13 +5208,13 @@ async def team_model_delete( ) # Get current models list - current_models: Final = team_obj.models or [] + current_models: Final[list[str]] = _STR_LIST_ADAPTER.validate_python(team_obj.models or []) # Remove specified models - updated_models: Final = [m for m in current_models if m not in data.models] + updated_models: Final[list[str]] = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team: Final = await TeamRepository(prisma_client).table.update( + updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, @@ -5419,15 +5467,14 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi BATCH_SIZE: Final = 500 while True: - find_args: dict = { - "take": BATCH_SIZE, - "order": {"team_id": "asc"}, - } - if cursor is not None: - find_args["cursor"] = {"team_id": cursor} - find_args["skip"] = 1 + cursor_arg = {"team_id": cursor} if cursor is not None else None - teams = await TeamRepository(prisma_client).table.find_many(**find_args) + teams = await _team_db(prisma_client).find_many( + take=BATCH_SIZE, + order={"team_id": "asc"}, + cursor=cursor_arg, + skip=1 if cursor is not None else None, + ) if not teams: break @@ -5522,11 +5569,13 @@ async def get_team_daily_activity( ) ## Fetch team aliases and check team admin status - where_condition: Final = {} + where_condition: Final[dict[str, object]] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases: Final = await TeamRepository(prisma_client).table.find_many(where=where_condition) - team_alias_metadata: Final = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} + team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition) + team_alias_metadata: Final[dict[str, dict[str, object]]] = { + t.team_id: {"team_alias": t.team_alias} for t in team_aliases + } # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 002ddd27e32..bb49b8b0d5b 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -32,12 +32,15 @@ from __future__ import annotations import json import re -from typing import Any, Final +from datetime import datetime +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, TypeVar, runtime_checkable from urllib.parse import quote, unquote from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.managed_resources.isolation import ( build_owner_filter, can_access_resource, @@ -51,6 +54,82 @@ from litellm.types.llms.openai import OpenAIFileObject from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +PrismaWhereValue: TypeAlias = ( + None | bool | int | float | str | datetime | list["PrismaWhereValue"] | dict[str, "PrismaWhereValue"] +) +PrismaWhere: TypeAlias = dict[str, PrismaWhereValue] + +_OBJECT_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) + + +class _OwnedRow(Protocol): + created_by: str | None + team_id: str | None + + +@runtime_checkable +class _FileIdLookupHook(Protocol): + async def get_unified_file_id( + self, file_id: str, litellm_parent_otel_span: object | None = None + ) -> _OwnedRow | None: ... + + +@runtime_checkable +class _FileIdStoreHook(Protocol): + async def store_unified_file_id( + self, + file_id: str, + file_object: OpenAIFileObject | None, + litellm_parent_otel_span: object | None, + model_mappings: dict[str, str], + user_api_key_dict: UserAPIKeyAuth, + ) -> None: ... + + +class _ManagedFileRow(BaseModel): + model_config = ConfigDict(from_attributes=True) + + unified_file_id: str + created_by: str | None = None + team_id: str | None = None + created_at: datetime | None = None + file_object: object = None + + +class _ManagedObjectRow(BaseModel): + model_config = ConfigDict(from_attributes=True) + + unified_object_id: str + created_by: str | None = None + team_id: str | None = None + created_at: datetime | None = None + file_object: object = None + + +class _CreatedAtRow(BaseModel): + model_config = ConfigDict(from_attributes=True) + + created_at: datetime | None = None + + +class _OwnershipRow(BaseModel): + model_config = ConfigDict(from_attributes=True) + + created_by: str | None = None + team_id: str | None = None + + +_FILE_ROW_ADAPTER: Final[TypeAdapter[_ManagedFileRow | None]] = TypeAdapter(_ManagedFileRow | None) +_FILE_ROWS_ADAPTER: Final[TypeAdapter[list[_ManagedFileRow]]] = TypeAdapter(list[_ManagedFileRow]) +_OBJECT_ROW_ADAPTER: Final[TypeAdapter[_ManagedObjectRow | None]] = TypeAdapter(_ManagedObjectRow | None) +_CREATED_AT_ADAPTER: Final[TypeAdapter[_CreatedAtRow | None]] = TypeAdapter(_CreatedAtRow | None) +_OWNERSHIP_ADAPTER: Final[TypeAdapter[_OwnershipRow | None]] = TypeAdapter(_OwnershipRow | None) + +_RowT: Final = TypeVar("_RowT", _ManagedFileRow, _ManagedObjectRow) + # --------------------------------------------------------------------------- # Field map # --------------------------------------------------------------------------- @@ -172,7 +251,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 @@ -197,7 +276,7 @@ class _RawIdGuardBudget: # --------------------------------------------------------------------------- # Maps (provider, canonical_path) -> "files" | "batches" -_LIST_ROUTE_TABLE: Final[dict[tuple[str, str], str]] = { +_LIST_ROUTE_TABLE: Final[dict[tuple[str, str], Literal["files", "batches"]]] = { ("openai", "/v1/files"): "files", ("openai", "/v1/batches"): "batches", ("azure", "/v1/files"): "files", @@ -263,8 +342,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. @@ -305,7 +384,7 @@ async def _resolve_one( # 2. DB lookup — pick table based on raw ID prefix if any(raw_id.startswith(p) for p in _FILE_PREFIXES): # File table — use hook's internal cache for speed when available - if managed_files_hook is not None: + if managed_files_hook is not None and isinstance(managed_files_hook, _FileIdLookupHook): try: file_row: Final = await managed_files_hook.get_unified_file_id( managed_id, @@ -322,8 +401,9 @@ async def _resolve_one( ) if not found and prisma_client is not None: try: - db_row: Final = await ManagedFileRepository(prisma_client).table.find_first( - where={"unified_file_id": managed_id} + db_row: Final = _OWNERSHIP_ADAPTER.validate_python( + await ManagedFileRepository(prisma_client).table.find_first(where={"unified_file_id": managed_id}), + from_attributes=True, ) if db_row is not None: row_created_by = db_row.created_by @@ -338,8 +418,11 @@ async def _resolve_one( # Object table (batches, responses) if prisma_client is not None: try: - obj_row: Final = await ManagedObjectRepository(prisma_client).table.find_first( - where={"unified_object_id": managed_id} + obj_row: Final = _OWNERSHIP_ADAPTER.validate_python( + await ManagedObjectRepository(prisma_client).table.find_first( + where={"unified_object_id": managed_id} + ), + from_attributes=True, ) if obj_row is not None: row_created_by = obj_row.created_by @@ -372,7 +455,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,14 +481,19 @@ 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: Final = await ManagedFileRepository(prisma_client).table.find_many( - where={"flat_model_file_ids": {"has": raw_id}}, + candidates: Final[tuple[_ManagedFileRow, ...]] = tuple( + _FILE_ROWS_ADAPTER.validate_python( + await ManagedFileRepository(prisma_client).table.find_many( + where={"flat_model_file_ids": {"has": raw_id}}, + ), + from_attributes=True, + ) ) except Exception: verbose_proxy_logger.debug("managed_id_rewriter: raw file-id guard lookup failed", exc_info=True) return provider_rows: Final = [ - row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider) + 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 +507,11 @@ 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: Final = await ManagedObjectRepository(prisma_client).table.find_first( - where={"model_object_id": f"passthrough:{provider}:{raw_id}"} + existing: Final = _OBJECT_ROW_ADAPTER.validate_python( + await ManagedObjectRepository(prisma_client).table.find_first( + where={"model_object_id": f"passthrough:{provider}:{raw_id}"} + ), + from_attributes=True, ) except Exception: verbose_proxy_logger.debug("managed_id_rewriter: raw object-id guard lookup failed", exc_info=True) @@ -434,7 +525,7 @@ async def _guard_raw_provider_id( # --------------------------------------------------------------------------- -def _build_managed_file_object(snapshot: dict[str, Any] | None, managed_id: str) -> OpenAIFileObject | None: +def _build_managed_file_object(snapshot: dict[str, object] | None, managed_id: str) -> OpenAIFileObject | None: """Build an ``OpenAIFileObject`` (with the managed ID swapped in) from an upstream file response so the DB-served list returns the same metadata as a direct file GET. Returns ``None`` when no usable snapshot is available, in @@ -442,7 +533,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 +546,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: dict[str, object] | None = None, is_create_route: bool = True, ) -> str: """Return an existing managed file ID or mint + store a new one.""" @@ -479,15 +570,22 @@ 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: tuple[_ManagedFileRow, ...] = ( + tuple( # rebind-ok: except branch below assigns the same name once + _FILE_ROWS_ADAPTER.validate_python( + await ManagedFileRepository(prisma_client).table.find_many( + where={"flat_model_file_ids": {"has": raw_id}}, + order={"created_at": "asc"}, + ), + from_attributes=True, + ) + ) ) except Exception: - candidates = [] + candidates = () # rebind-ok: try branch above assigns the same name once verbose_proxy_logger.debug("managed_id_rewriter: file dedup lookup failed", exc_info=True) provider_rows: Final = [ - row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider) + row for row in candidates if _managed_id_matches_provider(row.unified_file_id, provider) ] owned_row: Final = next( (row for row in provider_rows if can_access_resource(user_api_key_dict, row.created_by, row.team_id)), @@ -523,7 +621,7 @@ async def _mint_or_reuse_file( "managed_id_rewriter: minted new managed file id for raw prefix=%s", raw_id.split("-", 1)[0], ) - if managed_files_hook is not None: + if managed_files_hook is not None and isinstance(managed_files_hook, _FileIdStoreHook): try: await managed_files_hook.store_unified_file_id( file_id=managed_id, @@ -551,9 +649,9 @@ async def _mint_or_reuse_object( raw_id: str, provider: str, file_purpose: str, - body_snapshot: dict, + body_snapshot: dict[str, object], user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: PrismaClient | None, is_create_route: bool, ) -> str: """Return an existing managed object ID (batch/response) or mint + store one.""" @@ -569,7 +667,7 @@ async def _mint_or_reuse_object( # f"{purpose}:{provider}:{raw_id}" for the same reason. namespaced_model_object_id: Final = f"passthrough:{provider}:{raw_id}" - async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str: + async def _reuse_existing(existing: _ManagedObjectRow, 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): @@ -618,12 +716,15 @@ 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 = _OBJECT_ROW_ADAPTER.validate_python( # rebind-ok: except branch below assigns the same name once + await ManagedObjectRepository(prisma_client).table.find_first( + where={"model_object_id": namespaced_model_object_id} + ), + from_attributes=True, ) except Exception: verbose_proxy_logger.debug("managed_id_rewriter: object dedup lookup failed", exc_info=True) - existing = None + existing = None # rebind-ok: try branch above assigns the same name once if existing is not None: return await _reuse_existing(existing, refresh_snapshot=True) @@ -659,11 +760,14 @@ 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 = _OBJECT_ROW_ADAPTER.validate_python( # rebind-ok: except branch below assigns the same name once + await ManagedObjectRepository(prisma_client).table.find_first( + where={"model_object_id": namespaced_model_object_id} + ), + from_attributes=True, ) except Exception: - raced = None + raced = None # rebind-ok: try branch above assigns the same name once if raced is not None: return await _reuse_existing(raced, refresh_snapshot=False) # No row backs the minted ID, so every later resolve would 404. Fall @@ -681,11 +785,11 @@ async def rewrite_response_ids( provider: str, method: str, route: str, - body: dict, + body: dict[str, object], user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, -) -> dict: + prisma_client: PrismaClient | None, + managed_files_hook: CustomLogger | None, +) -> dict[str, object]: """ Mint managed IDs for raw provider values listed in ``BUILTIN_OUTPUT_ID_FIELD_MAP`` and swap them into *body*. @@ -714,7 +818,7 @@ async def rewrite_response_ids( # creates may degrade to a raw id on a cross-owner collision. is_create_route: Final = "{" not in canonical - mutated: Final = dict(body) # shallow copy; only return if something changed + mutated: Final[dict[str, object]] = dict(body) # shallow copy; only return if something changed changed = False def _record(field_name: str, raw_value: str, managed_id: str) -> None: @@ -795,7 +899,7 @@ def is_passthrough_list_route(provider: str, method: str, route: str) -> bool: return (provider, canonical) in _LIST_ROUTE_TABLE -def _parse_file_object(file_object: Any) -> Any: +def _parse_file_object(file_object: object) -> object: """Prisma may return ``Json`` columns as either a parsed dict or the raw JSON string (depending on driver / row source). Mirror the handling used elsewhere (see ``openai_files_endpoints/common_utils.py``) so callers can @@ -803,13 +907,13 @@ def _parse_file_object(file_object: Any) -> Any: """ if isinstance(file_object, str): try: - return json.loads(file_object) + return _OBJECT_ADAPTER.validate_python(json.loads(file_object)) except (TypeError, ValueError): return None return file_object -def _empty_list_response() -> dict[str, Any]: +def _empty_list_response() -> dict[str, object]: return { "object": "list", "data": [], @@ -819,28 +923,28 @@ def _empty_list_response() -> dict[str, Any]: } -def _parse_list_limit(query_params: dict[str, Any] | None) -> tuple[int, int]: +def _parse_list_limit(query_params: dict[str, str] | None) -> tuple[int, int]: params: Final = query_params or {} try: - raw_limit = int(params.get("limit", 20)) + raw_limit = int(params.get("limit", "20")) # rebind-ok: except branch below assigns the same name once except (TypeError, ValueError): - raw_limit = 20 + raw_limit = 20 # rebind-ok: try branch above assigns the same name once # Fetch one extra to cheaply detect has_more. return raw_limit, min(raw_limit, 100) + 1 async def _build_list_where_with_cursor( - prisma_client: Any, - resource_kind: str, + prisma_client: PrismaClient, + resource_kind: Literal["files", "batches"], provider: str, - owner_filter: dict[str, Any], - query_params: dict[str, Any] | None, -) -> tuple[dict[str, Any], str]: + owner_filter: PrismaWhere, + query_params: dict[str, str] | None, +) -> tuple[PrismaWhere, str]: """Return a Prisma ``where`` clause and fetch order for a list query.""" params: Final = query_params or {} - after_id: Final[str | None] = params.get("after") - before_id: Final[str | None] = params.get("before") - where: dict[str, Any] = dict(owner_filter) + after_id: Final = params.get("after") + before_id: Final = params.get("before") + where: PrismaWhere = dict(owner_filter) # rebind-ok: narrowed with a cursor boundary further below fetch_order = "desc" cursor_id: Final = after_id or before_id @@ -850,14 +954,22 @@ 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: Final = ( - ManagedFileRepository(prisma_client).table - if resource_kind == "files" - else ManagedObjectRepository(prisma_client).table - ) cursor_field: Final = "unified_file_id" if resource_kind == "files" else "unified_object_id" try: - cursor_row: Final = await cursor_table.find_first(where={**owner_filter, cursor_field: cursor_id}) + if resource_kind == "files": + cursor_row = _CREATED_AT_ADAPTER.validate_python( # rebind-ok: else branch below assigns the same name once + await ManagedFileRepository(prisma_client).table.find_first( + where={**owner_filter, cursor_field: cursor_id} + ), + from_attributes=True, + ) + else: + cursor_row = _CREATED_AT_ADAPTER.validate_python( # rebind-ok: if branch above assigns the same name once + await ManagedObjectRepository(prisma_client).table.find_first( + where={**owner_filter, cursor_field: cursor_id} + ), + from_attributes=True, + ) if cursor_row is not None: if after_id: op = "lt" @@ -867,7 +979,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: Final = { + boundary: Final[PrismaWhere] = { "OR": [ {"created_at": {op: cursor_row.created_at}}, { @@ -884,75 +996,75 @@ 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_page( + prisma_client: PrismaClient, + provider: str, + where: PrismaWhere, 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. +) -> tuple[_ManagedFileRow, ...] | None: + scoped_where: Final[PrismaWhere] = { + **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, - order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}], - take=fetch_limit, + return tuple( + _FILE_ROWS_ADAPTER.validate_python( + await ManagedFileRepository(prisma_client).table.find_many( + where=scoped_where, + order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}], + take=fetch_limit, + ), + from_attributes=True, ) - 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_object_list_page( + prisma_client: PrismaClient, provider: str, - where: dict[str, Any], + where: PrismaWhere, 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. +) -> tuple[_ManagedObjectRow, ...] | None: + scoped_where: Final[PrismaWhere] = { + **where, + "model_object_id": {"startswith": f"passthrough:{provider}:"}, + "file_purpose": "batch", + } + try: + return tuple( + TypeAdapter(list[_ManagedObjectRow]).validate_python( + await ManagedObjectRepository(prisma_client).table.find_many( + where=scoped_where, + order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}], + take=fetch_limit, + ), + from_attributes=True, + ) + ) + except Exception: # noqa: BLE001 # DB/validation failures degrade to an empty page, matching every other lookup in this module + verbose_proxy_logger.warning("managed_id_rewriter: list DB query failed", exc_info=True) + return None - 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. - """ - scoped_where: Final = 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}:"} - - rows: Final = await _fetch_list_rows(prisma_client, resource_kind, scoped_where, fetch_order, fetch_limit) +def _paginate(rows: tuple[_RowT, ...] | None, raw_limit: int, fetch_order: str) -> tuple[tuple[_RowT, ...], bool]: + """Trim a fetched (over-fetched-by-one) page to the caller's limit and + detect has_more; ``asc``-ordered pages (built for a ``before`` cursor) are + flipped back to the client-facing newest-first order.""" if rows is None: - return [], False - + return (), False effective_limit: Final = min(raw_limit, 100) has_more: Final = len(rows) > effective_limit - page = rows[:effective_limit] - if fetch_order == "asc": - page = list(reversed(page)) - return page, has_more + page: Final = rows[:effective_limit] + return (tuple(reversed(page)) if fetch_order == "asc" else page), has_more -def _serialize_file_list_item(row: Any) -> dict[str, Any]: - item: Final[dict[str, Any]] = { +def _serialize_file_list_item(row: _ManagedFileRow) -> dict[str, object]: + item: Final[dict[str, object]] = { "id": row.unified_file_id, "object": "file", "created_at": int(row.created_at.timestamp()) if row.created_at else None, @@ -964,8 +1076,8 @@ def _serialize_file_list_item(row: Any) -> dict[str, Any]: return item -def _serialize_batch_list_item(row: Any) -> dict[str, Any]: - item: Final[dict[str, Any]] = {} +def _serialize_batch_list_item(row: _ManagedObjectRow) -> dict[str, object]: + item: Final[dict[str, object]] = {} file_object: Final = _parse_file_object(row.file_object) if isinstance(file_object, dict): item.update(file_object) @@ -974,20 +1086,13 @@ def _serialize_batch_list_item(row: Any) -> dict[str, Any]: 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: Final = "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: dict[str, str] | None = None, +) -> dict[str, object] | None: """Query the DB for managed IDs the caller owns and return an OpenAI-style paginated list response. @@ -1018,23 +1123,28 @@ async def list_passthrough_ids_from_db( raw_limit, fetch_limit = _parse_list_limit(query_params) where, fetch_order = await _build_list_where_with_cursor( - prisma_client, resource_kind, provider, owner_filter, query_params + prisma_client, resource_kind, provider, dict(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] - else: - data = [_serialize_batch_list_item(row) for row in page] - first_id, last_id = _list_boundary_ids(page, resource_kind) + if resource_kind == "files": + file_rows: Final = await _fetch_file_list_page(prisma_client, provider, where, fetch_order, fetch_limit) + file_page, has_more = _paginate(file_rows, raw_limit, fetch_order) + data = tuple( + _serialize_file_list_item(row) for row in file_page + ) # rebind-ok: else branch below assigns the same name once + first_id, last_id = (file_page[0].unified_file_id, file_page[-1].unified_file_id) if file_page else (None, None) + else: + object_rows: Final = await _fetch_object_list_page(prisma_client, provider, where, fetch_order, fetch_limit) + object_page, has_more = _paginate( + object_rows, raw_limit, fetch_order + ) # rebind-ok: has_more re-binds the if-branch's binding above + data = tuple( + _serialize_batch_list_item(row) for row in object_page + ) # rebind-ok: if branch above assigns the same name once + first_id, last_id = ( # rebind-ok: re-binds the if-branch's unpacked names above + (object_page[0].unified_object_id, object_page[-1].unified_object_id) if object_page else (None, None) + ) + verbose_proxy_logger.debug( "managed_id_rewriter: list served from DB provider=%s kind=%s count=%d admin=%s", provider, @@ -1044,7 +1154,7 @@ async def list_passthrough_ids_from_db( ) return { "object": "list", - "data": data, + "data": list(data), "first_id": first_id, "last_id": last_id, "has_more": has_more, @@ -1060,8 +1170,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 +1202,12 @@ async def rewrite_path_ids( async def rewrite_query_ids( - params: dict[str, Any] | None, + params: dict[str, str] | 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, str] | None: """ Walk query param values and resolve any passthrough managed IDs. Returns *params* unchanged (same object) when nothing is resolved. @@ -1108,12 +1218,11 @@ async def rewrite_query_ids( mutated: Final = dict(params) rewritten_keys: Final[list[str]] = [] for key, val in list(mutated.items()): - if isinstance(val, str): - if is_managed(val): - mutated[key] = await _resolve_one(val, provider, user_api_key_dict, prisma_client, managed_files_hook) - rewritten_keys.append(key) - else: - await _guard_raw_provider_id(val, provider, user_api_key_dict, prisma_client, budget) + if is_managed(val): + mutated[key] = await _resolve_one(val, provider, user_api_key_dict, prisma_client, managed_files_hook) + rewritten_keys.append(key) + else: + await _guard_raw_provider_id(val, provider, user_api_key_dict, prisma_client, budget) if rewritten_keys: verbose_proxy_logger.debug( "managed_id_rewriter: query ids rewritten provider=%s keys=%s", @@ -1124,12 +1233,12 @@ async def rewrite_query_ids( async def rewrite_body_ids( - body: dict[str, Any] | None, + body: dict[str, object] | None, provider: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, - managed_files_hook: Any, -) -> dict[str, Any] | None: + prisma_client: PrismaClient | None, + managed_files_hook: CustomLogger | None, +) -> dict[str, object] | None: """ Recursively walk a request body dict/list and resolve any passthrough managed IDs. Skips litellm internal keys (``litellm_*``). @@ -1140,11 +1249,11 @@ async def rewrite_body_ids( budget: Final = _RawIdGuardBudget() - async def _walk(node: Any, depth: int) -> Any: + async def _walk(node: object, depth: int) -> object: if depth >= _MAX_BODY_REWRITE_DEPTH: return node if isinstance(node, dict): - result: Final[dict[str, Any]] = {} + result: Final[dict[str, object]] = {} changed_inner = False for k, v in node.items(): # Skip litellm internal injection keys (e.g. litellm_logging_obj) @@ -1169,6 +1278,7 @@ async def rewrite_body_ids( return node rewritten: Final = 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/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f54023836e5..c7635b8dc97 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,8 +4,8 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence -from typing import Any, Final, Literal, cast +from collections.abc import Iterable, Mapping, Sequence +from typing import Final, Literal, TypeAlias, cast from openai.types.chat.chat_completion_named_tool_choice_param import ( ChatCompletionNamedToolChoiceParam, @@ -33,6 +33,7 @@ from litellm.types.llms.openai import ( ChatCompletionImageUrlObject, ChatCompletionResponseMessage, ChatCompletionSystemMessage, + ChatCompletionTextObject, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolMessage, @@ -82,6 +83,12 @@ from .custom_tools import ( TOOL_CALLS_CACHE: Final = InMemoryCache() +ChatCompletionOutputMessage: TypeAlias = ( + AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage +) +ChatCompletionSessionMessage: TypeAlias = ChatCompletionOutputMessage | Message + + class ChatCompletionSession(TypedDict, total=False): messages: list[ AllMessageValues @@ -122,8 +129,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. @@ -147,37 +154,39 @@ class LiteLLMCompletionResponsesConfig: if isinstance(tool_choice, str): return tool_choice - if isinstance(tool_choice, dict): - tool_choice_type: Final = tool_choice.get("type") + tool_choice_type: Final = 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"): - return tool_choice + # If it has a function with name, it's standard OpenAI format - pass through + function_value: Final = tool_choice.get("function") + if isinstance(function_value, Mapping) and function_value.get("name"): + return tool_choice - # Handle Cursor IDE dict formats without function name - if tool_choice_type == "auto": - return "auto" - elif tool_choice_type == "none": - return "none" - elif tool_choice_type in ["required", "tool", "any"]: - # "tool" without a specific function name means "use any tool" - # which is equivalent to "required" in OpenAI format - return "required" - elif tool_choice_type == "function": - function_name: Final = tool_choice.get("name") - if function_name: - return ChatCompletionNamedToolChoiceParam( - type="function", function=NamedToolChoiceFunction(name=function_name) - ) - return "required" - elif tool_choice_type == "custom": - custom: Final = tool_choice.get("custom") - custom_name = tool_choice.get("name") or (custom.get("name") if isinstance(custom, dict) else None) - if custom_name: - return ChatCompletionNamedToolChoiceParam( - type="function", function=NamedToolChoiceFunction(name=custom_name) - ) - return "required" + # Handle Cursor IDE dict formats without function name + if tool_choice_type == "auto": + return "auto" + elif tool_choice_type == "none": + return "none" + elif tool_choice_type in ["required", "tool", "any"]: + # "tool" without a specific function name means "use any tool" + # which is equivalent to "required" in OpenAI format + return "required" + elif tool_choice_type == "function": + function_name: Final = tool_choice.get("name") + if isinstance(function_name, str) and function_name: + return ChatCompletionNamedToolChoiceParam( + type="function", function=NamedToolChoiceFunction(name=function_name) + ) + return "required" + elif tool_choice_type == "custom": + custom: Final = tool_choice.get("custom") + custom_name: Final = tool_choice.get("name") or ( + custom.get("name") if isinstance(custom, Mapping) else None + ) + if isinstance(custom_name, str) and custom_name: + return ChatCompletionNamedToolChoiceParam( + type="function", function=NamedToolChoiceFunction(name=custom_name) + ) + return "required" # Return as-is for unknown formats return tool_choice @@ -205,7 +214,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None = None, stream: bool | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, **kwargs, ) -> dict: """ @@ -462,7 +471,7 @@ class LiteLLMCompletionResponsesConfig: if not chat_completion_messages: continue - deduped_in_place: list[Any] = [] + deduped_in_place: list[ChatCompletionOutputMessage] = [] for m in chat_completion_messages: role = "" if isinstance(m, dict): @@ -472,7 +481,7 @@ class LiteLLMCompletionResponsesConfig: # Drop assistant tool_calls wrappers if we already have this call_id if role == "assistant": - tool_calls: Any = ( + tool_calls: object = ( m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None) ) call_id = "" @@ -534,7 +543,7 @@ class LiteLLMCompletionResponsesConfig: call_id = "" if role == "assistant": - tool_calls: Any = None + tool_calls: object = None if isinstance(tool_call_message, dict): tool_calls = tool_call_message.get("tool_calls") else: @@ -578,7 +587,7 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None: + def _find_previous_assistant_idx(messages: list[ChatCompletionSessionMessage], 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": @@ -586,9 +595,9 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str: + def _recover_tool_call_id_from_assistant(assistant_message: ChatCompletionSessionMessage) -> str: """Try to recover empty tool_call_id from assistant message's tool_calls.""" - tool_calls_raw: Final = ( + tool_calls_raw: Final[object] = ( assistant_message.get("tool_calls") if isinstance(assistant_message, dict) else getattr(assistant_message, "tool_calls", None) @@ -604,9 +613,9 @@ class LiteLLMCompletionResponsesConfig: return "" @staticmethod - def _get_tool_calls_list(assistant_message: Any) -> list[Any]: + def _get_tool_calls_list(assistant_message: ChatCompletionSessionMessage) -> Sequence[object]: """Extract tool_calls as a list from assistant message.""" - tool_calls_raw: Final = ( + tool_calls_raw: Final[object] = ( assistant_message.get("tool_calls") if isinstance(assistant_message, dict) else getattr(assistant_message, "tool_calls", None) @@ -615,12 +624,12 @@ class LiteLLMCompletionResponsesConfig: 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)): + if isinstance(tool_calls_raw, Iterable) and not isinstance(tool_calls_raw, (str, bytes)): return list(tool_calls_raw) return [] @staticmethod - def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool: + def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool: """Check if a tool_call with the given ID exists in the list.""" for tool_call in tool_calls: tool_call_id_to_check: str | None = None @@ -633,7 +642,7 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None: + def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[object]) -> dict[str, object] | None: """Reconstruct a minimal tool_call definition from tools list.""" for tool in tools: if isinstance(tool, dict): @@ -651,7 +660,7 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: + def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object: """ Safely read a field from dict-like or attribute-based objects. """ @@ -672,13 +681,13 @@ class LiteLLMCompletionResponsesConfig: @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: Final = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments") - function: Final[dict[str, Any]] = { + function: Final[dict[str, object]] = { "name": function_name_raw or "", "arguments": function_arguments_raw or "{}", } @@ -697,7 +706,7 @@ class LiteLLMCompletionResponsesConfig: ) @staticmethod - def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None: + def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[str, object] | None: """ Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. """ @@ -705,7 +714,9 @@ class LiteLLMCompletionResponsesConfig: return None if isinstance(tool_use_definition, dict): - normalized_definition: dict[str, Any] = dict(tool_use_definition) + normalized_definition: dict[str, object] = dict( + tool_use_definition + ) # rebind-ok: else branch below assigns the same name once 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") @@ -738,20 +749,30 @@ 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: Final = cast(dict[str, Any], assistant_message) - if "tool_calls" not in prev_assistant_dict: - prev_assistant_dict["tool_calls"] = [] - tool_calls_list: Final = prev_assistant_dict["tool_calls"] + if "tool_calls" not in assistant_message: + assistant_message["tool_calls"] = [] + tool_calls_list: Final = assistant_message["tool_calls"] if isinstance(tool_calls_list, list): 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) + current_tool_calls = getattr( + assistant_message, "tool_calls", None + ) # rebind-ok: conditionally overridden below + if current_tool_calls is None: + current_tool_calls = [] # rebind-ok: conditionally overrides the default set above + setattr(assistant_message, "tool_calls", current_tool_calls) # noqa: B010 # assistant_message is typed `object` (duck-typed SDK message) + if isinstance(current_tool_calls, list): + current_tool_calls.append(tool_call_chunk) + + @staticmethod + def _set_tool_call_id(message: object, tool_call_id: str) -> None: + if isinstance(message, dict): + message["tool_call_id"] = tool_call_id + elif hasattr(message, "tool_call_id"): + setattr(message, "tool_call_id", tool_call_id) # noqa: B010 # message is typed `object` (duck-typed SDK message) @staticmethod def _ensure_tool_results_have_corresponding_tool_calls( @@ -762,7 +783,7 @@ class LiteLLMCompletionResponsesConfig: | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: list[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -819,17 +840,9 @@ class LiteLLMCompletionResponsesConfig: # Try to recover empty tool_call_id from previous assistant message if not tool_call_id and prev_assistant_idx is not None: prev_assistant = fixed_messages[prev_assistant_idx] - tool_call_id = LiteLLMCompletionResponsesConfig._recover_tool_call_id_from_assistant( - prev_assistant, message - ) + tool_call_id = LiteLLMCompletionResponsesConfig._recover_tool_call_id_from_assistant(prev_assistant) 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["tool_call_id"] = tool_call_id - elif hasattr(message, "tool_call_id"): - setattr(message, "tool_call_id", tool_call_id) + LiteLLMCompletionResponsesConfig._set_tool_call_id(message, tool_call_id) # Only remove messages with empty tool_call_id if we have other non-tool messages # This prevents ending up with an empty messages list when using previous_response_id @@ -878,7 +891,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 @@ -915,7 +928,7 @@ class LiteLLMCompletionResponsesConfig: return [] return [ GenericChatCompletionMessage( - role=input_item.get("role") or "user", + role=str(input_item.get("role") or "user"), content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( content ), @@ -923,7 +936,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 """ @@ -936,7 +949,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 @@ -946,7 +959,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 @@ -958,8 +971,8 @@ class LiteLLMCompletionResponsesConfig: return [] def _normalize_function_call_output_to_tool_content( - output: Any, - ) -> Any: + output: object, + ) -> str | list[ChatCompletionTextObject | ChatCompletionImageObject]: """ Normalize Responses API function_call_output.output into a shape that downstream chat adapters (esp. Gemini) can reliably consume. @@ -981,7 +994,7 @@ class LiteLLMCompletionResponsesConfig: # Some adapters represent tool output as a list of "input_*" parts if isinstance(output, list): - normalized_blocks: Final[list[dict[str, Any]]] = [] + normalized_blocks: Final[list[ChatCompletionTextObject | ChatCompletionImageObject]] = [] text_acc: Final[list[str]] = [] for part in output: if not isinstance(part, dict): @@ -991,19 +1004,23 @@ class LiteLLMCompletionResponsesConfig: txt = part.get("text") if isinstance(txt, str) and txt: text_acc.append(txt) - normalized_blocks.append({"type": "text", "text": txt}) + normalized_blocks.append(ChatCompletionTextObject(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") if isinstance(url, str) and url: - normalized_blocks.append({"type": "image_url", "image_url": {"url": url}}) + normalized_blocks.append( + ChatCompletionImageObject( + type="image_url", image_url=ChatCompletionImageUrlObject(url=url) + ) + ) elif isinstance(image_url_val, str) and image_url_val: normalized_blocks.append( - { - "type": "image_url", - "image_url": {"url": image_url_val}, - } + ChatCompletionImageObject( + type="image_url", + image_url=ChatCompletionImageUrlObject(url=image_url_val), + ) ) # Prefer structured blocks if we have images; otherwise return a string. @@ -1082,7 +1099,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 @@ -1108,10 +1125,10 @@ class LiteLLMCompletionResponsesConfig: raw_input: Final = function_call.get("input") or "" raw_arguments = json.dumps({"content": raw_input}) if raw_input else "" tool_call: Final = ChatCompletionToolCallChunk( - id=function_call.get("call_id") or function_call.get("id") or "", + id=str(function_call.get("call_id") or function_call.get("id") or ""), type="function", function=ChatCompletionToolCallFunctionChunk( - name=function_call.get("name") or "", + name=str(function_call.get("name") or ""), arguments=str(raw_arguments or ""), ), index=0, @@ -1127,16 +1144,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 + file_id: Final = item.get("file_id") or item.get("file_url") + return str(file_id) if file_id 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 @@ -1146,35 +1164,35 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary with transformed file structure for Chat Completion """ - file_dict: Final[dict[str, Any]] = {} + file_dict: Final[dict[str, object]] = {} file_id: Final = 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: Final[dict[str, Any]] = {"type": "file", "file": file_dict} + new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: new_item["cache_control"] = item["cache_control"] return new_item @staticmethod def _transform_input_image_item_to_image_item( - item: dict[str, Any], + item: Mapping[str, object], ) -> ChatCompletionImageObject: """ Transform a Responses API input_image item to a Chat Completion image item """ image_url_obj: Final = ChatCompletionImageUrlObject( - url=item.get("image_url") or "", detail=item.get("detail") or "auto" + url=str(item.get("image_url") or ""), detail=str(item.get("detail") or "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 @@ -1188,7 +1206,7 @@ class LiteLLMCompletionResponsesConfig: elif isinstance(content, str): return content elif isinstance(content, list): - content_list: Final[list[str | dict[str, Any]]] = [] + content_list: Final[list[str | dict[str, object]]] = [] for item in content: if isinstance(item, str): content_list.append(item) @@ -1198,7 +1216,7 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) ) elif item.get("type") == "input_image": - image_block = dict( + image_block: dict[str, object] = dict( LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) ) if "cache_control" in item: @@ -1209,7 +1227,7 @@ class LiteLLMCompletionResponsesConfig: text_value = item.get("text") if text_value is None: continue - content_block: dict[str, Any] = { + content_block: dict[str, object] = { "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( item.get("type") or "text" ), @@ -1299,7 +1317,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 "", @@ -1340,7 +1358,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 @@ -1348,17 +1366,19 @@ class LiteLLMCompletionResponsesConfig: """ if chat_completion_tools is None or not chat_completion_tools: return [] - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: if not isinstance(tool, dict): result.append(tool) continue if tool.get("type") == "function": - fn = cast(dict[str, Any], tool.get("function") or {}) - parameters = dict(fn.get("parameters", {}) or {}) + fn_value = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool, "function", {}) + fn: dict[str, object] = fn_value if isinstance(fn_value, dict) else {} + parameters_value = fn.get("parameters") or {} + parameters: dict[str, object] = dict(parameters_value) if isinstance(parameters_value, dict) 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 "", @@ -1431,20 +1451,22 @@ class LiteLLMCompletionResponsesConfig: 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) 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) if hasattr(provider_specific_fields, "__dict__") else {} + tool_provider_fields = getattr(tool, "provider_specific_fields", None) + function_provider_fields = getattr(function_definition, "provider_specific_fields", None) + if tool_provider_fields: + provider_specific_fields = ( + tool_provider_fields + if isinstance(tool_provider_fields, dict) + else (dict(tool_provider_fields) if hasattr(tool_provider_fields, "__dict__") else {}) + ) + elif function_provider_fields: + provider_specific_fields = ( + function_provider_fields + if isinstance(function_provider_fields, dict) + else ( + dict(function_provider_fields) if hasattr(function_provider_fields, "__dict__") else {} ) + ) output_tool_call: ResponseFunctionToolCall = ResponseFunctionToolCall( name=tool_name, @@ -1508,9 +1530,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. @@ -1527,8 +1549,9 @@ class LiteLLMCompletionResponsesConfig: 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): - provider_fields: Final = tool_call_item.get("provider_specific_fields") + else: + getter: Final = getattr(tool_call_item, "get", None) + provider_fields: Final = getter("provider_specific_fields") if callable(getter) else None if provider_fields: provider_specific_fields = ( provider_fields @@ -1536,15 +1559,15 @@ class LiteLLMCompletionResponsesConfig: else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {}) ) - function_dict: Final[dict[str, Any]] = { - "name": tool_call_item.name, - "arguments": tool_call_item.arguments, + function_dict: Final[dict[str, object]] = { + "name": getattr(tool_call_item, "name"), # noqa: B009 # tool_call_item is typed `object` (duck-typed SDK tool call) + "arguments": getattr(tool_call_item, "arguments"), # noqa: B009 # tool_call_item is typed `object` (duck-typed SDK tool call) } if provider_specific_fields: function_dict["provider_specific_fields"] = provider_specific_fields - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item( getattr(tool_call_item, "id", None), getattr(tool_call_item, "call_id", None), @@ -1561,9 +1584,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_apply_patch_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 ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. @@ -1580,9 +1603,9 @@ class LiteLLMCompletionResponsesConfig: """ import json - operation_dict: Final = tool_call_item.operation.model_dump() - tool_call_dict: Final[dict[str, Any]] = { - "id": tool_call_item.call_id, + operation_dict: Final = getattr(tool_call_item, "operation").model_dump() # noqa: B009 # tool_call_item is typed `object` (duck-typed SDK tool call) + tool_call_dict: Final[dict[str, object]] = { + "id": getattr(tool_call_item, "call_id"), # noqa: B009 # tool_call_item is typed `object` (duck-typed SDK tool call) "function": { "name": "apply_patch", "arguments": json.dumps(operation_dict), @@ -1719,11 +1742,8 @@ class LiteLLMCompletionResponsesConfig: """ output_items: Final[list] = [] for choice in chat_completion_response.choices or []: - message = getattr(choice, "message", None) - if not message: - continue - psf = getattr(message, "provider_specific_fields", None) - if not psf or not isinstance(psf, dict): + psf = choice.message.provider_specific_fields + if not psf: continue results = psf.get("code_interpreter_results") if results and isinstance(results, list): @@ -1791,13 +1811,13 @@ class LiteLLMCompletionResponsesConfig: """ image_generation_items: Final[list[OutputImageGenerationCall]] = [] - images: Final = getattr(choice.message, "images", []) + images: Final = 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 = image_item["image_url"].get("url", "") base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: @@ -2048,8 +2068,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. diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 94cfe0fcb1f..8951237099d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -9,10 +9,11 @@ from collections.abc import Mapping from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol import httpx from openai._streaming import SSEDecoder +from pydantic import BaseModel import litellm from litellm.constants import ( @@ -34,6 +35,22 @@ from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook +if TYPE_CHECKING: + from fastapi import WebSocket + from websockets.asyncio.client import ClientConnection + + from litellm.types.llms.openai import ( + ContentPartDoneEvent, + ContentPartDonePartOutputText, + ContentPartDonePartReasoningText, + ContentPartDonePartRefusal, + ResponseAPIUsage, + ResponseCreatedEvent, + ResponseInProgressEvent, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, + ) + @lru_cache(maxsize=1) def _get_openai_response_types(): @@ -42,7 +59,7 @@ def _get_openai_response_types(): return openai_types -def _log_background_task_failure(task: asyncio.Task[Any], *, task_name: str) -> None: +def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str) -> None: if task.cancelled(): return exception: Final = task.exception() @@ -131,7 +148,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 @@ -176,7 +193,7 @@ class BaseResponsesAPIStreamingIterator: llm_provider=self.custom_llm_provider or "", ) - def _process_chunk(self, chunk) -> Any | None: + def _process_chunk(self, chunk) -> ResponsesAPIStreamingResponse | None: """Process a single chunk of data from the stream""" if not chunk: return None @@ -222,9 +239,9 @@ class BaseResponsesAPIStreamingIterator: setattr(openai_responses_api_chunk, "response", response) # Encode container_id on streaming events so proxy/UI follow-ups route correctly - _event_type: Final = getattr(openai_responses_api_chunk, "type", None) + _event_type: Final[object] = getattr(openai_responses_api_chunk, "type", None) if _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: - _delta: Final = getattr(openai_responses_api_chunk, "delta", None) + _delta: Final[object] = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_delta, str): self._generated_content += _delta _stream_model_id: Final = ( @@ -234,7 +251,7 @@ class BaseResponsesAPIStreamingIterator: ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - _item: Final = getattr(openai_responses_api_chunk, "item", None) + _item: Final[object] = getattr(openai_responses_api_chunk, "item", None) if _item is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_item, @@ -242,7 +259,7 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED: - _annotation: Final = getattr(openai_responses_api_chunk, "annotation", None) + _annotation: Final[object] = getattr(openai_responses_api_chunk, "annotation", None) if _annotation is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_annotation, @@ -250,7 +267,7 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: - _part: Final = getattr(openai_responses_api_chunk, "part", None) + _part: Final[object] = getattr(openai_responses_api_chunk, "part", None) if _part is not None: if isinstance(_part, dict): ResponsesAPIRequestUtils._encode_container_ids_in_annotations( @@ -268,14 +285,14 @@ class BaseResponsesAPIStreamingIterator: # 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: Final = getattr(openai_responses_api_chunk, "type", None) + event_type: Final[object] = getattr(openai_responses_api_chunk, "type", None) if event_type in ( openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - item: Final = getattr(openai_responses_api_chunk, "item", None) + item: Final[object] = getattr(openai_responses_api_chunk, "item", None) if item: - encrypted_content: Final = getattr(item, "encrypted_content", None) + encrypted_content: Final[object] = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): model_id: Final = ( self.litellm_metadata.get("model_info", {}).get("id") @@ -289,7 +306,7 @@ class BaseResponsesAPIStreamingIterator: setattr(item, "encrypted_content", wrapped_content) # Store the completed response (also for incomplete/failed so logging still fires) - _chunk_type: Final = getattr(openai_responses_api_chunk, "type", None) + _chunk_type: Final[object] = 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, @@ -299,9 +316,11 @@ class BaseResponsesAPIStreamingIterator: self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Final[Any | None] = getattr(openai_responses_api_chunk, "response", None) + response_obj: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_obj: - usage_obj: Final[Any | None] = getattr(response_obj, "usage", None) + usage_obj: Final[object | None] = getattr(response_obj, "usage", None) if usage_obj is not None: try: cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) @@ -389,8 +408,10 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None - error_info: Final = getattr(response_obj, "error", None) if response_obj else None + response_obj: Final[ResponsesAPIResponse | None] = ( + getattr(self.completed_response, "response", None) if self.completed_response else None + ) + error_info: Final[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: Final = litellm.APIError( @@ -401,10 +422,10 @@ class BaseResponsesAPIStreamingIterator: ) self._handle_failure(exception) - def _record_failed_response_usage(self, response_obj: Any | None) -> None: + def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is None: return try: @@ -451,7 +472,7 @@ class BaseResponsesAPIStreamingIterator: is_pre_first_chunk=not self._yielded_first_chunk, ) - def _get_completed_response_object(self) -> Any | None: + def _get_completed_response_object(self) -> ResponsesAPIResponse | None: openai_types: Final = _get_openai_response_types() completed_response: Final = self.completed_response if isinstance(completed_response, openai_types.ResponsesAPIResponse): @@ -880,7 +901,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _set_events_from_response( self, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events = _build_synthetic_response_events( @@ -894,7 +915,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __aiter__(self): return self - async def __anext__(self) -> Any: + async def __anext__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopAsyncIteration evt: Final = self._events[self._idx] @@ -908,7 +929,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __iter__(self): return self - def __next__(self) -> Any: + def __next__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopIteration evt: Final = self._events[self._idx] @@ -923,7 +944,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __init__( self, - response: Any, + response: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, request_data: dict[str, Any] | None = None, call_type: str | None = None, @@ -941,13 +962,13 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): ) self._completed_response_cache_hit = True self._persist_completed_response_before_logging = False - self._events: list[Any] = [] + self._events: list[ResponsesAPIStreamingResponse] = [] self._idx = 0 self._set_events_from_response(transformed=response, logging_obj=logging_obj) def _set_events_from_response( self, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: self._events = _build_synthetic_response_events( @@ -961,7 +982,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __aiter__(self): return self - async def __anext__(self) -> Any: + async def __anext__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopAsyncIteration evt: Final = self._events[self._idx] @@ -975,7 +996,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __iter__(self): return self - def __next__(self) -> Any: + def __next__(self) -> ResponsesAPIStreamingResponse: if self._idx >= len(self._events): raise StopIteration evt: Final = self._events[self._idx] @@ -987,8 +1008,8 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): return evt -def _dump_response_object(obj: Any) -> dict[str, Any]: - if hasattr(obj, "model_dump"): +def _dump_response_object(obj: object) -> dict[str, Any]: + if isinstance(obj, BaseModel): return obj.model_dump() if isinstance(obj, dict): return obj @@ -1000,8 +1021,8 @@ def _build_response_status_event( "response.created", "response.in_progress", ], - transformed: Any, -) -> Any: + transformed: ResponsesAPIResponse, +) -> ResponseCreatedEvent | ResponseInProgressEvent: openai_types: Final = _get_openai_response_types() in_progress_response: Final = transformed.model_copy( deep=True, @@ -1018,10 +1039,10 @@ def _build_content_part_done_event( output_index: int, content_index: int, part_payload: dict[str, Any], -) -> Any | None: +) -> ContentPartDoneEvent | None: openai_types: Final = _get_openai_response_types() part_type: Final = part_payload.get("type") - part: Any + part: ContentPartDonePartOutputText | ContentPartDonePartRefusal | ContentPartDonePartReasoningText if part_type == "output_text": annotations: Final = [ openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) @@ -1057,7 +1078,7 @@ def _build_content_part_done_event( def _add_text_like_part_events( *, - events: list[Any], + events: list[ResponsesAPIStreamingResponse], item_id: str, output_index: int, content_index: int, @@ -1123,13 +1144,13 @@ def _add_text_like_part_events( def _build_synthetic_response_events( *, - transformed: Any, + transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, chunk_size: int, -) -> list[Any]: +) -> list[ResponsesAPIStreamingResponse]: openai_types: Final = _get_openai_response_types() if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Final[Any | None] = getattr(transformed, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(transformed, "usage", None) if usage_obj is not None: try: cost: Final[float | None] = logging_obj._response_cost_calculator(result=transformed) @@ -1138,10 +1159,13 @@ def _build_synthetic_response_events( except Exception: pass - events: Final[list[Any]] = [ - _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), - _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), - ] + created_event: Final[ResponsesAPIStreamingResponse] = _build_response_status_event( + openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed + ) + in_progress_event: Final[ResponsesAPIStreamingResponse] = _build_response_status_event( + openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed + ) + events: Final[list[ResponsesAPIStreamingResponse]] = [created_event, in_progress_event] sequence_number = 0 for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): @@ -1277,6 +1301,21 @@ RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"}) +class PIIGuardrailCallback(Protocol): + """Structural shape relied on for Presidio-style PII guardrail callbacks.""" + + async def check_pii( + self, + *, + text: str, + output_parse_pii: bool, + presidio_config: object, + request_data: dict[str, Any], + ) -> str: ... + + def get_presidio_settings_from_request_data(self, data: dict[str, Any]) -> object: ... + + class ResponsesWebSocketStreaming: """ Manages bidirectional WebSocket forwarding for the Responses API @@ -1292,26 +1331,26 @@ class ResponsesWebSocketStreaming: def __init__( self, - websocket: Any, - backend_ws: Any, + websocket: WebSocket, + backend_ws: ClientConnection, logging_obj: LiteLLMLoggingObj, user_api_key_dict: Any | None = None, - request_data: dict | None = None, + request_data: dict[str, Any] | None = None, first_message: str | None = None, - guardrail_callbacks: list[Any] | None = None, - output_guardrail_callbacks: list[Any] | None = None, + guardrail_callbacks: list[PIIGuardrailCallback] | None = None, + output_guardrail_callbacks: list[PIIGuardrailCallback] | 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.request_data: dict[str, Any] = request_data or {} self.messages: list[dict] = [] self.input_messages: list[dict[str, str]] = [] self.first_message = first_message - self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] - self.output_guardrail_callbacks: list[Any] = output_guardrail_callbacks or [] + self.guardrail_callbacks: list[PIIGuardrailCallback] = guardrail_callbacks or [] + self.output_guardrail_callbacks: list[PIIGuardrailCallback] = 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 @@ -1536,10 +1575,10 @@ class ResponsesWebSocketStreaming: if ( isinstance(block, dict) and block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES - and isinstance(block.get("text"), str) + and isinstance(block_text := block.get("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, @@ -1618,7 +1657,7 @@ class ResponsesWebSocketStreaming: continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text(text, pii_tokens) + unmasked: str = getattr(cb, "_unmask_pii_text")(text, pii_tokens) # noqa: B009 # _unmask_pii_text is a protected method accessed across a duck-typed guardrail callback if unmasked != text: content_block["text"] = unmasked modified = True @@ -1627,9 +1666,9 @@ class ResponsesWebSocketStreaming: if event_type in self._DELTA_EVENT_TYPES: delta: Final = evt_obj.get("delta") if isinstance(delta, str): - unmasked = cb._unmask_pii_text(delta, pii_tokens) - if unmasked != delta: - evt_obj["delta"] = unmasked + unmasked_delta: Final[str] = getattr(cb, "_unmask_pii_text")(delta, pii_tokens) # noqa: B009 # _unmask_pii_text is a protected method accessed across a duck-typed guardrail callback + if unmasked_delta != delta: + evt_obj["delta"] = unmasked_delta return json.dumps(evt_obj) return response_str @@ -1793,7 +1832,7 @@ class ManagedResponsesWebSocketHandler: def __init__( self, - websocket: Any, + websocket: WebSocket, model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: Any | None = None, @@ -1803,7 +1842,7 @@ class ManagedResponsesWebSocketHandler: timeout: float | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, - **kwargs: Any, + **kwargs: object, ) -> None: self.websocket = websocket self.model = model diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 523255c0064..bf2e882fefb 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -746,7 +746,7 @@ class ChatCompletionAssistantMessage(OpenAIChatCompletionAssistantMessage, total class ChatCompletionToolMessage(TypedDict): role: Literal["tool"] - content: Union[str, Iterable[ChatCompletionTextObject]] + content: Union[str, Iterable[Union[ChatCompletionTextObject, ChatCompletionImageObject]]] tool_call_id: str diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index d03881416f0..761a6a5b00d 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 776 }, "ANN201": { - "limit": 1961 + "limit": 1958 }, "ANN202": { "limit": 868 @@ -24,7 +24,7 @@ "limit": 121 }, "ANN401": { - "limit": 1848 + "limit": 1569 }, "ASYNC230": { "limit": 11 @@ -39,10 +39,10 @@ "limit": 505 }, "B009": { - "limit": 81 + "limit": 77 }, "B010": { - "limit": 194 + "limit": 193 }, "B018": { "limit": 2 @@ -72,7 +72,7 @@ "limit": 19 }, "C408": { - "limit": 11 + "limit": 8 }, "C414": { "limit": 4 @@ -81,7 +81,7 @@ "limit": 1 }, "C901": { - "limit": 310 + "limit": 307 }, "D419": { "limit": 6 @@ -264,7 +264,7 @@ "limit": 58 }, "SIM102": { - "limit": 314 + "limit": 296 }, "SIM103": { "limit": 119 @@ -312,7 +312,7 @@ "limit": 526 }, "TRY004": { - "limit": 96 + "limit": 95 }, "TRY201": { "limit": 407 diff --git a/ruff.toml b/ruff.toml index b652e206f41..255cc93ca89 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,7 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", "TID251", + "C901", "TID251", "B009", "B010", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index c3c5a5fb24a..941ae60bb16 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23346 + "limit": 23319 }, "LIT002": { - "limit": 27227 + "limit": 27185 }, "LIT003": { "limit": 286 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1103 + "limit": 1082 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16828 + "limit": 16762 }, "LIT011": { - "limit": 5603 + "limit": 5596 } } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cf7a9d3907e..fce75583a99 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22748,7 +22748,7 @@ export interface components { /** ChatCompletionToolMessage */ ChatCompletionToolMessage: { /** Content */ - content: string | components["schemas"]["ChatCompletionTextObject"][]; + content: string | (components["schemas"]["ChatCompletionTextObject"] | components["schemas"]["ChatCompletionImageObject"])[]; /** * Role * @constant