diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 3f89ff179eb..a23028c079c 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,15 +1,15 @@ { "reportAny": { - "limit": 33171 + "limit": 32185 }, "reportArgumentType": { - "limit": 2645 + "limit": 2635 }, "reportAssignmentType": { "limit": 329 }, "reportAttributeAccessIssue": { - "limit": 516 + "limit": 514 }, "reportCallIssue": { "limit": 123 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 10228 + "limit": 10118 }, "reportFunctionMemberAccess": { "limit": 11 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5869 + "limit": 5867 }, "reportMissingTypeArgument": { - "limit": 15864 + "limit": 15863 }, "reportMissingTypeStubs": { "limit": 41 @@ -84,7 +84,7 @@ "limit": 77 }, "reportPrivateUsage": { - "limit": 2437 + "limit": 2415 }, "reportRedeclaration": { "limit": 12 @@ -99,25 +99,25 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45522 + "limit": 45254 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40479 + "limit": 40440 }, "reportUnknownParameterType": { - "limit": 20341 + "limit": 20338 }, "reportUnknownVariableType": { - "limit": 32052 + "limit": 32048 }, "reportUnnecessaryCast": { "limit": 177 }, "reportUnnecessaryComparison": { - "limit": 1023 + "limit": 1015 }, "reportUnnecessaryContains": { "limit": 7 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 60dbf7c644a..563f35a0665 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -30,7 +30,6 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) -from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.types.llms.openai import OpenAIChatCompletionChunk from litellm.types.router import GenericLiteLLMParams @@ -125,14 +124,17 @@ class CustomStreamWrapper: self.model = model self.make_call = make_call self.custom_llm_provider = custom_llm_provider - self.logging_obj: LiteLLMLoggingObject = logging_obj + self.logging_obj: Any = logging_obj self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False self._stream_created_time: float = time.time() - litellm_params: GenericLiteLLMParams = GenericLiteLLMParams( - **self.logging_obj.model_call_details.get("litellm_params", {}) + _litellm_params_kwargs = self.logging_obj.model_call_details.get("litellm_params", {}) + litellm_params: GenericLiteLLMParams = ( + GenericLiteLLMParams.model_validate(_litellm_params_kwargs) + if isinstance(_litellm_params_kwargs, dict) + else GenericLiteLLMParams() ) self.merge_reasoning_content_in_choices: bool = litellm_params.merge_reasoning_content_in_choices or False self.sent_first_thinking_block = False @@ -677,7 +679,7 @@ class CustomStreamWrapper: if chunk: args.update({k: v for k, v in chunk.items() if k != "stream"}) - model_response = ModelResponseStream(**args) + model_response = ModelResponseStream.model_validate(args) if self.response_id is not None: model_response.id = self.response_id if self.system_fingerprint is not None: @@ -817,7 +819,7 @@ class CustomStreamWrapper: _initial_delta = model_response.choices[0].delta.model_dump() _initial_delta.pop("role", None) - model_response.choices[0].delta = Delta(**_initial_delta) + model_response.choices[0].delta = Delta.model_validate(_initial_delta) return model_response def _has_special_delta_content(self, model_response: ModelResponseStream) -> bool: @@ -915,7 +917,7 @@ class CustomStreamWrapper: choice_json.pop( "finish_reason", None ) # for mistral etc. which return a value in their last chunk (not-openai compatible). - choices.append(StreamingChoices(**choice_json)) + choices.append(StreamingChoices.model_validate(choice_json)) except Exception: choices.append(StreamingChoices()) setattr(model_response, "choices", choices) @@ -946,7 +948,7 @@ class CustomStreamWrapper: self.sent_first_chunk = True if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] - model_response.choices[0].delta = Delta(**completion_obj) + model_response.choices[0].delta = Delta.model_validate(completion_obj) _index: Optional[int] = completion_obj.get("index") if _index is not None: model_response.choices[0].index = _index diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ec1301e5923..b0a2f225400 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -160,8 +160,10 @@ def _rust_responses_websocket_enabled( from .http_handler import get_shared_realtime_ssl_context if TYPE_CHECKING: + import tiktoken from aiohttp import ClientSession + from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( AnthropicMessagesStreamingResponse, @@ -210,7 +212,7 @@ def _responses_api_optional_request_param_names() -> frozenset[str]: return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys()) -def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: +def _custom_logger_callbacks(logging_obj: LiteLLMLoggingObj) -> tuple["CustomLogger", ...]: from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import ( get_custom_logger_compatible_class, @@ -221,7 +223,7 @@ def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: if isinstance(dynamic_success_callbacks, (list, tuple)): callbacks.extend(dynamic_success_callbacks) - custom_loggers: list[Any] = [] + custom_loggers = [] for cb in callbacks: if isinstance(cb, str): resolved = get_custom_logger_compatible_class(cb) # type: ignore[arg-type] @@ -230,10 +232,10 @@ def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: cb = resolved if isinstance(cb, CustomLogger): custom_loggers.append(cb) - return custom_loggers + return tuple(custom_loggers) -def _has_pre_call_deployment_hook(logging_obj: Any) -> bool: +def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: from litellm.integrations.custom_logger import CustomLogger base_func = CustomLogger.async_pre_call_deployment_hook @@ -359,7 +361,7 @@ class BaseLLMHTTPHandler: messages: list, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: Optional["tiktoken.Encoding"], api_key: Optional[str] = None, client: Optional[AsyncHTTPHandler] = None, json_mode: bool = False, @@ -425,7 +427,7 @@ class BaseLLMHTTPHandler: api_base: Optional[str], custom_llm_provider: str, model_response: ModelResponse, - encoding, + encoding: Optional["tiktoken.Encoding"], logging_obj: LiteLLMLoggingObj, optional_params: dict, timeout: Union[float, httpx.Timeout], @@ -5008,7 +5010,7 @@ class BaseLLMHTTPHandler: return depth, max(max_loops, 1), fingerprints @staticmethod - def _has_agentic_completion_hook(logging_obj: Any) -> bool: + def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool: """ True if any registered callback actually overrides ``async_should_run_agentic_loop`` (the gate every agentic hook goes @@ -5248,7 +5250,7 @@ class BaseLLMHTTPHandler: self, result: Any, model: str, - responses_api_provider_config: Any, + responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, ) -> Any: @@ -5769,7 +5771,7 @@ class BaseLLMHTTPHandler: websockets_module: Any, url: str, headers: dict, - ssl_context: Any, + ssl_context: bool | str | ssl.SSLContext | None, *, open_timeout: float = 8.0, max_attempts: int = 3, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a94a75fdfa3..2565bd5ae7d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -490,7 +490,7 @@ _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS = frozenset({"llm_api_routes", "info_rout def _validate_caller_can_change_key_ownership( data: Optional[BaseModel], - existing_key_row: Any, + existing_key_row: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, ) -> None: """ @@ -523,7 +523,7 @@ def _validate_caller_can_change_key_ownership( status_code=403, detail="Non-admin users cannot remove the user_id from a key.", ) - existing_user_id = getattr(existing_key_row, "user_id", None) + existing_user_id = existing_key_row.user_id if incoming_user_id != existing_user_id: raise HTTPException( status_code=403, @@ -2112,7 +2112,7 @@ async def _process_single_key_update( litellm_changed_by: Optional[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: Any, + proxy_logging_obj: ProxyLogging, llm_router: Optional[Router], user_custom_key_update: Optional[Callable] = None, existing_key_row: Optional[LiteLLM_VerificationToken] = None, @@ -2265,9 +2265,9 @@ async def _process_single_key_update( async def _validate_mcp_servers_for_key_update( data: "UpdateKeyRequest", team_obj: Optional["LiteLLM_TeamTableCachedObj"], - existing_key_row: Any, - prisma_client: Any, - user_api_key_cache: Any, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, is_proxy_admin: bool, ) -> Optional[ObjectPermissionDict]: """Validate MCP servers in object_permission against the effective team.""" @@ -2302,14 +2302,20 @@ async def _validate_mcp_servers_for_key_update( async def _validate_update_key_data( data: UpdateKeyRequest, - existing_key_row: Any, + existing_key_row: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, - llm_router: Any, + llm_router: Router | None, premium_user: bool, - prisma_client: Any, - user_api_key_cache: Any, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, ) -> None: """Validate permissions and constraints for key update.""" + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected"}, + ) + # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(data.spend) @@ -2414,11 +2420,12 @@ async def _validate_update_key_data( # _check_key_admin_access that would otherwise require team/org admin status. _key_is_team_key = getattr(existing_key_row, "team_id", None) is not None can_skip_admin_check = (caller_is_creator or _key_is_team_key) and not _is_budget_change - if (not _is_proxy_admin) and prisma_client is not None and not can_skip_admin_check: - hashed_key = existing_key_row.token + if (not _is_proxy_admin) and not can_skip_admin_check: + if existing_key_row.token is None: + raise HTTPException(status_code=500, detail={"error": "Existing key is missing its token"}) await _check_key_admin_access( user_api_key_dict=user_api_key_dict, - hashed_token=hashed_key, + hashed_token=existing_key_row.token, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"), @@ -5421,7 +5428,7 @@ async def list_keys( async def _apply_non_admin_alias_scope( user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: PrismaClient | None, query_params: List[Any], where_parts: List[str], ) -> None: @@ -5852,9 +5859,7 @@ async def _list_key_helper( where=where # type: ignore ) else: - total_count = await VerificationTokenRepository(prisma_client).table.count( - where=where # type: ignore - ) + total_count = await VerificationTokenRepository(prisma_client).count(where=where) verbose_proxy_logger.debug(f"Total count of keys: {total_count}") @@ -5932,7 +5937,7 @@ def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]: async def _check_key_admin_access( user_api_key_dict: UserAPIKeyAuth, hashed_token: str, - prisma_client: Any, + prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, route: str, ) -> None: @@ -6443,7 +6448,7 @@ def _validate_key_alias_format(key_alias: Optional[str]) -> None: async def _enforce_unique_key_alias( key_alias: Optional[str], - prisma_client: Any, + prisma_client: PrismaClient | None, existing_key_token: Optional[str] = None, ) -> None: """ @@ -6451,7 +6456,7 @@ async def _enforce_unique_key_alias( Args: key_alias (Optional[str]): The key alias to check - prisma_client (Any): Prisma client instance + prisma_client (Optional[PrismaClient]): Prisma client instance existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check (The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c35c17aa359..7f8a10f8ac8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -14,7 +14,7 @@ import json import math import traceback from datetime import datetime, timezone -from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple, Union, cast +from typing import Annotated, Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -31,6 +31,7 @@ from litellm.proxy._types import ( CommonProxyErrors, DeleteTeamRequest, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, LiteLLM_DeletedTeamTable, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, @@ -78,6 +79,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.common_utils import ( _check_passthrough_routes_caller_permission, _is_user_org_admin_for_team, @@ -106,7 +108,7 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) -from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy +from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( @@ -151,9 +153,9 @@ def _sanitize_for_log(value: object) -> str: async def _refresh_cached_team( - team_row: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + team_row: LiteLLM_TeamTable, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> None: """ Refresh the in-memory cached team object after a DB write. @@ -274,7 +276,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - team_member_budget_table = await new_budget( + team_member_budget_table: LiteLLM_BudgetTable = await new_budget( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -322,7 +324,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - budget_row = await update_budget( + budget_row: LiteLLM_BudgetTable = await update_budget( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -395,7 +397,7 @@ class TeamMemberBudgetHandler: @staticmethod async def backfill_team_member_budget_entries( team_id: str, - members_with_roles: List[Union[Member, dict]], + members_with_roles: Sequence[Union[Member, dict]], team_member_budget_id: str, prisma_client: PrismaClient, ) -> None: @@ -414,7 +416,9 @@ class TeamMemberBudgetHandler: return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = await TeamMembershipRepository(prisma_client).table.find_many(where={"team_id": team_id}) + existing_memberships: Sequence[LiteLLM_TeamMembership] = await TeamMembershipRepository( + prisma_client + ).table.find_many(where={"team_id": team_id}) existing_user_ids = {m.user_id for m in existing_memberships} # Identify members with no existing membership row. @@ -447,7 +451,7 @@ class TeamMemberBudgetHandler: # Heal existing membership rows that predate the team_member_budget # configuration: populate budget_id where it is currently NULL. # Rows with an explicit budget_id (per-member override) are left alone. - updated = await TeamMembershipRepository(prisma_client).table.update_many( + updated: int = await TeamMembershipRepository(prisma_client).table.update_many( where={"team_id": team_id, "budget_id": None}, data={"budget_id": team_member_budget_id}, ) @@ -503,7 +507,7 @@ async def get_all_team_memberships( # else: # where_obj = {"user_id": str(user_id), "team_id": {"in": team_id}} - team_memberships = await TeamMembershipRepository(prisma_client).table.find_many( + team_memberships: Sequence[LiteLLM_TeamMembership] = await TeamMembershipRepository(prisma_client).table.find_many( where=where_obj, include={"litellm_budget_table": True}, ) @@ -765,7 +769,7 @@ async def _check_org_team_limits( # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - teams = await TeamRepository(prisma_client).table.find_many( + teams: Sequence[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.find_many( where={"organization_id": org_table.organization_id}, ) @@ -790,7 +794,7 @@ async def _check_user_team_limits( data: Union[NewTeamRequest, UpdateTeamRequest], user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, - user_api_key_cache: Any, + user_api_key_cache: UserApiKeyCache, ) -> None: """ Enforce the caller's personal limits when CREATING a standalone team. @@ -1394,7 +1398,7 @@ async def _update_model_table( async def _auto_add_team_members_to_organization( team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, - prisma_client: Any, + prisma_client: PrismaClient, ) -> None: """ When moving a team to an org, ensure all team members are also org members. @@ -1432,11 +1436,11 @@ async def _auto_add_team_members_to_organization( async def fetch_and_validate_organization( organization_id: str, - existing_team_row: Any, + existing_team_row: LiteLLM_TeamTable, llm_router: Optional[Router], - prisma_client: Any, + prisma_client: PrismaClient, user_api_key_dict: Optional[UserAPIKeyAuth] = None, -) -> Any: +) -> LiteLLM_OrganizationTable | None: """ Fetch and validate an organization for team update operations. @@ -1455,7 +1459,7 @@ async def fetch_and_validate_organization( if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - organization_row = await OrganizationRepository(prisma_client).table.find_unique( + organization_row: LiteLLM_OrganizationTable | None = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": organization_id}, include={"litellm_budget_table": True, "members": True, "teams": True}, ) @@ -1704,7 +1708,9 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": data.team_id} + ) if existing_team_row is None: raise HTTPException( @@ -1792,7 +1798,9 @@ async def update_team( # check org team limits - if updating team that belongs to an org org_id_to_check = ( - data.organization_id if data.organization_id is not None else existing_team_row.organization_id + data.organization_id + if data.organization_id is not None + else getattr(existing_team_row, "organization_id", None) ) if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None: org_table = await get_org_object( @@ -1895,7 +1903,7 @@ async def update_team( updated_kv.pop("model_aliases") _model_id = await _update_model_table( data=data, - model_id=existing_team_row.model_id, + model_id=getattr(existing_team_row, "model_id", None), prisma_client=prisma_client, user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, @@ -2004,13 +2012,16 @@ async def patch_team( patch_fields = data.model_dump(exclude_unset=True, exclude={"team_id"}) if "metadata" in patch_fields: - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + existing_team_row: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": team_id} + ) if existing_team_row is None: raise HTTPException( status_code=404, detail={"error": f"Team not found, passed team_id={team_id}"}, ) - existing_metadata = existing_team_row.metadata if isinstance(existing_team_row.metadata, dict) else {} + _existing_metadata = getattr(existing_team_row, "metadata", None) + existing_metadata = _existing_metadata if isinstance(_existing_metadata, dict) else {} patch_fields["metadata"] = apply_json_merge_patch(existing_metadata, patch_fields["metadata"]) update_request = UpdateTeamRequest.model_validate({"team_id": team_id, **patch_fields}) @@ -2706,7 +2717,9 @@ async def team_member_delete( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + _existing_team_row: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": data.team_id} + ) if _existing_team_row is None: raise HTTPException( @@ -2910,7 +2923,9 @@ async def team_member_update( _validate_budget_duration(data.budget_duration) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + _existing_team_row: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": data.team_id} + ) if _existing_team_row is None: raise HTTPException( @@ -3818,7 +3833,9 @@ async def block_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": data.team_id} + ) if existing_team is None: raise HTTPException( status_code=404, @@ -3831,7 +3848,7 @@ async def block_team( user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": True}, # type: ignore ) @@ -3867,7 +3884,9 @@ async def unblock_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": data.team_id} + ) if existing_team is None: raise HTTPException( status_code=404, @@ -3880,7 +3899,7 @@ async def unblock_team( user_api_key_dict=user_api_key_dict, ) - record = await TeamRepository(prisma_client).table.update( + record: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": False}, # type: ignore ) @@ -3914,7 +3933,9 @@ async def list_available_teams( return [] # filter out teams that the user is already a member of - user_info = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_api_key_dict.user_id}) + user_info: LiteLLM_UserTable | None = await UserRepository(prisma_client).table.find_unique( + where={"user_id": user_api_key_dict.user_id} + ) if user_info is None: raise HTTPException( status_code=404, @@ -3924,7 +3945,9 @@ async def list_available_teams( available_teams = [team for team in available_teams if team not in user_info_correct_type.teams] - available_teams_db = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": available_teams}}) + available_teams_db: Sequence[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.find_many( + where={"team_id": {"in": available_teams}} + ) available_teams_correct_type = [LiteLLM_TeamTable.model_validate(team.model_dump()) for team in available_teams_db] @@ -3933,9 +3956,9 @@ async def list_available_teams( async def _get_org_admin_org_ids( user_id: str, - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> Optional[List[str]]: """ Return the list of organization IDs where the user is an org admin. @@ -4113,7 +4136,7 @@ def _convert_teams_to_response_models( async def _get_keys_count_by_team( - prisma_client: Any, + prisma_client: PrismaClient, teams: list, ) -> Dict[str, int]: """Aggregate virtual-key counts per team for the given page of teams. @@ -4138,9 +4161,9 @@ async def _enforce_list_team_v2_access( user_api_key_dict: UserAPIKeyAuth, user_id: Optional[str], organization_id: Optional[str], - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[List[str]]]: """Enforce access control for list_team_v2. @@ -4391,9 +4414,9 @@ async def list_team_v2( async def _authorize_and_filter_teams( user_api_key_dict: UserAPIKeyAuth, user_id: Optional[str], - prisma_client: Any, - user_api_key_cache: Any, - proxy_logging_obj: Any, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, ) -> list: """ Authorize the /team/list request and return filtered teams. @@ -4568,7 +4591,9 @@ async def get_paginated_teams( total_count = await TeamRepository(prisma_client).table.count() # Get paginated teams - teams = await TeamRepository(prisma_client).table.find_many( + teams: list[LiteLLM_TeamTable] = await TeamRepository( # mutable-ok: matches this function's list return type + prisma_client + ).table.find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4633,7 +4658,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams = await TeamRepository(prisma_client).table.find_many( + teams: Sequence[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -4701,7 +4726,9 @@ async def team_model_add( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + team_row: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": data.team_id} + ) if team_row is None: raise HTTPException( @@ -4747,11 +4774,16 @@ async def team_model_add( # the writer and lets Prisma bump updated_at. # `include` mirrors the relations the auth path consumes off the cached # team object so that `_refresh_cached_team` doesn't null them out. - updated_team = await TeamRepository(prisma_client).table.update( + updated_team: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"updated_at": datetime.now(timezone.utc)}, include={"object_permission": True}, # type: ignore ) + if updated_team is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) await _refresh_cached_team( team_row=updated_team, @@ -4801,7 +4833,9 @@ async def team_model_delete( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + team_row: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": data.team_id} + ) if team_row is None: raise HTTPException( @@ -4829,11 +4863,16 @@ async def team_model_delete( updated_models = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team = await TeamRepository(prisma_client).table.update( + updated_team: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, # type: ignore ) + if updated_team is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) await _refresh_cached_team( team_row=updated_team, @@ -4964,10 +5003,15 @@ async def update_team_member_permissions( }, ) # Update the team member permissions - updated_team = await TeamRepository(prisma_client).table.update( + updated_team: LiteLLM_TeamTable | None = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"team_member_permissions": data.team_member_permissions}, ) + if updated_team is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) return updated_team @@ -5058,7 +5102,7 @@ async def _compute_and_batch_updates(prisma_client, teams, permissions_to_add: s async def _append_permissions_to_specific_teams(prisma_client, team_ids: List[str], permissions_to_add: set) -> int: """Fetch specific teams by ID and append permissions.""" - teams = await TeamRepository(prisma_client).table.find_many( + teams: Sequence[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": team_ids}}, ) @@ -5088,7 +5132,7 @@ async def _append_permissions_to_all_teams(prisma_client, permissions_to_add: se find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await TeamRepository(prisma_client).table.find_many(**find_args) + teams: Sequence[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.find_many(**find_args) if not teams: break @@ -5188,8 +5232,12 @@ async def get_team_daily_activity( where_condition = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases = await TeamRepository(prisma_client).table.find_many(where=where_condition) - team_alias_metadata = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} + team_aliases: Sequence[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.find_many( + where=where_condition + ) + team_alias_metadata: Mapping[str, dict[str, object]] = { + t.team_id: {"team_alias": t.team_alias} for t in team_aliases + } # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. @@ -5219,9 +5267,9 @@ async def get_team_daily_activity( # If user does not have full team view, filter by their API keys if not has_full_team_view: # Get all API keys for this user - user_keys = await VerificationTokenRepository(prisma_client).table.find_many( - where={"user_id": user_api_key_dict.user_id} - ) + user_keys: Sequence[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"user_id": user_api_key_dict.user_id}) user_api_keys = [key.token for key in user_keys if key.token] # If user has no API keys, return empty result if not user_api_keys: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 18a927e7a44..00d00ff2f2f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,6 +26,7 @@ from typing import ( List, Literal, Optional, + Sequence, Set, Tuple, TypedDict, @@ -128,7 +129,7 @@ from litellm.utils import ( ) if TYPE_CHECKING: - from aiohttp import ClientSession + from aiohttp import ClientSession, SocketFactoryType from opentelemetry.trace import Span as _Span from litellm.integrations.opentelemetry import OpenTelemetry @@ -648,7 +649,7 @@ from fastapi.responses import ( RedirectResponse, StreamingResponse, ) -from fastapi.routing import APIRouter +from fastapi.routing import APIRoute, APIRouter from fastapi.security import OAuth2PasswordBearer from fastapi.security.api_key import APIKeyHeader from fastapi.staticfiles import StaticFiles @@ -830,6 +831,15 @@ async def proxy_shutdown_event(): cleanup_router_config_variables() +class _AiohttpConnectorKwargs(TypedDict, total=False): + keepalive_timeout: int + ttl_dns_cache: int + enable_cleanup_closed: bool + limit: int + limit_per_host: int + socket_factory: "SocketFactoryType" + + async def _initialize_shared_aiohttp_session(): """Initialize shared aiohttp session for connection reuse with connection limits.""" try: @@ -839,7 +849,7 @@ async def _initialize_shared_aiohttp_session(): _build_aiohttp_keepalive_socket_factory, ) - connector_kwargs: Dict[str, Any] = { + connector_kwargs: _AiohttpConnectorKwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, } @@ -1153,7 +1163,7 @@ async def proxy_startup_event(app: FastAPI): await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] -def _generate_stable_operation_id(route: Any) -> str: +def _generate_stable_operation_id(route: APIRoute) -> str: operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") route_methods = sorted(route.methods or []) if len(route_methods) == 1: @@ -1267,7 +1277,10 @@ vertex_live_passthrough_vertex_base = VertexBase() from fastapi.routing import APIWebSocketRoute -def _inject_websocket_stubs_into_openapi_schema(openapi_schema: dict, websocket_routes: list) -> dict: +def _inject_websocket_stubs_into_openapi_schema( + openapi_schema: dict, # mutable-ok: mutated in place and returned to the caller + websocket_routes: Sequence[APIWebSocketRoute], +) -> dict: # mutable-ok: same object as the openapi_schema argument, mutated in place """ Add a synthetic GET stub for each WebSocket route so it appears in Swagger UI. @@ -1281,7 +1294,7 @@ def _inject_websocket_stubs_into_openapi_schema(openapi_schema: dict, websocket_ parameters = [] try: - if hasattr(route, "dependant") and route.dependant is not None: + if hasattr(route, "dependant"): # Handle both FastAPI <0.120 and >=0.120 query_params = getattr(route.dependant, "query_params", []) if query_params: @@ -3518,7 +3531,7 @@ _DB_OVERLAY_REMOTE_MODULE_LIST_FIELDS: Dict[str, Tuple[str, ...]] = { } -def _is_remote_module_url(value: Any) -> bool: +def _is_remote_module_url(value: object) -> bool: return isinstance(value, str) and (value.startswith("s3://") or value.startswith("gcs://")) @@ -5398,7 +5411,7 @@ class ProxyConfig: # decrypt values for k, v in _litellm_params.items(): _litellm_params[k] = self._resolve_db_litellm_param(key=k, value=v) - _litellm_params = LiteLLM_Params(**_litellm_params) + _litellm_params = LiteLLM_Params.model_validate(_litellm_params) else: verbose_proxy_logger.error( @@ -5429,7 +5442,7 @@ class ProxyConfig: # decrypt values for k, v in _litellm_params.items(): _litellm_params[k] = self._resolve_db_litellm_param(key=k, value=v) - _litellm_params = LiteLLM_Params(**_litellm_params) + _litellm_params = LiteLLM_Params.model_validate(_litellm_params) else: verbose_proxy_logger.error( f"Invalid model added to proxy db. Invalid litellm params. litellm_params={_litellm_params}" @@ -6897,7 +6910,7 @@ class ProxyConfig: if isinstance(credential, dict): credential_object = CredentialItem(**credential) elif isinstance(credential, BaseModel): - credential_object = CredentialItem(**credential.model_dump()) + credential_object = CredentialItem.model_validate(credential.model_dump()) decrypted_credential_values = {} for k, v in credential_object.credential_values.items(): @@ -13063,7 +13076,7 @@ def _get_model_group_info( _model_group_info = llm_router.get_model_group_info(model_group=model) if _model_group_info is not None: - model_groups.append(ModelGroupInfoProxy(**_model_group_info.model_dump())) + model_groups.append(ModelGroupInfoProxy.model_validate(_model_group_info.model_dump())) else: model_group_info = ModelGroupInfoProxy( model_group=model, @@ -14782,7 +14795,7 @@ async def update_config_general_settings( ) try: - ConfigGeneralSettings(**{data.field_name: data.field_value}) + ConfigGeneralSettings.model_validate({data.field_name: data.field_value}) except Exception: raise HTTPException( status_code=400, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 924189fed4b..222103fe6f7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -345,7 +345,7 @@ def _accepts_litellm_call_info(cb: CustomLogger) -> bool: return _CALLBACK_ACCEPTS_CALL_INFO[key] -def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: Any) -> None: +def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: CustomLogger) -> None: """ If `exc` is an HTTPException with a dict `detail`, mutate it in place to add `guardrail_name` and `guardrail_mode` taken from the callback instance. @@ -394,11 +394,11 @@ class _CallbackCapabilities: # Tuple[(resolved_callback, "override" | "apply_guardrail"), ...] # Ordered the same as ``litellm.callbacks``; used to build the streaming # iterator chain without re-scanning per request. - iterator_overrides: Tuple[Tuple[Any, str], ...] = field(default_factory=tuple) + iterator_overrides: tuple[tuple[CustomLogger, str], ...] = field(default_factory=tuple) # Resolved CustomLogger callbacks in original order. Pre-resolving once # avoids the per-request ``get_custom_logger_compatible_class`` walk for # every string entry in ``litellm.callbacks``. - resolved_callbacks: Tuple[Any, ...] = field(default_factory=tuple) + resolved_callbacks: tuple[CustomLogger, ...] = field(default_factory=tuple) class ProxyLogging: @@ -676,7 +676,7 @@ class ProxyLogging: return synthetic_data - def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> Optional[Any]: + def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None: """ Convert LLM guardrail result back to MCP response format. """ @@ -802,7 +802,9 @@ class ProxyLogging: verbose_proxy_logger.error(f"Error in manual argument parsing: {e}") return None - def _convert_llm_result_to_mcp_during_response(self, llm_result, request_obj) -> Optional[Any]: + def _convert_llm_result_to_mcp_during_response( + self, llm_result, request_obj + ) -> MCPDuringCallResponseObject | None: """ Convert LLM guardrail result back to MCP during call response format. """ @@ -1142,8 +1144,8 @@ class ProxyLogging: self, data: dict, litellm_logging_obj: Any, - prompt_id: Any, - prompt_version: Any, + prompt_id: str, + prompt_version: int | None, call_type: CallTypesLiteral, ) -> None: """Process prompt template if applicable.""" @@ -1438,9 +1440,7 @@ class ProxyLogging: data = result elif ( - _callback is not None - and isinstance(_callback, CustomLogger) - and "async_pre_call_hook" in vars(_callback.__class__) + "async_pre_call_hook" in vars(_callback.__class__) and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook ): if call_type == "call_mcp_tool" and user_api_key_dict is None: @@ -1614,7 +1614,7 @@ class ProxyLogging: break @staticmethod - async def _run_guardrail_with_metrics(callback: Any, coro: Awaitable[Any], hook_type: str) -> Any: + async def _run_guardrail_with_metrics(callback: CustomLogger, coro: Awaitable[Any], hook_type: str) -> Any: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and @@ -1646,7 +1646,7 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( - callback: Any, gen: AsyncGenerator[Any, None] + callback: CustomLogger, gen: AsyncGenerator[Any, None] ) -> AsyncGenerator[Any, None]: """ Yield from `gen`; if iteration raises an HTTPException with dict detail, @@ -1691,8 +1691,11 @@ class ProxyLogging: has_streaming_chunk_override = False has_guardrail = False has_pre_call_override = False - iterator_overrides: List[Tuple[Any, str]] = [] # (callback, kind) - resolved_callbacks: List[Any] = [] + # (callback, kind) pairs; built alongside the 4 flags above and the + # resolved_callbacks list below in one pass over callbacks, since a + # functional per-field rewrite would re-walk this perf-sensitive scan. + iterator_overrides: list[tuple[CustomLogger, str]] = [] # mutable-ok: see comment above + resolved_callbacks: list[CustomLogger] = [] # mutable-ok: see comment above for callback in callbacks: if isinstance(callback, str): @@ -2795,7 +2798,7 @@ _DEPRECATED_KEY_CACHE_TTL_SECONDS = 60 async def _lookup_deprecated_key( - db: Any, + db: Union[PrismaWrapper, RoutingPrismaWrapper], hashed_token: str, ) -> Optional[str]: """ @@ -2875,7 +2878,7 @@ def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]: return None -async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any]: +async def get_config_param(prisma_client: "PrismaClient", param_name: str) -> Any | None: """Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None.""" cache_key = _config_cache_key(param_name) cached = await litellm_config_cache.async_get_cache(cache_key) @@ -2893,7 +2896,7 @@ async def invalidate_config_param(param_name: str) -> None: await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) -async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: +async def prefetch_config_params(prisma_client: "PrismaClient", param_names: Sequence[str]) -> None: """Batch-load LiteLLM_Config rows into the cache with one find_many.""" if not param_names: return @@ -5932,7 +5935,10 @@ def _check_and_merge_model_level_guardrails( return _merge_guardrails_with_existing(data, model_level_guardrails) -def _merge_guardrails_with_existing(data: dict, model_level_guardrails: Any) -> dict: +def _merge_guardrails_with_existing( + data: dict, # mutable-ok: request payload dict, needs .copy()/.setdefault() + model_level_guardrails: Sequence[str] | str | None, +) -> dict: # mutable-ok: returns a merged copy of the request dict, not a read-only view """ Merge model-level guardrails with any existing guardrails in the request data. @@ -5952,11 +5958,13 @@ def _merge_guardrails_with_existing(data: dict, model_level_guardrails: Any) -> existing_guardrails = [existing_guardrails] if existing_guardrails else [] # Ensure model_level_guardrails is a list - if not isinstance(model_level_guardrails, list): - model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] + if isinstance(model_level_guardrails, list): + normalized_model_level_guardrails = model_level_guardrails + else: + normalized_model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] # Combine existing and model-level guardrails - metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) + metadata["guardrails"] = list(set(existing_guardrails + normalized_model_level_guardrails)) return modified_data diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 3b5ec5b0dee..367c1ff34f0 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3118 + "limit": 3116 }, "ANN002": { "limit": 69 @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 2013 + "limit": 1915 }, "ASYNC230": { "limit": 14 @@ -222,7 +222,7 @@ "limit": 38 }, "RET504": { - "limit": 716 + "limit": 710 }, "RUF010": { "limit": 874 @@ -324,7 +324,7 @@ "limit": 883 }, "UP006": { - "limit": 12146 + "limit": 12137 }, "UP007": { "limit": 2526 @@ -363,6 +363,6 @@ "limit": 105 }, "UP045": { - "limit": 17806 + "limit": 17803 } } diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 25cc1621d54..c42e889c451 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23267 + "limit": 23256 }, "LIT002": { "limit": 27434 @@ -24,6 +24,6 @@ "limit": 1004 }, "LIT009": { - "limit": 2474 + "limit": 2472 } }