From f747bc67048f2b38e6e99bd57ca284fe6f77bdb5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:14:07 +0000 Subject: [PATCH] refactor(types): replace Any with real types across 29 more backend files Second batch of the fifth basedpyright Any reduction round. Every change is typing-only and leaves runtime behavior identical. Guardrail hooks and file, vector store and usage endpoints move their payload, header and response annotations from Any to object, Mapping[str, object] or the concrete response model the call site already produces. Two private aggregation helpers in the guardrail usage endpoints take a key accessor function instead of an attribute name string, so the key they read is checked against the row type. The verification token repository reaches its two tables through Protocols that name the handles it calls, rather than reading them off an untyped prisma client, and the Azure AD credential wrapper describes the azure-identity credential it wraps the same way. --- .../litellm_core_utils/llm_cost_calc/utils.py | 15 +++-- .../azure/text_to_speech/transformation.py | 9 +-- .../image_edit/stability_transformation.py | 4 +- .../responses/transformation.py | 14 ++--- .../llms/cohere/embed/v1_transformation.py | 15 +++-- litellm/llms/gdc/chat/transformation.py | 22 ++++++- .../llama3/transformation.py | 4 +- litellm/llms/voyage/rerank/transformation.py | 6 +- litellm/proxy/client/users.py | 3 +- .../proxy/common_utils/http_parsing_utils.py | 14 +++-- litellm/proxy/db/prisma_client.py | 9 ++- .../block_code_execution.py | 9 ++- .../cato_networks/cato_networks.py | 10 ++-- .../guardrail_hooks/compresr/compresr.py | 11 ++-- .../llm_as_a_judge/__init__.py | 23 +++++++- .../guardrails/guardrail_hooks/noma/noma.py | 15 +++-- .../panw_prisma_airs/panw_prisma_airs.py | 4 +- litellm/proxy/guardrails/usage_endpoints.py | 25 ++++---- .../proxy/hooks/proxy_track_cost_callback.py | 28 +++++++-- ...model_access_group_management_endpoints.py | 25 ++++---- .../usage_endpoints/ai_usage_chat.py | 11 ++-- .../openai_files_endpoints/files_endpoints.py | 4 +- .../storage_backend_service.py | 9 ++- .../proxy/vector_store_endpoints/endpoints.py | 2 +- .../vector_store_files_endpoints/endpoints.py | 10 ++-- litellm/proxy_auth/credentials.py | 58 ++++++++++++++----- litellm/rag/rag_query.py | 4 +- .../verification_token_repository.py | 46 ++++++++++++--- litellm/router_utils/search_api_router.py | 4 +- 29 files changed, 274 insertions(+), 139 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 68dc27ec25e..52dac92ee22 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone, tzinfo from types import MappingProxyType -from typing import Any, Final, Literal, TypedDict, cast +from typing import Final, Literal, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import litellm @@ -89,7 +89,7 @@ def _requested_image_size(optional_params: Mapping[str, object] | None) -> str | return value if value is not None and _IMAGE_SIZE_PATTERN.fullmatch(value) else None -def get_web_search_requests(server_tool_use: Any) -> int | None: +def get_web_search_requests(server_tool_use: object) -> int | None: """ Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, @@ -1494,7 +1494,7 @@ def calculate_image_response_cost_from_usage( if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0: return None - input_tokens_details: Final = getattr(usage, "input_tokens_details", None) + input_tokens_details: Final[object] = getattr(usage, "input_tokens_details", None) prompt_tokens_details: PromptTokensDetailsWrapper | None = None if input_tokens_details is not None: # input_tokens_details may be a dict (e.g. OpenAI image edit responses) @@ -1507,9 +1507,12 @@ def calculate_image_response_cost_from_usage( cached_tokens=0, ) - output_tokens_details = getattr(usage, "completion_tokens_details", None) - if output_tokens_details is None: - output_tokens_details = getattr(usage, "output_tokens_details", None) + completion_tokens_details_attr: Final[object] = getattr(usage, "completion_tokens_details", None) + output_tokens_details: Final[object] = ( + getattr(usage, "output_tokens_details", None) + if completion_tokens_details_attr is None + else completion_tokens_details_attr + ) if output_tokens_details is None: completion_tokens_details = CompletionTokensDetailsWrapper( diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index d8ccf26ce60..eed7a3178ca 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -19,6 +19,7 @@ from litellm.secret_managers.main import get_secret_str if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.openai import HttpxBinaryResponseContent else: LiteLLMLoggingObj = Any @@ -67,15 +68,15 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): litellm_params_dict: dict, logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout, - extra_headers: dict[str, Any] | None, - base_llm_http_handler: Any, + extra_headers: dict[str, object] | None, + base_llm_http_handler: "BaseLLMHTTPHandler", aspeech: bool, api_base: str | None, api_key: str | None, - **kwargs: Any, + **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[object, object, "HttpxBinaryResponseContent"], ]: """ Dispatch method to handle Azure AVA TTS requests diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index bc9a64f587a..01e25f4671e 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -125,7 +125,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): } # Create a copy to not mutate original - convert TypedDict to regular dict - mapped_params: Final[dict[str, Any]] = dict(image_edit_optional_params) + mapped_params: Final[dict[str, object]] = dict(image_edit_optional_params) for k, v in image_edit_optional_params.items(): if k in param_mapping: @@ -172,7 +172,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): Returns the request body dict that will be JSON-encoded by the handler. """ # Build Bedrock Stability request - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "output_format": "png", # Default to PNG } diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index bbbda4d14b6..3d2eab8fcee 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -16,8 +16,8 @@ BaseAWSLLM._sign_request after the request body is finalized. """ import json -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final import httpx from typing_extensions import ReadOnly, TypedDict @@ -142,9 +142,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI return False @staticmethod - def _filter_unsupported_tools(tools: list[Any]) -> list[Any]: + def _filter_unsupported_tools(tools: "Sequence[object]") -> "list[object]": """Keep only tool types Mantle's Responses API accepts.""" - kept: Final[list[Any]] = [] + kept: Final[list[object]] = [] dropped_types: Final[list[str]] = [] for tool in tools: if not isinstance(tool, dict): @@ -217,11 +217,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) @staticmethod - def _is_codex_additional_tools_item(item: Any) -> bool: + def _is_codex_additional_tools_item(item: object) -> bool: return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE @staticmethod - def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]": + def _tools_of_additional_tools_item(item: "Mapping[str, object]") -> "list[object]": tools: Final = item.get("tools") return tools if isinstance(tools, list) else [] @@ -229,7 +229,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def _hoist_codex_additional_tools( cls, input: "str | ResponseInputParam", - ) -> "tuple[str | ResponseInputParam, list[Any]]": + ) -> "tuple[str | ResponseInputParam, list[object]]": """Codex's "responses lite" wire mode ships tool definitions inside `input` as {"type": "additional_tools", "role": "developer", "tools": [...]} items. api.openai.com accepts that item type; Mantle diff --git a/litellm/llms/cohere/embed/v1_transformation.py b/litellm/llms/cohere/embed/v1_transformation.py index ee40464362d..b35fae5a1ac 100644 --- a/litellm/llms/cohere/embed/v1_transformation.py +++ b/litellm/llms/cohere/embed/v1_transformation.py @@ -2,7 +2,8 @@ Legacy /v1/embedding transformation logic for Bedrock Cohere. """ -from typing import Any, Final +from collections.abc import Sized +from typing import Final, Protocol import httpx @@ -16,6 +17,12 @@ from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, U from litellm.utils import is_base64_encoded +class _SupportsEncode(Protocol): + """Tokenizer handle: the embedding usage path only encodes text to measure its token length.""" + + def encode(self, text: str, /) -> Sized: ... + + class CohereEmbeddingConfig: """ Reference: https://docs.cohere.com/v2/reference/embed @@ -61,7 +68,7 @@ class CohereEmbeddingConfig: return transformed_request - def _calculate_usage(self, input: list[str], encoding: Any, meta: dict) -> Usage: + def _calculate_usage(self, input: list[str], encoding: _SupportsEncode, meta: dict) -> Usage: input_tokens = 0 text_tokens: Final[int | None] = meta.get("billed_units", {}).get("input_tokens") @@ -97,7 +104,7 @@ class CohereEmbeddingConfig: data: dict | CohereEmbeddingRequest, model_response: EmbeddingResponse, model: str, - encoding: Any, + encoding: _SupportsEncode, input: list, ) -> EmbeddingResponse: response_json: Final = response.json() @@ -121,7 +128,7 @@ class CohereEmbeddingConfig: response_json: dict, model_response: EmbeddingResponse, model: str, - encoding: Any, + encoding: _SupportsEncode, input: list, ) -> EmbeddingResponse: """ diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 03037512551..6eac3ac79cd 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -7,14 +7,32 @@ import os import re import threading from collections.abc import Callable -from typing import Any, Final, Protocol +from typing import Final, Protocol from urllib.parse import urlsplit +from typing_extensions import ReadOnly, TypedDict, Unpack + import litellm from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues +class _OpenAIGPTConfigOptions(TypedDict, total=False): + """The sampling defaults ``OpenAIGPTConfig.__init__`` accepts and stashes on the class.""" + + frequency_penalty: ReadOnly[int | None] + function_call: ReadOnly[str | dict[str, object] | None] + functions: ReadOnly[list[object] | None] + logit_bias: ReadOnly[dict[str, object] | None] + max_tokens: ReadOnly[int | None] + n: ReadOnly[int | None] + presence_penalty: ReadOnly[int | None] + stop: ReadOnly[str | list[object] | None] + temperature: ReadOnly[int | None] + top_p: ReadOnly[int | None] + response_format: ReadOnly[dict[str, object] | None] + + class _GDCHAudienceCredentials(Protocol): """A GDCH service account credential already bound to an audience, ready to mint a bearer token.""" @@ -32,7 +50,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): _GDCH_CREDENTIAL_TYPE: Final[str] = "gdch_service_account" _PATH_ID_PATTERN: Final[re.Pattern[str]] = re.compile(r"^[a-zA-Z0-9_-]+$") - def __init__(self, **kwargs: Any) -> None: + def __init__(self, **kwargs: Unpack[_OpenAIGPTConfigOptions]) -> None: super().__init__(**kwargs) self._creds_lock = threading.Lock() self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {} diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index 279035c455d..89a5b8a570e 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -1,6 +1,6 @@ import types from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -95,7 +95,7 @@ class VertexAILlama3Config(OpenAIGPTConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "VertexAILlama3StreamingHandler": return VertexAILlama3StreamingHandler( streaming_response=streaming_response, sync_stream=sync_stream, diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index fea8452d934..6f8d024f0b1 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -4,8 +4,8 @@ Transformation logic for Voyage AI's /v1/rerank endpoint. Docs - https://docs.voyageai.com/docs/reranker """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final import httpx @@ -33,7 +33,7 @@ class VoyageRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: Sequence[str | Mapping[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index 3f11fe94043..503c92228a8 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final import requests @@ -50,7 +51,7 @@ class UsersManagementClient: response.raise_for_status() return response.json() - def create_user(self, user_data: dict[str, Any]) -> dict[str, Any]: + def create_user(self, user_data: Mapping[str, object]) -> dict[str, Any]: """Create a new user (POST /user/new)""" url: Final = f"{self.base_url}/user/new" response: Final = requests.post(url, headers=self._get_headers(), json=user_data, timeout=self.timeout) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 54a0f18fd63..2bae3e946f7 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -52,14 +52,18 @@ def _unqualified(annotation: object) -> object: return _unqualified(qualified[0]) +def _union_members(annotation: object) -> tuple[object, ...]: + """The non-``None`` members of a union annotation, or the annotation itself when it is not a union.""" + if get_origin(annotation) not in (Union, UnionType): + return (annotation,) + members: Final[tuple[object, ...]] = get_args(annotation) + return tuple(arg for arg in members if arg is not type(None)) + + def _numeric_form_type(annotation: object) -> type[int] | type[float] | None: """The scalar to parse an ``int``/``float``-typed field as, else ``None``.""" unwrapped: Final = _unqualified(annotation) - candidates: Final = ( - tuple(arg for arg in get_args(unwrapped) if arg is not type(None)) - if get_origin(unwrapped) in (Union, UnionType) - else (unwrapped,) - ) + candidates: Final = _union_members(unwrapped) if len(candidates) != 1: return None if candidates[0] is int: diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 2190ae55fd2..2f2ebfdf2bb 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -13,7 +13,7 @@ import urllib import urllib.parse from collections.abc import Callable from datetime import datetime, timedelta -from typing import Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_proxy_logger from litellm.proxy.db.token_auth import ( @@ -27,6 +27,9 @@ from litellm.proxy.db.token_auth import ( ) from litellm.secret_managers.main import str_to_bool +if TYPE_CHECKING: + from prisma import Prisma + __all__ = ( "IAMEndpoint", "PrismaManager", @@ -242,7 +245,7 @@ class PrismaWrapper: def _write_engine(prisma_client: _PrismaClient, engine: _PrismaEngine) -> None: prisma_client._Prisma__engine = engine - def _instrument_prisma_client(self, prisma_client: _PrismaClient) -> _PrismaDrainTracker | None: + def _instrument_prisma_client(self, prisma_client: "Prisma | _PrismaClient") -> _PrismaDrainTracker | None: from prisma.errors import ClientNotConnectedError try: @@ -255,7 +258,7 @@ class PrismaWrapper: self._write_engine(prisma_client, _TrackedPrismaEngine(engine, tracker)) return tracker - def _get_engine_pid(self, prisma_client: _PrismaClient | None = None) -> int: + def _get_engine_pid(self, prisma_client: "Prisma | _PrismaClient | None" = None) -> int: """Get the PID of the current Prisma engine subprocess, or 0 if unavailable. Must never raise: it runs inside the reconnect path, where the client diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py index bf2aa1f76e0..2b697671eda 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -8,9 +8,10 @@ confidence scoring and a tunable threshold (only block when confidence >= thresh import re from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, cast from fastapi import HTTPException +from typing_extensions import TypedDict, Unpack from litellm.integrations.custom_guardrail import ( CustomGuardrail, @@ -314,6 +315,10 @@ def _confidence_for_block( return 0.0 +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + class BlockCodeExecutionGuardrail(CustomGuardrail): """ Guardrail that detects fenced code blocks (markdown ```) and blocks or masks them @@ -332,7 +337,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): detect_execution_intent: bool = True, event_hook: Literal["pre_call", "post_call", "during_call"] | list[str] | None = None, default_on: bool = False, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: # Normalize to type expected by CustomGuardrail _event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None 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 176c308eda6..2d203c31974 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -264,7 +264,7 @@ class CatoNetworksGuardrail(CustomGuardrail): stack.extend(reversed(node)) @classmethod - def _extra_inspection_sources(cls, data: Mapping[str, Any]) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]: + def _extra_inspection_sources(cls, data: Mapping[str, object]) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]: """Text the proxy forwards to the model outside chat ``messages``: Responses-API ``input`` and ``instructions``, legacy completion ``prompt`` and tool/function/``response_format`` schema strings. Returned @@ -336,7 +336,7 @@ class CatoNetworksGuardrail(CustomGuardrail): ) raise HTTPException(status_code=400, detail=detection_message) - def _anonymize_request(self, res: Any, data: dict) -> dict: + def _anonymize_request(self, res: _CatoAnalyzeResponse, data: dict) -> dict: verbose_proxy_logger.info("Cato: anonymize action") redacted_chat: Final = res.get("redacted_chat") if not redacted_chat: @@ -379,7 +379,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return data @classmethod - def _apply_extra_redaction(cls, data: dict, field: str, redacted: list) -> bool: + def _apply_extra_redaction(cls, data: dict, field: str, redacted: Sequence[Mapping[str, object]]) -> bool: if field == "input": input_only: Final = {"input": data["input"]} if not redacted: @@ -400,7 +400,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return True @classmethod - def _apply_schema_string_redaction(cls, data: dict, redacted: list) -> None: + def _apply_schema_string_redaction(cls, data: dict, redacted: Sequence[Mapping[str, object]]) -> None: redactions: Final = iter(redacted) for container, key in cls._iter_schema_string_refs(data): replacement = next(redactions, None) @@ -408,7 +408,7 @@ class CatoNetworksGuardrail(CustomGuardrail): container[key] = replacement["content"] @staticmethod - def _apply_prompt_redaction(data: dict, redacted: list) -> None: + def _apply_prompt_redaction(data: dict, redacted: Sequence[Mapping[str, object]]) -> None: contents: Final = [m.get("content") for m in redacted if isinstance(m, dict)] prompt: Final = data.get("prompt") if isinstance(prompt, str): diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 93d859066b0..1ecdb1b0f63 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -22,7 +22,7 @@ import json import time from collections import Counter, OrderedDict from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Final, Literal, TypeGuard +from typing import TYPE_CHECKING, Final, Literal, TypeGuard from urllib.parse import urlparse import httpx @@ -64,6 +64,9 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) from litellm.types.proxy.guardrails.guardrail_hooks.base import ( GuardrailConfigModel, ) @@ -1049,7 +1052,7 @@ class CompresrGuardrail(CustomGuardrail): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -1069,8 +1072,8 @@ class CompresrGuardrail(CustomGuardrail): tools: dict, model: str, messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None, anthropic_messages_optional_request_params: dict, logging_obj: LiteLLMLoggingObj | None, stream: bool, diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index 172b1440ca3..806ab5161e8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -2,10 +2,10 @@ from collections.abc import Callable, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar +from typing import TYPE_CHECKING, Final, Generic, Literal, Optional, TypeVar from fastapi import HTTPException -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -106,6 +106,23 @@ def _build_judge_prompt( ) +class _CustomGuardrailOptions(TypedDict, total=False): + """The ``CustomGuardrail`` options this guardrail accepts and forwards untouched.""" + + mask_request_content: ReadOnly[bool] + mask_response_content: ReadOnly[bool] + violation_message_template: ReadOnly[str | None] + end_session_after_n_fails: ReadOnly[int | None] + on_violation: ReadOnly[str | None] + realtime_violation_message: ReadOnly[str | None] + on_sensitive_data: ReadOnly[str | None] + sensitive_data_route_to_model: ReadOnly[str | None] + sticky_session_routing: ReadOnly[bool] + run_in_parallel: ReadOnly[bool] + scan_raw_request: ReadOnly[bool] + only_scan_new_messages: ReadOnly[bool] + + class LLMAsAJudgeGuardrail(CustomGuardrail): """Post-call guardrail that judges response quality via an LLM.""" @@ -119,7 +136,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, default_on: bool = False, router_provider: "Callable[[], Router | None] | None" = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: _event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None if event_hook is not None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 7ef0a9f73f3..edd78e0bbc6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -9,13 +9,13 @@ import asyncio import json import os import warnings -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterable from datetime import datetime from typing import ( TYPE_CHECKING, - Any, Final, Literal, + TypeVar, ) from urllib.parse import urljoin @@ -39,9 +39,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( CallTypes, CallTypesLiteral, - EmbeddingResponse, GuardrailStatus, - ImageResponse, ModelResponseStream, TextCompletionResponse, ) @@ -53,7 +51,8 @@ SENSITIVE_DATA_DETECTOR_KEYS: Final[list[str]] = ["sensitiveData", "dataDetector # Type aliases MessageRole = Literal["user", "assistant"] -LLMResponse = Any | ModelResponse | EmbeddingResponse | ImageResponse +LLMResponse = object +_LLMResponseT: Final = TypeVar("_LLMResponseT") _LEGACY_NOMA_DEPRECATION_WARNED = False if TYPE_CHECKING: @@ -709,10 +708,10 @@ class NomaGuardrail(CustomGuardrail): async def _check_llm_response( self, request_data: dict, - response: LLMResponse, + response: _LLMResponseT, user_auth: UserAPIKeyAuth, event_type: GuardrailEventHooks | None = None, - ) -> Any: + ) -> _LLMResponseT: """Check LLM response for policy violations""" content: Final = await self._process_llm_response_check(request_data, response, user_auth, event_type) if not content: @@ -798,7 +797,7 @@ class NomaGuardrail(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[ModelResponseStream], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: """Process streaming response chunks with Noma guardrail.""" 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 b73d3adb99e..86ad2f9db5f 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 @@ -793,7 +793,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): }, ) - def _prepare_metadata_from_request(self, data: dict[str, Any]) -> dict[str, Any]: + def _prepare_metadata_from_request(self, data: dict[str, Any]) -> dict[str, object]: """ Extract and prepare metadata from request data for PANW API call. @@ -809,7 +809,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): """ user_metadata: Final = data.get("metadata", {}) or {} requester_meta: Final = user_metadata.get("requester_metadata", {}) or {} - metadata: Final = { + metadata: Final[dict[str, object]] = { "user": data.get("user") or "litellm_user", "model": data.get("model") or "unknown", } diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 6259efb6654..3f7c36bbbf2 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -45,6 +45,7 @@ _EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({}) _ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2}) _T = TypeVar("_T") +_MetricsRowT = TypeVar("_MetricsRowT", bound="_DailyMetricsRow") _USAGE_MAX_RANGE_DAYS: Final = 366 @@ -360,10 +361,12 @@ def _trend_from_comparison(current_fail: float, previous_fail: float) -> str: return "stable" -def _aggregate_daily_metrics(metrics: "Sequence[_DailyMetricsRow]", id_attr: str) -> Mapping[str, _MetricTotals]: +def _aggregate_daily_metrics( + metrics: "Sequence[_MetricsRowT]", id_of: "Callable[[_MetricsRowT], str]" +) -> Mapping[str, _MetricTotals]: agg: Final[dict[str, _MetricTotals]] = {} for m in metrics: - gid: str = getattr(m, id_attr) + gid: str = id_of(m) if gid not in agg: agg[gid] = {"requests": 0, "passed": 0, "blocked": 0, "flagged": 0} agg[gid]["requests"] += int(m.requests_evaluated or 0) @@ -373,10 +376,12 @@ def _aggregate_daily_metrics(metrics: "Sequence[_DailyMetricsRow]", id_attr: str return agg -def _prev_fail_rates(metrics_prev: "Sequence[_DailyMetricsRow]", id_attr: str) -> Mapping[str, float]: +def _prev_fail_rates( + metrics_prev: "Sequence[_MetricsRowT]", id_of: "Callable[[_MetricsRowT], str]" +) -> Mapping[str, float]: prev_agg_raw: Final[dict[str, _PrevPeriodCounts]] = {} for m in metrics_prev: - gid: str = getattr(m, id_attr) + gid: str = id_of(m) r, b = int(m.requests_evaluated or 0), int(m.blocked_count or 0) if gid not in prev_agg_raw: prev_agg_raw[gid] = {"req": 0, "blocked": 0} @@ -429,7 +434,7 @@ def _field_str(mapping: Mapping[str, object], key: str, default: str) -> str: return str(mapping.get(key, default)) -def _get_guardrail_attrs(g: "_DbOrConfigGuardrail") -> tuple[Any, str]: +def _get_guardrail_attrs(g: "_DbOrConfigGuardrail") -> tuple[str | None, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" gid: Final = _get_guardrail_field(g, "guardrail_id") name: Final = _get_guardrail_field(g, "guardrail_name") @@ -592,8 +597,8 @@ async def guardrails_usage_overview( Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits] ] = await _find_daily_guardrail_usage_units(prisma_client, where=units_where) - agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id") - prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id") + agg: Final = _aggregate_daily_metrics(metrics, lambda m: m.guardrail_id) + prev_agg: Final = _prev_fail_rates(metrics_prev, lambda m: m.guardrail_id) units_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_counter_units) cost_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_tracked_cost) untracked_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_untracked_units) @@ -811,7 +816,7 @@ def _usage_log_entry_from_row( ) -def _snippet(text: Any, max_len: int = 200) -> str | None: +def _snippet(text: object, max_len: int = 200) -> str | None: if text is None: return None if isinstance(text, str): @@ -964,8 +969,8 @@ async def policies_usage_overview( } }, ) - agg: Final = _aggregate_daily_metrics(metrics, "policy_id") - prev_agg: Final = _prev_fail_rates(metrics_prev, "policy_id") + agg: Final = _aggregate_daily_metrics(metrics, lambda m: m.policy_id) + prev_agg: Final = _prev_fail_rates(metrics_prev, lambda m: m.policy_id) chart: Final = _chart_from_metrics(metrics) total_requests: Final = sum(a["requests"] for a in agg.values()) total_blocked: Final = sum(a["blocked"] for a in agg.values()) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c4fba8ecf9e..e6ada41f062 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -2,7 +2,7 @@ import asyncio import traceback from collections.abc import Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm from litellm._logging import verbose_proxy_logger @@ -578,18 +578,36 @@ def _get_request_tags_for_cost_tracking( return None +class _IncrementSpendCounters(Protocol): + """The ``increment_spend_counters`` coroutine :func:`_update_database_and_spend_counters` awaits.""" + + async def __call__( + self, + token: str | None, + team_id: str | None, + user_id: str | None, + response_cost: float | None, + org_id: str | None = None, + budget_reservation: dict[str, object] | None = None, + end_user_id: str | None = None, + tags: list[str] | None = None, + request_started_at: datetime | None = None, + model_access_groups: Sequence[str] | None = None, + ) -> None: ... + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", - increment_spend_counters: Any, + increment_spend_counters: _IncrementSpendCounters, user_api_key: str | None, user_id: str | None, end_user_id: str | None, team_id: str | None, org_id: str | None, kwargs: dict, - completion_response: litellm.ModelResponse | Any | None, - start_time: Any, - end_time: Any, + completion_response: object, + start_time: datetime | None, + end_time: datetime | None, response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index a48130a4f22..e960bdfe337 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -540,7 +540,7 @@ async def get_all_access_groups_from_db( deployments: Final = await ModelRepository(prisma_client).table.find_many() # Build access group map - access_group_map: Final[dict[str, dict[str, Any]]] = {} + model_names_by_group: Final[dict[str, list[str]]] = {} for deployment in deployments: model_info = deployment.model_info or {} @@ -550,25 +550,20 @@ async def get_all_access_groups_from_db( model_name = deployment.model_name for access_group in access_groups: - if access_group not in access_group_map: - access_group_map[access_group] = { - "model_names": set(), - "deployment_count": 0, - } + if access_group not in model_names_by_group: + model_names_by_group[access_group] = [] - access_group_map[access_group]["model_names"].add(model_name) - access_group_map[access_group]["deployment_count"] += 1 + model_names_by_group[access_group].append(model_name) # Convert to AccessGroupInfo objects - result: Final = {} - for access_group, data in access_group_map.items(): - result[access_group] = AccessGroupInfo( + return { + access_group: AccessGroupInfo( access_group=access_group, - model_names=sorted(list(data["model_names"])), - deployment_count=data["deployment_count"], + model_names=sorted(frozenset(model_names)), + deployment_count=len(model_names), ) - - return result + for access_group, model_names in model_names_by_group.items() + } @router.post( diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index da4ddbd0aac..1265da99d89 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -6,7 +6,7 @@ usage/spend data by querying the aggregated daily activity endpoints. import json from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import date -from typing import Any, Final, Literal, NamedTuple, Protocol, cast, overload +from typing import Final, Literal, NamedTuple, Protocol, cast, overload from typing_extensions import ReadOnly, TypedDict @@ -16,6 +16,7 @@ from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) +from litellm.types.utils import ChatCompletionMessageToolCall # --------------------------------------------------------------------------- # Constants @@ -489,19 +490,19 @@ async def _execute_tool_call( async def _process_tool_call( - tc: Any, + tc: ChatCompletionMessageToolCall, chat_messages: list[Mapping[str, object]], user_id: str | None, is_admin: bool, ) -> AsyncIterator[str]: """Execute a single tool call, yielding SSE events for status.""" - fn_name: Final[str] = tc.function.name + fn_name: Final = tc.function.name fn_args: Final[Mapping[str, str]] = json.loads(tc.function.arguments) allowed_names: Final = {t["function"]["name"] for t in get_tools_for_role(is_admin)} - handler: Final = TOOL_HANDLERS.get(fn_name) + handler: Final = TOOL_HANDLERS.get(fn_name) if fn_name is not None else None - if fn_name not in allowed_names or not handler: + if fn_name is None or fn_name not in allowed_names or not handler: chat_messages.append( { "role": "tool", diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index bf07f4748ef..4d3a397e519 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -485,8 +485,8 @@ async def create_file( # Parse expires_after if provided expires_after: FileExpiresAfter | None = None form_data_raw: Final = await request.form() - form_data_dict: Final[dict[str, Any]] = dict(form_data_raw) - extracted_litellm_metadata: Final[dict[str, Any] | None] = extract_nested_form_metadata( + form_data_dict: Final[Mapping[str, object]] = dict(form_data_raw) + extracted_litellm_metadata: Final[Mapping[str, object] | None] = extract_nested_form_metadata( form_data=form_data_dict, prefix="litellm_metadata[" ) expires_after_anchor: Final = form_data_raw.get("expires_after[anchor]") diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index e766f335071..0407499cbcc 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -7,7 +7,6 @@ storage backends (e.g., Azure Blob Storage) and managing associated metadata. import base64 import time -from collections.abc import Mapping from typing import Any, Final, cast from litellm._logging import verbose_proxy_logger @@ -17,7 +16,7 @@ from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.llms.openai import OpenAIFileObject, OpenAIFilesPurpose -from litellm.types.utils import SpecialEnums +from litellm.types.utils import ExtractedFileData, SpecialEnums class StorageBackendFileService: @@ -33,7 +32,7 @@ class StorageBackendFileService: @staticmethod async def upload_file_to_storage_backend( - file_data: Mapping[str, Any], + file_data: ExtractedFileData, target_storage: str, target_model_names: list[str], purpose: OpenAIFilesPurpose, @@ -163,7 +162,7 @@ class StorageBackendFileService: @staticmethod def _create_unified_file_id( - file_type: str, + file_type: str | None, target_model_names: list[str], file_id: str, ) -> str: @@ -193,7 +192,7 @@ class StorageBackendFileService: @staticmethod async def _store_in_managed_files( file_object: OpenAIFileObject, - file_data: Mapping[str, Any], + file_data: ExtractedFileData, target_model_names: list[str], target_storage: str, storage_url: str, diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 1feda0b0bb5..f21c294e5a2 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -244,7 +244,7 @@ async def vector_store_create( ) # Create vector store across multiple models - response: Final = await managed_vector_stores.acreate_vector_store( + response: Final[object] = await managed_vector_stores.acreate_vector_store( create_request=data, llm_router=llm_router, target_model_names_list=target_model_names_list, diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 957ed9fd0b9..3ddea288ab8 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -553,7 +553,7 @@ async def vector_store_file_create( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -756,7 +756,7 @@ async def vector_store_file_retrieve( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -863,7 +863,7 @@ async def vector_store_file_content( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -973,7 +973,7 @@ async def vector_store_file_update( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -1080,7 +1080,7 @@ async def vector_store_file_delete( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy_auth/credentials.py b/litellm/proxy_auth/credentials.py index a4e29241959..f8814954a7a 100644 --- a/litellm/proxy_auth/credentials.py +++ b/litellm/proxy_auth/credentials.py @@ -7,7 +7,7 @@ It follows the same TokenCredential protocol used by Azure SDK. import time from dataclasses import dataclass -from typing import Any, Final, Protocol, runtime_checkable +from typing import Final, Protocol, runtime_checkable @dataclass @@ -50,6 +50,22 @@ class TokenCredential(Protocol): ... +class _AzureAccessToken(Protocol): + """The two attributes :class:`AzureADCredential` reads off an azure-identity token.""" + + @property + def token(self) -> str: ... + + @property + def expires_on(self) -> int: ... + + +class _AzureTokenCredential(Protocol): + """The single method :class:`AzureADCredential` calls on the credential it wraps.""" + + def get_token(self, *scopes: str) -> _AzureAccessToken: ... + + class AzureADCredential: """ Wrapper for Azure Identity credentials. @@ -71,7 +87,7 @@ class AzureADCredential: cred = AzureADCredential(credential=azure_cred) """ - def __init__(self, credential: Any | None = None): + def __init__(self, credential: _AzureTokenCredential | None = None): """ Initialize with an optional Azure credential. @@ -79,7 +95,7 @@ class AzureADCredential: credential: An azure-identity credential object. If None, DefaultAzureCredential will be used on first token request. """ - self._credential: Any = credential + self._credential: _AzureTokenCredential | None = credential self._initialized = credential is not None def get_token(self, scope: str) -> AccessToken: @@ -95,20 +111,30 @@ class AzureADCredential: Raises: ImportError: If azure-identity is not installed. """ - if not self._initialized: - try: - from azure.identity import DefaultAzureCredential - - self._credential = DefaultAzureCredential() - self._initialized = True - except ImportError: - raise ImportError( - "azure-identity is required for AzureADCredential. Install it with: pip install azure-identity" - ) - - result: Final = self._credential.get_token(scope) + result: Final = self._resolve_credential().get_token(scope) return AccessToken(token=result.token, expires_on=result.expires_on) + def _resolve_credential(self) -> _AzureTokenCredential: + """Return the wrapped credential, building the Azure default chain on first use. + + Raises: + ImportError: If azure-identity is not installed. + """ + existing: Final = self._credential + if existing is not None: + return existing + try: + from azure.identity import DefaultAzureCredential + + created: Final = DefaultAzureCredential() + except ImportError: + raise ImportError( + "azure-identity is required for AzureADCredential. Install it with: pip install azure-identity" + ) + self._credential = created + self._initialized = True + return created + class GenericOAuth2Credential: """ @@ -228,7 +254,7 @@ class ProxyAuthHandler: self._cached_token = self.credential.get_token(self.scope) return self._cached_token - def get_auth_headers(self) -> dict: + def get_auth_headers(self) -> dict[str, str]: """ Get HTTP headers for authentication. diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 255faf94402..9325547c17d 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -124,9 +124,9 @@ class RAGQuery: @staticmethod def extract_documents_from_search( search_response: Any, - ) -> list[str | dict[str, Any]]: + ) -> list[str | dict[str, object]]: """Extract text documents from vector store search response.""" - documents: Final[list[str | dict[str, Any]]] = [] + documents: Final[list[str | dict[str, object]]] = [] search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} for result in search_data["results"]: content_list = result.get("content", []) diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index c0e59f9b975..d02c2114136 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -5,7 +5,8 @@ VerificationToken repository for database operations on LiteLLM_VerificationToke import json from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Final +from types import TracebackType +from typing import TYPE_CHECKING, Final, Protocol from litellm.models.verification_token import ( LiteLLM_VerificationToken, @@ -25,7 +26,36 @@ if TYPE_CHECKING: LiteLLM_VerificationToken as PrismaVerificationToken, ) - from litellm.proxy.utils import PrismaClient + +class _VerificationTokenTables(Protocol): + """The two verification token tables this repository reads and writes.""" + + @property + def litellm_verificationtoken(self) -> TableActions["PrismaVerificationToken"]: ... + + @property + def litellm_deletedverificationtoken(self) -> TableActions["PrismaDeletedVerificationToken"]: ... + + +class _VerificationTokenTransactionManager(Protocol): + async def __aenter__(self) -> _VerificationTokenTables: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> bool | None: ... + + +class _PrismaVerificationTokenDb(_VerificationTokenTables, Protocol): + def tx(self) -> _VerificationTokenTransactionManager: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _PrismaVerificationTokenDb: ... + _JSON_ENCODED_TOKEN_FIELDS: Final = ( "aliases", @@ -44,17 +74,17 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): """Repository for verification token (API key) database operations.""" @property - def prisma_client(self) -> "PrismaClient": - prisma_client: Final[PrismaClient] = super().prisma_client - return prisma_client + def _db(self) -> _PrismaVerificationTokenDb: + client: Final[_PrismaClientView] = self.prisma_client + return client.db @property def table(self) -> TableActions["PrismaVerificationToken"]: - return self.prisma_client.db.litellm_verificationtoken + return self._db.litellm_verificationtoken @property def deleted_table(self) -> TableActions["PrismaDeletedVerificationToken"]: - return self.prisma_client.db.litellm_deletedverificationtoken + return self._db.litellm_deletedverificationtoken @property def model_class(self) -> type[LiteLLM_VerificationToken]: @@ -325,7 +355,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): archive_data["litellm_changed_by"] = litellm_changed_by archive_data["deleted_at"] = datetime.utcnow() - async with self.prisma_client.db.tx() as tx: + async with self._db.tx() as tx: await tx.litellm_deletedverificationtoken.create(data=archive_data) await tx.litellm_verificationtoken.delete(where={"token": token}) diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index 309894957ea..76e833563ba 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -15,7 +15,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_router_logger if TYPE_CHECKING: - from litellm.types.router import SearchToolTypedDict + from litellm.types.router import SearchToolLiteLLMParams, SearchToolTypedDict class _SearchToolsRouter(Protocol): @@ -34,7 +34,7 @@ class SearchAPIRouter: @staticmethod def _resolve_search_provider_credentials( *, - tool_litellm_params: dict[str, Any], + tool_litellm_params: "SearchToolLiteLLMParams", ) -> tuple[str | None, str | None]: """ Resolve search provider credentials from tool configuration ONLY.