From 9c0cd6e02e24a48a97215761cb098cb28e7aea7f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 11 Mar 2026 14:34:43 -0700 Subject: [PATCH] Fix ruff lint errors: unused imports, undefined names, and statement counts Remove unused imports (F401), add missing TYPE_CHECKING imports for forward references (F821), and extract helpers to reduce statement counts below the 50-statement limit (PLR0915). Co-Authored-By: Claude Opus 4.6 --- litellm/litellm_core_utils/core_helpers.py | 2 +- litellm/main.py | 1 + .../guardrails/guardrail_hooks/presidio.py | 21 +++++--- .../key_management_endpoints.py | 51 ++++++++++++------- .../tool_management_endpoints.py | 1 - .../object_permission_utils.py | 8 ++- 6 files changed, 57 insertions(+), 27 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 85ed955af4a..ee111f35929 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1,6 +1,6 @@ # What is this? ## Helper utilities -from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union, get_args +from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union import httpx diff --git a/litellm/main.py b/litellm/main.py index 2b210c79a56..4e4ce976ac4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -63,6 +63,7 @@ from litellm.utils import exception_type, get_litellm_params, get_optional_param # Logging is imported lazily when needed to avoid loading litellm_logging at import time if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import TokenCountResponse from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index ddeba2100c3..4ce0f3ef5e8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1228,13 +1228,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # --- PRESERVE USAGE METADATA --- # stream_chunk_builder might miss usage if it's only in the last chunk - if ( - not getattr(assembled_model_response, "usage", None) - ) and remaining_chunks: - last_chunk = remaining_chunks[-1] - last_chunk_usage = getattr(last_chunk, "usage", None) - if last_chunk_usage: - setattr(assembled_model_response, "usage", last_chunk_usage) + self._preserve_usage_from_last_chunk( + assembled_model_response, remaining_chunks + ) # Apply PII unmasking to assembled content (unmasking tokens back to original text) await self._process_response_for_pii( @@ -1253,6 +1249,17 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for chunk in remaining_chunks: yield chunk + @staticmethod + def _preserve_usage_from_last_chunk( + assembled_model_response: Any, + chunks: List[Any], + ) -> None: + """Copy usage metadata from the last chunk when stream_chunk_builder misses it.""" + if not getattr(assembled_model_response, "usage", None) and chunks: + last_chunk_usage = getattr(chunks[-1], "usage", None) + if last_chunk_usage: + setattr(assembled_model_response, "usage", last_chunk_usage) + def get_presidio_settings_from_request_data( self, data: dict ) -> Optional[PresidioPerRequestConfig]: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b9dcc514d29..654205252b6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1756,6 +1756,34 @@ async def _process_single_key_update( return updated_key_info +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, +) -> None: + """Validate MCP servers in object_permission against the effective team.""" + effective_team_obj = team_obj + # If team_id isn't being changed, resolve the existing key's team + if effective_team_obj is None and existing_key_row.team_id: + effective_team_obj = await get_team_object( + team_id=existing_key_row.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + object_permission_dict = ( + data.object_permission.model_dump() + if hasattr(data.object_permission, "model_dump") + else data.object_permission + ) + await validate_key_mcp_servers_against_team( + object_permission=object_permission_dict, + team_obj=effective_team_obj, + ) + + @router.post( "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) @@ -1956,23 +1984,12 @@ async def update_key_fn( # Validate MCP servers in object_permission against the effective team if data.object_permission is not None: - effective_team_obj = team_obj - # If team_id isn't being changed, resolve the existing key's team - if effective_team_obj is None and existing_key_row.team_id: - effective_team_obj = await get_team_object( - team_id=existing_key_row.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - object_permission_dict = ( - data.object_permission.model_dump() - if hasattr(data.object_permission, "model_dump") - else data.object_permission - ) - await validate_key_mcp_servers_against_team( - object_permission=object_permission_dict, - team_obj=effective_team_obj, + await _validate_mcp_servers_for_key_update( + data=data, + team_obj=team_obj, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, ) non_default_values = await prepare_key_update_data( diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 7fdd3475c04..19ca2c9f6be 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -26,7 +26,6 @@ from litellm.types.tool_management import ( ToolDetailResponse, ToolInputPolicy, ToolListResponse, - ToolOutputPolicy, ToolPolicyOption, ToolPolicyOptionsResponse, ToolPolicyUpdateRequest, diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 319a0b5eb79..0f426bf6045 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -4,7 +4,7 @@ organizations, teams, and keys. """ import json -from typing import Dict, List, Optional, Set, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union from fastapi import HTTPException, status @@ -12,6 +12,12 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient + +if TYPE_CHECKING: + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTableCachedObj, + )