From 638cdc0c82d702a4b14d70fdeed4ff97c8f96d24 Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 18 Oct 2026 10:57:03 +0000 Subject: [PATCH] refactor(types): replace Any with proven types in 39 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../handler.py | 18 ++----- .../azure_sentinel/azure_sentinel.py | 13 ++++- .../integrations/datadog/datadog_llm_obs.py | 24 +++++++-- litellm/integrations/email_alerting.py | 26 ++++++---- litellm/litellm_core_utils/core_helpers.py | 6 ++- litellm/llms/anthropic/chat/transformation.py | 3 +- litellm/llms/azure_ai/embed/handler.py | 7 ++- litellm/llms/base_llm/chat/transformation.py | 2 +- ...n_nova_canvas_image_edit_transformation.py | 21 ++++++-- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- .../llms/gemini/image_edit/transformation.py | 45 +++++++++++++--- litellm/llms/gigachat/file_handler.py | 15 +++++- litellm/llms/github_copilot/authenticator.py | 12 +++-- .../litellm_proxy/skills/code_execution.py | 10 +++- litellm/llms/oci/chat/transformation.py | 15 +++--- litellm/llms/ollama/common_utils.py | 49 +++++++++++------- litellm/llms/ollama/completion/handler.py | 20 ++++---- litellm/llms/openai_like/json_loader.py | 7 +-- .../opensandbox/sandbox/transformation.py | 12 +++-- litellm/llms/sap/chat/handler.py | 25 ++++----- .../stability/image_edit/transformations.py | 17 ++++++- .../llms/watsonx/completion/transformation.py | 42 +++++++++++---- litellm/models/organization_membership.py | 3 +- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/ip_address_utils.py | 2 +- litellm/proxy/auth/litellm_license.py | 13 ++++- litellm/proxy/client/chat.py | 5 +- litellm/proxy/client/credentials.py | 17 ++++--- litellm/proxy/client/keys.py | 13 +++-- litellm/proxy/client/teams.py | 10 ++-- .../cato_networks/cato_networks.py | 51 ++++++++++++------- .../generic_guardrail_api.py | 5 +- .../guardrail_hooks/microsoft_purview/base.py | 17 +++++-- .../panw_prisma_airs/panw_prisma_airs.py | 4 +- .../mcp_management_endpoints.py | 38 ++++++++++---- .../workflow_management_endpoints.py | 36 ++++++++++--- .../complexity_router/complexity_router.py | 19 +++++-- litellm/router_utils/search_api_router.py | 2 +- litellm/secret_managers/main.py | 14 ++++- 39 files changed, 457 insertions(+), 187 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 3f7017eccc2..94af75de249 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -2,12 +2,10 @@ Handler for transforming /chat/completions api requests to litellm.responses requests """ -from collections.abc import AsyncIterable, Coroutine, Iterable, Mapping +from collections.abc import AsyncIterable, Iterable, Mapping from typing import ( TYPE_CHECKING, - Any, Final, - Union, cast, # noqa: TID251 # optional_params reaches the bridge as an untyped dict ) @@ -160,13 +158,7 @@ class ResponsesToCompletionBridgeHandler: encoding=typed_kwargs.get("encoding"), ) - def completion( - self, *args, **kwargs - ) -> Union[ - Coroutine[None, None, Union["ModelResponse", "CustomStreamWrapper"]], - "ModelResponse", - "CustomStreamWrapper", - ]: + def completion(self, *args, **kwargs) -> object: if kwargs.get("acompletion") is True: return self.acompletion(**kwargs) @@ -269,7 +261,7 @@ class ResponsesToCompletionBridgeHandler: ) return self._apply_post_stream_processing(streamwrapper, model, custom_llm_provider) - async def acompletion(self, *args, **kwargs) -> Union["ModelResponse", "CustomStreamWrapper"]: + async def acompletion(self, *args, **kwargs) -> object: from litellm import aresponses from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -379,7 +371,7 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider: str, logging_obj: "LiteLLMLoggingObj", json_mode: bool | None, - ) -> "CustomStreamWrapper": + ) -> object: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.base_model_iterator import MockResponseIterator @@ -396,7 +388,7 @@ class ResponsesToCompletionBridgeHandler: stream: "CustomStreamWrapper", model: str, custom_llm_provider: str, - ) -> Any: + ) -> object: """Apply provider-specific post-stream processing if available.""" from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 84dfc770e8f..7f900d693a4 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -22,6 +22,8 @@ from typing import Final, TypeVar from urllib.parse import urlparse import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.integrations.batch_utils import ( @@ -52,6 +54,15 @@ MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType( ) +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _AzureOAuthTokenResponse(TypedDict, total=False): + access_token: ReadOnly[NotRequired[str | None]] + expires_in: ReadOnly[NotRequired[int | float]] + + +_AZURE_OAUTH_TOKEN_RESPONSE_ADAPTER: Final = TypeAdapter(_AzureOAuthTokenResponse) + + class AzureSentinelLogger(CustomBatchLogger): """ Logger that sends LiteLLM logs to Azure Sentinel via Azure Monitor Logs Ingestion API @@ -226,7 +237,7 @@ class AzureSentinelLogger(CustomBatchLogger): if response.status_code != 200: raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}") - token_response: Final = response.json() + token_response: Final = _AZURE_OAUTH_TOKEN_RESPONSE_ADAPTER.validate_python(response.json()) self.oauth_token = token_response.get("access_token") expires_in: Final = token_response.get("expires_in", 3600) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 270d799742c..a72023eaedf 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -12,7 +12,7 @@ import os from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType -from typing import Any, Final, Literal +from typing import Final, Literal import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -54,6 +54,7 @@ from litellm.types.utils import ( PROMPT_CARRYING_GUARDRAIL_FIELDS, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CallTypes, + StandardLoggingMetadata, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) @@ -93,8 +94,12 @@ _ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "ro _COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS) -def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]: - metadata: Final = standard_logging_payload.get("metadata") +def _logging_metadata_of(standard_logging_payload: StandardLoggingPayload) -> StandardLoggingMetadata | None: + return standard_logging_payload.get("metadata") + + +def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, object]: + metadata: Final = _logging_metadata_of(standard_logging_payload) return metadata or _EMPTY_MAPPING @@ -1057,14 +1062,23 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Get budget information from metadata metadata: Final = _metadata_of(standard_logging_payload) + logging_metadata: Final = _logging_metadata_of(standard_logging_payload) # API key max budget - user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget") + user_api_key_max_budget: Final = ( + logging_metadata["user_api_key_max_budget"] + if logging_metadata and "user_api_key_max_budget" in logging_metadata + else None + ) if user_api_key_max_budget is not None: spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget) # API key spend - user_api_key_spend: Final = metadata.get("user_api_key_spend") + user_api_key_spend: Final = ( + logging_metadata["user_api_key_spend"] + if logging_metadata and "user_api_key_spend" in logging_metadata + else None + ) if user_api_key_spend is not None: try: spend_metrics["user_api_key_spend"] = float(user_api_key_spend) diff --git a/litellm/integrations/email_alerting.py b/litellm/integrations/email_alerting.py index 27ee6bb60b4..938ea7a605c 100644 --- a/litellm/integrations/email_alerting.py +++ b/litellm/integrations/email_alerting.py @@ -5,16 +5,21 @@ Functions for sending Email Alerts import os from typing import Final +from pydantic import ConfigDict, TypeAdapter + from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.repositories.team_repository import TeamRepository from litellm.types.integrations.slack_alerting import WebhookEvent +_QUERY_RAW_RESULT: Final = TypeAdapter(object, config=ConfigDict(strict=True)) +_EMAIL_QUERY_RESULT: Final = TypeAdapter(list[object], config=ConfigDict(strict=True)) + # we use this for the email header, please send a test email if you change this. verify it looks good on email LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" LITELLM_SUPPORT_CONTACT: Final = "support@berri.ai" -async def get_all_team_member_emails(team_id: str | None = None) -> list: +async def get_all_team_member_emails(team_id: str | None = None) -> list[str]: verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id) if team_id is None: return [] @@ -51,18 +56,21 @@ async def get_all_team_member_emails(team_id: str | None = None) -> list: WHERE user_id = ANY($1::TEXT[]); """ - _result: Final = await prisma_client.db.query_raw(sql_query, _team_member_user_ids) + raw_result: Final = _QUERY_RAW_RESULT.validate_python( + await prisma_client.db.query_raw(sql_query, _team_member_user_ids) + ) + result: Final = None if raw_result is None else _EMAIL_QUERY_RESULT.validate_python(raw_result) - verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result) + verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", raw_result) - if _result is None: + if result is None: return [] - emails: Final = [] - for user in _result: - if user and isinstance(user, dict) and user.get("user_email", None) is not None: - emails.append(user.get("user_email")) - return emails + return [ + user_email + for user in result + if user and isinstance(user, dict) and isinstance((user_email := user.get("user_email")), str) + ] async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index e706bcadef8..4f72b45b3b7 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id @@ -25,6 +25,7 @@ else: _CODEX_CLIENT_PREFIX_RE: Final = re.compile(r"^codex[-_ /]", re.IGNORECASE) +_MODEL_DUMP: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True)) def is_codex_user_agent(user_agent: str) -> bool: @@ -559,7 +560,8 @@ def preserve_upstream_non_openai_attributes( """ # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings expected_keys: Final = set(type(model_response).model_fields.keys()).union({"usage"}) - for key, value in original_chunk.model_dump().items(): + original_attributes: Final = _MODEL_DUMP.validate_python(original_chunk.model_dump()) + for key, value in original_attributes.items(): if key not in expected_keys: setattr(model_response, key, value) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 38af6e7998b..859b3ddd2f0 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -3100,7 +3100,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## RESPONSE OBJECT try: - completion_response: Final = raw_response.json() + completion_response: Final = _COMPLETION_RESPONSE.validate_python(raw_response.json()) except Exception as e: response_headers: Final = getattr(raw_response, "headers", None) raise AnthropicError( @@ -3186,6 +3186,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _DECODED_JSON: Final = TypeAdapter(object) +_COMPLETION_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) _REQUEST_TOOLS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 57e6ca51ade..d820a403c42 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -3,6 +3,7 @@ from typing import Final from urllib.parse import urlsplit, urlunsplit from openai import OpenAI +from pydantic import TypeAdapter import litellm from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base @@ -18,6 +19,8 @@ from litellm.utils import convert_to_model_response_object from .cohere_transformation import AzureAICohereConfig +_EMBEDDING_RESPONSE: Final = TypeAdapter[dict[str, object]](dict[str, object]) + def _foundry_models_route_base(api_base: str | None) -> str | None: if api_base is None or not is_foundry_model_inference_base(api_base): @@ -88,7 +91,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): headers={"Authorization": f"Bearer {api_key}"}, ) - embedding_response: Final = response.json() + embedding_response: Final = _EMBEDDING_RESPONSE.validate_python(response.json()) embedding_headers: Final = dict(response.headers) returned_response: Final[EmbeddingResponse] = convert_to_model_response_object( response_object=embedding_response, @@ -131,7 +134,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): headers={"Authorization": f"Bearer {api_key}"}, ) - embedding_response: Final = response.json() + embedding_response: Final = _EMBEDDING_RESPONSE.validate_python(response.json()) embedding_headers: Final = dict(response.headers) returned_response: Final[EmbeddingResponse] = convert_to_model_response_object( response_object=embedding_response, diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 51c8e4b8bef..b2eb4cde5ce 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -428,7 +428,7 @@ class BaseConfig(ABC): streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse"], sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> object: pass def get_model_response_iterator_for_request( diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index 1cece8743de..4963a314cb9 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -14,10 +14,13 @@ from __future__ import annotations import base64 import os +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig @@ -39,6 +42,17 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _NovaCanvasSuccessResponse(TypedDict): + images: NotRequired[ReadOnly[Sequence[str] | None]] + errors: NotRequired[ReadOnly[object]] + message: NotRequired[ReadOnly[object]] + error: NotRequired[ReadOnly[object]] + + +_NOVA_CANVAS_SUCCESS_RESPONSE: Final = TypeAdapter(_NovaCanvasSuccessResponse) + + def _nova_canvas_task_body( *, image_b64: str, @@ -401,7 +415,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): json_mode: bool | None = None, ) -> ImageResponse: try: - response_data: Final = raw_response.json() + raw_response_data: Final[object] = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error parsing Nova Canvas image edit response: {e}", @@ -411,12 +425,13 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): if raw_response.status_code not in (200,): raise self.get_error_class( - error_message=f"Nova Canvas image edit error: {response_data}", + error_message=f"Nova Canvas image edit error: {raw_response_data}", status_code=raw_response.status_code, headers=raw_response.headers, ) - images: Final[list[str]] = response_data.get("images") or [] + response_data: Final = _NOVA_CANVAS_SUCCESS_RESPONSE.validate_python(raw_response_data) + images: Final[Sequence[str]] = response_data.get("images") or [] if "errors" in response_data and not images: raise self.get_error_class( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b4dae7ba8b6..13082d639f7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -8487,7 +8487,7 @@ class BaseLLMHTTPHandler: prefetch_resp.raise_for_status() except Exception as e: raise self._handle_error(e=e, provider_config=video_provider_config) - prefetched_source_data = prefetch_resp.json() + prefetched_source_data = _JSON_OBJECT.validate_python(prefetch_resp.json()) try: url, data, files = video_provider_config.transform_video_edit_request( @@ -8585,7 +8585,7 @@ class BaseLLMHTTPHandler: prefetch_resp.raise_for_status() except Exception as e: raise self._handle_error(e=e, provider_config=video_provider_config) - prefetched_source_data = prefetch_resp.json() + prefetched_source_data = _JSON_OBJECT.validate_python(prefetch_resp.json()) try: url, data, files = video_provider_config.transform_video_edit_request( diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 31d3963c70c..b54e4ab1468 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -1,10 +1,12 @@ import base64 -from collections.abc import Mapping +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, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig @@ -33,6 +35,35 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True)) +class _GeminiInlineData(TypedDict): + data: NotRequired[ReadOnly[str | None]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _GeminiPart(TypedDict): + inlineData: NotRequired[ReadOnly[_GeminiInlineData | None]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _GeminiContent(TypedDict): + parts: NotRequired[ReadOnly[Sequence[_GeminiPart]]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _GeminiCandidate(TypedDict): + content: NotRequired[ReadOnly[_GeminiContent]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _GeminiImageEditResponse(TypedDict): + candidates: NotRequired[ReadOnly[Sequence[_GeminiCandidate]]] + usageMetadata: NotRequired[ReadOnly[Mapping[str, object]]] + + +_GEMINI_IMAGE_EDIT_RESPONSE: Final = TypeAdapter(_GeminiImageEditResponse) + + class GeminiImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" SUPPORTED_PARAMS: list[str] = ["n", "size", "imageConfig"] @@ -133,7 +164,8 @@ class GeminiImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) - candidates: Final = response_json.get("candidates", []) + response_body: Final = _GEMINI_IMAGE_EDIT_RESPONSE.validate_python(response_json) + candidates: Final = response_body.get("candidates", []) data_list: Final[list[ImageObject]] = [] for candidate in candidates: @@ -141,17 +173,18 @@ class GeminiImageEditConfig(BaseImageEditConfig): parts = content.get("parts", []) for part in parts: inline_data = part.get("inlineData") - if inline_data and inline_data.get("data"): + image_data = inline_data.get("data") if inline_data else None + if image_data: data_list.append( ImageObject( - b64_json=inline_data["data"], + b64_json=image_data, url=None, ) ) model_response.data = cast(list[OpenAIImage], data_list) - if "usageMetadata" in response_json: - model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"]) + if "usageMetadata" in response_body: + model_response.usage = transform_gemini_image_usage(dict(response_body["usageMetadata"])) return model_response def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[dict[str, object]]: diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index c051d693720..09dd09cb4bf 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -12,6 +12,9 @@ import uuid from collections.abc import Mapping from typing import Final +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -26,6 +29,14 @@ from .authenticator import get_access_token, get_access_token_async _file_cache: Final[dict[str, str]] = {} +@with_config(ConfigDict(extra="allow", strict=True)) +class _UploadedFile(TypedDict): + id: ReadOnly[NotRequired[str | None]] + + +_UPLOADED_FILE: Final = TypeAdapter(_UploadedFile) + + def _get_url_hash(url: str) -> str: """Generate hash for URL to use as cache key.""" return hashlib.sha256(url.encode()).hexdigest() @@ -137,7 +148,7 @@ def upload_file_sync( timeout=60, ) response.raise_for_status() - result: Final = response.json() + result: Final = _UPLOADED_FILE.validate_python(response.json()) file_id: Final = result.get("id") if file_id: @@ -206,7 +217,7 @@ async def upload_file_async( timeout=60, ) response.raise_for_status() - result: Final = response.json() + result: Final = _UPLOADED_FILE.validate_python(response.json()) file_id: Final = result.get("id") if file_id: diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index 59e1804f453..f8ab2b4b6ab 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -2,7 +2,7 @@ import json import os import time from datetime import datetime -from typing import Any, Final +from typing import Final import httpx from pydantic import ConfigDict, TypeAdapter, with_config @@ -40,9 +40,15 @@ class _AccessTokenPoll(TypedDict): error: ReadOnly[NotRequired[object]] +@with_config(ConfigDict(extra="allow", strict=True)) +class _ApiKeyInfo(TypedDict): + token: ReadOnly[NotRequired[str | None]] + + _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) _DEVICE_CODE: Final = TypeAdapter(_DeviceCode) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPoll) +_API_KEY_INFO: Final = TypeAdapter(_ApiKeyInfo) def github_api_headers( @@ -194,7 +200,7 @@ class Authenticator: verbose_logger.warning("Error reading API endpoint from file: %s", e) return None - def _refresh_api_key(self) -> dict[str, Any]: + def _refresh_api_key(self) -> _ApiKeyInfo: """ Refresh the API key using the access token. @@ -215,7 +221,7 @@ class Authenticator: response = sync_client.get(api_key_url, headers=headers) response.raise_for_status() - response_json = response.json() + response_json = _API_KEY_INFO.validate_python(response.json()) if "token" in response_json: return response_json diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index fbc287589b3..cc61a9add02 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -17,7 +17,9 @@ from collections.abc import Mapping, Sequence from enum import Enum from typing import Any, Final, Protocol, TypedDict +from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly +from typing_extensions import TypedDict as _TypedDict from litellm._logging import verbose_logger @@ -50,10 +52,14 @@ class _AnthropicToolSpec(TypedDict, total=False): input_schema: ReadOnly[_ToolArgumentSchema] -class _CodeExecutionArguments(TypedDict, total=False): +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _CodeExecutionArguments(_TypedDict, total=False): code: ReadOnly[str] +_CODE_EXECUTION_ARGUMENTS: Final = TypeAdapter(_CodeExecutionArguments) + + class _GeneratedFile(TypedDict, total=False): name: ReadOnly[str] mime_type: ReadOnly[str] @@ -115,7 +121,7 @@ class _CodeExecutionOutcome(TypedDict, total=False): def _parse_code_execution_arguments(serialized_arguments: str) -> _CodeExecutionArguments: - return json.loads(serialized_arguments) + return _CODE_EXECUTION_ARGUMENTS.validate_python(json.loads(serialized_arguments)) class LiteLLMInternalTools(str, Enum): diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index e62e5e63642..959858c1a33 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Callable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -188,6 +188,7 @@ def _normalize_tool_choice(selected_params: dict) -> None: _TOOL_CHOICE_MODES: Final = frozenset({"auto", "none", "required", "any"}) _TOOLS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) _JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OCI_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(strict=True)) def _as_json_object(value: object) -> Mapping[str, object] | None: @@ -714,11 +715,12 @@ class OCIChatConfig(BaseConfig): status_code=raw_response.status_code, ) + validated_response_json: Final = _OCI_JSON_OBJECT_ADAPTER.validate_python(response_json) vendor: Final = get_vendor_from_model(model) if vendor == OCIVendors.COHERE: - model_response = handle_cohere_response(response_json, model, model_response, raw_response) + model_response = handle_cohere_response(validated_response_json, model, model_response, raw_response) else: - model_response = handle_generic_response(response_json, model, model_response, raw_response) + model_response = handle_generic_response(validated_response_json, model, model_response, raw_response) model_response.hidden_params["additional_headers"] = raw_response.headers return model_response @@ -869,9 +871,10 @@ class OCIStreamWrapper(CustomStreamWrapper): message=f"Chunk cannot be parsed as JSON: {e}", ) - if dict_chunk.get("apiFormat") == "COHERE": + validated_dict_chunk: Final = _OCI_JSON_OBJECT_ADAPTER.validate_python(dict_chunk) + if validated_dict_chunk.get("apiFormat") == "COHERE": result: Final = handle_cohere_stream_chunk( - dict_chunk, + validated_dict_chunk, prior_tool_calls_emitted=self._cohere_tool_calls_emitted, prior_text_emitted=self._cohere_text_emitted, ) @@ -886,7 +889,7 @@ class OCIStreamWrapper(CustomStreamWrapper): self._cohere_text_emitted = True break return self._emit_chunk(result) - return self._emit_chunk(handle_generic_stream_chunk(dict_chunk)) + return self._emit_chunk(handle_generic_stream_chunk(validated_dict_chunk)) __all__ = [ diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 284a6ffea62..ec305eca847 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -12,6 +12,19 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ModelResponse, ModelResponseStream _PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object]) +_OLLAMA_JSON: Final = TypeAdapter[object](object) +_OLLAMA_FIELDS: Final = TypeAdapter[dict[str, object]](dict[str, object]) +_OLLAMA_ENTRIES: Final = TypeAdapter[list[object]](list[object]) + + +def _model_name_from_ollama_tag(entry: object) -> str | None: + if not isinstance(entry, dict): + return None + entry_fields: Final = _OLLAMA_FIELDS.validate_python(entry) + name: Final = entry_fields.get("name") or entry_fields.get("model") + if not isinstance(name, str): + return None + return name if name.startswith("ollama/") else f"ollama/{name}" class OllamaResponseIdNormalizer: @@ -147,24 +160,22 @@ class OllamaModelInfo(BaseLLMModelInfo): api_key = self.get_api_key(api_key) if passed_api_base is None or api_key else None headers: Final = {"Authorization": f"Bearer {api_key}"} if api_key else {} - names: Final[set[str]] = set() try: resp: Final = httpx.get(f"{base}/api/tags", headers=headers) resp.raise_for_status() - data: Final = resp.json() - # Expecting a dict with a 'models' list - models_list = [] - if isinstance(data, dict) and "models" in data and isinstance(data["models"], list): - models_list = data["models"] - elif isinstance(data, list): - models_list = data - # Extract model names - for entry in models_list: - if not isinstance(entry, dict): - continue - nm = entry.get("name") or entry.get("model") - if isinstance(nm, str): - names.add(nm if nm.startswith("ollama/") else f"ollama/{nm}") + data: Final = _OLLAMA_JSON.validate_python(resp.json()) + data_fields: Final = _OLLAMA_FIELDS.validate_python(data) if isinstance(data, dict) else None + model_field: Final = data_fields.get("models") if data_fields is not None else None + models_list: Final = ( + _OLLAMA_ENTRIES.validate_python(model_field) + if isinstance(model_field, list) + else _OLLAMA_ENTRIES.validate_python(data) + if isinstance(data, list) + else [] + ) + names: Final = frozenset( + name for entry in models_list if (name := _model_name_from_ollama_tag(entry)) is not None + ) except Exception as e: verbose_logger.warning("Error retrieving ollama tag endpoint: %s", e) # If tags endpoint fails, fall back to static list @@ -201,13 +212,13 @@ class OllamaModelInfo(BaseLLMModelInfo): return any(name.lower() in model_cost_keys for name in potential_model_names) @staticmethod - def _supports_function_calling(ollama_model_info: dict) -> bool: + def _supports_function_calling(ollama_model_info: dict[str, object]) -> bool: _template: Final[str] = str(ollama_model_info.get("template", "") or "") return "tools" in _template.lower() @staticmethod - def _get_max_tokens(ollama_model_info: dict) -> int | None: - _model_info: Final[dict] = ollama_model_info.get("model_info", {}) + def _get_max_tokens(ollama_model_info: dict[str, object]) -> object: + _model_info: Final = _OLLAMA_FIELDS.validate_python(ollama_model_info.get("model_info", {})) for key, value in _model_info.items(): if "context_length" in key: @@ -248,7 +259,7 @@ class OllamaModelInfo(BaseLLMModelInfo): "max_output_tokens": None, } - model_info: Final = response.json() + model_info: Final = _OLLAMA_FIELDS.validate_python(response.json()) max_tokens: Final = self._get_max_tokens(model_info) return { diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 816028499e8..33c00e7ff1d 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -5,9 +5,9 @@ Ollama /chat/completion calls handled in llm_http_handler.py """ from collections.abc import Mapping, Sequence -from typing import Any, Final, Protocol, TypedDict +from typing import Any, Final, Protocol -from typing_extensions import NotRequired, ReadOnly +from pydantic import BaseModel, ConfigDict import litellm from litellm.types.utils import EmbeddingResponse @@ -19,11 +19,13 @@ class TokenEncoder(Protocol): def encode(self, text: str, /) -> Sequence[int]: ... -class OllamaEmbeddingResponse(TypedDict): +class OllamaEmbeddingResponse(BaseModel): """Body of an Ollama ``/api/embed`` response.""" - embeddings: ReadOnly[list[list[float]]] - prompt_eval_count: ReadOnly[NotRequired[int]] + model_config = ConfigDict(extra="allow", frozen=True, strict=True) + + embeddings: list[list[int | float]] + prompt_eval_count: int | None = None def _ollama_embed_url(api_base: str) -> str: @@ -58,12 +60,12 @@ def _process_ollama_embedding_response( encoding: TokenEncoder | None, ) -> EmbeddingResponse: output_data: Final = [] - embeddings: Final[list[list[float]]] = response_json["embeddings"] + embeddings: Final[list[list[int | float]]] = response_json.embeddings for idx, emb in enumerate(embeddings): output_data.append({"object": "embedding", "index": idx, "embedding": emb}) - input_tokens = response_json.get("prompt_eval_count", None) + input_tokens: int | None = response_json.prompt_eval_count if input_tokens is None: if encoding is not None: @@ -100,7 +102,7 @@ async def ollama_aembeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = await litellm.module_level_aclient.post(url=_ollama_embed_url(api_base), json=data) - response_json: Final[OllamaEmbeddingResponse] = response.json() + response_json: Final = OllamaEmbeddingResponse.model_validate(response.json()) return _process_ollama_embedding_response( response_json=response_json, @@ -124,7 +126,7 @@ def ollama_embeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = litellm.module_level_client.post(url=_ollama_embed_url(api_base), json=data) - response_json: Final[OllamaEmbeddingResponse] = response.json() + response_json: Final = OllamaEmbeddingResponse.model_validate(response.json()) return _process_ollama_embedding_response( response_json=response_json, diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 23b63177137..79af4f19037 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -6,12 +6,13 @@ import json from pathlib import Path from typing import Final -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.types.utils import ProviderSpecificModelInfo _PROVIDER_INFO_ADAPTER: Final = TypeAdapter(ProviderSpecificModelInfo) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) class SimpleProviderConfig: @@ -53,10 +54,10 @@ class JSONProviderRegistry: try: with open(json_path) as f: - data: Final = json.load(f) + data: Final = _JSON_OBJECT.validate_python(json.load(f)) for slug, config in data.items(): - cls._providers[slug] = SimpleProviderConfig(slug, config) + cls._providers[slug] = SimpleProviderConfig(slug, _JSON_OBJECT.validate_python(config)) cls._loaded = True except Exception as e: diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py index 13c9d805638..ecf05ffc903 100644 --- a/litellm/llms/opensandbox/sandbox/transformation.py +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -41,6 +41,7 @@ MAX_OUTPUT_BYTES: Final = SANDBOX_MAX_OUTPUT_BYTES _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_VALUE: Final = TypeAdapter(object) +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True, strict=True)) class OpenSandboxSandboxConfig(BaseSandboxConfig): @@ -491,21 +492,22 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): if not data: return None try: - parsed: Final = json.loads(data) + parsed: Final = _JSON_VALUE.validate_python(json.loads(data)) except json.JSONDecodeError: return None if not isinstance(parsed, dict): return None - if "type" not in parsed and "code" in parsed and "message" in parsed: + parsed_object: Final = _JSON_DICT.validate_python(parsed) + if "type" not in parsed_object and "code" in parsed_object and "message" in parsed_object: return { "type": "error", "error": { - "ename": str(parsed["code"]), - "evalue": str(parsed["message"]), + "ename": str(parsed_object["code"]), + "evalue": str(parsed_object["message"]), "traceback": [], }, } - return parsed + return parsed_object @staticmethod def _normalize_result(message: dict[str, object]) -> dict[str, object]: diff --git a/litellm/llms/sap/chat/handler.py b/litellm/llms/sap/chat/handler.py index b4f4d4faddb..746c5920a4d 100755 --- a/litellm/llms/sap/chat/handler.py +++ b/litellm/llms/sap/chat/handler.py @@ -6,12 +6,15 @@ from collections.abc import AsyncIterator, Iterator from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import OpenAIChatCompletionChunk from ...custom_httpx.llm_http_handler import BaseLLMHTTPHandler +_STREAM_EVENT_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + # ------------------------------- # Errors @@ -107,6 +110,14 @@ class _StreamParser: return None +def _parse_stream_event(payload: str) -> OpenAIChatCompletionChunk | None: + try: + event_object: Final[object] = json.loads(payload) + except Exception: + return None + return _StreamParser.to_openai_chunk(_STREAM_EVENT_ADAPTER.validate_python(event_object)) + + # ------------------------------- # Iterators # ------------------------------- @@ -146,12 +157,7 @@ class SAPStreamIterator: raise StopIteration try: - obj = json.loads(payload) - except Exception: - continue - - try: - chunk = _StreamParser.to_openai_chunk(obj) + chunk = _parse_stream_event(payload) except ValueError as e: self._safe_close() raise e @@ -217,12 +223,7 @@ class AsyncSAPStreamIterator: await self._aclose() raise StopAsyncIteration try: - obj = json.loads(payload) - except Exception: - continue - - try: - chunk = _StreamParser.to_openai_chunk(obj) + chunk = _parse_stream_event(payload) except ValueError as e: await self._aclose() raise GenAIHubOrchestrationError(502, str(e)) diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 5d5fad6128e..c1b3390e91b 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -10,6 +10,8 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -30,6 +32,15 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True)) +class _StabilityImageEditResponse(TypedDict): + finish_reason: ReadOnly[NotRequired[object]] + image: ReadOnly[NotRequired[str | None]] + + +_STABILITY_IMAGE_EDIT_RESPONSE: Final = TypeAdapter(_StabilityImageEditResponse) + + class StabilityImageEditConfig(BaseImageEditConfig): """ Configuration for Stability AI image edit. @@ -275,8 +286,10 @@ class StabilityImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) + validated_response: Final = _STABILITY_IMAGE_EDIT_RESPONSE.validate_python(response_data) + # Check finish_reason - finish_reason: Final = response_data.get("finish_reason", "") + finish_reason: Final = validated_response.get("finish_reason", "") if finish_reason == "CONTENT_FILTERED": raise self.get_error_class( error_message="Content was filtered by Stability AI safety systems", @@ -289,7 +302,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): model_response.data = [] # Extract image from response - image_b64: Final = response_data.get("image") + image_b64: Final = validated_response.get("image") if image_b64: model_response.data.append( ImageObject( diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 5c0416bd6ca..188893b5b37 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -1,9 +1,11 @@ import time -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock @@ -29,6 +31,25 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True)) +class _WatsonXTextGenResponse(TypedDict): + results: ReadOnly[Sequence[object]] + model_id: NotRequired[ReadOnly[str | None]] + created_at: NotRequired[ReadOnly[str | None]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _WatsonXTextGenResult(TypedDict): + generated_text: ReadOnly[str] + input_token_count: ReadOnly[int] + generated_token_count: ReadOnly[int] + stop_reason: ReadOnly[str] + + +_WATSONX_TEXT_GEN_RESPONSE: Final = TypeAdapter(_WatsonXTextGenResponse) +_WATSONX_TEXT_GEN_RESULT: Final = TypeAdapter(_WatsonXTextGenResult) + + class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): """ Reference: https://cloud.ibm.com/apidocs/watsonx-ai#text-generation @@ -297,19 +318,22 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): status_code=500, message=f"Error: Invalid response from Watsonx.ai API: {json_resp}", ) + payload: Final = _WATSONX_TEXT_GEN_RESPONSE.validate_python(json_resp) if model_response is None: - model_response = ModelResponse(model=json_resp.get("model_id", None)) - generated_text: Final = json_resp["results"][0]["generated_text"] - prompt_tokens: Final = json_resp["results"][0]["input_token_count"] - completion_tokens: Final = json_resp["results"][0]["generated_token_count"] + model_response = ModelResponse(model=payload.get("model_id", None)) + first_result: Final = _WATSONX_TEXT_GEN_RESULT.validate_python(payload["results"][0]) + generated_text: Final = first_result["generated_text"] + prompt_tokens: Final = first_result["input_token_count"] + completion_tokens: Final = first_result["generated_token_count"] model_response.choices[0].message.content = generated_text - model_response.choices[0].finish_reason = map_finish_reason(json_resp["results"][0]["stop_reason"]) - if json_resp.get("created_at"): + model_response.choices[0].finish_reason = map_finish_reason(first_result["stop_reason"]) + created_at: Final = payload.get("created_at") + if created_at: try: - created_datetime = datetime.fromisoformat(json_resp["created_at"]) + created_datetime = datetime.fromisoformat(created_at) except ValueError: # datetime.fromisoformat cannot handle 'Z' in Python 3.10 - created_datetime = datetime.fromisoformat(f"{json_resp['created_at'].rstrip('Z')}+00:00") + created_datetime = datetime.fromisoformat(f"{created_at.rstrip('Z')}+00:00") model_response.created = int(created_datetime.timestamp()) else: model_response.created = int(time.time()) diff --git a/litellm/models/organization_membership.py b/litellm/models/organization_membership.py index e9697cf57b6..6a0d22cb63b 100644 --- a/litellm/models/organization_membership.py +++ b/litellm/models/organization_membership.py @@ -6,7 +6,6 @@ Canonical definition for ``litellm_organizationmembership``. Re-exported from """ from datetime import datetime -from typing import Any from pydantic import ConfigDict, model_validator @@ -24,7 +23,7 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: str | None = None created_at: datetime updated_at: datetime - user: Any | None = None + user: object | None = None litellm_budget_table: LiteLLM_BudgetTable | None = None user_email: str | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a04fba2d905..8bae8b9dc79 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2627,7 +2627,7 @@ async def _load_team_membership_on_cache_miss( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, ) -> LiteLLM_TeamMembership | None: - redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + redis_cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) redis_membership: Final = _membership_from_cached_payload(redis_cached) if not isinstance(redis_membership, _TeamMembershipCacheMiss): return redis_membership diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index f1ca2c70f3b..65c24f16fc4 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -371,4 +371,4 @@ class IPAddressUtils: return client_ip case _HopCountUnset(): pass - return get_request_ip_address(request, use_x_forwarded_for=use_xff) + return get_request_ip_address(request, use_x_forwarded_for=use_xff is True) diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index e1d6b2e090c..7ae04b69db8 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -7,7 +7,8 @@ from datetime import datetime from typing import TYPE_CHECKING, Final import httpx -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly, Required, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import NON_LLM_CONNECTION_TIMEOUT @@ -23,6 +24,14 @@ AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' fe _LICENSE_VERDICT: Final = TypeAdapter(object) +@with_config(ConfigDict(extra="allow", strict=True)) +class _AirgappedLicensePayload(TypedDict, total=False): + expiration_date: Required[ReadOnly[str]] + + +_AIRGAPPED_LICENSE_PAYLOAD_ADAPTER: Final = TypeAdapter(_AirgappedLicensePayload) + + class LicenseCheck: """ - Check if license in env @@ -201,7 +210,7 @@ class LicenseCheck: ) # Decode and parse the data - license_data: Final = json.loads(message.decode()) + license_data: Final = _AIRGAPPED_LICENSE_PAYLOAD_ADAPTER.validate_python(json.loads(message.decode())) # Check expiration date expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index c027bf2fe98..3a8f8eb7f21 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -7,6 +7,7 @@ from pydantic import ConfigDict, TypeAdapter from .exceptions import StreamError, UnauthorizedError +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) _SSE_LINE: Final[TypeAdapter[bytes | bytearray]] = TypeAdapter( bytes | bytearray, config=ConfigDict(arbitrary_types_allowed=True, strict=True, hide_input_in_errors=True) ) @@ -59,7 +60,7 @@ class ChatClient: frequency_penalty: float | None = None, user: str | None = None, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Create a chat completion. @@ -114,7 +115,7 @@ class ChatClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index ebcd79457a6..7771f16b856 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -3,9 +3,12 @@ from typing import Any, Final from urllib.parse import quote import requests +from pydantic import ConfigDict, TypeAdapter from .exceptions import UnauthorizedError +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class CredentialsManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -75,7 +78,7 @@ class CredentialsManagementClient: credential_values: Mapping[str, object], return_request: bool = False, display_name: str | None = None, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Create a new credential. @@ -111,7 +114,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -121,7 +124,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Delete a credential by name. @@ -148,7 +151,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -172,7 +175,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -182,7 +185,7 @@ class CredentialsManagementClient: self, credential_name: str, return_request: bool = False, - ) -> dict[str, Any] | requests.Request: + ) -> dict[str, object] | requests.Request: """ Get a credential by name. @@ -209,7 +212,7 @@ class CredentialsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index 028b338f412..cbeb283bef5 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -3,11 +3,14 @@ from collections.abc import Mapping from typing import Any, Final import requests +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.secret_redaction import redact_string from .exceptions import UnauthorizedError +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class KeysManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -179,7 +182,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -223,7 +226,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -239,7 +242,7 @@ class KeysManagementClient: key_alias: str | None = None, team_id: str | None = None, user_id: str | None = None, - ) -> dict[str, object] | requests.Request: + ) -> dict[str, object]: """ Update an existing API key's parameters. @@ -285,7 +288,7 @@ class KeysManagementClient: response: Final = session.send(request.prepare(), timeout=self._timeout) response_text = response.text response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except Exception: raise Exception(f"Error updating key: {response_text}") @@ -314,7 +317,7 @@ class KeysManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: redacted_message: Final = redact_string(str(e)) if e.response.status_code == 401: diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index dc6fa1e228a..bae97f2eb3d 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -4,9 +4,13 @@ import builtins from typing import Any, Final import requests +from pydantic import ConfigDict, TypeAdapter from .exceptions import UnauthorizedError +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT_LIST: Final = TypeAdapter(builtins.list[dict[str, object]], config=ConfigDict(strict=True)) + class TeamsManagementClient: """Client for managing teams in LiteLLM proxy.""" @@ -125,9 +129,9 @@ class TeamsManagementClient: raise UnauthorizedError("Authentication failed. Check your API key.") response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) - def get_available(self) -> builtins.list[dict[str, Any]]: + def get_available(self) -> builtins.list[dict[str, object]]: """ Get list of available teams that the user can join. @@ -146,4 +150,4 @@ class TeamsManagementClient: raise UnauthorizedError("Authentication failed. Check your API key.") response.raise_for_status() - return response.json() + return _JSON_OBJECT_LIST.validate_python(response.json()) diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 936f862b10b..b91c922b8bb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -14,8 +14,8 @@ from ssl import SSLContext from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException -from pydantic import BaseModel -from typing_extensions import NotRequired, TypedDict +from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ConnectionClosed @@ -58,38 +58,51 @@ class _WsSslKwargs(TypedDict, total=False): ssl: bool | str | SSLContext +_CATO_JSON_CONFIG: Final = ConfigDict(extra="allow", strict=True) + + +@with_config(_CATO_JSON_CONFIG) class _CatoRequiredAction(TypedDict, total=False): - action_type: str - detection_message: str + action_type: ReadOnly[str] + detection_message: ReadOnly[str] +@with_config(_CATO_JSON_CONFIG) class _CatoRedactedMessage(TypedDict): - role: NotRequired[str] - content: str | None + content: NotRequired[ReadOnly[str | None]] +@with_config(_CATO_JSON_CONFIG) class _CatoRedactedChat(TypedDict, total=False): - all_redacted_messages: Sequence[_CatoRedactedMessage] + all_redacted_messages: ReadOnly[Sequence[_CatoRedactedMessage]] +@with_config(_CATO_JSON_CONFIG) class _CatoAnalysisResult(TypedDict, total=False): - policy_drill_down: Mapping[str, object] + policy_drill_down: ReadOnly[Mapping[str, object]] +@with_config(_CATO_JSON_CONFIG) class _CatoAnalyzeResponse(TypedDict): - required_action: NotRequired[_CatoRequiredAction | None] - analysis_result: NotRequired[_CatoAnalysisResult] - redacted_chat: NotRequired[_CatoRedactedChat] + required_action: NotRequired[ReadOnly[_CatoRequiredAction | None]] + analysis_result: NotRequired[ReadOnly[_CatoAnalysisResult]] + redacted_chat: NotRequired[ReadOnly[_CatoRedactedChat]] +@with_config(_CATO_JSON_CONFIG) class _CatoOutputRedaction(TypedDict): - redacted_output: str + redacted_output: ReadOnly[str] +@with_config(_CATO_JSON_CONFIG) class _CatoStreamMessage(TypedDict, total=False): - verified_chunk: Mapping[str, object] - done: bool - blocking_message: str + verified_chunk: ReadOnly[Mapping[str, object]] + done: ReadOnly[object] + blocking_message: ReadOnly[str] + + +_CATO_ANALYZE_RESPONSE_ADAPTER: Final = TypeAdapter(_CatoAnalyzeResponse) +_CATO_STREAM_MESSAGE_ADAPTER: Final = TypeAdapter(_CatoStreamMessage) class CatoNetworksGuardrail(CustomGuardrail): @@ -308,7 +321,7 @@ class CatoNetworksGuardrail(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - res: Final[_CatoAnalyzeResponse] = response.json() + res: Final = _CATO_ANALYZE_RESPONSE_ADAPTER.validate_python(response.json()) required_action: Final = res.get("required_action") action_type: Final = required_action and required_action.get("action_type", None) if action_type is None: @@ -358,7 +371,7 @@ class CatoNetworksGuardrail(CustomGuardrail): if original_messages: data["messages"] = [ ( - {**original, "content": redacted_messages[idx]["content"]} + {**original, "content": redacted_messages[idx].get("content")} if idx < len(redacted_messages) and redacted_messages[idx].get("content") is not None else original ) @@ -449,7 +462,7 @@ class CatoNetworksGuardrail(CustomGuardrail): timeout=self.timeout, ) response.raise_for_status() - res: Final[_CatoAnalyzeResponse] = response.json() + res: Final = _CATO_ANALYZE_RESPONSE_ADAPTER.validate_python(response.json()) required_action: Final = res.get("required_action") action_type: Final = required_action and required_action.get("action_type", None) if action_type == "block_action" and required_action is not None: @@ -637,7 +650,7 @@ class CatoNetworksGuardrail(CustomGuardrail): try: while True: raw_message = await self._await_cato_message(websocket, sender) - result: _CatoStreamMessage = json.loads(raw_message) + result = _CATO_STREAM_MESSAGE_ADAPTER.validate_python(json.loads(raw_message)) if verified_chunk := result.get("verified_chunk"): yield ModelResponseStream.model_validate(verified_chunk) continue diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 199da7df30d..75a12b2df37 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -11,7 +11,7 @@ from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Literal, Optional import httpx -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError from pydantic_core import PydanticSerializationError from litellm._logging import verbose_proxy_logger @@ -45,6 +45,7 @@ _CHAT_ROW: Final[TypeAdapter[AllMessageValues]] = TypeAdapter(AllMessageValues) _SENT_ROW: Final = TypeAdapter(Mapping[str, object]) _SENT_ROWS: Final = TypeAdapter(Sequence[Mapping[str, object]]) _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) # Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*). _HEADER_VALUE_ALLOWLIST: Final = frozenset( @@ -600,7 +601,7 @@ class GenericGuardrailAPI(CustomGuardrail): ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final = _JSON_OBJECT.validate_python(response.json()) verbose_proxy_logger.debug("Generic Guardrail API response: %s", response_json) diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index f83314af548..8659ac135d0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -6,7 +6,8 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx -from typing_extensions import NotRequired, TypedDict +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -35,6 +36,16 @@ class GraphTokenResponse(TypedDict): expires_in: NotRequired[int] +@with_config(ConfigDict(extra="allow", strict=True)) +class _GraphTokenPayload(TypedDict): + access_token: ReadOnly[str] + expires_in: ReadOnly[NotRequired[int | float | str]] + + +_GRAPH_TOKEN_RESPONSE_ADAPTER: Final = TypeAdapter(_GraphTokenPayload) +_GRAPH_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + + class PurviewGuardrailBase: """ Base class for Microsoft Purview guardrails. @@ -112,7 +123,7 @@ class PurviewGuardrailBase: timeout=self.timeout, ) response.raise_for_status() - token_data: Final[GraphTokenResponse] = response.json() + token_data: Final = _GRAPH_TOKEN_RESPONSE_ADAPTER.validate_python(response.json()) access_token: Final = token_data["access_token"] expires_in: Final = int(token_data.get("expires_in", 3599)) # Recompute ``now`` after the await so the expiry reflects when the @@ -148,7 +159,7 @@ class PurviewGuardrailBase: verbose_proxy_logger.debug("Purview Graph POST %s", url) response: Final = await self.async_handler.post(url=url, headers=headers, json=json_body, timeout=self.timeout) response.raise_for_status() - response_json: Final[dict[str, object]] = response.json() + response_json: Final = _GRAPH_RESPONSE_ADAPTER.validate_python(response.json()) response_headers: Final = dict(response.headers) verbose_proxy_logger.debug("Purview Graph response: %s", response_json) return response_json, response_headers diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 8e03f7db4f5..8beda158ada 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -67,6 +67,8 @@ ToolCallLike: TypeAlias = ( | ChatCompletionToolCallChunk ) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) + class _ToolCallFunctionSlice(LiteLLMBaseModel): model_config = ConfigDict(from_attributes=True, extra="ignore") @@ -452,7 +454,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) response.raise_for_status() - result: Final[dict[str, object]] = response.json() + result: Final = _JSON_OBJECT.validate_python(response.json()) # Validate response format if "action" not in result: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a014b9cf8a5..40c5e46a813 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -46,7 +46,7 @@ from fastapi import ( status, ) from fastapi.responses import JSONResponse -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import ReadOnly, TypedDict from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, public_catalog_operation @@ -123,6 +123,22 @@ def does_mcp_server_exist(mcp_server_records: Iterable[_HasServerId], mcp_server DEFAULT_MCP_REGISTRY_VERSION: Final = "1.0.0" + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _MCPRegistryServer(TypedDict, total=False): + name: ReadOnly[str] + title: ReadOnly[str] + description: ReadOnly[str] + category: ReadOnly[str] + + +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _MCPRegistry(TypedDict, total=False): + servers: ReadOnly[Sequence[_MCPRegistryServer]] + + +_MCP_REGISTRY_ADAPTER: Final = TypeAdapter(_MCPRegistry) + if TYPE_CHECKING: from prisma import models as prisma_models @@ -3417,19 +3433,23 @@ if MCP_AVAILABLE: "mcp_registry.json", ) - _mcp_registry_cache: Mapping[str, Sequence[Mapping[str, str]]] | None = None + _mcp_registry_cache: _MCPRegistry | None = None - def _load_mcp_registry() -> Mapping[str, Sequence[Mapping[str, str]]]: + def _load_mcp_registry() -> _MCPRegistry: """Load the curated MCP registry from disk. Cached after first read.""" global _mcp_registry_cache if _mcp_registry_cache is not None: return _mcp_registry_cache - try: - with open(_MCP_REGISTRY_PATH, "r") as f: - data: Mapping[str, Sequence[Mapping[str, str]]] = json.load(f) - except Exception as e: - verbose_proxy_logger.warning("Failed to load MCP registry from %s: %s", _MCP_REGISTRY_PATH, e) - data = {"servers": []} + + def _read_registry() -> _MCPRegistry: + try: + with open(_MCP_REGISTRY_PATH, "r") as f: + return _MCP_REGISTRY_ADAPTER.validate_python(json.load(f)) + except Exception as e: + verbose_proxy_logger.warning("Failed to load MCP registry from %s: %s", _MCP_REGISTRY_PATH, e) + return {"servers": []} + + data: Final = _read_registry() _mcp_registry_cache = data return data diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 890a344dfe0..868ff86b442 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -147,6 +147,35 @@ class _MessageCreateData(TypedDict, total=False): session_id: str +class _EventTable(Protocol): + async def create(self, *, data: _EventCreateData) -> object: ... + + +class _RunTable(Protocol): + async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> object: ... + + +class _WorkflowTx(Protocol): + @property + def litellm_workflowevent(self) -> _EventTable: ... + + @property + def litellm_workflowrun(self) -> _RunTable: ... + + +async def _append_event_in_tx( + tx: _WorkflowTx, run_id: str, event_data: _EventCreateData, new_status: str | None +) -> object: + """Create the event and, when the event type maps to a status, move the run to it in the same transaction.""" + event: Final = await tx.litellm_workflowevent.create(data=event_data) + if new_status: + await tx.litellm_workflowrun.update( + where={"run_id": run_id}, + data={"status": new_status}, + ) + return event + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -387,12 +416,7 @@ async def append_workflow_event( event_data["data"] = _json(data.data) async with prisma_client.db.tx() as tx: - event: object = await tx.litellm_workflowevent.create(data=event_data) - if new_status: - await tx.litellm_workflowrun.update( - where={"run_id": run_id}, - data={"status": new_status}, - ) + event: object = await _append_event_in_tx(tx, run_id, event_data, new_status) return event diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 1638f20c8e3..a475358e1f0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -22,7 +22,7 @@ import json import random import re import time -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType @@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, TypeAdapter, ValidationError, create_model from pydantic_core import ErrorDetails +from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger @@ -411,9 +412,19 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str] return [*base_keywords, *deduped_custom.values()] -def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, Any]: +class _ParentSessionKwargs(TypedDict, total=False): + litellm_session_id: ReadOnly[object] + litellm_trace_id: ReadOnly[object] + + +def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> _ParentSessionKwargs: kwargs: Final = request_kwargs or {} - return {k: kwargs[k] for k in ("litellm_session_id", "litellm_trace_id") if kwargs.get(k) is not None} + session_id: Final = kwargs.get("litellm_session_id") + trace_id: Final = kwargs.get("litellm_trace_id") + return _ParentSessionKwargs( + **({"litellm_session_id": session_id} if session_id is not None else {}), + **({"litellm_trace_id": trace_id} if trace_id is not None else {}), + ) def _response_cost_or_none(response: object) -> float | None: @@ -2828,7 +2839,7 @@ class ComplexityRouter(CustomLogger): "originating_request_masked": masked_originating_request(request_kwargs), "body": {"model": llm_config.model, **payload}, } - classify: Final = ( + classify: Final[Callable[..., Awaitable[ModelResponse | ResponsesAPIResponse]]] = ( self.litellm_router_instance.aresponses if encrypted_task is not None else self.litellm_router_instance.acompletion diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index 1cfb311d796..390e499e7ce 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -90,7 +90,7 @@ class SearchAPIRouter: def get_matching_search_tools( router_instance: _SearchToolsRouter, search_tool_name: str, - ) -> list: + ) -> list[SearchToolTypedDict]: """ Get all search tools matching the given name. diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 716407d6def..072fabda905 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -6,7 +6,8 @@ import traceback from typing import Final import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -30,6 +31,14 @@ class _OidcTokenClaims(LiteLLMBaseModel): exp: float | None = None +@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True)) +class _GitHubOidcTokenResponse(TypedDict): + value: NotRequired[ReadOnly[str | None]] + + +_GITHUB_OIDC_TOKEN_RESPONSE: Final = TypeAdapter(_GitHubOidcTokenResponse) + + def _oidc_token_cache_ttl(oidc_token: str, max_ttl: int) -> int: """Cache TTL for a fetched OIDC token: ``max_ttl``, capped so the cache entry never outlives the token's own ``exp`` claim (minus a safety margin). @@ -271,7 +280,8 @@ def get_secret( }, ) if response.status_code == 200: - oidc_token = response.json().get("value", None) + response_data: Final = _GITHUB_OIDC_TOKEN_RESPONSE.validate_python(response.json()) + oidc_token = response_data.get("value", None) oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=300 - 5) return oidc_token else: