diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f5b8210096..7ee0dc87a88 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from enum import Enum from typing import Any, Final, Literal, cast -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm from litellm._internal_context import current_service_target, service_target @@ -44,6 +44,7 @@ from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" _ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) _DECODED_LITERAL: Final = TypeAdapter(object) +_REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter(list[dict[str, object]], config=ConfigDict(strict=True)) @contextmanager @@ -211,7 +212,7 @@ class Cache: if not redis_startup_nodes: _env_cluster_nodes: Final = litellm.get_secret("REDIS_CLUSTER_NODES") if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str): - redis_startup_nodes = json.loads(_env_cluster_nodes) + redis_startup_nodes = _REDIS_STARTUP_NODES_ADAPTER.validate_python(json.loads(_env_cluster_nodes)) if redis_startup_nodes: # Only pass GCP parameters if they are provided diff --git a/litellm/images/utils.py b/litellm/images/utils.py index 24454954714..96e41cc390e 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -17,7 +17,7 @@ class ImageEditRequestUtils: image_edit_optional_params: ImageEditOptionalRequestParams, drop_params: bool | None = None, additional_drop_params: list[str] | None = None, - ) -> dict: + ) -> dict[str, object]: """ Get optional parameters for the image edit API. diff --git a/litellm/integrations/arize/arize_phoenix_client.py b/litellm/integrations/arize/arize_phoenix_client.py index 6e8ebb99a9e..3c1acfd36c3 100644 --- a/litellm/integrations/arize/arize_phoenix_client.py +++ b/litellm/integrations/arize/arize_phoenix_client.py @@ -5,9 +5,20 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix. import urllib.parse from typing import Final +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict + from litellm.llms.custom_httpx.http_handler import HTTPHandler +@with_config(ConfigDict(extra="allow", strict=True)) +class _PromptVersionResponse(TypedDict): + data: ReadOnly[NotRequired[dict[str, object] | None]] + + +_PROMPT_VERSION_RESPONSE: Final = TypeAdapter(_PromptVersionResponse) + + def _sanitize_id(identifier: str) -> str: """Reject path traversal characters and URL-encode the identifier.""" if any(c in identifier for c in ("/", "\\", "#", "?")): @@ -71,7 +82,7 @@ class ArizePhoenixClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() + data: Final = _PROMPT_VERSION_RESPONSE.validate_python(response.json()) return data.get("data") except Exception as e: diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 0d0e309034a..a9bd93e6420 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -72,7 +72,7 @@ if TYPE_CHECKING: ) from litellm.proxy.proxy_server import UserAPIKeyAuth as _UserAPIKeyAuth - Span = _Span | Any + Span = _Span Tracer = _Tracer Context = _Context SpanExporter = _SpanExporter @@ -1440,7 +1440,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): span = None # Only set attributes if the span is still recording (not closed) # Note: parent_span is guaranteed to be not None here - if hasattr(parent_span, "set_status"): + if parent_span is not None and hasattr(parent_span, "set_status"): parent_span.set_status(Status(StatusCode.OK)) self.set_attributes(parent_span, kwargs, response_obj) # Raw-request as direct child of parent_span @@ -2233,7 +2233,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # record error on parent span (keeps hierarchy shallow) # Only set attributes if the span is still recording (not closed) # Note: parent_otel_span is guaranteed to be not None here - if parent_otel_span.is_recording(): + if parent_otel_span is not None and parent_otel_span.is_recording(): self.set_attributes(parent_otel_span, kwargs, response_obj) parent_error_message: Final = self._record_exception_on_span(span=parent_otel_span, kwargs=kwargs) parent_otel_span.set_status(Status(StatusCode.ERROR, parent_error_message)) diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index afd9ba84c96..e68f8699c4c 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -14,6 +14,7 @@ import json from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Optional +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -28,6 +29,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class _A2ATextPart(TypedDict, total=False): """The subset of an A2A message part this handler reads text from.""" @@ -333,7 +336,7 @@ class A2AGuardrailHandler(BaseTranslation): obj = item elif isinstance(item, str): try: - obj = json.loads(item.strip()) + obj = _JSON_OBJECT.validate_python(json.loads(item.strip())) except (json.JSONDecodeError, TypeError): continue else: diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 1818b406823..91db4e266b7 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -15,6 +15,14 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues +def _get_a2a_message_role(msg: AllMessageValues) -> str: + if isinstance(msg, BaseModel): + return f"{msg.model_dump().get('role', 'user')}" + if isinstance(msg, dict): + return f"{msg.get('role', 'user')}" + return f"{dict(msg).get('role', 'user')}" + + class A2AError(BaseLLMException): """Base exception for A2A protocol errors""" @@ -50,12 +58,7 @@ def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: content_text = convert_content_list_to_str(message=msg) # Get role - if isinstance(msg, BaseModel): - role = msg.model_dump().get("role", "user") - elif isinstance(msg, dict): - role = msg.get("role", "user") - else: - role = dict(msg).get("role", "user") + role = _get_a2a_message_role(msg) if content_text: conversation_parts.append(f"{role}: {content_text}") diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index aa2bee1f31d..415f39260f7 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -20,7 +20,7 @@ from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger @@ -293,14 +293,29 @@ def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ... ) -class _AnthropicSSEDelta(TypedDict, total=False): - type: ReadOnly[str] +@with_config(ConfigDict(extra="allow", strict=True)) +class _SSETextDelta(TypedDict, total=False): + type: ReadOnly[object] text: ReadOnly[str] - stop_reason: ReadOnly[str | None] -class _AnthropicSSEEvent(TypedDict, total=False): - delta: ReadOnly[_AnthropicSSEDelta] +@with_config(ConfigDict(extra="allow", strict=True)) +class _SSETextDeltaEvent(TypedDict, total=False): + delta: ReadOnly[_SSETextDelta] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _SSEStopDelta(TypedDict, total=False): + stop_reason: ReadOnly[object] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _SSEStopDeltaEvent(TypedDict, total=False): + delta: ReadOnly[_SSEStopDelta] + + +_SSE_TEXT_DELTA_EVENT: Final = TypeAdapter(_SSETextDeltaEvent) +_SSE_STOP_DELTA_EVENT: Final = TypeAdapter(_SSEStopDeltaEvent) class AnthropicMessagesHandler(BaseTranslation): @@ -1629,14 +1644,13 @@ class AnthropicMessagesHandler(BaseTranslation): if not line.startswith("data:"): return line try: - data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads( - line[len("data:") :].strip() - ) + data: Final = _DECODED_SSE_DATA.validate_python(json.loads(line[len("data:") :].strip())) except json.JSONDecodeError: return line - if not isinstance(data, dict): + event: Final = as_str_mapping(data) + if event is None: return line - rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event) + rewritten: Final = _rewritten_event(_as_str_mapping(event), rewrite_event) return line if rewritten is data else "data: " + json.dumps(rewritten) def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: @@ -1737,8 +1751,8 @@ class AnthropicMessagesHandler(BaseTranslation): # Only process content_block_delta events if event_type == "content_block_delta" and data_line: try: - data: _AnthropicSSEEvent = json.loads(data_line) - delta: _AnthropicSSEDelta = data.get("delta", {}) + data: _SSETextDeltaEvent = _SSE_TEXT_DELTA_EVENT.validate_python(json.loads(data_line)) + delta: _SSETextDelta = data.get("delta", {}) if delta.get("type") == "text_delta": text += delta.get("text", "") except json.JSONDecodeError: @@ -1800,9 +1814,9 @@ class AnthropicMessagesHandler(BaseTranslation): # Check for message_delta event with stop_reason if event_type == "message_delta" and data_line: try: - data: _AnthropicSSEEvent = json.loads(data_line) - delta: _AnthropicSSEDelta = data.get("delta", {}) - stop_reason: str | None = delta.get("stop_reason") + data: _SSEStopDeltaEvent = _SSE_STOP_DELTA_EVENT.validate_python(json.loads(data_line)) + delta: _SSEStopDelta = data.get("delta", {}) + stop_reason: object = delta.get("stop_reason") if stop_reason is not None: return True except json.JSONDecodeError: diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 9624dbb582b..fa795d36e0f 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -161,6 +161,7 @@ else: _ANTHROPIC_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 _DEFER_LOADING_ADAPTER: Final = TypeAdapter(bool) +_COMPLETION_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object]) _TOOL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _TOOLS_ADAPTER: Final = TypeAdapter(list[object]) _NATIVE_OUTPUT_STRING_FORMATS: Final = frozenset( @@ -3115,7 +3116,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## RESPONSE OBJECT try: - completion_response: Final = raw_response.json() + completion_response: Final = _COMPLETION_RESPONSE_ADAPTER.validate_python(raw_response.json(), strict=True) except Exception as e: response_headers: Final = getattr(raw_response, "headers", None) raise AnthropicError( diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 18b4b6f456a..ebbb086eb96 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -1,6 +1,9 @@ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.exceptions import UnsupportedParamsError from litellm.llms.base_llm.image_generation.transformation import ( @@ -17,6 +20,34 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +@with_config(ConfigDict(extra="allow", strict=True)) +class _MAIImageUsageDetails(TypedDict, total=False): + text_tokens: ReadOnly[NotRequired[int | float | None]] + image_tokens: ReadOnly[NotRequired[int | float | None]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _MAIImageUsage(TypedDict, total=False): + input_tokens_details: ReadOnly[NotRequired[object]] + num_input_text_tokens: ReadOnly[NotRequired[int | float | None]] + num_input_image_tokens: ReadOnly[NotRequired[int | float | None]] + input_tokens: ReadOnly[NotRequired[int | float | None]] + output_tokens: ReadOnly[NotRequired[int | float | None]] + num_output_tokens: ReadOnly[NotRequired[int | float | None]] + output_image_tokens: ReadOnly[NotRequired[int | float | None]] + total_tokens: ReadOnly[NotRequired[int | float | None]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _MAIImageResponse(TypedDict, total=False): + usage: ReadOnly[NotRequired[object]] + + +_MAI_IMAGE_USAGE_DETAILS_ADAPTER: Final[TypeAdapter[_MAIImageUsageDetails]] = TypeAdapter(_MAIImageUsageDetails) +_MAI_IMAGE_USAGE_ADAPTER: Final[TypeAdapter[_MAIImageUsage]] = TypeAdapter(_MAIImageUsage) +_MAI_IMAGE_RESPONSE_ADAPTER: Final[TypeAdapter[_MAIImageResponse]] = TypeAdapter(_MAIImageResponse) + + class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): """Azure AI Foundry MAI image generation (e.g. MAI-Image-2.5).""" @@ -77,7 +108,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): return "maiimage" in model_normalized @staticmethod - def normalize_mai_image_usage(usage: dict[str, Any] | None) -> dict[str, Any]: + def normalize_mai_image_usage(usage: object) -> Mapping[str, object]: """Map Azure MAI usage fields to OpenAI ImageUsage schema.""" if usage is None: return { @@ -87,10 +118,13 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): "total_tokens": 0, } - normalized_usage: Final = dict(usage) - input_tokens_details = normalized_usage.get("input_tokens_details") - if not isinstance(input_tokens_details, dict): - input_tokens_details = {} + normalized_usage: Final[_MAIImageUsage] = _MAI_IMAGE_USAGE_ADAPTER.validate_python(usage) + input_tokens_details_value: Final = normalized_usage.get("input_tokens_details") + input_tokens_details: Final[_MAIImageUsageDetails] = ( + _MAI_IMAGE_USAGE_DETAILS_ADAPTER.validate_python(input_tokens_details_value) + if isinstance(input_tokens_details_value, dict) + else {} + ) text_tokens = normalized_usage.get("num_input_text_tokens") if text_tokens is None: @@ -120,18 +154,16 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if total_tokens is None: total_tokens = input_tokens + output_tokens - normalized_usage.update( - { - "input_tokens": input_tokens, - "input_tokens_details": { - "image_tokens": image_tokens, - "text_tokens": text_tokens, - }, - "output_tokens": output_tokens, - "total_tokens": total_tokens, - } - ) - return normalized_usage + return { + **normalized_usage, + "input_tokens": input_tokens, + "input_tokens_details": { + "image_tokens": image_tokens, + "text_tokens": text_tokens, + }, + "output_tokens": output_tokens, + "total_tokens": total_tokens, + } def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: return ["n", "size"] @@ -251,12 +283,18 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): json_mode: bool | None = None, ) -> ImageResponse: try: - response: Final = raw_response.json() + response_json: Final = _MAI_IMAGE_RESPONSE_ADAPTER.validate_python(raw_response.json()) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) - if "usage" in response: - response["usage"] = self.normalize_mai_image_usage(response.get("usage")) + response: Final = ( + { + **response_json, + "usage": self.normalize_mai_image_usage(response_json["usage"]), + } + if "usage" in response_json + else dict(response_json) + ) logging_obj.post_call( input=request_data.get("prompt", ""), diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 58cfb6cb95e..a88734ea26c 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -9,6 +9,7 @@ from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final, get_args, overload import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL @@ -45,6 +46,8 @@ from .amazon_titan_v2_transformation import AmazonTitanV2Config from .cohere_transformation import BedrockCohereEmbeddingConfig from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig, drop_params_enabled +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -140,7 +143,7 @@ class BedrockEmbedding(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return response.json() + return _JSON_OBJECT.validate_python(response.json()) async def _make_async_call( self, @@ -174,7 +177,7 @@ class BedrockEmbedding(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return response.json() + return _JSON_OBJECT.validate_python(response.json()) def _transform_response( self, @@ -663,7 +666,7 @@ class BedrockEmbedding(BaseAWSLLM): # Parse response if response.status_code == 200: - return response.json() + return _JSON_OBJECT.validate_python(response.json()) else: raise Exception(f"Failed to get async invoke status: {response.status_code} - {response.text}") diff --git a/litellm/llms/custom_httpx/mock_transport.py b/litellm/llms/custom_httpx/mock_transport.py index 0322cceec0f..96ca41eb47c 100644 --- a/litellm/llms/custom_httpx/mock_transport.py +++ b/litellm/llms/custom_httpx/mock_transport.py @@ -12,6 +12,8 @@ import uuid from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict # --------------------------------------------------------------------------- # Pre-built response templates @@ -56,6 +58,14 @@ _JSON_HEADERS: Final = { } +@with_config(ConfigDict(extra="allow", strict=True)) +class _MockChatRequest(TypedDict): + model: ReadOnly[NotRequired[str]] + + +_MOCK_CHAT_REQUEST: Final = TypeAdapter(_MockChatRequest) + + class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): """ httpx transport that returns canned OpenAI ChatCompletion responses. @@ -70,7 +80,7 @@ class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): body: Final = json.loads(request.content) except (json.JSONDecodeError, ValueError): return ("mock-model", False) - model: Final = body.get("model", "mock-model") + model: Final = _MOCK_CHAT_REQUEST.validate_python(body).get("model", "mock-model") return (model, False) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: diff --git a/litellm/llms/edenai/audio_transcription/transformation.py b/litellm/llms/edenai/audio_transcription/transformation.py index ee574a4f406..8430e2c3cdc 100644 --- a/litellm/llms/edenai/audio_transcription/transformation.py +++ b/litellm/llms/edenai/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params @@ -21,6 +22,8 @@ from litellm.utils import convert_to_model_response_object from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + def _form_fields(model: str, optional_params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: httpx form data """LiteLLM parks non-OpenAI params, `model` included, under `extra_body` for the OpenAI SDK; a @@ -75,7 +78,7 @@ class EdenAIAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig): def transform_audio_transcription_response(self, raw_response: httpx.Response) -> TranscriptionResponse: if "application/json" not in raw_response.headers.get("content-type", ""): return TranscriptionResponse(text=raw_response.text) - body: Final = raw_response.json() + body: Final = _JSON_OBJECT.validate_python(raw_response.json()) response: Final[TranscriptionResponse] = convert_to_model_response_object( response_object=body, model_response_object=TranscriptionResponse(), response_type="audio_transcription" ) diff --git a/litellm/llms/edenai/image_generation/transformation.py b/litellm/llms/edenai/image_generation/transformation.py index e2a0b8684f6..22e95fa505b 100644 --- a/litellm/llms/edenai/image_generation/transformation.py +++ b/litellm/llms/edenai/image_generation/transformation.py @@ -8,6 +8,7 @@ Docs: https://www.edenai.co/docs/v3/llms/image-generation from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -35,6 +36,8 @@ _SUPPORTED_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( "user", ) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class EdenAIImageGenerationConfig(BaseImageGenerationConfig): def get_supported_openai_params( @@ -97,7 +100,7 @@ class EdenAIImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - body: Final = raw_response.json() + body: Final = _JSON_OBJECT.validate_python(raw_response.json()) logging_obj.post_call(original_response=body) response: Final[ImageResponse] = convert_to_model_response_object( response_object=body, model_response_object=model_response, response_type="image_generation" diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 34bd9660c43..f140fd7fd2a 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -9,6 +9,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig @@ -37,6 +38,7 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiRetrievedContext(TypedDict, total=False): """Passage Gemini retrieved from a File Search store.""" @@ -45,26 +47,30 @@ class GeminiRetrievedContext(TypedDict, total=False): title: ReadOnly[str] +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiGroundingChunk(TypedDict, total=False): """One source Gemini grounded its answer on.""" retrievedContext: ReadOnly[GeminiRetrievedContext] +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiGroundingSegment(TypedDict, total=False): """Span of the generated answer a grounding support refers to.""" text: ReadOnly[str] +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiGroundingSupport(TypedDict, total=False): """Citation linking an answer span to the grounding chunks that back it.""" segment: ReadOnly[GeminiGroundingSegment] groundingChunkIndices: ReadOnly[Sequence[int]] - confidenceScores: ReadOnly[Sequence[float]] + confidenceScores: ReadOnly[Sequence[int | float]] +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiFileSearchGroundingMetadata(TypedDict, total=False): """Grounding metadata Gemini returns for a File Search candidate.""" @@ -72,18 +78,21 @@ class GeminiFileSearchGroundingMetadata(TypedDict, total=False): groundingSupports: ReadOnly[Sequence[GeminiGroundingSupport]] +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiFileSearchCandidate(TypedDict, total=False): """One candidate of a Gemini File Search ``generateContent`` response.""" groundingMetadata: ReadOnly[GeminiFileSearchGroundingMetadata] +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiFileSearchResponse(TypedDict, total=False): """Body of a ``generateContent`` call made with the File Search tool.""" candidates: ReadOnly[Sequence[GeminiFileSearchCandidate]] +@with_config(ConfigDict(extra="allow", strict=True)) class GeminiFileSearchStore(TypedDict, total=False): """Body of a Gemini ``fileSearchStores`` create response.""" @@ -92,6 +101,10 @@ class GeminiFileSearchStore(TypedDict, total=False): createTime: ReadOnly[str] +_FILE_SEARCH_RESPONSE: Final = TypeAdapter(GeminiFileSearchResponse) +_FILE_SEARCH_STORE: Final = TypeAdapter(GeminiFileSearchStore) + + class GeminiVectorStoreConfig(BaseVectorStoreConfig): """ Vector store configuration for Google Gemini File Search. @@ -237,7 +250,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): Extracts grounding metadata and citations from the response. """ try: - response_data: Final[GeminiFileSearchResponse] = response.json() + response_data: Final = _FILE_SEARCH_RESPONSE.validate_python(response.json()) results: Final[list[VectorStoreSearchResult]] = [] # Extract candidates and grounding metadata @@ -346,7 +359,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): Transform Gemini's fileSearchStore response to standard format. """ try: - response_data: Final[GeminiFileSearchStore] = response.json() + response_data: Final = _FILE_SEARCH_STORE.validate_python(response.json()) # Extract store name (format: fileSearchStores/xxxxxxx) store_name: Final = response_data.get("name", "") diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 9553327db97..2ee3a12996d 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # spread dict loses the role literal without a cast import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -19,6 +20,8 @@ from ..common_utils import ( ) from ..per_user_auth import require_github_copilot_user_session +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding @@ -309,7 +312,7 @@ class GithubCopilotConfig(OpenAIConfig): json_mode: bool | None = None, ) -> "ModelResponse": try: - response_json = raw_response.json() + raw_json: Final = raw_response.json() except Exception: return super().transform_response( model=model, @@ -325,12 +328,12 @@ class GithubCopilotConfig(OpenAIConfig): json_mode=json_mode, ) + response_json: Final = _JSON_OBJECT.validate_python(raw_json) if not response_json.get("choices"): - response_json = self._synthesize_choices_for_anthropic_native(response_json) raw_response = httpx.Response( status_code=raw_response.status_code, headers=raw_response.headers, - content=json.dumps(response_json).encode(), + content=json.dumps(self._synthesize_choices_for_anthropic_native(response_json)).encode(), ) return super().transform_response( diff --git a/litellm/llms/langflow/chat/transformation.py b/litellm/llms/langflow/chat/transformation.py index 1781082d36f..f451ca40cf3 100644 --- a/litellm/llms/langflow/chat/transformation.py +++ b/litellm/llms/langflow/chat/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final from urllib.parse import quote import httpx +from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -13,6 +14,8 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMExcepti from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer @@ -236,12 +239,13 @@ class LangFlowConfig(BaseConfig): json_mode: bool | None = None, ) -> ModelResponse: try: - response_json: Final = raw_response.json() + raw_json: Final = raw_response.json() except Exception as e: raise LangFlowError( message=f"LangFlow returned a non-JSON response: {e}", status_code=raw_response.status_code, ) + response_json: Final = _JSON_OBJECT.validate_python(raw_json) verbose_logger.debug("LangFlow response: %s", response_json) diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index 4f93f435957..891490b8598 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -13,6 +13,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional, Union import httpx +from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -23,6 +24,8 @@ from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer @@ -436,7 +439,7 @@ class LangGraphConfig(BaseConfig): Transform the LangGraph response to LiteLLM ModelResponse format. """ try: - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) verbose_logger.debug("LangGraph response: %s", response_json) content: Final = self._extract_content_from_response(response_json) diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py index ffc6f1d5fe9..47908f20d22 100644 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_2_transformation.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -9,6 +10,8 @@ from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + if TYPE_CHECKING: from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj @@ -57,7 +60,7 @@ class DallE2ImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - response: Final = raw_response.json() + response: Final = _JSON_OBJECT.validate_python(raw_response.json()) stringified_response: Final = response ## LOGGING diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py index 90b7eaedf2f..dbab449ca25 100644 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_3_transformation.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -9,6 +10,8 @@ from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + if TYPE_CHECKING: from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj @@ -57,7 +60,7 @@ class DallE3ImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - response: Final = raw_response.json() + response: Final = _JSON_OBJECT.validate_python(raw_response.json()) stringified_response: Final = response ## LOGGING diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index c3a826616ed..88a0b8cc01f 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -9,6 +10,8 @@ from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + if TYPE_CHECKING: from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj @@ -66,7 +69,7 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - response: Final = raw_response.json() + response: Final = _JSON_OBJECT.validate_python(raw_response.json()) stringified_response: Final = response ## LOGGING diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index 3633f7b6127..ece41a7de82 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -3,6 +3,7 @@ import os from typing import Final, cast import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.llms.custom_httpx.http_handler import ( @@ -11,6 +12,8 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +_TOKEN_INFO_RESPONSE: Final = TypeAdapter(dict[str, object]) + class Oauth2Handler: """ @@ -179,7 +182,7 @@ class Oauth2Handler: response.raise_for_status() # If we get here, the request was successful - data = response.json() + data = _TOKEN_INFO_RESPONSE.validate_python(response.json(), strict=True) verbose_proxy_logger.debug( "Oauth2 token validation for token=%s, response from endpoint=%s", diff --git a/litellm/proxy/client/cli/commands/encryption.py b/litellm/proxy/client/cli/commands/encryption.py index ef73b4f0d2d..5d36daa69c9 100644 --- a/litellm/proxy/client/cli/commands/encryption.py +++ b/litellm/proxy/client/cli/commands/encryption.py @@ -1,14 +1,30 @@ """CLI commands for the at-rest credential encryption migration.""" -from typing import Final +from typing import Final, TypedDict import click import rich +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly from ...http_client import HTTPClient from .auth import cli_context_obj +@with_config(ConfigDict(extra="allow", strict=True)) +class _MigrationReport(TypedDict, total=False): + residual_legacy: ReadOnly[int | float | None] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _MigrationResponse(TypedDict, total=False): + report: ReadOnly[_MigrationReport] + + +_JSON_OBJECT: Final = TypeAdapter(object) +_MIGRATION_RESPONSE: Final = TypeAdapter(_MigrationResponse) + + @click.group() def encryption(): """Migrate at-rest credentials to AES-256-GCM and attest residual state.""" @@ -40,23 +56,35 @@ def migrate(ctx: click.Context, check_only: bool, dry_run: bool): lite encryption migrate --check # attestation scan, no writes lite encryption migrate # perform the migration """ - client: Final = HTTPClient(cli_context_obj(ctx)["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_obj(ctx) + base_url: Final = context["base_url"] + if "api_key" not in context: + raise KeyError("api_key") + client: Final = HTTPClient(base_url, context["api_key"]) - if check_only: - response = client.request("GET", "/credentials/migrate-encryption/check") - else: - response = client.request( + response: Final[object] = _JSON_OBJECT.validate_python( + client.request("GET", "/credentials/migrate-encryption/check") + if check_only + else client.request( "POST", "/credentials/migrate-encryption", json={}, params={"dry_run": "true"} if dry_run else None, ) + ) rich.print_json(data=response) - report: Final = response.get("report", {}) if isinstance(response, dict) else {} - residual: Final = report.get("residual_legacy") + residual: Final = _residual_legacy(response) if residual is not None and residual > 0: rich.print(f"[yellow]Residual legacy values remaining: {residual}[/yellow]") elif residual == 0: rich.print("[green]No legacy values remaining (residual_legacy == 0).[/green]") + + +def _residual_legacy(response: object) -> int | float | None: + if not isinstance(response, dict): + return None + payload: Final = _MIGRATION_RESPONSE.validate_python(response) + report: Final = payload.get("report", {}) + return report.get("residual_legacy") diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index ac7d5ee45ea..538b42a8c83 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -1,7 +1,7 @@ # stdlib imports import re from collections import defaultdict -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -61,6 +61,10 @@ def _get_model_info_obj_from_yaml(model: dict[str, Any]) -> ModelYamlInfo: ) +def _str_field(model: Mapping[str, object], key: str, default: str) -> str: + return str(model.get(key, default)) + + def format_iso_datetime_str(iso_datetime_str: str | None) -> str: """Format an ISO format datetime string to human-readable date with minute resolution.""" if not iso_datetime_str: @@ -142,10 +146,10 @@ def list_models(ctx: click.Context, output_format: Literal["table", "json"]) -> created = int(created) table.add_row( - str(model.get("id", "")), - str(model.get("object", "model")), + _str_field(model, "id", ""), + _str_field(model, "object", "model"), (format_timestamp(created) if isinstance(created, int) else format_iso_datetime_str(created)), - str(model.get("owned_by", "")), + _str_field(model, "owned_by", ""), ) rich.print(table) diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index 7b4f36dbc2f..7593cf809d0 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -2,9 +2,12 @@ from collections.abc import Iterable, Mapping from typing import Any, Final import requests +from pydantic import ConfigDict, TypeAdapter from .exceptions import NotFoundError, UnauthorizedError +_JSON_VALUE: Final = TypeAdapter(object, config=ConfigDict(strict=True)) + class UsersManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -29,7 +32,7 @@ class UsersManagementClient: response.raise_for_status() return response.json().get("users", response.json()) - def get_user(self, user_id: str | None = None) -> dict[str, object]: + def get_user(self, user_id: str | None = None) -> object: """Get user info (GET /user/info)""" url: Final = f"{self.base_url}/user/info" params: Final = {"user_id": user_id} if user_id else {} @@ -39,9 +42,9 @@ class UsersManagementClient: if response.status_code == 404: raise NotFoundError(response.text) response.raise_for_status() - return response.json() + return _JSON_VALUE.validate_python(response.json()) - def get_user_v2(self, user_id: str | None = None) -> dict[str, object]: + def get_user_v2(self, user_id: str | None = None) -> object: """Get user info v2 - lightweight, returns only user object (GET /v2/user/info)""" url: Final = f"{self.base_url}/v2/user/info" params: Final = {"user_id": user_id} if user_id else {} @@ -51,18 +54,18 @@ class UsersManagementClient: if response.status_code == 404: raise NotFoundError(response.text) response.raise_for_status() - return response.json() + return _JSON_VALUE.validate_python(response.json()) - def create_user(self, user_data: Mapping[str, object]) -> dict[str, object]: + def create_user(self, user_data: Mapping[str, object]) -> object: """Create a new user (POST /user/new)""" url: Final = f"{self.base_url}/user/new" response: Final = requests.post(url, headers=self._get_headers(), json=user_data, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() - return response.json() + return _JSON_VALUE.validate_python(response.json()) - def delete_user(self, user_ids: list[str]) -> dict[str, object]: + def delete_user(self, user_ids: list[str]) -> object: """Delete users (POST /user/delete)""" url: Final = f"{self.base_url}/user/delete" response: Final = requests.post( @@ -71,4 +74,4 @@ class UsersManagementClient: if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() - return response.json() + return _JSON_VALUE.validate_python(response.json()) diff --git a/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py index af4a59efb06..fe61f8c5b01 100644 --- a/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py +++ b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py @@ -16,6 +16,10 @@ unknown impl value raises, which the proxy converts to the fail-closed 503. """ import os +from typing import Final + +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.management_helpers.team_metadata_validation import ( @@ -23,12 +27,20 @@ from litellm.proxy.management_helpers.team_metadata_validation import ( TeamMetadataValidationResult, ) from litellm.types.llms.custom_http import httpxSpecialProvider -from typing import Final ALLOWED_COST_CENTERS: Final = frozenset({"CC-1001", "CC-1002"}) CLOSED_PORT_URL: Final = "http://127.0.0.1:9/validate" +@with_config(ConfigDict(extra="allow", strict=True)) +class _CostCenterReply(TypedDict): + ok: NotRequired[ReadOnly[object]] + reason: NotRequired[ReadOnly[str | None]] + + +_COST_CENTER_REPLY: Final = TypeAdapter(_CostCenterReply) + + async def _validate_allowlist(payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult: cost_center: Final = payload.metadata.get("cost_center") if cost_center is None: @@ -54,7 +66,7 @@ async def _validate_via_http(payload: TeamMetadataValidationPayload, service_url }, timeout=2.0, ) - body: Final = response.json() + body: Final = _COST_CENTER_REPLY.validate_python(response.json()) if body.get("ok") is True: return TeamMetadataValidationResult(valid=True) return TeamMetadataValidationResult( diff --git a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py index 611738ede8a..701b7bbb3c6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py @@ -300,7 +300,7 @@ class VigilGuardGuardrail(CustomGuardrail): def _build_output( inputs: GenericGuardrailAPIInputs, final_texts: list[str], - final_tool_calls: Any, + final_tool_calls: _ToolCalls | None, ) -> GenericGuardrailAPIInputs: # When nothing was changed, return the input shape verbatim so the guardrail # logs "allow" rather than "mask". When a text or a tool-call argument was @@ -315,7 +315,7 @@ class VigilGuardGuardrail(CustomGuardrail): guardrailed["images"] = inputs["images"] if "tools" in inputs: guardrailed["tools"] = inputs["tools"] - if tool_calls_changed: + if tool_calls_changed and final_tool_calls is not None: guardrailed["tool_calls"] = final_tool_calls return guardrailed diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index ef60a800cdd..e07045c0145 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -452,7 +452,7 @@ class KeyManagementEventHooks: except ImportError: return None - prisma_client: Final = getattr(proxy_server_module, "prisma_client", None) + prisma_client: Final = proxy_server_module.prisma_client user_api_key_cache: Final = getattr(proxy_server_module, "user_api_key_cache", None) if prisma_client is None or user_api_key_cache is None: diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index cf3025d8c7b..002aeb145a5 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -18,6 +18,8 @@ import logging import os from typing import TYPE_CHECKING, Any, Final +from pydantic import ConfigDict, TypeAdapter + from litellm import DualCache from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -36,6 +38,14 @@ else: InternalUsageCache = Any +_CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float]] = TypeAdapter( + str | bytes | int | float, config=ConfigDict(strict=True, hide_input_in_errors=True) +) +_OPTIONAL_CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float | None]] = TypeAdapter( + str | bytes | int | float | None, config=ConfigDict(strict=True, hide_input_in_errors=True) +) + + # Redis Lua script for atomic float increment with TTL. # INCRBYFLOAT returns the new value as a string. # Only sets EXPIRE on first call (when prior value was nil). @@ -216,9 +226,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): """Read current accumulated spend for a session.""" if self.internal_usage_cache.dual_cache.redis_cache is not None: try: - result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key) - if result is not None: - return float(result) + redis_result: Final[str | bytes | int | float | None] = _OPTIONAL_CACHE_AMOUNT_ADAPTER.validate_python( + await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key) + ) + if redis_result is not None: + return float(redis_result) return 0.0 except Exception as e: log_redis_failure( @@ -228,24 +240,28 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): e, ) - result = await self.internal_usage_cache.async_get_cache( - key=cache_key, - litellm_parent_otel_span=None, - local_only=True, + cache_result: Final[str | bytes | int | float | None] = _OPTIONAL_CACHE_AMOUNT_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) ) - if result is not None: - return float(result) + if cache_result is not None: + return float(cache_result) return 0.0 async def _increment_spend(self, cache_key: str, amount: float) -> float: """Atomically increment the session spend and return the new value.""" if self.increment_script is not None: try: - result: Final = await self.increment_script( - keys=[cache_key], - args=[str(amount), self.ttl], + script_result: Final[str | bytes | int | float] = _CACHE_AMOUNT_ADAPTER.validate_python( + await self.increment_script( + keys=[cache_key], + args=[str(amount), self.ttl], + ) ) - return float(result) + return float(script_result) except Exception as e: log_redis_failure( verbose_proxy_logger, @@ -257,12 +273,15 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return await self._in_memory_increment_spend(cache_key, amount) async def _in_memory_increment_spend(self, cache_key: str, amount: float) -> float: - current: Final = await self.internal_usage_cache.async_get_cache( - key=cache_key, - litellm_parent_otel_span=None, - local_only=True, + current: Final[str | bytes | int | float | None] = _OPTIONAL_CACHE_AMOUNT_ADAPTER.validate_python( + await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) ) - new_value: Final = (float(current) if current is not None else 0.0) + amount + current_spend: Final = float(current) if current is not None else 0.0 + new_value: Final = current_spend + amount await self.internal_usage_cache.async_set_cache( key=cache_key, value=new_value, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index b563dbc5c4c..46c3a3ed13a 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -22,7 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, TypedDict import litellm @@ -123,6 +123,15 @@ _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) _USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 +@with_config(ConfigDict(extra="allow", strict=True)) +class _AdminTeamAndKeyRows(TypedDict): + keys: ReadOnly[Sequence[Mapping[str, object]] | None] + teams: ReadOnly[Sequence[Mapping[str, object]] | None] + + +_ADMIN_TEAM_AND_KEY_ROWS_ADAPTER: Final = TypeAdapter(list[_AdminTeamAndKeyRows]) + + def _user_table( prisma_client: "PrismaClient | None", ) -> "TableActions[prisma_models.LiteLLM_UserTable]": @@ -1169,8 +1178,8 @@ async def user_info_v2( async def _fetch_admin_teams_and_keys_rows( prisma_client: "PrismaClient", sql_query: str -) -> Sequence[Mapping[str, Sequence[Mapping[str, object]] | None]]: - return await prisma_client.db.query_raw(sql_query) +) -> Sequence[_AdminTeamAndKeyRows]: + return _ADMIN_TEAM_AND_KEY_ROWS_ADAPTER.validate_python(await prisma_client.db.query_raw(sql_query)) async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index b93c05c555f..a611151a36a 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -5,6 +5,7 @@ from functools import partial from typing import Annotated, Final, Literal from fastapi import APIRouter, Depends, Query, Request +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth @@ -31,6 +32,7 @@ router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) # SPEND_LOGS_PAGINATION_COUNT_CAP, the bound ui_view_spend_logs puts on its count # query, so both reads of the same table stop at the same depth. SPEND_LOGS_FACET_SCAN_CAP: Final = 10000 +_FACET_ROWS: Final = TypeAdapter(list[dict[str, str | None]]) def _as_utc(value: datetime) -> datetime: @@ -125,8 +127,8 @@ async def _list_spend_log_facet( f" ORDER BY {column_sql} ASC" f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" ) - rows: Final = await prisma_client.db.query_raw(facet_sql, *params) - values: Final[list[str]] = [row[column] for row in rows if row.get(column)] + rows: Final = _FACET_ROWS.validate_python(await prisma_client.db.query_raw(facet_sql, *params), strict=True) + values: Final[list[str]] = [value for row in rows if (value := row.get(column))] has_more: Final = len(values) > page_size return FacetListResponse( diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index 9c0f3c4b81c..9dbf6bcc43d 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -3,11 +3,13 @@ import json import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi import status as http_status +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException, UserAPIKeyAuth @@ -44,6 +46,20 @@ _DEFAULT_TRANSCRIPTION_MODEL: Final = "gpt-realtime-whisper" _ALLOWED_SESSION_TYPES: Final = ("realtime", "transcription") +@with_config(ConfigDict(strict=True, extra="allow")) +class _RealtimeTokenPayload(TypedDict, total=False): + v: NotRequired[ReadOnly[object]] + ephemeral_key: NotRequired[ReadOnly[str]] + model_id: NotRequired[ReadOnly[str]] + user_id: NotRequired[ReadOnly[str | None]] + team_id: NotRequired[ReadOnly[str | None]] + expires_at: NotRequired[ReadOnly[object]] + session_type: NotRequired[ReadOnly[str | None]] + + +_REALTIME_TOKEN_PAYLOAD_ADAPTER: Final = TypeAdapter(_RealtimeTokenPayload) + + def _coerce_realtime_session_type(session_type: str | None) -> str: if session_type in _ALLOWED_SESSION_TYPES: return session_type @@ -201,17 +217,14 @@ def _header_value(headers: Mapping[str, str], name: str, default: str) -> str: def _decode_realtime_token_payload( decrypted_value: str, -) -> dict[str, Any] | None: +) -> _RealtimeTokenPayload | None: """ Decode realtime token payload; returns None for legacy/raw ephemeral tokens. """ try: - decoded: Final = json.loads(decrypted_value) + decoded: Final = _REALTIME_TOKEN_PAYLOAD_ADAPTER.validate_python(json.loads(decrypted_value)) except Exception: return None - - if not isinstance(decoded, dict): - return None if decoded.get("v") != _REALTIME_TOKEN_VERSION: return None if not isinstance(decoded.get("ephemeral_key"), str): diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index 1d119375778..b6e7941dbbc 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -47,13 +47,13 @@ class QualityRouter(CustomLogger): model_name: str, litellm_router_instance: "Router", default_model: str | None = None, - quality_router_config: dict[str, Any] | None = None, + quality_router_config: Mapping[str, object] | None = None, ): self.model_name = model_name self.litellm_router_instance = litellm_router_instance if quality_router_config: - self.config = QualityRouterConfig(**quality_router_config) + self.config = QualityRouterConfig.model_validate(quality_router_config) else: self.config = QualityRouterConfig() diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py b/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py index ae2df8069a9..4cbd252f65c 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py @@ -2,7 +2,7 @@ import enum from collections.abc import Mapping -from typing import Any, Literal, TypedDict +from typing import Literal, TypedDict from pydantic import Field @@ -42,7 +42,7 @@ class PolicyApplicableTo(str, enum.Enum): class DynamoAIRequest(TypedDict, total=False): """Request structure for DynamoAI /moderation/analyze endpoint""" - messages: list[dict[str, Any]] + messages: list[dict[str, object]] textType: DynamoTextType | None policyIds: list[str] modelId: str | None @@ -58,8 +58,8 @@ class PolicyInfo(TypedDict, total=False): description: str method: PolicyMethod action: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"] - methodParams: dict[str, Any] - decisionParams: dict[str, Any] + methodParams: dict[str, object] + decisionParams: dict[str, object] applicableTo: PolicyApplicableTo created_at: str creatorId: str @@ -76,7 +76,7 @@ class AppliedPolicyDto(TypedDict, total=False): """Applied policy details from DynamoAI response""" policy: PolicyInfo - outputs: dict[str, Any] | None + outputs: dict[str, object] | None action: str | None diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py index 992815bf454..30b0fb80970 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py @@ -89,7 +89,7 @@ class NomaV2StreamingParams(BaseModel): cls, value: object, handler: ValidatorFunctionWrapHandler, info: ValidationInfo ) -> object: try: - validated: Final = handler(value) + validated: Final[object] = handler(value) except ValidationError: verbose_proxy_logger.warning( "Noma v2 guardrail %s: ignoring %s=%.100r, which is not a valid value; the default applies", diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 6d1ab85012b..84e787f4ee1 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -95,11 +95,11 @@ class RerankResponse(LiteLLMBaseModel): def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation self._hidden_params = hidden_params - def __getitem__(self, key): + def __getitem__(self, key: str) -> object: return self.__dict__[key] - def get(self, key, default=None): + def get(self, key: str, default: object = None) -> object: return self.__dict__.get(key, default) - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: return key in self.__dict__