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.