From fa966ca2d0bc9bca963112c8befbda5997f306ef Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 08:10:30 +0000 Subject: [PATCH] refactor(types): replace Any with real types across 20 backend files Fifth round of basedpyright Any reduction. Every change is typing-only and leaves runtime behavior identical. Prisma table access now goes through the PrismaTableRepository and TableActions protocols the repo already has, instead of reading untyped attributes off prisma_client.db. Payload and parameter annotations move from dict[str, Any] to dict[str, object] or Mapping[str, object]. Guardrail constructors that took **kwargs: Any now take Unpack of a PEP 728 TypedDict, the same _CustomGuardrailOptions shape three other guardrails already use. Calls into the OpenAI and Azure assistants SDKs pass explicit keywords rather than splatting a dict, so the arguments are checked against the real SDK signatures. --- litellm/llms/azure/assistants.py | 9 +-- litellm/llms/ollama/chat/transformation.py | 2 +- .../llms/ollama/completion/transformation.py | 2 +- litellm/llms/openai/openai.py | 61 ++++++++++++++----- litellm/llms/vertex_ai/fine_tuning/handler.py | 5 +- .../vertex_imagen_transformation.py | 25 ++++---- .../vertex_gemini_transformation.py | 6 +- litellm/proxy/_experimental/mcp_server/db.py | 8 ++- .../proxy/common_utils/reset_budget_job.py | 15 +++-- .../cisco_ai_defense/cisco_ai_defense.py | 7 ++- .../guardrail_hooks/deepkeep/deepkeep.py | 14 +++-- .../guardrail_hooks/headroom/headroom.py | 13 ++-- .../model_armor/model_armor.py | 6 +- .../auto_router_endpoints.py | 10 +-- .../spend_tracking/ptu_flat_cost_rollup.py | 30 +++++++-- litellm/proxy/utils.py | 11 ++-- litellm/realtime_api/main.py | 2 +- litellm/repositories/model_repository.py | 20 ++---- litellm/repositories/team_repository.py | 42 ++++++++++++- .../router_utils/fallback_event_handlers.py | 11 ++-- 20 files changed, 199 insertions(+), 100 deletions(-) diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index f7b419405ac..a4742a25a87 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -1,5 +1,5 @@ from collections.abc import Coroutine, Iterable -from typing import Any, Final, Literal, TypedDict +from typing import Final, Literal, TypedDict import httpx from openai import AsyncAzureOpenAI, AzureOpenAI @@ -715,7 +715,8 @@ class AzureAssistantsAPI(BaseAzureLLM): event_handler: AssistantEventHandler | None, litellm_params: dict | None = None, ) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]: - data: Final[dict[str, Any]] = { + stream_fn: Final = client.beta.threads.runs.stream + base_data: Final[_RunThreadStreamData] = { "thread_id": thread_id, "assistant_id": assistant_id, "additional_instructions": additional_instructions, @@ -725,8 +726,8 @@ class AzureAssistantsAPI(BaseAzureLLM): "tools": tools, } if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return stream_fn(**base_data, event_handler=event_handler) + return stream_fn(**base_data) def run_thread_stream( self, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 181894646e3..257ebf921d9 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -124,7 +124,7 @@ class OllamaChatConfig(BaseConfig): setattr(self.__class__, key, value) @classmethod - def get_config(cls): + def get_config(cls) -> dict[str, object]: return super().get_config() def get_supported_openai_params(self, model: str): diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index dccc83efed4..a9bbfedafd6 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -227,7 +227,7 @@ class OllamaConfig(BaseConfig): model: str, api_base: str | None = None, api_key: str | None = None, - ) -> Any: + ) -> dict[str, object] | None: """ curl http://localhost:11434/api/show -d '{ "name": "mistral" diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index edc8d64d9c2..15c543b94f0 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1,7 +1,7 @@ import time import types from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, cast import httpx @@ -2754,7 +2754,12 @@ class OpenAIAssistantsAPI(BaseLLM): message_thread: Final = await openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread( + id=message_thread.id, + created_at=message_thread.created_at, + metadata=message_thread.metadata, + object=message_thread.object, + ) # fmt: off @@ -2840,7 +2845,12 @@ class OpenAIAssistantsAPI(BaseLLM): message_thread: Final = openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread( + id=message_thread.id, + created_at=message_thread.created_at, + metadata=message_thread.metadata, + object=message_thread.object, + ) async def async_get_thread( self, @@ -2863,7 +2873,12 @@ class OpenAIAssistantsAPI(BaseLLM): response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread( + id=response.id, + created_at=response.created_at, + metadata=response.metadata, + object=response.object, + ) # fmt: off @@ -2929,7 +2944,12 @@ class OpenAIAssistantsAPI(BaseLLM): response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread( + id=response.id, + created_at=response.created_at, + metadata=response.metadata, + object=response.object, + ) def delete_thread(self): pass @@ -2986,18 +3006,27 @@ class OpenAIAssistantsAPI(BaseLLM): tools: Iterable[AssistantToolParam] | None, event_handler: AssistantEventHandler | None, ) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]: - data: Final[dict[str, Any]] = { - "thread_id": thread_id, - "assistant_id": assistant_id, - "additional_instructions": additional_instructions, - "instructions": instructions, - "metadata": metadata, - "model": model, - "tools": tools, - } + runs_stream: Final = client.beta.threads.runs.stream if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return runs_stream( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + event_handler=event_handler, + ) + return runs_stream( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + ) def run_thread_stream( self, diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index 7ecc5e8ff3d..c79b6ffce43 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -280,7 +280,7 @@ class VertexFineTuningAPI(VertexLLM): vertex_location: str, vertex_credentials: str, request_route: str, - ): + ) -> object: _auth_header, vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, @@ -341,5 +341,4 @@ class VertexFineTuningAPI(VertexLLM): f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}" ) - response_json: Final = response.json() - return response_json + return response.json() diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index c6ad5928b74..fddc075bfc6 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -1,6 +1,7 @@ import base64 import json import os +from collections.abc import Mapping from io import BufferedRandom, BufferedReader, BytesIO from pathlib import Path from typing import TYPE_CHECKING, Any, Final, cast @@ -47,11 +48,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): image_edit_optional_params: ImageEditOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, object]: supported_params: Final = self.get_supported_openai_params(model) filtered_params = {key: value for key, value in image_edit_optional_params.items() if key in supported_params} - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Map OpenAI parameters to Imagen format if "n" in filtered_params: @@ -148,10 +149,10 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): model: str, prompt: str | None, image: FileTypes | None, - image_edit_optional_request_params: dict[str, Any], + image_edit_optional_request_params: Mapping[str, object], litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[dict[str, Any], RequestFiles | None]: + ) -> tuple[dict[str, object], RequestFiles | None]: # Prepare reference images in the correct Imagen format if image is None: raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") @@ -182,14 +183,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): parameters["guidanceScale"] = 7.5 # Default guidance scale parameters["seed"] = None # Let Vertex AI choose random seed - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "instances": instances, "parameters": parameters, } - payload: Final[Any] = json.dumps(request_body) + payload: Final = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) - return cast(tuple[dict[str, Any], RequestFiles | None], (payload, empty_files)) + return cast(tuple[dict[str, object], RequestFiles | None], (payload, empty_files)) def transform_image_edit_response( self, @@ -237,8 +238,8 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): def _prepare_reference_images( self, image: FileTypes | list[FileTypes], - image_edit_optional_request_params: dict[str, Any], - ) -> list[dict[str, Any]]: + image_edit_optional_request_params: Mapping[str, object], + ) -> list[dict[str, object]]: """ Prepare reference images in the correct Imagen API format """ @@ -248,7 +249,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): else: images = [image] - reference_images: Final[list[dict[str, Any]]] = [] + reference_images: Final[list[dict[str, object]]] = [] for idx, img in enumerate(images): if img is None: @@ -258,7 +259,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): base64_data = base64.b64encode(image_bytes).decode("utf-8") # Create reference image structure - reference_image = { + reference_image: dict[str, object] = { "referenceType": "REFERENCE_TYPE_RAW", "referenceId": idx + 1, "referenceImage": {"bytesBase64Encoded": base64_data}, @@ -272,7 +273,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): mask_bytes: Final = self._read_all_bytes(mask_image) mask_base64: Final = base64.b64encode(mask_bytes).decode("utf-8") - mask_reference: Final = { + mask_reference: Final[dict[str, object]] = { "referenceType": "REFERENCE_TYPE_MASK", "referenceId": len(reference_images) + 1, "referenceImage": {"bytesBase64Encoded": mask_base64}, diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index d7a2491c04a..b2c52c53580 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -218,10 +218,10 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): contents: Final = [{"role": "user", "parts": [{"text": prompt}]}] # Prepare generation config - generation_config: Final[dict[str, Any]] = {"responseModalities": ["IMAGE"]} + generation_config: Final[dict[str, object]] = {"responseModalities": ["IMAGE"]} # Seed from user-supplied imageConfig dict; flat params are overlaid for backward compat. - image_config: Final[dict[str, Any]] = dict(optional_params.get("imageConfig") or {}) + image_config: Final[dict[str, object]] = dict(optional_params.get("imageConfig") or {}) if "aspectRatio" in optional_params: image_config["aspectRatio"] = optional_params["aspectRatio"] @@ -242,7 +242,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): elif "n" in optional_params: generation_config["candidateCount"] = optional_params["n"] - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "contents": contents, "generationConfig": generation_config, } diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 082a90fdcfb..980a67b30d6 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -37,6 +37,7 @@ from litellm.repositories.table_repositories import ( MCPServerOAuthClientRepository, MCPServerRepository, MCPUserCredentialsRepository, + PrismaTableRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.verification_token_repository import ( @@ -522,11 +523,14 @@ def _user_credential_actions( return table +class _MCPUserEnvVarsRepository(PrismaTableRepository["prisma_db_models.LiteLLM_MCPUserEnvVars"]): + table_name = "litellm_mcpuserenvvars" + + def _user_env_var_actions( prisma_client: PrismaClient, ) -> "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]": - table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars - return table + return _MCPUserEnvVarsRepository(prisma_client).table async def _db_find_user_credential_row( diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index f2648c8466e..dc29a8377c9 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -48,7 +48,7 @@ from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManage from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository -from litellm.repositories.prisma_protocols import SpendLinkedTable +from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable from litellm.repositories.table_repositories import ( EndUserRepository, ModelAccessGroupBudgetRepository, @@ -435,6 +435,11 @@ class ResetBudgetJob: self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings() self.pod_lock_manager: PodLockManager | None = pod_lock_manager + @property + def _new_batch(self) -> Callable[[], PrismaBatch]: + new_batch: Final[Callable[[], PrismaBatch]] = self.prisma_client.db.batch_ + return new_batch + async def _lease_is_held(self, lock_manager: PodLockManager) -> bool: """True only when the lease is readable and someone holds it. @@ -721,7 +726,7 @@ class ResetBudgetJob: ) async def _commit_budget_cascade_once(self, cascade: _BudgetCascade) -> None: - async with budget_cascade_unit_of_work(self.prisma_client.db.batch_) as uow: + async with budget_cascade_unit_of_work(self._new_batch) as uow: _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) @@ -861,7 +866,7 @@ class ResetBudgetJob: ) async def _write_key_reset_updates_once(self, updated_keys: list[LiteLLM_VerificationToken]) -> None: - async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: + async with spend_reset_unit_of_work(self._new_batch) as uow: for k in updated_keys: if k.token is None: continue @@ -885,7 +890,7 @@ class ResetBudgetJob: ) async def _write_user_reset_updates_once(self, updated_users: list[LiteLLM_UserTable]) -> None: - async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: + async with spend_reset_unit_of_work(self._new_batch) as uow: for u in updated_users: uow.users.queue_spend_reset( user_id=u.user_id, @@ -907,7 +912,7 @@ class ResetBudgetJob: ) async def _write_team_reset_updates_once(self, updated_teams: list[LiteLLM_TeamTable]) -> None: - async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: + async with spend_reset_unit_of_work(self._new_batch) as uow: for t in updated_teams: uow.teams.queue_spend_reset( team_id=t.team_id, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index facb822d00d..017ef6e09f6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal import httpx from fastapi import HTTPException +from typing_extensions import TypedDict, Unpack from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -111,6 +112,10 @@ class CiscoAIDefenseGuardrailAPIError(Exception): """Raised when there is an error talking to the Cisco AI Defense API.""" +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): """ Cisco AI Defense guardrail integration. @@ -144,7 +149,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): on_flagged_action: str | None = None, fallback_on_error: str | None = None, timeout: float | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: resolved_api_key: Final = api_key or os.environ.get("CISCO_AI_DEFENSE_API_KEY") if not resolved_api_key: diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index 214d4b486d4..539dc1ea1e9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -7,10 +7,10 @@ import os from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol import httpx -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version @@ -56,7 +56,13 @@ class DeepKeepFirewallResponse(TypedDict): class _DeepKeepInitKwargsView(TypedDict): """Typed read of the guardrail name carried in the untyped base-guardrail kwargs.""" - guardrail_name: ReadOnly[str] + guardrail_name: ReadOnly[str | None] + + +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + guardrail_name: ReadOnly[str | None] class _DeepKeepMetadataSource(TypedDict, total=False): @@ -110,7 +116,7 @@ class DeepKeepGuardrail(CustomGuardrail): firewall_id: str | None = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", extra_headers: Mapping[str, str] | list[str] | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index fa113aa4d33..eb62b896784 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -7,7 +7,7 @@ import time import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard +from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypeGuard import httpx from fastapi import HTTPException @@ -50,6 +50,7 @@ from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs 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.guardrails import LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -878,9 +879,9 @@ class HeadroomGuardrail(CustomGuardrail): async def async_pre_call_deployment_hook( self, - kwargs: dict[str, Any], + kwargs: dict[str, object], call_type: CallTypes | None, - ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict + ) -> dict[str, object] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type) effective: Final = base_result if base_result is not None else kwargs if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES: @@ -897,7 +898,7 @@ class HeadroomGuardrail(CustomGuardrail): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -919,8 +920,8 @@ class HeadroomGuardrail(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/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index fde40111d49..88c7f21ad72 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -360,7 +360,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else: return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}} - def _should_block_content(self, armor_response: Mapping[str, Any], allow_sanitization: bool = False) -> bool: + def _should_block_content(self, armor_response: Mapping[str, object], allow_sanitization: bool = False) -> bool: """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult.""" for filt in self._filter_result_items(armor_response): # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before @@ -429,7 +429,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return filter_results return [] - def _has_deidentify_match(self, armor_response: Mapping[str, Any]) -> bool: + def _has_deidentify_match(self, armor_response: Mapping[str, object]) -> bool: """Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction.""" for filter_entry in self._filter_result_items(armor_response): sdp = filter_entry.get("sdpFilterResult") @@ -439,7 +439,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): def _resolve_streaming_outcome( self, - armor_response: Mapping[str, Any], + armor_response: Mapping[str, object], assembled_response: object, content: str, ) -> tuple[bool, str | None]: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index bbc914a772a..12d23550cb9 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -8,7 +8,6 @@ POST /auto_router/validate_complexity_router_config - Dry-run the complexity-rou from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from itertools import chain, groupby -from operator import attrgetter from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Protocol from uuid import uuid4 @@ -1094,6 +1093,10 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]: ) +def _leg_group_id(leg: "_LegRow") -> str: + return leg.group_id + + class _LegRow(BaseModel): """One LiteLLM_ShadowEvalJob row, validated off the untyped prisma record. A row is one target's leg of a job; the legs of a job share group_id and identical config, @@ -1598,10 +1601,7 @@ async def list_shadow_eval_jobs( or () ) by_group: Final[Mapping[str, tuple[_LegRow, ...]]] = MappingProxyType( - { - group_id: tuple(group) - for group_id, group in groupby(sorted(legs, key=attrgetter("group_id")), key=attrgetter("group_id")) - } + {group_id: tuple(group) for group_id, group in groupby(sorted(legs, key=_leg_group_id), key=_leg_group_id)} ) newest_first: Final = sorted( by_group, key=lambda group_id: max(leg.created_at for leg in by_group[group_id]), reverse=True diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 0e6412a2c64..b8af432029f 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -31,11 +31,31 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: + from prisma import models as prisma_models + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient + +class _DailyTeamSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyTeamSpend"]): + table_name = "litellm_dailyteamspend" + + +def _daily_team_spend_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_DailyTeamSpend]": + """The sentinel rows this rollup writes, reads back and prunes.""" + return _DailyTeamSpendRepository(prisma_client).table + + +def _proxy_model_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_ProxyModelTable]": + """The stored deployments the rollup scans for PTU config.""" + return ModelRepository(prisma_client).table + + _HOURS_PER_DAY: Final = 24 _PRUNE_ID_CHUNK_SIZE: Final = 5_000 _UPSERT_ATTEMPTS: Final = 3 @@ -97,7 +117,7 @@ def _decode_model_info(raw: object) -> "Mapping[str, object] | None": """ if isinstance(raw, str): try: - decoded: Final = json.loads(raw) + decoded: Final[object] = json.loads(raw) except (TypeError, ValueError): return None return decoded if isinstance(decoded, dict) else None @@ -240,7 +260,7 @@ async def _upsert_ptu_daily_row( } } now: Final = datetime.now(timezone.utc) - await prisma_client.db.litellm_dailyteamspend.upsert( + await _daily_team_spend_table(prisma_client).upsert( where=where, data={ # mutable-ok: prisma upsert data payload "create": { # mutable-ok: prisma create payload @@ -353,7 +373,7 @@ async def _load_ptu_models(prisma_client: "PrismaClient", *, router: object | No The router is handed in rather than read off the proxy module, so a run prices exactly the deployments its caller declares and nothing a co-resident process left behind. """ - rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many() + rows: Final = await _proxy_model_table(prisma_client).find_many() db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or ""))) config_records: Final = _config_deployments(router, owned_by_db=db_ids) models: Final = tuple( @@ -503,7 +523,7 @@ async def _existing_sentinel_keys( survives a rename. Nothing here reads the display name. """ date_range: Final = {"gte": start.isoformat(), "lte": end.isoformat()} # mutable-ok: prisma range filter - rows: Final = await prisma_client.db.litellm_dailyteamspend.find_many( + rows: Final = await _daily_team_spend_table(prisma_client).find_many( where={"api_key": PTU_SENTINEL_API_KEY, "date": date_range} # mutable-ok: prisma find filter ) return frozenset( @@ -771,7 +791,7 @@ async def _prune_unrefreshed_sentinel_rows( ) filters: Final = tuple(_prune_filter(date_str=date_str, cutoff=cutoff, chunk=chunk) for chunk in chunks) deletions: Final = tuple( - [await prisma_client.db.litellm_dailyteamspend.delete_many(where=where) for where in filters] + [await _daily_team_spend_table(prisma_client).delete_many(where=where) for where in filters] ) deleted: Final = sum(deletions) if deleted: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddf31cb1d8a..21328a4962a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3401,7 +3401,7 @@ class _ConfigRow: __slots__ = ("param_name", "param_value") - def __init__(self, param_name: str, param_value: Any) -> None: + def __init__(self, param_name: str, param_value: object) -> None: self.param_name = param_name self.param_value = param_value @@ -3414,7 +3414,7 @@ def _pack_config_row(row: Any) -> dict[str, object]: return {"param_name": row.param_name, "param_value": row.param_value} -def _unpack_config_row(cached: Any) -> _ConfigRow | None: +def _unpack_config_row(cached: object) -> _ConfigRow | None: if cached is None or cached == _CONFIG_CACHE_MISS: return None if isinstance(cached, dict): @@ -3557,6 +3557,7 @@ class PrismaClient: verbose_proxy_logger.debug("Creating Prisma Client..") try: from prisma import Prisma + from prisma.types import DatasourceOverride except Exception as e: verbose_proxy_logger.error("Failed to import Prisma client: %s", e) verbose_proxy_logger.error("This usually means 'prisma generate' hasn't been run yet.") @@ -3607,11 +3608,11 @@ class PrismaClient: reader_token: Final = mint_database_token(token_auth, reader_iam_endpoint) read_replica_url = reader_iam_endpoint.build_url(reader_token) os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url - reader_kwargs: Final[dict[str, Any]] = {"datasource": {"url": read_replica_url}} + reader_datasource: Final = DatasourceOverride(url=read_replica_url) if http_client is not None: - reader_prisma = Prisma(http=http_client, **reader_kwargs) + reader_prisma = Prisma(http=http_client, datasource=reader_datasource) else: - reader_prisma = Prisma(**reader_kwargs) + reader_prisma = Prisma(datasource=reader_datasource) reader_wrapper: Final = PrismaWrapper( original_prisma=reader_prisma, token_auth=token_auth, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index b824a5928c6..9d81c725c70 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -53,7 +53,7 @@ bedrock_realtime: Final = BedrockRealtime() xai_realtime: Final = XAIRealtime() vertex_llm_base: Final = VertexBase() base_llm_http_handler = BaseLLMHTTPHandler() -_EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MODEL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index d24eb8ffc62..8ee76b93923 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -4,29 +4,23 @@ Model repository for database operations on LiteLLM_ProxyModelTable. import json from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final from litellm.models.model import LiteLLM_ProxyModelTable -from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _PrismaModelDb(Protocol): - @property - def litellm_proxymodeltable(self) -> TableActions["prisma_models.LiteLLM_ProxyModelTable"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _PrismaModelDb: ... +class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ProxyModelTable"]): + table_name = "litellm_proxymodeltable" class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): @@ -38,11 +32,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_ProxyModelTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return wrap_table_actions_for_config_sync( - actions=client.db.litellm_proxymodeltable, - table_name="litellm_proxymodeltable", - ) + return _ProxyModelTableRepository(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_ProxyModelTable]: diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 5ff07d76b5d..cbe263699c9 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -5,6 +5,7 @@ Team repository for database operations on LiteLLM_TeamTable. import json from collections.abc import Mapping, Sequence from datetime import datetime +from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter @@ -40,6 +41,36 @@ def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays: return team +class _TeamTables(Protocol): + """The two team tables this repository reads and writes.""" + + @property + def litellm_teamtable(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: ... + + @property + def litellm_deletedteamtable(self) -> TableActions["prisma_models.LiteLLM_DeletedTeamTable"]: ... + + +class _TeamTransactionManager(Protocol): + async def __aenter__(self) -> _TeamTables: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> bool | None: ... + + +class _PrismaTeamDb(_TeamTables, Protocol): + def tx(self) -> _TeamTransactionManager: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _PrismaTeamDb: ... + + _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) _JSON_ENCODED_TEAM_FIELDS: Final = ( "metadata", @@ -54,13 +85,18 @@ _JSON_ENCODED_TEAM_FIELDS: Final = ( class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" + @property + def _db(self) -> _PrismaTeamDb: + client: Final[_PrismaClientView] = self.prisma_client + return client.db + @property def table(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: - return self.prisma_client.db.litellm_teamtable + return self._db.litellm_teamtable @property def deleted_table(self) -> TableActions["prisma_models.LiteLLM_DeletedTeamTable"]: - return self.prisma_client.db.litellm_deletedteamtable + return self._db.litellm_deletedteamtable @property def model_class(self) -> type[LiteLLM_TeamTable]: @@ -256,7 +292,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): 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_deletedteamtable.create(data=archive_data) await tx.litellm_teamtable.delete(where={"team_id": team_id}) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 7fda5d96fb0..bf1838e50f7 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -266,7 +266,7 @@ def get_pre_routing_selection(kwargs: Mapping[str, object]) -> str | None: DISABLE_FALLBACKS_METADATA_KEY: Final = "_disable_fallbacks" -def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bool) -> None: +def record_disable_fallbacks(request_kwargs: Mapping[str, object] | None, disabled: bool) -> None: """ Write-or-clear the request's disable_fallbacks verdict into the router-internal metadata bucket. The wrapper pops the raw kwarg before any downstream frame runs, so the refusal @@ -286,7 +286,7 @@ def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bucket.pop(DISABLE_FALLBACKS_METADATA_KEY, None) -def fallbacks_disabled_for_request(kwargs: Mapping[str, Any]) -> bool: +def fallbacks_disabled_for_request(kwargs: Mapping[str, object]) -> bool: """True when this request opted out of fallbacks, read from the raw kwarg (pre-pop snapshots keep it) or the router-internal bucket the wrapper stamps after popping it.""" if kwargs.get("disable_fallbacks") is True: @@ -639,7 +639,7 @@ async def log_failure_fallback_event(original_model_group: str, kwargs: dict, or verbose_router_logger.error("Error in log_failure_fallback_event: %s", e) -def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool: +def _check_non_standard_fallback_format(fallbacks: Sequence[object] | None) -> bool: """ Checks if the fallbacks list is a list of strings or a list of dictionaries. @@ -653,8 +653,9 @@ def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool: return False if all(isinstance(item, str) for item in fallbacks): return True - elif all(isinstance(item, dict) for item in fallbacks): - for item in fallbacks: + dict_entries: Final = tuple(item for item in fallbacks if isinstance(item, dict)) + if len(dict_entries) == len(fallbacks): + for item in dict_entries: for key in LiteLLMParamsTypedDict.__annotations__: if key in item: # If the value is a list, it's likely a standard fallback model group mapping