diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 22dd4170963..e78c4bd24c3 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -91,13 +91,13 @@ class SensitiveDataMasker: def _mask_sequence( self, - values: list[Any], + values: Sequence[object], depth: int, max_depth: int, excluded_keys: set[str] | None, key_is_sensitive: bool, - ) -> list[Any]: - masked_items: Final[list[Any]] = [] + ) -> Sequence[object]: + masked_items: Final[list[object]] = [] if depth >= max_depth: return values @@ -197,7 +197,7 @@ def _walk_payload(node: object, key_is_sensitive: bool, depth: int) -> object: return node -def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dict[str, Any]: +def mask_sensitive_keys(data: Mapping[str, object], sensitive_fields: set[str]) -> dict[str, object]: """Return a new dict with values masked for keys listed in ``sensitive_fields``. Unlike :meth:`SensitiveDataMasker.mask_dict`, this does exact key-name @@ -209,7 +209,7 @@ def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dic range and are replaced with a fixed-length all-mask string, so a short credential is never returned verbatim. """ - masked: Final[dict[str, Any]] = {} + masked: Final[dict[str, object]] = {} mask_char: Final = _default_masker.mask_char min_visible: Final = _default_masker.visible_prefix + _default_masker.visible_suffix for key, value in data.items(): diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index d188fac8704..bea77a6761c 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -27,7 +27,6 @@ without the optional STT extras installed. import asyncio import inspect from collections.abc import Callable, Iterable -from types import ModuleType from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.litellm_core_utils.audio_utils.utils import ( @@ -95,11 +94,37 @@ class _AudioEncoding(Protocol): def LINEAR_PCM(self) -> object: ... -def _auth_factory(riva_module: ModuleType) -> Callable[..., _RivaAuth]: +class _RivaClientModule(Protocol): + """The ``riva.client`` entry points this handler calls.""" + + @property + def Auth(self) -> Callable[..., _RivaAuth]: ... + + @property + def ASRService(self) -> Callable[[_RivaAuth], _AsrService]: ... + + +class _RivaAsrModule(Protocol): + """The protobuf constructors this handler calls, from whichever module exposes them.""" + + @property + def AudioEncoding(self) -> _AudioEncoding: ... + + @property + def RecognitionConfig(self) -> Callable[..., _RecognitionConfig]: ... + + @property + def StreamingRecognitionConfig(self) -> Callable[..., _StreamingRecognitionConfig]: ... + + @property + def EndpointingConfig(self) -> Callable[..., _EndpointingConfig]: ... + + +def _auth_factory(riva_module: _RivaClientModule) -> Callable[..., _RivaAuth]: return riva_module.Auth -def _audio_encoding(riva_asr_module: ModuleType) -> _AudioEncoding: +def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding: return riva_asr_module.AudioEncoding @@ -317,7 +342,7 @@ class NvidiaRivaAudioTranscription: def _construct_auth( self, - riva_module: ModuleType, + riva_module: _RivaClientModule, api_base: str, api_key: str | None, optional_params: dict, @@ -349,7 +374,7 @@ class NvidiaRivaAudioTranscription: return _auth_factory(riva_module)(None, use_ssl, api_base, metadata) def _build_recognition_config_proto( - self, riva_asr_module: ModuleType, recognition_config_dict: dict[str, Any] + self, riva_asr_module: _RivaAsrModule, recognition_config_dict: dict[str, Any] ) -> _RecognitionConfig: encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper() encoding_enum: Final[object] = getattr( @@ -436,7 +461,7 @@ class NvidiaRivaAudioTranscription: return final_results -def _import_riva() -> tuple[ModuleType, ModuleType]: +def _import_riva() -> tuple[_RivaClientModule, _RivaAsrModule]: """ Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``. diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 58cf7c7e702..67b01c2dc43 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: import tiktoken from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator class VertexGemmaConfig(OpenAIGPTConfig): @@ -56,7 +57,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): self, model_response: ModelResponse, stream: bool, - ) -> ModelResponse | Any: + ) -> "ModelResponse | MockResponseIterator": """ Helper method to return fake stream iterator if streaming is requested. @@ -138,7 +139,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client: HTTPHandler | httpx.Client | None, api_base: str, headers: dict[str, str], # mutable-ok: forwarded to post(headers: dict | None) - request_data: dict[str, Any], # mutable-ok: forwarded to post(json: dict | ...) + request_data: dict[str, object], # mutable-ok: forwarded to post(json: dict | ...) timeout: float | httpx.Timeout | None, ) -> httpx.Response: if isinstance(client, HTTPHandler): @@ -173,7 +174,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client: AsyncHTTPHandler | httpx.AsyncClient | None, api_base: str, headers: dict[str, str], # mutable-ok: forwarded to post(headers: dict | None) - request_data: dict[str, Any], # mutable-ok: forwarded to post(json: dict | ...) + request_data: dict[str, object], # mutable-ok: forwarded to post(json: dict | ...) timeout: float | httpx.Timeout | None, ) -> httpx.Response: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 16f58ef5b76..4cf63b1fceb 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -453,7 +453,7 @@ def _raise_for_upstream_failure( if response.status_code == 401 and relays_upstream_auth: raise MCPUpstreamAuthError( status_code=response.status_code, - www_authenticate=response.headers.get("www-authenticate"), + www_authenticate=dict(response.headers).get("www-authenticate"), server_name=upstream, ) raise MCPOpenApiUpstreamError(response.status_code, upstream) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index f7b92df5ba3..5503d19211b 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -18,6 +18,7 @@ TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against st from __future__ import annotations import json +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, Protocol @@ -29,6 +30,8 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS if TYPE_CHECKING: + from prisma.models import LiteLLM_SSOIdentityAssertion + from litellm.proxy.utils import PrismaClient _ASSERTION_DECRYPT_LOG_KEY: Final = "sso_identity_assertion" @@ -36,6 +39,34 @@ _STR_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) _MAYBE_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) +class _SSOAssertionTable(Protocol): + """The ``LiteLLM_SSOIdentityAssertion`` table operations this store calls.""" + + async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_SSOIdentityAssertion | None: ... + + async def find_many(self) -> Sequence[LiteLLM_SSOIdentityAssertion]: ... + + async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> object: ... + + async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> object: ... + + +class _MCPServerTable(Protocol): + """The ``LiteLLM_MCPServerTable`` lookup the retention gate calls.""" + + async def find_first(self, *, where: Mapping[str, str]) -> object | None: ... + + +def _assertion_table(prisma_client: PrismaClient) -> _SSOAssertionTable: + """The SSO assertion table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_ssoidentityassertion + + +def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: + """The MCP server table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpservertable + + class SSOIdentityAssertion(BaseModel): """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token, ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login.""" @@ -163,9 +194,7 @@ async def ema_assertion_retention_enabled() -> bool: return True if prisma_client is None: return False - row: Final = await prisma_client.db.litellm_mcpservertable.find_first( - where={"auth_type": MCPAuth.oauth2_id_jag.value} - ) + row: Final = await _mcp_server_table(prisma_client).find_first(where={"auth_type": MCPAuth.oauth2_id_jag.value}) return row is not None @@ -184,7 +213,7 @@ async def persist_sso_identity_assertion( **({"expires_at": assertion.expires_at.isoformat()} if assertion.expires_at else {}), } encoded: Final = _STR_ADAPTER.validate_python(encrypt_value_helper(json.dumps(payload))) - await prisma_client.db.litellm_ssoidentityassertion.upsert( + await _assertion_table(prisma_client).upsert( where={"user_id": user_id}, data={ "create": {"user_id": user_id, "assertion_b64": encoded}, @@ -200,7 +229,7 @@ async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None: if prisma_client is None: return None - row: Final = await prisma_client.db.litellm_ssoidentityassertion.find_unique(where={"user_id": user_id}) + row: Final = await _assertion_table(prisma_client).find_unique(where={"user_id": user_id}) if row is None: return None raw: Final = _MAYBE_STR_ADAPTER.validate_python( @@ -310,13 +339,13 @@ async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, re_encrypted: Final = _STR_ADAPTER.validate_python( encrypt_value_helper(plaintext, new_encryption_key=new_master_key) ) - await prisma_client.db.litellm_ssoidentityassertion.update( + await _assertion_table(prisma_client).update( where={"user_id": row.user_id}, data={"assertion_b64": re_encrypted}, ) return True - rows: Final = await prisma_client.db.litellm_ssoidentityassertion.find_many() + rows: Final = await _assertion_table(prisma_client).find_many() outcomes: Final = [await _rotate_row(row) for row in rows] verbose_proxy_logger.info( "rotate_sso_identity_assertions_master_key: rotated %d row(s), skipped %d", diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 50e0a961a49..17743e153f2 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -13,7 +13,7 @@ from collections.abc import Set as AbstractSet from dataclasses import dataclass, field from typing import TYPE_CHECKING, Final -from starlette.types import Receive, Scope, Send +from starlette.types import ASGIApp, Receive, Scope, Send from litellm._logging import verbose_proxy_logger @@ -266,7 +266,7 @@ class LazyFeatureMiddleware: def __init__( self, - app, + app: ASGIApp, fastapi_app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES, ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index b31ed4b0f4a..c69b24c0553 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -10,6 +10,7 @@ import os from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from typing_extensions import ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException @@ -33,6 +34,12 @@ BLOCKED_BY_OVALIX_FALLBACK_MESSAGE: Final = "This message was blocked by Ovalix" BLOCKED_ACTION_TYPE: Final = "block" +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + supported_event_hooks: ReadOnly[list[GuardrailEventHooks]] + + class OvalixGuardrailMissingSecrets(Exception): """Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing.""" @@ -80,7 +87,7 @@ class OvalixGuardrail(CustomGuardrail): application_id: str | None = None, pre_checkpoint_id: str | None = None, post_checkpoint_id: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ): self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE") self._tracker_api_key = tracker_api_key or os.environ.get("OVALIX_TRACKER_API_KEY") @@ -88,10 +95,9 @@ class OvalixGuardrail(CustomGuardrail): self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get("OVALIX_PRE_CHECKPOINT_ID") self._post_checkpoint_id = post_checkpoint_id or os.environ.get("OVALIX_POST_CHECKPOINT_ID") - if "supported_event_hooks" not in kwargs: - kwargs["supported_event_hooks"] = [] + supported_event_hooks: Final = kwargs.get("supported_event_hooks", []) - self._validate_config(kwargs["supported_event_hooks"]) + self._validate_config(supported_event_hooks) self._tracker_headers = httpx.Headers( { @@ -103,7 +109,8 @@ class OvalixGuardrail(CustomGuardrail): self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) - super().__init__(**kwargs) + forwarded: Final[_CustomGuardrailOptions] = {**kwargs, "supported_event_hooks": supported_event_hooks} + super().__init__(**forwarded) verbose_proxy_logger.debug( "Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s", self._tracker_api_base, diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index fe658a13c24..b7fcd4ac7a9 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -473,7 +473,7 @@ async def delete_team_callback( raise _callback_error(404, f"callback_name = {callback_name} is not registered for team_id = {team_id}.") updated_metadata: Final = {**team_metadata, "logging": remaining_callbacks} # mutable-ok: persisted as JSON - encrypted_metadata: Final = encrypt_callback_vars(updated_metadata) + encrypted_metadata: Final[object] = encrypt_callback_vars(updated_metadata) team_metadata_json: Final = json.dumps(encrypted_metadata) updated_team: Final = await TeamRepository(prisma_client).table.update( @@ -610,8 +610,8 @@ async def disable_team_logging( # _get_dynamic_logging_metadata stops at metadata["logging"], where the API # and Admin UI register callbacks, without ever reading callback_settings. team_metadata["logging"] = [] # mutable-ok: the disabled state is persisted as an empty JSON array - team_metadata = encrypt_callback_vars(team_metadata) - team_metadata_json: Final = json.dumps(team_metadata) + encrypted_metadata: Final[object] = encrypt_callback_vars(team_metadata) + team_metadata_json: Final = json.dumps(encrypted_metadata) # Update team in database updated_team: Final = await TeamRepository(prisma_client).table.update( @@ -643,7 +643,7 @@ async def disable_team_logging( await _emit_team_callback_audit_log( team_id=team_id, before_metadata=before_metadata, - after_metadata=team_metadata, + after_metadata=encrypted_metadata, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, ) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 0ca2c4c8865..2fb6813a471 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -88,7 +88,7 @@ def _redact_sensitive_litellm_params(litellm_params: object, _depth: int = 0) -> return None if isinstance(litellm_params, str): try: - parsed: Final = json.loads(litellm_params) + parsed: Final[object] = json.loads(litellm_params) except (TypeError, ValueError): return REDACTED_BY_LITELM_STRING return json.dumps(_redact_sensitive_litellm_params(parsed, _depth + 1)) @@ -589,7 +589,8 @@ async def update_vector_store( try: update_data: Final = data.model_dump(exclude_unset=True) - vector_store_id: Final[str] = update_data.pop("vector_store_id") + vector_store_id: Final[str] = data.vector_store_id + update_data.pop("vector_store_id") # Per-store access control: anyone authenticated who passes the # premium-feature gate could otherwise update *any* vector store — diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 73a0159fc9f..b81c2cc0ebe 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -277,7 +277,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): raise Exception(error_msg) verbose_logger.debug("Initiate resumable upload response: %s", response.headers) # Extract upload URL from response headers - upload_url: Final = response.headers.get("x-goog-upload-url") + upload_url: Final = dict(response.headers).get("x-goog-upload-url") if not upload_url: raise Exception("No upload URL returned in response headers") diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index e86c8e7c919..9f5bf783958 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -16,7 +16,7 @@ Requires: import json import os -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -35,6 +35,9 @@ from litellm.types.secret_managers.main import KeyManagementSettings from .base_secret_manager import BaseSecretManager +if TYPE_CHECKING: + from botocore.awsrequest import HTTPHeaders + class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): def __init__( @@ -530,7 +533,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): secret_value: str | None = None, optional_params: dict | None = None, request_data: dict | None = None, - ) -> tuple[str, Any, bytes]: + ) -> tuple[str, "HTTPHeaders", bytes]: """Prepare the AWS Secrets Manager request""" try: from botocore.auth import SigV4Auth