diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 715e492b159..d6c97b120c9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2507,7 +2507,7 @@ class Logging(LiteLLMLoggingBaseClass): self.callback_duration_ms += (time.time() - _start) * 1000 return payload - def _transform_usage_objects(self, result): + def _transform_usage_objects(self, result: object) -> object: if isinstance(result, ResponsesAPIResponse): result = result.model_copy() transformed_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(result.usage) @@ -2545,13 +2545,13 @@ class Logging(LiteLLMLoggingBaseClass): def _success_handler_helper_fn( self, - result=None, - start_time=None, - end_time=None, - cache_hit=None, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, standard_logging_object: StandardLoggingPayload | None = None, build_logging_payload: bool = True, - ): + ) -> tuple[datetime.datetime, datetime.datetime, object]: try: if start_time is None: start_time = self.start_time diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 2070c9f29ec..2ceb0042f1f 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,9 +1,10 @@ import base64 import time from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, replace from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict, Union, cast +from typing import TYPE_CHECKING, Final, Optional, TypeAlias, TypedDict, Union, cast from pydantic import BaseModel from typing_extensions import ReadOnly, Required @@ -154,6 +155,15 @@ class _ToolCallChoice(TypedDict, total=False): _ToolCallKey: TypeAlias = tuple[int, int] +@dataclass(frozen=True, slots=True) +class _ToolCallAccumulator: + id: str | None = None + name: str | None = None + type: str | None = None + custom_name: str | None = None + provider_specific_fields: Mapping[str, object] | None = None + + class _ToolCallChunk(TypedDict): choices: ReadOnly[Sequence[_ToolCallChoice]] @@ -503,7 +513,7 @@ class ChunkProcessor: self, tool_call_chunks: Sequence["_ToolCallChunk"] ) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]: tool_calls_list: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] = [] - tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {} + tool_call_map: Final[dict[_ToolCallKey, _ToolCallAccumulator]] = {} for chunk in tool_call_chunks: choices = chunk["choices"] @@ -537,49 +547,43 @@ class ChunkProcessor: index = (choice_index, getattr(tool_call, "index", 0)) if index not in tool_call_map: - tool_call_map[index] = { - "id": None, - "name": None, - "type": None, - "custom_name": None, - "provider_specific_fields": None, - } + tool_call_map[index] = _ToolCallAccumulator() # Extract id, type, and function data (handle both dict and object) if isinstance(tool_call, dict): if fragment_id := tool_call.get("id"): - tool_call_map[index]["id"] = fragment_id + tool_call_map[index] = replace(tool_call_map[index], id=fragment_id) if fragment_type := tool_call.get("type"): - tool_call_map[index]["type"] = fragment_type + tool_call_map[index] = replace(tool_call_map[index], type=fragment_type) function = tool_call.get("function", {}) if isinstance(function, dict): if fragment_name := function.get("name"): - tool_call_map[index]["name"] = fragment_name + tool_call_map[index] = replace(tool_call_map[index], name=fragment_name) else: # function is an object if function_name := getattr(function, "name", None): - tool_call_map[index]["name"] = function_name + tool_call_map[index] = replace(tool_call_map[index], name=function_name) custom = tool_call.get("custom") if isinstance(custom, dict): if custom_name := custom.get("name"): - tool_call_map[index]["custom_name"] = custom_name + tool_call_map[index] = replace(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 + tool_call_map[index] = replace(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_map[index] = replace(tool_call_map[index], type=tool_call.type) if object_function_name := getattr(getattr(tool_call, "function", None), "name", None): - tool_call_map[index]["name"] = object_function_name + tool_call_map[index] = replace(tool_call_map[index], name=object_function_name) object_custom: ChatCompletionDeltaCustomToolCallPayload | None = getattr( tool_call, "custom", None ) if object_custom is not None: if getattr(object_custom, "name", None): - tool_call_map[index]["custom_name"] = object_custom.name + tool_call_map[index] = replace(tool_call_map[index], custom_name=object_custom.name) # Preserve provider_specific_fields from streaming chunks provider_fields: object = None @@ -602,12 +606,14 @@ class ChunkProcessor: if provider_fields: # Merge provider_specific_fields if multiple chunks have them - merged_provider_fields = tool_call_map[index]["provider_specific_fields"] - if merged_provider_fields is None: - merged_provider_fields = {} - tool_call_map[index]["provider_specific_fields"] = merged_provider_fields - if isinstance(provider_fields, dict): - merged_provider_fields.update(provider_fields) + tool_call_map[index] = replace( + tool_call_map[index], + provider_specific_fields=( + {**(tool_call_map[index].provider_specific_fields or {}), **provider_fields} + if isinstance(provider_fields, dict) + else dict(tool_call_map[index].provider_specific_fields or {}) + ), + ) joined_fragments: Final = self._join_fragments_by_key_and_field( self._iter_tool_call_fragments(tool_call_chunks) @@ -616,37 +622,37 @@ class ChunkProcessor: # Convert the map to a list of tool calls for index in sorted(tool_call_map.keys()): tool_call_data = tool_call_map[index] - if tool_call_data["id"] and tool_call_data["custom_name"]: + if tool_call_data.id and tool_call_data.custom_name: tool_calls_list.append( ChatCompletionMessageCustomToolCall( - id=tool_call_data["id"], + id=tool_call_data.id, custom=ChatCompletionCustomToolCallPayload( - name=tool_call_data["custom_name"], + name=tool_call_data.custom_name, input=joined_fragments.get((index, "custom_input"), ""), ), ) ) - elif tool_call_data["id"] and tool_call_data["name"]: + elif tool_call_data.id and tool_call_data.name: combined_arguments = joined_fragments.get((index, "arguments"), "") or "{}" # Build function - provider_specific_fields should be on tool_call level, not function level function = Function( arguments=combined_arguments, - name=tool_call_data["name"], + name=tool_call_data.name, ) # Prepare params for ChatCompletionMessageToolCall - tool_call_params = { - "id": tool_call_data["id"], - "function": function, - "type": tool_call_data["type"] or "function", - } - # Add provider_specific_fields if present (for thought signatures in Gemini 3) - if tool_call_data.get("provider_specific_fields"): - tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"] - - tool_call = ChatCompletionMessageToolCall(**tool_call_params) + tool_call = ChatCompletionMessageToolCall( + id=tool_call_data.id, + function=function, + type=tool_call_data.type or "function", + **( + {"provider_specific_fields": dict(tool_call_data.provider_specific_fields)} + if tool_call_data.provider_specific_fields + else {} + ), + ) tool_calls_list.append(tool_call) return tool_calls_list diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 207ed249866..0813ef58528 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -474,7 +474,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): Returns: AgentCoreParsedResponse: Parsed response data """ - content_type: Final = raw_response.headers.get("content-type", "").lower() + content_type: Final = ", ".join(httpx.Headers(raw_response.headers).get_list("content-type")).lower() verbose_logger.debug("AgentCore response Content-Type: %s", content_type) # Parse response based on content type @@ -698,7 +698,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) # Check if response is JSON (agent used sync return) instead of SSE - content_type: Final = response.headers.get("content-type", "").lower() + content_type: Final = ", ".join(httpx.Headers(response.headers).get_list("content-type")).lower() if "application/json" in content_type: verbose_logger.debug( "AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream" @@ -920,7 +920,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) # Check if response is JSON (agent used sync return) instead of SSE - content_type: Final = response.headers.get("content-type", "").lower() + content_type: Final = ", ".join(httpx.Headers(response.headers).get_list("content-type")).lower() if "application/json" in content_type: verbose_logger.debug( "AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream" diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index 67e84f6da24..ca47ac67957 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -1,6 +1,9 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, RootModel, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.llms.custom_httpx.http_handler import get_async_httpx_client @@ -12,6 +15,16 @@ else: GenerateContentContentListUnionDict = Any +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _GeminiCountTokensResponse(TypedDict): + totalTokens: NotRequired[ReadOnly[int | float]] + tokenizer_used: NotRequired[ReadOnly[str]] + + +class _GeminiCountTokensResponseRoot(RootModel[_GeminiCountTokensResponse]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + class GoogleAIStudioTokenCounter: def _clean_contents_for_gemini_api(self, contents: Any) -> object: """ @@ -79,7 +92,7 @@ class GoogleAIStudioTokenCounter: async def acount_tokens( self, - contents: Any, + contents: Sequence[Mapping[str, object]] | None, model: str, api_key: str | None = None, api_base: str | None = None, @@ -143,8 +156,8 @@ class GoogleAIStudioTokenCounter: response.raise_for_status() # Parse response - result: Final = response.json() - return result + result: Final = _GeminiCountTokensResponseRoot.model_validate(response.json()).root + return dict(result) except httpx.HTTPStatusError as e: error_msg = f"Google Gen AI Studio API error: {e.response.status_code} - {e.response.text}" diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index 96b16e3e510..a73a2fe4192 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -41,11 +41,14 @@ Response format: """ import base64 +from collections.abc import Mapping, Sequence from io import BufferedReader, BytesIO from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from pydantic import ConfigDict, RootModel, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.images.utils import ImageEditRequestUtils @@ -71,6 +74,52 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _ImageUrl(TypedDict): + url: ReadOnly[NotRequired[str | None]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _Image(TypedDict): + image_url: ReadOnly[NotRequired[_ImageUrl]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _Message(TypedDict): + images: ReadOnly[NotRequired[Sequence[_Image]]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _Choice(TypedDict): + message: ReadOnly[NotRequired[_Message]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _TokensDetails(TypedDict): + image_tokens: ReadOnly[NotRequired[int | float]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _Usage(TypedDict): + prompt_tokens: ReadOnly[NotRequired[int | float]] + total_tokens: ReadOnly[NotRequired[int | float]] + completion_tokens_details: ReadOnly[NotRequired[_TokensDetails]] + prompt_tokens_details: ReadOnly[NotRequired[_TokensDetails | None]] + cost: ReadOnly[NotRequired[int | float | str | None]] + cost_details: ReadOnly[NotRequired[Mapping[str, object] | None]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _OpenRouterImageEditResponse(TypedDict): + choices: ReadOnly[NotRequired[Sequence[_Choice]]] + usage: ReadOnly[NotRequired[_Usage | None]] + model: ReadOnly[NotRequired[object]] + + +class _OpenRouterImageEditResponseRoot(RootModel[_OpenRouterImageEditResponse]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + class OpenRouterImageEditConfig(BaseImageEditConfig): """ Configuration for OpenRouter image editing via chat completions. @@ -201,7 +250,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): logging_obj: LiteLLMLoggingObj, ) -> ImageResponse: try: - response_json: Final = raw_response.json() + response_json: Final[object] = raw_response.json() except Exception as e: raise OpenRouterException( message=f"Error parsing OpenRouter response: {e}", @@ -213,7 +262,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): model_response.data = [] try: - choices: Final = response_json.get("choices", []) + payload: Final = _OpenRouterImageEditResponseRoot.model_validate(response_json).root + choices: Final = payload.get("choices", []) for choice in choices: message = choice.get("message", {}) @@ -252,7 +302,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): headers={}, ) - self._set_usage_and_cost(model_response, response_json, model) + self._set_usage_and_cost(model_response, payload, model) return model_response def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: @@ -303,11 +353,11 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): def _set_usage_and_cost( self, model_response: ImageResponse, - response_json: dict, + payload: _OpenRouterImageEditResponse, model: str, ) -> None: """Extract and set usage and cost information from OpenRouter response.""" - usage_data: Final = response_json.get("usage", {}) + usage_data: Final = payload.get("usage", {}) if usage_data: prompt_tokens: Final = usage_data.get("prompt_tokens", 0) total_tokens: Final = usage_data.get("total_tokens", 0) @@ -347,7 +397,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): ) response_cost_details.update(cost_details) - model_response.hidden_params["model"] = response_json.get("model", model) + model_response.hidden_params["model"] = payload.get("model", model) def _read_image_bytes(self, image: FileTypes) -> bytes: """Read raw bytes from various image input types.""" diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index 38a06a37f7e..007b047cc1b 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -2,6 +2,8 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, RootModel, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.secret_managers.main import get_secret_str @@ -22,6 +24,25 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _RAGFlowDataset(TypedDict): + id: ReadOnly[NotRequired[str | None]] + name: ReadOnly[NotRequired[str | None]] + create_time: ReadOnly[NotRequired[int | float | None]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _RAGFlowCreateResponse(TypedDict): + data: ReadOnly[NotRequired[_RAGFlowDataset]] + + +class _RAGFlowCreateResponseRoot(RootModel[_RAGFlowCreateResponse]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + +_RAGFLOW_EMPTY_DATA: Final[_RAGFlowDataset] = {} + + class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): """Vector store configuration for RAGFlow datasets.""" @@ -185,7 +206,8 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): headers=response.headers, ) - data: Final = response_json.get("data", {}) + create_response: Final = _RAGFlowCreateResponseRoot.model_validate(response_json).root + data: Final = create_response.get("data", _RAGFLOW_EMPTY_DATA) # Extract dataset ID dataset_id: Final = data.get("id") diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 1f7fba3c4c4..7772ca6869a 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -4,7 +4,8 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import Final import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, RootModel, TypeAdapter, with_config +from typing_extensions import ReadOnly, Required, TypedDict import litellm from litellm import verbose_logger @@ -15,6 +16,38 @@ from litellm.types.utils import StreamingChatCompletionChunk _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _SageMakerToken(TypedDict, total=False): + text: ReadOnly[str | None] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _SageMakerChunkData(TypedDict, total=False): + token: ReadOnly[_SageMakerToken | None] + index: ReadOnly[int | None] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _SageMakerMessagesChunkData(TypedDict, total=False): + choices: ReadOnly[Required[list[Mapping[str, object]]]] + + +class _SageMakerChunkDataRoot(RootModel[_SageMakerChunkData]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + +class _SageMakerMessagesChunkDataRoot(RootModel[_SageMakerMessagesChunkData]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + +def _parse_sagemaker_chunk_data(data: str) -> _SageMakerChunkData: + return _SageMakerChunkDataRoot.model_validate(json.loads(data)).root + + +def _parse_sagemaker_messages_chunk_data(data: str) -> _SageMakerMessagesChunkData: + return _SageMakerMessagesChunkDataRoot.model_validate(json.loads(data)).root + + def _load_sagemaker_response_stream_shape(): try: from botocore.loaders import Loader @@ -62,12 +95,12 @@ class AWSEventStreamDecoder: self.content_blocks: list = [] self.is_messages_api = is_messages_api - def _chunk_parser_messages_api(self, chunk_data: dict) -> StreamingChatCompletionChunk: + def _chunk_parser_messages_api(self, chunk_data: _SageMakerMessagesChunkData) -> StreamingChatCompletionChunk: openai_chunk: Final = StreamingChatCompletionChunk(**{"model": self.model, **chunk_data}) return openai_chunk - def _chunk_parser(self, chunk_data: dict) -> GChunk: + def _chunk_parser(self, chunk_data: _SageMakerChunkData) -> GChunk: verbose_logger.debug("in sagemaker chunk parser, chunk_data %s", chunk_data) _token: Final = chunk_data.get("token", {}) or {} _index: Final = chunk_data.get("index", None) or 0 @@ -113,11 +146,12 @@ class AWSEventStreamDecoder: # Try to parse the accumulated JSON try: - _data = json.loads(accumulated_json) if self.is_messages_api: - yield self._chunk_parser_messages_api(chunk_data=_data) + yield self._chunk_parser_messages_api( + chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) + ) else: - yield self._chunk_parser(chunk_data=_data) + yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) # Reset accumulated_json after successful parsing accumulated_json = "" except json.JSONDecodeError: @@ -127,11 +161,12 @@ class AWSEventStreamDecoder: # Handle any remaining data after the iterator is exhausted if accumulated_json: try: - _data = json.loads(accumulated_json) if self.is_messages_api: - yield self._chunk_parser_messages_api(chunk_data=_data) + yield self._chunk_parser_messages_api( + chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) + ) else: - yield self._chunk_parser(chunk_data=_data) + yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) except json.JSONDecodeError: # Handle or log any unparseable data at the end verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json) @@ -161,11 +196,12 @@ class AWSEventStreamDecoder: accumulated_json += message # Try to parse the accumulated JSON - _data = json.loads(accumulated_json) if self.is_messages_api: - yield self._chunk_parser_messages_api(chunk_data=_data) + yield self._chunk_parser_messages_api( + chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) + ) else: - yield self._chunk_parser(chunk_data=_data) + yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) # Reset accumulated_json after successful parsing accumulated_json = "" except json.JSONDecodeError: @@ -181,11 +217,12 @@ class AWSEventStreamDecoder: # Handle any remaining data after the iterator is exhausted if accumulated_json: try: - _data = json.loads(accumulated_json) if self.is_messages_api: - yield self._chunk_parser_messages_api(chunk_data=_data) + yield self._chunk_parser_messages_api( + chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json) + ) else: - yield self._chunk_parser(chunk_data=_data) + yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json)) except json.JSONDecodeError: # Handle or log any unparseable data at the end verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json) diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 2d35dd9b480..ebbe2c59886 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -1,6 +1,9 @@ +from collections.abc import Sequence from typing import Final, Literal import httpx +from pydantic import ConfigDict, RootModel, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -13,10 +16,6 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, ) from litellm.llms.openai.openai import AllMessageValues -from litellm.types.llms.vertex_ai import ( - CachedContentListAllResponseBody, - VertexAICachedContentResponseObject, -) from litellm.utils import is_prompt_caching_valid_prompt from ..common_utils import VertexAIError, get_vertex_base_url @@ -32,6 +31,31 @@ local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for cal MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _CachedContentItem(TypedDict): + displayName: ReadOnly[NotRequired[object]] + name: ReadOnly[NotRequired[str | None]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _CachedContentListBody(TypedDict): + cachedContents: ReadOnly[NotRequired[Sequence[_CachedContentItem]]] + nextPageToken: ReadOnly[NotRequired[str | None]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _CreatedCachedContent(TypedDict): + name: ReadOnly[NotRequired[str | None]] + + +class _CachedContentListBodyRoot(RootModel[_CachedContentListBody]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + +class _CreatedCachedContentRoot(RootModel[_CreatedCachedContent]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + class ContextCachingEndpoints(VertexBase): """ Covers context caching endpoints for Vertex AI + Google AI Studio @@ -157,25 +181,20 @@ class ContextCachingEndpoints(VertexBase): except Exception as e: raise VertexAIError(status_code=500, message=str(e)) - raw_response = resp.json() + raw_response = _CachedContentListBodyRoot.model_validate(resp.json()).root logging_obj.post_call(original_response=raw_response) if "cachedContents" not in raw_response: return None - all_cached_items = CachedContentListAllResponseBody(**raw_response) - - if "cachedContents" not in all_cached_items: - return None - # Check current page for matching cache_key - for cached_item in all_cached_items["cachedContents"]: + for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") if display_name is not None and display_name == cache_key: return cached_item.get("name") # Check if there are more pages - page_token = all_cached_items.get("nextPageToken") + page_token = raw_response.get("nextPageToken") if not page_token: # No more pages, cache not found break @@ -249,25 +268,20 @@ class ContextCachingEndpoints(VertexBase): except Exception as e: raise VertexAIError(status_code=500, message=str(e)) - raw_response = resp.json() + raw_response = _CachedContentListBodyRoot.model_validate(resp.json()).root logging_obj.post_call(original_response=raw_response) if "cachedContents" not in raw_response: return None - all_cached_items = CachedContentListAllResponseBody(**raw_response) - - if "cachedContents" not in all_cached_items: - return None - # Check current page for matching cache_key - for cached_item in all_cached_items["cachedContents"]: + for cached_item in raw_response["cachedContents"]: display_name = cached_item.get("displayName") if display_name is not None and display_name == cache_key: return cached_item.get("name") # Check if there are more pages - page_token = all_cached_items.get("nextPageToken") + page_token = raw_response.get("nextPageToken") if not page_token: # No more pages, cache not found break @@ -424,14 +438,11 @@ class ContextCachingEndpoints(VertexBase): except httpx.TimeoutException: raise VertexAIError(status_code=408, message="Timeout error occurred.") - raw_response_cached: Final = response.json() - cached_content_response_obj: Final = VertexAICachedContentResponseObject( - name=raw_response_cached.get("name"), model=raw_response_cached.get("model") - ) + raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root return ( non_cached_messages, optional_params, - cached_content_response_obj["name"], + raw_response_cached.get("name"), ) async def async_check_and_create_cache( @@ -580,14 +591,11 @@ class ContextCachingEndpoints(VertexBase): except httpx.TimeoutException: raise VertexAIError(status_code=408, message="Timeout error occurred.") - raw_response_cached: Final = response.json() - cached_content_response_obj: Final = VertexAICachedContentResponseObject( - name=raw_response_cached.get("name"), model=raw_response_cached.get("model") - ) + raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root return ( non_cached_messages, optional_params, - cached_content_response_obj["name"], + raw_response_cached.get("name"), ) def get_cache(self): diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index fb846d6b196..93e0a183a2b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2209,10 +2209,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _set_stream_metadata_on_response( model_response: Any, - grounding_metadata: list[dict], - url_context_metadata: list[dict], - safety_ratings: list[dict], - citation_metadata: list[dict], + grounding_metadata: Sequence[Mapping[str, object]], + url_context_metadata: Sequence[Mapping[str, object]], + safety_ratings: Sequence[Mapping[str, object]], + citation_metadata: Sequence[Mapping[str, object]], ) -> None: setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) if grounding_metadata: diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 1d4a5d48a65..21c94b76d21 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal import httpx from fastapi import HTTPException +from pydantic import ConfigDict, RootModel, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -65,6 +67,55 @@ from litellm.types.utils import ( GUARDRAIL_NAME: Final = "model_armor" + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _MatchResult(TypedDict): + matchState: NotRequired[ReadOnly[str]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _DeidentifyData(TypedDict): + text: NotRequired[ReadOnly[str]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _DeidentifyResult(TypedDict): + matchState: NotRequired[ReadOnly[str]] + data: NotRequired[ReadOnly[_DeidentifyData]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _SdpResult(TypedDict): + inspectResult: NotRequired[ReadOnly[_MatchResult]] + deidentifyResult: NotRequired[ReadOnly[_DeidentifyResult]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _Filter(TypedDict): + raiFilterResult: NotRequired[ReadOnly[_MatchResult]] + piAndJailbreakFilterResult: NotRequired[ReadOnly[_MatchResult]] + maliciousUriFilterResult: NotRequired[ReadOnly[_MatchResult]] + csamFilterFilterResult: NotRequired[ReadOnly[_MatchResult]] + virusScanFilterResult: NotRequired[ReadOnly[_MatchResult]] + sdpFilterResult: NotRequired[ReadOnly[_SdpResult | None]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _SanitizationResult(TypedDict): + filterResults: NotRequired[ReadOnly[Mapping[str, _Filter] | Sequence[_Filter] | str | int | float | bool | None]] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _ModelArmorResponse(TypedDict): + sanitizationResult: NotRequired[ReadOnly[_SanitizationResult]] + sanitizedText: NotRequired[ReadOnly[str | None]] + text: NotRequired[ReadOnly[str | None]] + + +class _ModelArmorResponseRoot(RootModel[_ModelArmorResponse]): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + # Only these carry the finished output; response.created carries an empty body _RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) @@ -99,10 +150,11 @@ class ModelArmorAPIError(Exception): _SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"}) + RedactablePayload = dict | list | str | int | float | bool | None -def _redact_scanned_content(payload: RedactablePayload, depth: int = 0) -> RedactablePayload: +def _redact_scanned_content(payload: object, depth: int = 0) -> object: if depth >= DEFAULT_MAX_RECURSE_DEPTH: return "[REDACTED]" if isinstance(payload, dict): @@ -223,12 +275,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return f"Model Armor API error (upstream {status_code})" return f"Model Armor API error (upstream {status_code}): {response_text}" - def _build_block_error_detail(self, message: str, armor_response: RedactablePayload) -> dict: + def _build_block_error_detail(self, message: str, armor_response: Mapping[str, object]) -> dict: if self.sanitize_error_detail: return {"error": message} return {"error": message, "model_armor_response": armor_response} - def _build_logging_response(self, armor_response: RedactablePayload) -> RedactablePayload: + def _build_logging_response(self, armor_response: Mapping[str, object]) -> object: if self.sanitize_error_detail: return _redact_scanned_content(armor_response) return armor_response @@ -286,7 +338,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): request_data: dict | None = None, file_bytes: bytes | None = None, file_type: str | None = None, - ) -> dict: + ) -> _ModelArmorResponse: """ Make request to Model Armor API. Supports both text and file prompt sanitization. If file_bytes and file_type are provided, file prompt sanitization is performed. @@ -361,8 +413,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): json_response: Final = response.json() if hasattr(json_response, "__await__"): - return await json_response - return json_response + return _ModelArmorResponseRoot.model_validate(await json_response).root + return _ModelArmorResponseRoot.model_validate(json_response).root def sanitize_file_prompt(self, file_bytes: bytes, file_type: str, source: str = "user_prompt") -> dict: """ @@ -378,7 +430,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else: return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}} - def _should_block_content(self, armor_response: Mapping[str, object], allow_sanitization: bool = False) -> bool: + def _should_block_content(self, armor_response: _ModelArmorResponse, allow_sanitization: bool = False) -> bool: """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult.""" for filt in self._filter_result_items(armor_response): # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before @@ -404,7 +456,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Fallback dict code removed; all cases handled above return False - def _get_sanitized_content(self, armor_response: Mapping[str, Any]) -> str | None: + def _get_sanitized_content(self, armor_response: _ModelArmorResponse) -> str | None: """ Get the sanitized content from a Model Armor response, if available. Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found. @@ -434,7 +486,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return armor_response.get("sanitizedText") or armor_response.get("text") @staticmethod - def _filter_result_items(armor_response: Mapping[str, Any]) -> Sequence[Any]: + def _filter_result_items(armor_response: _ModelArmorResponse) -> Sequence[_Filter]: """Every filter result in a scan response. filterResults is a dict of named filters on most templates and a list on some, so both @@ -447,7 +499,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return filter_results return [] - def _has_deidentify_match(self, armor_response: Mapping[str, object]) -> bool: + def _has_deidentify_match(self, armor_response: _ModelArmorResponse) -> bool: """Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction.""" for filter_entry in self._filter_result_items(armor_response): sdp = filter_entry.get("sdpFilterResult") @@ -457,7 +509,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): def _resolve_streaming_outcome( self, - armor_response: Mapping[str, object], + armor_response: _ModelArmorResponse, assembled_response: object, content: str, ) -> tuple[bool, str | None]: @@ -489,7 +541,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return False, sanitized_content @staticmethod - def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object: + def _append_armor_response(existing: object, armor_response: object) -> object: """Accumulate scan responses so a later text scan does not drop an earlier file scan. Returns the single response on its own (backward compatible) and a list once a request diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 5a313b14c4a..4e46a070ec2 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -245,7 +245,7 @@ class _TransactionFactory(Protocol): class _ModelTransactionClient(LiteLLMBaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True) + model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True, hide_input_in_errors=True) tx: _TransactionFactory @@ -1856,8 +1856,8 @@ async def delete_team_models( Returns the model_ids that were deleted. """ deleted_model_ids: Final[list[str]] = [] - async with prisma_client.db.tx() as tx_ctx: - tx: Final[_TxModelTables] = tx_ctx + transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db) + async with transaction_client.tx() as tx: for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index ed5843f4bce..a4d2977a971 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -41,7 +41,7 @@ if TYPE_CHECKING: import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm._internal_context import with_service_target @@ -237,6 +237,12 @@ class CliSsoTeamDetail(LiteLLMBaseModel): _CLI_SSO_TEAM_DETAILS_ADAPTER: Final = TypeAdapter(tuple[CliSsoTeamDetail, ...]) +_CLI_SSO_FLOW_JSON_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter( + dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True) +) +_SSO_USERINFO_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter( + dict[str, object] | None, config=ConfigDict(strict=True, hide_input_in_errors=True) +) _TEAMLESS_CLI_SSO_TEAM_DETAIL: Final = CliSsoTeamDetail(team_models=()) @@ -371,7 +377,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: flow = _read_cli_sso_flow(cache, _get_cli_sso_flow_cache_key(cast(str, login_id))) if isinstance(flow, str): try: - flow = _as_object(json.loads(flow)) + flow = _CLI_SSO_FLOW_JSON_ADAPTER.validate_python(json.loads(flow)) except ValueError: flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: @@ -4199,7 +4205,7 @@ class SSOAuthenticationHandler: ) if resp.status_code == 200: try: - userinfo_raw: Final[dict[str, object] | None] = resp.json() + userinfo_raw: Final = _SSO_USERINFO_RESPONSE_ADAPTER.validate_python(resp.json()) if not userinfo_raw: # JSON null (None) or empty dict ({}) — no identity claims. # Treat as failure so id_token fallback can be attempted.