From bf9437780b7fd0aa74827239cfe9ea2e19bbacc7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:09:22 +0000 Subject: [PATCH] refactor(types): replace Any with real types across 11 more backend files Final batch of the fifth basedpyright Any reduction round. Every change is typing-only and leaves runtime behavior identical. These are the densest remaining files, so the yield per file is small and most of the batch was left alone deliberately. The Riva transcription handler describes the SDK module attributes it reads with Protocols instead of a bare ModuleType, the AWS secret manager stops hiding a botocore header object behind Any, and the sensitive data masker, MCP SSO assertion store and Ovalix guardrail move payload and option annotations to object and Mapping[str, object]. --- .../sensitive_data_masker.py | 10 ++--- .../audio_transcription/handler.py | 37 +++++++++++++--- .../vertex_gemma_models/transformation.py | 7 +-- .../mcp_server/openapi_to_mcp_generator.py | 2 +- .../sso_assertion_store.py | 43 ++++++++++++++++--- litellm/proxy/_lazy_features.py | 4 +- .../guardrail_hooks/ovalix/ovalix.py | 17 +++++--- .../team_callback_endpoints.py | 8 ++-- .../management_endpoints.py | 5 ++- litellm/rag/ingestion/gemini_ingestion.py | 2 +- .../secret_managers/aws_secret_manager_v2.py | 7 ++- 11 files changed, 104 insertions(+), 38 deletions(-) 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