From 5a62806fdc0902850a858cc0525b15a37d76e4ac Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 16 Jun 2026 16:52:49 -0700 Subject: [PATCH 1/3] chore(lint): remove PLR0915 too-many-statements ruff rule (#30574) Drops PLR0915 from ruff's extend-select along with its per-file-ignores, and strips the now-unused `# noqa: PLR0915` directives across the codebase (RUF100 would otherwise flag them as unused). The C901 suppression that shared a directive with PLR0915 in streaming_handler.py is preserved. --- db_scripts/create_views.py | 8 +++----- .../proxy/hooks/managed_files.py | 2 +- .../management_endpoints/project_endpoints.py | 2 +- litellm/_redis.py | 2 +- litellm/a2a_protocol/main.py | 2 +- litellm/batches/main.py | 2 +- litellm/caching/caching_handler.py | 2 +- litellm/caching/qdrant_semantic_cache.py | 2 +- .../transformation.py | 2 +- litellm/cost_calculator.py | 4 ++-- litellm/images/main.py | 4 ++-- .../SlackAlerting/slack_alerting.py | 4 ++-- litellm/integrations/braintrust_logging.py | 8 ++------ litellm/integrations/langfuse/langfuse.py | 2 +- litellm/integrations/mock_client_factory.py | 2 +- litellm/integrations/opentelemetry.py | 4 +--- litellm/integrations/prometheus.py | 4 ++-- .../exception_mapping_utils.py | 2 +- .../get_llm_provider_logic.py | 4 ++-- .../get_supported_openai_params.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 14 +++++++------- .../litellm_core_utils/llm_cost_calc/utils.py | 2 +- .../convert_dict_to_response.py | 2 +- .../prompt_templates/factory.py | 12 ++++++------ .../litellm_core_utils/realtime_streaming.py | 2 +- .../streaming_chunk_builder_utils.py | 2 +- .../litellm_core_utils/streaming_handler.py | 8 ++++---- litellm/llms/anthropic/chat/handler.py | 2 +- litellm/llms/anthropic/chat/transformation.py | 4 ++-- .../adapters/streaming_iterator.py | 4 ++-- .../adapters/transformation.py | 2 +- .../context_management/editors/compact.py | 4 ++-- .../responses_adapters/streaming_iterator.py | 2 +- .../responses_adapters/transformation.py | 2 +- litellm/llms/bedrock/chat/converse_handler.py | 2 +- .../bedrock/chat/converse_transformation.py | 2 +- litellm/llms/bedrock/chat/invoke_handler.py | 4 ++-- litellm/llms/bedrock/embed/embedding.py | 2 +- .../image_edit/stability_transformation.py | 2 +- .../guardrail_translation/handler.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 2 +- litellm/llms/gemini/realtime/transformation.py | 2 +- .../huggingface/embedding/transformation.py | 2 +- litellm/llms/openai/openai.py | 2 +- litellm/llms/predibase/chat/transformation.py | 2 +- .../llms/vertex_ai/gemini/transformation.py | 4 ++-- .../vertex_and_google_ai_studio_gemini.py | 10 ++++------ .../batch_embed_content_handler.py | 2 +- litellm/llms/vertex_ai/vertex_ai_non_gemini.py | 4 ++-- litellm/main.py | 12 ++++++------ .../mcp_server/auth/user_api_key_auth_mcp.py | 2 +- .../mcp_server/mcp_server_manager.py | 2 +- .../mcp_server/sampling_handler.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 8 ++++---- litellm/proxy/agent_endpoints/a2a_endpoints.py | 2 +- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/handle_jwt.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 4 ++-- litellm/proxy/batches_endpoints/endpoints.py | 4 ++-- litellm/proxy/common_request_processing.py | 4 ++-- litellm/proxy/common_utils/callback_utils.py | 2 +- litellm/proxy/db/create_views.py | 2 +- litellm/proxy/db/db_spend_update_writer.py | 2 +- .../guardrails/guardrail_hooks/lakera_ai.py | 2 +- .../panw_prisma_airs/panw_prisma_airs.py | 4 ++-- .../unified_guardrail/unified_guardrail.py | 2 +- .../health_endpoints/_health_endpoints.py | 2 +- .../proxy/hooks/parallel_request_limiter.py | 6 ++---- litellm/proxy/litellm_pre_call_utils.py | 2 +- .../common_daily_activity.py | 2 +- .../key_management_endpoints.py | 6 +++--- .../management_endpoints/team_endpoints.py | 4 ++-- litellm/proxy/management_endpoints/ui_sso.py | 2 +- .../openai_files_endpoints/files_endpoints.py | 4 ++-- .../anthropic_passthrough_logging_handler.py | 4 ++-- .../openai_passthrough_logging_handler.py | 2 +- .../vertex_passthrough_logging_handler.py | 2 +- .../pass_through_endpoints.py | 8 ++++---- litellm/proxy/proxy_cli.py | 2 +- litellm/proxy/proxy_server.py | 18 +++++++++--------- .../response_polling/background_streaming.py | 2 +- litellm/proxy/route_llm_request.py | 2 +- .../spend_management_endpoints.py | 4 ++-- .../spend_tracking/spend_tracking_utils.py | 4 +--- litellm/realtime_api/main.py | 2 +- litellm/rerank_api/main.py | 2 +- .../responses/mcp/chat_completions_handler.py | 2 +- .../responses/mcp/litellm_proxy_mcp_handler.py | 2 +- litellm/router.py | 14 +++++++------- litellm/router_strategy/lowest_cost.py | 2 +- litellm/router_strategy/lowest_latency.py | 10 +++------- litellm/router_strategy/lowest_tpm_rpm.py | 2 +- litellm/secret_managers/main.py | 2 +- .../secret_managers/secret_manager_handler.py | 2 +- litellm/types/utils.py | 4 ++-- litellm/utils.py | 18 +++++++++--------- ruff.toml | 9 ++------- .../test_pass_through_endpoints.py | 2 +- 98 files changed, 177 insertions(+), 200 deletions(-) diff --git a/db_scripts/create_views.py b/db_scripts/create_views.py index 3027b38958d..2b34664452d 100644 --- a/db_scripts/create_views.py +++ b/db_scripts/create_views.py @@ -15,7 +15,7 @@ db = Prisma( ) -async def check_view_exists(): # noqa: PLR0915 +async def check_view_exists(): """ Checks if the LiteLLM_VerificationTokenView and MonthlyGlobalSpend exists in the user's db. @@ -34,8 +34,7 @@ async def check_view_exists(): # noqa: PLR0915 print("LiteLLM_VerificationTokenView Exists!") # noqa except Exception: # If an error occurs, the view does not exist, so create it - await db.execute_raw( - """ + await db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -45,8 +44,7 @@ async def check_view_exists(): # noqa: PLR0915 t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """ - ) + """) print("LiteLLM_VerificationTokenView Created!") # noqa diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index a1f63f388b4..6830147116d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -412,7 +412,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", ) - async def async_pre_call_hook( # noqa: PLR0915 + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 75229bacc8f..a057df65500 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -483,7 +483,7 @@ async def new_project( response_model=LiteLLM_ProjectTable, ) @management_endpoint_wrapper -async def update_project( # noqa: PLR0915 +async def update_project( data: UpdateProjectRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/litellm/_redis.py b/litellm/_redis.py index e2b04f795cb..1b6e1a5e4b0 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -311,7 +311,7 @@ def get_redis_url_from_environment(): return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}" -def _get_redis_client_logic(**env_overrides): # noqa: PLR0915 +def _get_redis_client_logic(**env_overrides): """ Common functionality across sync + async redis client implementations """ diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index dcb5cb74ec4..2b6f2cd12b4 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -436,7 +436,7 @@ def _build_streaming_logging_obj( return logging_obj -async def asend_message_streaming( # noqa: PLR0915 +async def asend_message_streaming( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendStreamingMessageRequest"] = None, api_base: Optional[str] = None, diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 15ee9303969..f124882b5a4 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -157,7 +157,7 @@ async def acreate_batch( @client -def create_batch( # noqa: PLR0915 +def create_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 48691335b40..2a8bd856040 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -394,7 +394,7 @@ class LLMCachingHandler: return cr["model"] return None - def _process_async_embedding_cached_response( # noqa: PLR0915 + def _process_async_embedding_cached_response( self, final_embedding_cached_response: Optional[EmbeddingResponse], cached_result: List[Optional[CachedEmbedding]], diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index cb521efca05..68d3b8c20b3 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -28,7 +28,7 @@ from .base_cache import BaseCache class QdrantSemanticCache(BaseCache): CACHE_KEY_FIELD_NAME = "litellm_cache_key" - def __init__( # noqa: PLR0915 + def __init__( self, qdrant_api_base=None, qdrant_api_key=None, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index dabf09f8b2a..3fa6b983e5f 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1211,7 +1211,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): return self.chunk_parser(json.loads(str_line)) @staticmethod - def translate_responses_chunk_to_openai_stream( # noqa: PLR0915 + def translate_responses_chunk_to_openai_stream( parsed_chunk: Union[dict, BaseModel], ) -> "ModelResponseStream": """ diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e934c6a6f83..5c77400651b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -288,7 +288,7 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) -def cost_per_token( # noqa: PLR0915 +def cost_per_token( model: str = "", prompt_tokens: int = 0, completion_tokens: int = 0, @@ -1136,7 +1136,7 @@ def _store_cost_breakdown_in_logging_obj( pass -def completion_cost( # noqa: PLR0915 +def completion_cost( completion_response=None, model: Optional[str] = None, prompt="", diff --git a/litellm/images/main.py b/litellm/images/main.py index d95b7287d20..8b108ded4c9 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -195,7 +195,7 @@ def image_generation( @client -def image_generation( # noqa: PLR0915 +def image_generation( prompt: str, model: Optional[str] = None, n: Optional[int] = None, @@ -738,7 +738,7 @@ def image_variation( @client -def image_edit( # noqa: PLR0915 +def image_edit( image: Optional[Union[FileTypes, List[FileTypes]]] = None, prompt: Optional[str] = None, model: Optional[str] = None, diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index e7be004e62e..2108ebae312 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -351,7 +351,7 @@ class SlackAlerting(CustomBatchLogger): except Exception: return 0 - async def send_daily_reports(self, router) -> bool: # noqa: PLR0915 + async def send_daily_reports(self, router) -> bool: """ Send a daily report on: - Top 5 deployments with most failed requests @@ -1373,7 +1373,7 @@ Model Info: return False - async def send_alert( # noqa: PLR0915 + async def send_alert( self, message: str, level: Literal["Low", "Medium", "High"], diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 9b1c5077882..6a6313f72e1 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -133,9 +133,7 @@ class BraintrustLogger(CustomLogger): self.default_project_id = project_dict["id"] - def log_success_event( # noqa: PLR0915 - self, kwargs, response_obj, start_time, end_time - ): + def log_success_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") try: litellm_call_id = kwargs.get("litellm_call_id") @@ -271,9 +269,7 @@ class BraintrustLogger(CustomLogger): except Exception as e: raise e # don't use verbose_logger.exception, if exception is raised - async def async_log_success_event( # noqa: PLR0915 - self, kwargs, response_obj, start_time, end_time - ): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") try: litellm_call_id = kwargs.get("litellm_call_id") diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 0efc7d66876..b1c6956a16c 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -549,7 +549,7 @@ class LangFuseLogger: ) ) - def _log_langfuse_v2( # noqa: PLR0915 + def _log_langfuse_v2( self, user_id: Optional[str], metadata: dict, diff --git a/litellm/integrations/mock_client_factory.py b/litellm/integrations/mock_client_factory.py index 02a927fe64f..9b912ce70c8 100644 --- a/litellm/integrations/mock_client_factory.py +++ b/litellm/integrations/mock_client_factory.py @@ -107,7 +107,7 @@ def _is_url_match(url, matchers: List[str]) -> bool: return False -def create_mock_client_factory(config: MockClientConfig): # noqa: PLR0915 +def create_mock_client_factory(config: MockClientConfig): """ Factory function that creates mock client functions based on configuration. diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index fc37b6a34d8..6b50ef49b49 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -2198,9 +2198,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return kv_pairs - def set_attributes( # noqa: PLR0915 - self, span: Span, kwargs, response_obj: Optional[Any] - ): + def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): try: if self.callback_name == "langtrace": from litellm.integrations.langtrace import LangtraceAttributes diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2119527a8e5..c63f114514a 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -75,7 +75,7 @@ class PrometheusLogger(CustomLogger): return cb return None - def __init__( # noqa: PLR0915 + def __init__( self, **kwargs, ): @@ -2255,7 +2255,7 @@ class PrometheusLogger(CustomLogger): or _litellm_params_metadata.get("user_agent"), } - def set_llm_deployment_failure_metrics(self, request_kwargs: dict): # noqa: PLR0915 + def set_llm_deployment_failure_metrics(self, request_kwargs: dict): """ Sets Failure metrics when an LLM API call fails diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 0d35da9fa1a..6087e55b136 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -234,7 +234,7 @@ def extract_and_raise_litellm_exception( ) -def exception_type( # type: ignore # noqa: PLR0915 +def exception_type( # type: ignore model, original_exception, custom_llm_provider, diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 182c5117a3e..4941d52d7d6 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -154,7 +154,7 @@ def handle_anthropic_text_model_custom_llm_provider( return model, custom_llm_provider -def get_llm_provider( # noqa: PLR0915 +def get_llm_provider( model: str, custom_llm_provider: Optional[str] = None, api_base: Optional[str] = None, @@ -568,7 +568,7 @@ def get_llm_provider( # noqa: PLR0915 ) -def _get_openai_compatible_provider_info( # noqa: PLR0915 +def _get_openai_compatible_provider_info( model: str, api_base: Optional[str], api_key: Optional[str], diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 65c238344e9..e87042b9101 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -5,7 +5,7 @@ from litellm.exceptions import BadRequestError from litellm.types.utils import LlmProviders, LlmProvidersSet -def get_supported_openai_params( # noqa: PLR0915 +def get_supported_openai_params( model: str, custom_llm_provider: Optional[str] = None, request_type: Literal[ diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3c525f743ed..0e9c3783316 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -986,7 +986,7 @@ class Logging(LiteLLMLoggingBaseClass): self._get_masked_api_base(additional_args.get("api_base", "")) ) - def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 + def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API litellm.error_logs["PRE_CALL"] = locals() try: @@ -2119,7 +2119,7 @@ class Logging(LiteLLMLoggingBaseClass): await self.async_success_handler(result=complete_streaming_response) return - def success_handler( # noqa: PLR0915 + def success_handler( self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs ): verbose_logger.debug( @@ -2584,7 +2584,7 @@ class Logging(LiteLLMLoggingBaseClass): ), ) - async def async_success_handler( # noqa: PLR0915 + async def async_success_handler( self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs ): """ @@ -3036,7 +3036,7 @@ class Logging(LiteLLMLoggingBaseClass): kwargs=self.model_call_details, ) # type: ignore - def failure_handler( # noqa: PLR0915 + def failure_handler( self, exception, traceback_exception, start_time=None, end_time=None ): verbose_logger.debug( @@ -3753,7 +3753,7 @@ def _get_masked_values( } -def set_callbacks(callback_list, function_id=None): # noqa: PLR0915 +def set_callbacks(callback_list, function_id=None): """ Globally sets the callback client """ @@ -3854,7 +3854,7 @@ def set_callbacks(callback_list, function_id=None): # noqa: PLR0915 return None -def _init_custom_logger_compatible_class( # noqa: PLR0915 +def _init_custom_logger_compatible_class( logging_integration: _custom_logger_compatible_callbacks_literal, internal_usage_cache: Optional[DualCache], llm_router: Optional[ @@ -4611,7 +4611,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: ) -def get_custom_logger_compatible_class( # noqa: PLR0915 +def get_custom_logger_compatible_class( logging_integration: _custom_logger_compatible_callbacks_literal, ) -> Optional[CustomLogger]: try: diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index d75850984a9..a7ac5b53349 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -683,7 +683,7 @@ def _get_regional_uplift_multiplier( return 1.0 -def generic_cost_per_token( # noqa: PLR0915 +def generic_cost_per_token( model: str, usage: Usage, custom_llm_provider: str, diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 4e5b53a13d7..016bb6b1e22 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -471,7 +471,7 @@ def _should_convert_tool_call_to_json_mode( return False -def convert_to_model_response_object( # noqa: PLR0915 +def convert_to_model_response_object( response_object: Optional[dict] = None, model_response_object: Optional[ Union[ diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 5059e612f2f..b95b73398ac 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1475,7 +1475,7 @@ def convert_to_gemini_tool_call_invoke( ) -def convert_to_gemini_tool_call_result( # noqa: PLR0915 +def convert_to_gemini_tool_call_result( message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage], last_message_with_tool_calls: Optional[dict], model: Optional[str] = None, @@ -2227,7 +2227,7 @@ def _sanitize_empty_text_content( return message -def _add_missing_tool_results( # noqa: PLR0915 +def _add_missing_tool_results( current_message: AllMessageValues, messages: List[AllMessageValues], current_index: int, @@ -2484,7 +2484,7 @@ def sanitize_messages_for_tool_calling( return sanitized_messages -def anthropic_messages_pt( # noqa: PLR0915 +def anthropic_messages_pt( messages: List[AllMessageValues], model: str, llm_provider: str, @@ -3278,7 +3278,7 @@ def convert_to_cohere_tool_invoke(tool_calls: list) -> List[ToolCallObject]: return cohere_tool_invoke -def cohere_messages_pt_v2( # noqa: PLR0915 +def cohere_messages_pt_v2( messages: List, model: str, llm_provider: str, @@ -4703,7 +4703,7 @@ class BedrockConverseMessagesProcessor: return messages @staticmethod - async def _bedrock_converse_messages_pt_async( # noqa: PLR0915 + async def _bedrock_converse_messages_pt_async( messages: List, model: str, llm_provider: str, @@ -5133,7 +5133,7 @@ class BedrockConverseMessagesProcessor: return assistant_parts -def _bedrock_converse_messages_pt( # noqa: PLR0915 +def _bedrock_converse_messages_pt( messages: List, model: str, llm_provider: str, diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c8f87d96e2f..c56a70177bf 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1198,7 +1198,7 @@ class RealTimeStreaming: item["content"] = new_content return item - async def client_ack_messages(self): # noqa: PLR0915 + async def client_ack_messages(self): try: while True: message = await self.websocket.receive_text() diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d51b937d434..04f6b1241c3 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -209,7 +209,7 @@ class ChunkProcessor: ) return response - def get_combined_tool_content( # noqa: PLR0915 + def get_combined_tool_content( self, tool_call_chunks: List[Dict[str, Any]] ) -> List[ChatCompletionMessageToolCall]: tool_calls_list: List[ChatCompletionMessageToolCall] = [] diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7e4bf895a79..888a9658396 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -967,7 +967,7 @@ class CustomStreamWrapper: delta, model_response.choices[0].delta, attribute ) - def return_processed_chunk_logic( # noqa: PLR0915, C901 + def return_processed_chunk_logic( # noqa: C901 self, completion_obj: Dict[str, Any], model_response: ModelResponseStream, @@ -1145,7 +1145,7 @@ class CustomStreamWrapper: del model_response.choices[0].delta.reasoning_content return - def chunk_creator(self, chunk: Any): # type: ignore # noqa: PLR0915 + def chunk_creator(self, chunk: Any): # type: ignore if hasattr(chunk, "id"): self.response_id = chunk.id model_response = self.model_response_creator() @@ -1887,7 +1887,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = "tool_calls" return model_response - def __next__(self) -> "ModelResponseStream": # noqa: PLR0915 + def __next__(self) -> "ModelResponseStream": cache_hit = False if ( self.custom_llm_provider is not None @@ -2077,7 +2077,7 @@ class CustomStreamWrapper: return self.completion_stream - async def __anext__(self) -> "ModelResponseStream": # noqa: PLR0915 + async def __anext__(self) -> "ModelResponseStream": cache_hit = False if ( self.custom_llm_provider is not None diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 2fb29b32a61..5d14f3cc4ae 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -772,7 +772,7 @@ class ModelResponseIterator: ) return results - def chunk_parser(self, chunk: dict) -> ModelResponseStream: # noqa: PLR0915 + def chunk_parser(self, chunk: dict) -> ModelResponseStream: try: type_chunk = chunk.get("type", "") or "" diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e8c1e659e9f..cf97c946f1c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -605,7 +605,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) return _tool_choice - def _map_tool_helper( # noqa: PLR0915 + def _map_tool_helper( self, tool: ChatCompletionToolParam, ) -> Tuple[Optional[AllAnthropicToolsValues], Optional[AnthropicMcpServerTool]]: @@ -1399,7 +1399,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return None - def map_openai_params( # noqa: PLR0915 + def map_openai_params( self, non_default_params: dict, optional_params: dict, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index f049abcf47f..a8e2fceb4ee 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -372,7 +372,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): cache_read_input_tokens=0, ) - def __next__(self): # noqa: PLR0915 + def __next__(self): from .transformation import LiteLLMAnthropicMessagesAdapter try: @@ -618,7 +618,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): ) raise StopIteration - async def __anext__(self): # noqa: PLR0915 + async def __anext__(self): from .transformation import LiteLLMAnthropicMessagesAdapter try: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index af868051f4d..bf425637b56 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -383,7 +383,7 @@ class LiteLLMAnthropicMessagesAdapter: isinstance(tool_type, str) and tool_type.startswith("web_search") ) or tool_name == "web_search" - def translate_anthropic_messages_to_openai( # noqa: PLR0915 + def translate_anthropic_messages_to_openai( self, messages: List[ Union[ diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index 4aae85b17fe..6479ee999b0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -97,7 +97,7 @@ def _read_summary_max_tokens_setting() -> int: return COMPACT_SUMMARY_MAX_TOKENS -async def _check_summary_model_access( # noqa: PLR0915 +async def _check_summary_model_access( user_api_key_auth: Any, summary_model: str, llm_router: Any, @@ -970,7 +970,7 @@ def apply_client_compaction_block_history( ) -async def apply_compact_20260112( # noqa: PLR0915 +async def apply_compact_20260112( *, model: str, messages: List[Dict[str, Any]], diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 5f1362e259f..04819a416a2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -66,7 +66,7 @@ class AnthropicResponsesStreamWrapper: self._current_block_index += 1 return self._current_block_index - def _process_event(self, event: Any) -> None: # noqa: PLR0915 + def _process_event(self, event: Any) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) if event_type is None and isinstance(event, dict): diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 2badc2a3276..4fb1ddf5c46 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -51,7 +51,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return source.get("url") return None - def translate_messages_to_responses_input( # noqa: PLR0915 + def translate_messages_to_responses_input( self, messages: List[ Union[ diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7e1020000f4..7b1064ccef9 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -248,7 +248,7 @@ class BedrockConverseLLM(BaseAWSLLM): encoding=encoding, ) - def completion( # noqa: PLR0915 + def completion( self, model: str, messages: list, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b5e5e4de6fc..bb261ec85b2 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2189,7 +2189,7 @@ class AmazonConverseConfig(BaseConfig): real_tools = [t for i, t in enumerate(tools) if i not in json_tool_indices] return real_tools if real_tools else None - def _transform_response( # noqa: PLR0915 + def _transform_response( self, model: str, response: httpx.Response, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 0a1322a751e..75b560b4d6d 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -473,7 +473,7 @@ class BedrockLLM(BaseAWSLLM): prompt += f"{message['content']}" return prompt, chat_history # type: ignore - def process_response( # noqa: PLR0915 + def process_response( self, model: str, response: httpx.Response, @@ -765,7 +765,7 @@ class BedrockLLM(BaseAWSLLM): return model_response - def completion( # noqa: PLR0915 + def completion( self, model: str, messages: list, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 27dc785bf57..b6aa99842d7 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -388,7 +388,7 @@ class BedrockEmbedding(BaseAWSLLM): batch_data=batch_data, ) - def embeddings( # noqa: PLR0915 + def embeddings( self, model: str, input: List[str], diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index 2d73e47003d..d00d62a8530 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -149,7 +149,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): return mapped_params - def transform_image_edit_request( # noqa: PLR0915 + def transform_image_edit_request( self, model: str, prompt: Optional[str], diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 2d6bdb5298a..0522bb249e1 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -226,7 +226,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): return _is_converse_endpoint(endpoint) @staticmethod - async def de_anonymize_event_stream( # noqa: PLR0915 + async def de_anonymize_event_stream( body_bytes: bytes, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 5575385fb28..8ac5b47c6e7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5676,7 +5676,7 @@ class BaseLLMHTTPHandler: ) raise - async def async_responses_websocket( # noqa: PLR0915 + async def async_responses_websocket( self, model: str, websocket: Any, diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 51fa395d899..74f6cd4d831 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -1378,7 +1378,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): raise ValueError(f"Unknown openai event: {key}, value: {value}") return openai_event - def transform_realtime_response( # noqa: PLR0915 + def transform_realtime_response( self, message: Union[str, bytes], model: str, diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 88d42cfcdcc..7cddda617a9 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -404,7 +404,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): ) return completion_response - def convert_to_model_response_object( # noqa: PLR0915 + def convert_to_model_response_object( self, completion_response: Union[List[Dict[str, Any]], Dict[str, Any]], model_response: ModelResponse, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 194f29648c4..ea905d8ebca 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -608,7 +608,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): return streaming_response - def completion( # type: ignore # noqa: PLR0915 + def completion( # type: ignore self, model_response: ModelResponse, timeout: Union[float, httpx.Timeout], diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 3d251d24b0d..ce004f60bfc 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -129,7 +129,7 @@ class PredibaseConfig(BaseConfig): optional_params["response_format"] = value return optional_params - def transform_response( # noqa: PLR0915 + def transform_response( self, model: str, raw_response: Response, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index c578d6cd28b..f5a2b268263 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -678,7 +678,7 @@ def check_if_part_exists_in_parts( return False -def _gemini_convert_messages_with_history( # noqa: PLR0915 +def _gemini_convert_messages_with_history( messages: List[AllMessageValues], model: Optional[str] = None, litellm_params: Optional[dict] = None, @@ -1176,7 +1176,7 @@ def _rewrite_google_maps_response_format(data: RequestBody) -> None: _rewrite_mime_type_to_response_format(generation_config) -def _transform_request_body( # noqa: PLR0915 +def _transform_request_body( messages: List[AllMessageValues], model: str, optional_params: dict, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 3ec7b0814dd..dab21e2ce8e 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -614,9 +614,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return googleSearch, googleSearchRetrieval, enterpriseWebSearch, urlContext - def _map_function( # noqa: PLR0915 - self, value: List[dict], optional_params: dict - ) -> List[Tools]: + def _map_function(self, value: List[dict], optional_params: dict) -> List[Tools]: """ Map OpenAI-style tools/functions to Vertex AI format. @@ -1173,7 +1171,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params["include_server_side_tool_invocations"] = True return - def map_openai_params( # noqa: PLR0915 + def map_openai_params( self, non_default_params: Dict, optional_params: Dict, @@ -1904,7 +1902,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return False @staticmethod - def _calculate_usage( # noqa: PLR0915 + def _calculate_usage( completion_response: Union[ GenerateContentResponseBody, BidiGenerateContentServerMessage ], @@ -2380,7 +2378,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return annotations @staticmethod - def _process_candidates( # noqa: PLR0915 + def _process_candidates( _candidates: List[Candidates], model_response: Union[ModelResponse, "ModelResponseStream"], standard_optional_params: dict, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 99165c37c93..165dac24903 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -124,7 +124,7 @@ class GoogleBatchEmbeddings(VertexLLM): return resolved_files - def batch_embeddings( # noqa: PLR0915 + def batch_embeddings( self, model: str, input: GeminiEmbeddingInput, diff --git a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py index 222820d7ee5..c134dee7ad4 100644 --- a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py +++ b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py @@ -77,7 +77,7 @@ def _set_client_in_cache(client_cache_key: str, vertex_llm_model: Any): ) -def completion( # noqa: PLR0915 +def completion( model: str, messages: list, model_response: ModelResponse, @@ -485,7 +485,7 @@ def completion( # noqa: PLR0915 ) -async def async_completion( # noqa: PLR0915 +async def async_completion( llm_model, mode: str, prompt: str, diff --git a/litellm/main.py b/litellm/main.py index a1bade9bb1f..80176cc8b16 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1086,7 +1086,7 @@ def _build_custom_pricing_entry( @tracer.wrap() @client -def completion( # type: ignore # noqa: PLR0915 +def completion( # type: ignore model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create messages: List = [], @@ -4878,7 +4878,7 @@ def embedding( @client -def embedding( # noqa: PLR0915 +def embedding( model, input=[], # Optional params @@ -6125,7 +6125,7 @@ async def atext_completion( @client -def text_completion( # noqa: PLR0915 +def text_completion( prompt: Union[ str, List[Union[str, List[Union[str, List[int]]]]] ], # Required: The prompt(s) to generate completions for. @@ -6664,7 +6664,7 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: @client -def transcription( # noqa: PLR0915 +def transcription( model: str, file: FileTypes, ## OPTIONAL OPENAI PARAMS ## @@ -6971,7 +6971,7 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent: @client -def speech( # noqa: PLR0915 +def speech( model: str, input: str, voice: Optional[Union[str, dict]] = None, @@ -7662,7 +7662,7 @@ def stream_chunk_builder_text_completion( return TextCompletionResponse(**response) -def stream_chunk_builder( # noqa: PLR0915 +def stream_chunk_builder( chunks: list, messages: Optional[list] = None, start_time=None, diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 1535daeb01d..e47fc84b533 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -125,7 +125,7 @@ class MCPRequestHandler: LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME = SpecialHeaders.mcp_access_groups.value @staticmethod - async def process_mcp_request( # noqa: PLR0915 + async def process_mcp_request( scope: Scope, ) -> Tuple[ UserAPIKeyAuth, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 5e419b5c0a3..3c3f2afad6d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3355,7 +3355,7 @@ class MCPServerManager: ) ) - async def _call_regular_mcp_tool( # noqa: PLR0915 + async def _call_regular_mcp_tool( self, mcp_server: MCPServer, original_tool_name: str, diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index 1637c9eb0b9..b659ba6f813 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -661,7 +661,7 @@ def _convert_openai_response_to_mcp_result( ) -async def _check_model_access( # noqa: PLR0915 +async def _check_model_access( model: str, user_api_key_auth: Any ) -> Optional["ErrorData"]: """Enforce model-permission checks for MCP sampling requests. diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 746fc4e7d3f..1d9a4479f05 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -617,7 +617,7 @@ if MCP_AVAILABLE: active_mcp_session_var.reset(_session_reset_token) @server.call_tool() - async def mcp_server_tool_call( # noqa: PLR0915 + async def mcp_server_tool_call( name: str, arguments: Dict[str, Any] | None ) -> CallToolResult: """ @@ -1591,7 +1591,7 @@ if MCP_AVAILABLE: _mcp_gateway_initialize_instructions.reset(instructions_token) _mcp_gateway_server_name.reset(server_name_token) - async def _get_tools_from_mcp_servers( # noqa: PLR0915 + async def _get_tools_from_mcp_servers( user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str], mcp_servers: Optional[List[str]], @@ -2435,7 +2435,7 @@ if MCP_AVAILABLE: }, ) - async def execute_mcp_tool( # noqa: PLR0915 + async def execute_mcp_tool( name: str, arguments: Dict[str, Any], allowed_mcp_servers: List[MCPServer], @@ -3642,7 +3642,7 @@ if MCP_AVAILABLE: detail="Forbidden", ) - async def handle_streamable_http_mcp( # noqa: PLR0915 + async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send ) -> None: """Handle MCP requests through StreamableHTTP.""" diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 7b2f75e1cff..7446f61ad1c 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -509,7 +509,7 @@ async def get_agent_card( tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], ) -async def invoke_agent_a2a( # noqa: PLR0915 +async def invoke_agent_a2a( agent_id: str, request: Request, fastapi_response: Response, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index aa967732a90..814346eddf8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -519,7 +519,7 @@ MODEL_DISCOVERY_ROUTES = frozenset( ) -async def common_checks( # noqa: PLR0915 +async def common_checks( request_body: dict, team_object: Optional[LiteLLM_TeamTable], user_object: Optional[LiteLLM_UserTable], diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index fd6ff2ada7f..90845dfd824 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1954,7 +1954,7 @@ class JWTAuthManager: return None, None, None @staticmethod - async def auth_builder( # noqa: PLR0915 + async def auth_builder( api_key: str, jwt_handler: JWTHandler, request_data: dict, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 666c01562b5..6f359e52eeb 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -979,7 +979,7 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None: request.state.parent_otel_span = parent_otel_span -async def _user_api_key_auth_builder( # noqa: PLR0915 +async def _user_api_key_auth_builder( request: Request, api_key: str, azure_api_key_header: str, @@ -2126,7 +2126,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached @tracer.wrap() -async def _run_centralized_common_checks( # noqa: PLR0915 +async def _run_centralized_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, request: Request, request_data: dict, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index ea479a5721b..344f90aa144 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -58,7 +58,7 @@ router = APIRouter() dependencies=[Depends(user_api_key_auth)], tags=["batch"], ) -async def create_batch( # noqa: PLR0915 +async def create_batch( request: Request, fastapi_response: Response, provider: Optional[str] = None, @@ -343,7 +343,7 @@ async def create_batch( # noqa: PLR0915 dependencies=[Depends(user_api_key_auth)], tags=["batch"], ) -async def retrieve_batch( # noqa: PLR0915 +async def retrieve_batch( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8f330ada7d3..41cadd5bbc3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -300,7 +300,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse): ) -async def create_response( # noqa: PLR0915 +async def create_response( generator: AsyncGenerator[str, None], media_type: str, headers: dict, @@ -1148,7 +1148,7 @@ class ProxyBaseLLMRequestProcessing: _payload_str, ) - async def base_process_llm_request( # noqa: PLR0915 + async def base_process_llm_request( self, request: Request, fastapi_response: Response, diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index ab1eeaf1646..71dce163b78 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -36,7 +36,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -def initialize_callbacks_on_proxy( # noqa: PLR0915 +def initialize_callbacks_on_proxy( value: Any, premium_user: bool, config_file_path: str, diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index 97525a528d0..d9e21fc5d2a 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -11,7 +11,7 @@ _db = Any _VIEW_NOT_FOUND_MARKERS = ("does not exist", "no such table", "undefined table") -async def create_missing_views(db: _db): # noqa: PLR0915 +async def create_missing_views(db: _db): """ -------------------------------------------------- NOTE: Copy of `litellm/db_scripts/create_views.py`. diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e7f14df5294..4b7b20d75d0 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1128,7 +1128,7 @@ class DBSpendUpdateWriter: "_flush_tool_discovery_queue error (non-blocking): %s", e ) - async def _commit_spend_updates_to_db( # noqa: PLR0915 + async def _commit_spend_updates_to_db( self, prisma_client: PrismaClient, n_retry_times: int, diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index ff802223f21..72b9b7dc3c1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -121,7 +121,7 @@ class lakeraAI_Moderation(CustomGuardrail): return None - async def _check( # noqa: PLR0915 + async def _check( self, data: dict, user_api_key_dict: UserAPIKeyAuth, 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 e5200394b55..e8887fa712a 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 @@ -261,7 +261,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) return "" - async def _call_panw_api( # noqa: PLR0915 + async def _call_panw_api( self, content: str = "", is_response: bool = False, @@ -1762,7 +1762,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): return rd.get("name") if ("arguments" in rd or "mcp_arguments" in rd) else None @log_guardrail_information - async def apply_guardrail( # noqa: PLR0915 + async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, request_data: dict, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 2a2c758fa8a..09fff71062b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -288,7 +288,7 @@ class UnifiedLLMGuardrails(CustomLogger): return response - async def async_post_call_streaming_iterator_hook( # noqa: PLR0915 + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, response: Any, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index e0d018d4344..8a432eb2f42 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -167,7 +167,7 @@ async def test_endpoint(request: Request): tags=["health"], dependencies=[Depends(user_api_key_auth)], ) -async def health_services_endpoint( # noqa: PLR0915 +async def health_services_endpoint( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), service: services = fastapi.Query(description="Specify the service being hit."), ): diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 23af23e78bd..d36e9858b5a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -239,7 +239,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): request_count_end_user_id=results[5], ) - async def async_pre_call_hook( # noqa: PLR0915 + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, @@ -506,9 +506,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): return - async def async_log_success_event( # noqa: PLR0915 - self, kwargs, response_obj, start_time, end_time - ): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0587ce1cc29..0e21fd8e1f0 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1317,7 +1317,7 @@ class LiteLLMProxyRequestSetup: ) -async def add_litellm_data_to_request( # noqa: PLR0915 +async def add_litellm_data_to_request( data: dict, request: Request, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 13107b68864..e258ddc0410 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -726,7 +726,7 @@ def _key_metadata( return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id")) -def _aggregate_grouping_sets_records_sync( # noqa: PLR0915 +def _aggregate_grouping_sets_records_sync( *, records: List[Any], api_key_metadata: Dict[str, Dict[str, Any]], diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b0210e1123f..132060be76b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -675,7 +675,7 @@ def _enforce_upperbound_key_params( ) -async def _common_key_generation_helper( # noqa: PLR0915 +async def _common_key_generation_helper( data: GenerateKeyRequest, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str], @@ -3419,7 +3419,7 @@ def _check_model_access_group( return True -async def generate_key_helper_fn( # noqa: PLR0915 +async def generate_key_helper_fn( request_type: Literal[ "user", "key" ], # identifies if this request is from /user/new or /key/generate @@ -4070,7 +4070,7 @@ async def delete_key_aliases( ) -async def _rotate_master_key( # noqa: PLR0915 +async def _rotate_master_key( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, current_master_key: str, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 85b640d6b6c..4d4d1ef2774 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -933,7 +933,7 @@ def _check_team_budget_update_authority( response_model=LiteLLM_TeamTable, ) @management_endpoint_wrapper -async def new_team( # noqa: PLR0915 +async def new_team( data: NewTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -1637,7 +1637,7 @@ def validate_team_org_change( "/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) @management_endpoint_wrapper -async def update_team( # noqa: PLR0915 +async def update_team( data: UpdateTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 5af1dd321ed..2bf12880a75 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -3145,7 +3145,7 @@ class SSOAuthenticationHandler: ) @staticmethod - async def get_redirect_response_from_openid( # noqa: PLR0915 + async def get_redirect_response_from_openid( result: Union[OpenID, dict, CustomOpenID], request: Request, received_response: Optional[dict] = None, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 3e5873c2655..f43e876d111 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -284,7 +284,7 @@ async def route_create_file( dependencies=[Depends(user_api_key_auth)], tags=["files"], ) -async def create_file( # noqa: PLR0915 +async def create_file( request: Request, fastapi_response: Response, purpose: str = Form(...), @@ -589,7 +589,7 @@ async def create_file( # noqa: PLR0915 dependencies=[Depends(user_api_key_auth)], tags=["files"], ) -async def get_file_content( # noqa: PLR0915 +async def get_file_content( request: Request, fastapi_response: Response, file_id: str, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a912a88a993..6feb4e36bf9 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -366,7 +366,7 @@ class AnthropicPassthroughLoggingHandler: ) @staticmethod - def _collapse_pure_text_chunks( # noqa: PLR0915 + def _collapse_pure_text_chunks( all_chunks: Sequence[Union[str, bytes]], ) -> Optional[List[str]]: """ @@ -551,7 +551,7 @@ class AnthropicPassthroughLoggingHandler: return complete_streaming_response @staticmethod - def batch_creation_handler( # noqa: PLR0915 + def batch_creation_handler( httpx_response: httpx.Response, logging_obj: LiteLLMLoggingObj, url_route: str, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 9f353226dd0..b77c6e2f655 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -275,7 +275,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return litellm_model_response, response_cost @staticmethod - def openai_passthrough_handler( # noqa: PLR0915 + def openai_passthrough_handler( httpx_response: httpx.Response, response_body: dict, logging_obj: LiteLLMLoggingObj, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 6a138532617..73d4245670a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -645,7 +645,7 @@ class VertexPassthroughLoggingHandler: return kwargs @staticmethod - def batch_prediction_jobs_handler( # noqa: PLR0915 + def batch_prediction_jobs_handler( httpx_response: httpx.Response, logging_obj: LiteLLMLoggingObj, url_route: str, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e3cb9dec884..b84746758fb 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -143,7 +143,7 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona return headers -async def chat_completion_pass_through_endpoint( # noqa: PLR0915 +async def chat_completion_pass_through_endpoint( fastapi_response: Response, request: Request, adapter_id: str, @@ -701,7 +701,7 @@ from litellm.passthrough.timeout_utils import ( ) -async def pass_through_request( # noqa: PLR0915 +async def pass_through_request( request: Request, target: str, custom_headers: dict, @@ -1540,7 +1540,7 @@ async def _parse_request_data_by_content_type( return query_params_data, custom_body_data, file_data, stream -def create_pass_through_route( # noqa: PLR0915 +def create_pass_through_route( endpoint, target: str, custom_headers: Optional[Mapping[str, Any]] = None, @@ -1776,7 +1776,7 @@ def create_websocket_passthrough_route( return websocket_endpoint_func -async def websocket_passthrough_request( # noqa: PLR0915 +async def websocket_passthrough_request( websocket: WebSocket, target: str, custom_headers: dict, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index bd9746d2a47..e1fb65074cd 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -814,7 +814,7 @@ class ProxyInitializationHelpers: default=False, help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.", ) -def run_server( # noqa: PLR0915 +def run_server( cli_args, host, port, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1f765aa8d63..a48e9c58861 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -745,7 +745,7 @@ async def _initialize_shared_aiohttp_session(): @asynccontextmanager -async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 +async def proxy_startup_event(app: FastAPI): global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session import json @@ -2496,7 +2496,7 @@ async def _invalidate_spend_counter(counter_key: str): ) -async def update_cache( # noqa: PLR0915 +async def update_cache( token: Optional[str], user_id: Optional[str], end_user_id: Optional[str], @@ -3900,7 +3900,7 @@ class ProxyConfig: premium_user = _license_check.is_premium() return - async def load_config( # noqa: PLR0915 + async def load_config( self, router: Optional[litellm.Router], config_file_path: str ): """ @@ -6631,7 +6631,7 @@ def save_worker_config(**data): os.environ["WORKER_CONFIG"] = json.dumps(data) -async def initialize( # noqa: PLR0915 +async def initialize( model=None, alias=None, api_base=None, @@ -7022,7 +7022,7 @@ def _format_streaming_sse_chunk(chunk: Union[str, bytes]) -> Union[str, bytes]: return f"data: {chunk}\n\n" -async def async_data_generator( # noqa: PLR0915 +async def async_data_generator( response, user_api_key_dict: UserAPIKeyAuth, request_data: dict ): verbose_proxy_logger.debug("inside generator") @@ -7470,7 +7470,7 @@ class ProxyStartupEvent: ) @classmethod - async def initialize_scheduled_background_jobs( # noqa: PLR0915 + async def initialize_scheduled_background_jobs( cls, general_settings: dict, prisma_client: PrismaClient, @@ -8681,7 +8681,7 @@ async def chat_completion( dependencies=[Depends(user_api_key_auth)], tags=["completions"], ) -async def completion( # noqa: PLR0915 +async def completion( request: Request, fastapi_response: Response, model: Optional[str] = None, @@ -14411,7 +14411,7 @@ async def invitation_delete( dependencies=[Depends(user_api_key_auth)], include_in_schema=False, ) -async def update_config( # noqa: PLR0915 +async def update_config( config_info: ConfigYAML, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -15028,7 +15028,7 @@ async def delete_callback( include_in_schema=False, dependencies=[Depends(user_api_key_auth)], ) -async def get_config(): # noqa: PLR0915 +async def get_config(): """ For Admin UI - allows admin to view config via UI # return the callbacks and the env variables for the callback diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 03039d4f441..a69e6734d71 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -21,7 +21,7 @@ from litellm.proxy.response_polling.polling_handler import ResponsePollingHandle from litellm.types.llms.openai import ResponsesAPIStatus -async def background_streaming_task( # noqa: PLR0915 +async def background_streaming_task( polling_id: str, data: dict, polling_handler: ResponsePollingHandler, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 3626a21516d..bbd8b75fdd6 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -264,7 +264,7 @@ async def add_shared_session_to_data(data: dict) -> None: pass -async def route_request( # noqa: PLR0915 - Complex routing function, refactoring tracked separately +async def route_request( data: dict, llm_router: Optional[LitellmRouter], user_model: Optional[str], diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index ef06adb27fc..0ba77dcd2f0 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1734,7 +1734,7 @@ async def calculate_spend(request: SpendCalculateRequest): 200: {"model": List[LiteLLM_SpendLogs]}, }, ) -async def ui_view_spend_logs( # noqa: PLR0915 +async def ui_view_spend_logs( request: Request, api_key: Optional[str] = fastapi.Query( default=None, @@ -2273,7 +2273,7 @@ async def ui_view_request_response_for_request_id( 200: {"model": List[LiteLLM_SpendLogs]}, }, ) -async def view_spend_logs( # noqa: PLR0915 +async def view_spend_logs( api_key: Optional[str] = fastapi.Query( default=None, description="Get spend logs based on api key", diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index d215294fd04..aef06a3c668 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -228,9 +228,7 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d return {} -def get_logging_payload( # noqa: PLR0915 - kwargs, response_obj, start_time, end_time -) -> SpendLogsPayload: +def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogsPayload: if kwargs is None: kwargs = {} diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 7031ecaa1a0..f6f0a92def0 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -284,7 +284,7 @@ async def arealtime_calls( @wrapper_client -async def _arealtime( # noqa: PLR0915 +async def _arealtime( model: str, websocket: Any, # fastapi websocket api_base: Optional[str] = None, diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index e27585116ce..e40e12e9197 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -75,7 +75,7 @@ async def arerank( @client -def rerank( # noqa: PLR0915 +def rerank( model: str, query: str, documents: List[Union[str, Dict[str, Any]]], diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 24b5db28571..acb7487f430 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -77,7 +77,7 @@ def _add_mcp_metadata_to_response( setattr(message, "provider_specific_fields", provider_fields) -async def acompletion_with_mcp( # noqa: PLR0915 +async def acompletion_with_mcp( model: str, messages: List, tools: Optional[List] = None, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 94cff6922b5..df5de205d45 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -644,7 +644,7 @@ class LiteLLM_Proxy_MCP_Handler: return result_text or "Tool executed successfully" @staticmethod - async def _execute_tool_calls( # noqa: PLR0915 + async def _execute_tool_calls( tool_server_map: dict[str, str], tool_calls: List[Any], user_api_key_auth: Any, diff --git a/litellm/router.py b/litellm/router.py index 5c4dc3eb943..5f26097443f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -241,7 +241,7 @@ class Router: lowesttpm_logger: Optional[LowestTPMLoggingHandler] = None optional_callbacks: Optional[List[Union[CustomLogger, Callable, str]]] = None - def __init__( # noqa: PLR0915 + def __init__( self, model_list: Optional[ Union[List[DeploymentTypedDict], List[Dict[str, Any]]] @@ -2887,7 +2887,7 @@ class Router: f"Silent experiment failed for model {silent_model}: {str(e)}" ) - async def _acompletion( # noqa: PLR0915 + async def _acompletion( self, model: str, messages: List[Dict[str, str]], **kwargs ) -> Union[ ModelResponse, @@ -5158,7 +5158,7 @@ class Router: ) raise e - async def _acreate_file( # noqa: PLR0915 + async def _acreate_file( self, model: str, **kwargs, @@ -6467,7 +6467,7 @@ class Router: # propagate so they remain visible. return None - async def async_function_with_fallbacks_common_utils( # noqa: PLR0915 + async def async_function_with_fallbacks_common_utils( self, e: Exception, disable_fallbacks: Optional[bool], @@ -6843,7 +6843,7 @@ class Router: ) @tracer.wrap() - async def async_function_with_retries(self, *args, **kwargs): # noqa: PLR0915 + async def async_function_with_retries(self, *args, **kwargs): verbose_router_logger.debug("Inside async function with retries.") original_function = kwargs.pop("original_function") fallbacks = kwargs.pop("fallbacks", self.fallbacks) @@ -9324,7 +9324,7 @@ class Router: return model_info - def _set_model_group_info( # noqa: PLR0915 + def _set_model_group_info( self, model_group: str, user_facing_model_group_name: str ) -> Optional[ModelGroupInfo]: """ @@ -10566,7 +10566,7 @@ class Router: ) return client - def _pre_call_checks( # noqa: PLR0915 + def _pre_call_checks( self, model: str, healthy_deployments: List, diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 54498363f51..3f641d4f0fb 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -190,7 +190,7 @@ class LowestCostLoggingHandler(CustomLogger): ) pass - async def async_get_available_deployments( # noqa: PLR0915 + async def async_get_available_deployments( self, model_group: str, healthy_deployments: list, diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 870b3f29d48..3adb8d43920 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -35,9 +35,7 @@ class LowestLatencyLoggingHandler(CustomLogger): self.router_cache = router_cache self.routing_args = RoutingArgs(**routing_args) - def log_success_event( # noqa: PLR0915 - self, kwargs, response_obj, start_time, end_time - ): + def log_success_event(self, kwargs, response_obj, start_time, end_time): try: """ Update latency usage on success @@ -259,9 +257,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ) pass - async def async_log_success_event( # noqa: PLR0915 - self, kwargs, response_obj, start_time, end_time - ): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: """ Update latency usage on success @@ -413,7 +409,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ) pass - def _get_available_deployments( # noqa: PLR0915 + def _get_available_deployments( self, model_group: str, healthy_deployments: list, diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 488f8450941..f807ba7232a 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -158,7 +158,7 @@ class LowestTPMLoggingHandler(CustomLogger): verbose_router_logger.debug(traceback.format_exc()) pass - def get_available_deployments( # noqa: PLR0915 + def get_available_deployments( self, model_group: str, healthy_deployments: list, diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 5c31d81f04c..f4b1d4a1b69 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -156,7 +156,7 @@ def get_secret_bool( return str_to_bool(_secret_value) -def get_secret( # noqa: PLR0915 +def get_secret( secret_name: str, default_value: Optional[Union[str, bool]] = None, ): diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index 4ff94d18eff..3a3cf6272dc 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -23,7 +23,7 @@ def _is_base64(s): return False -def get_secret_from_manager( # noqa: PLR0915 +def get_secret_from_manager( client: Any, key_manager: str, secret_name: str, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a5032942011..f2152577b4d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1572,7 +1572,7 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None """Breakdown of tokens used in the prompt.""" - def __init__( # noqa: PLR0915 + def __init__( self, prompt_tokens: Optional[int] = None, completion_tokens: Optional[int] = None, @@ -1908,7 +1908,7 @@ class ModelResponse(ModelResponseBase): choices: List[Choices] """The list of completion choices the model generated for the input prompt.""" - def __init__( # noqa: PLR0915 + def __init__( self, id=None, choices=None, diff --git a/litellm/utils.py b/litellm/utils.py index ee7d952008e..9c5989a11d3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -760,7 +760,7 @@ def _remove_thought_signatures_from_messages( return processed_messages -def function_setup( # noqa: PLR0915 +def function_setup( original_function: str, rules_obj, start_time, *args, **kwargs ): # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc. ### NOTICES ### @@ -1422,12 +1422,12 @@ def post_call_processing( raise e -def client(original_function): # noqa: PLR0915 +def client(original_function): Rules = getattr(sys.modules[__name__], "Rules") rules_obj = Rules() @wraps(original_function) - def wrapper(*args, **kwargs): # noqa: PLR0915 + def wrapper(*args, **kwargs): # DO NOT MOVE THIS. It always needs to run first # Check if this is an async function. If so only execute the async function call_type = original_function.__name__ @@ -1775,7 +1775,7 @@ def client(original_function): # noqa: PLR0915 raise e @wraps(original_function) - async def wrapper_async(*args, **kwargs): # noqa: PLR0915 + async def wrapper_async(*args, **kwargs): print_args_passed_to_litellm(original_function, args, kwargs) start_time = datetime.datetime.now() result = None @@ -2942,7 +2942,7 @@ def _resolve_builtin_model_cost_entry( return None -def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 +def register_model(model_cost: Union[str, dict]): """ Register new / Override existing models (and their pricing) to specific providers. Provide EITHER a model cost dictionary or a url to a hosted json blob @@ -3365,7 +3365,7 @@ def get_optional_params_image_gen( return optional_params -def get_optional_params_embeddings( # noqa: PLR0915 +def get_optional_params_embeddings( # 2 optional params model: str, user: Optional[str] = None, @@ -4112,7 +4112,7 @@ def pre_process_optional_params( return optional_params -def get_optional_params( # noqa: PLR0915 +def get_optional_params( # use the openai defaults # https://platform.openai.com/docs/api-reference/chat/create model: str, @@ -5842,7 +5842,7 @@ def _is_potential_model_name_in_model_cost( ) -def _get_model_info_helper( # noqa: PLR0915 +def _get_model_info_helper( model: str, custom_llm_provider: Optional[str] = None, api_base: Optional[str] = None, @@ -6566,7 +6566,7 @@ def create_proxy_transport_and_mounts(): return sync_proxy_mounts, async_proxy_mounts -def validate_environment( # noqa: PLR0915 +def validate_environment( model: Optional[str] = None, api_key: Optional[str] = None, api_base: Optional[str] = None, diff --git a/ruff.toml b/ruff.toml index 7baa1c5f92d..2db4122a30e 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,5 +1,5 @@ lint.ignore = ["F405", "E402", "E501", "F403"] -lint.extend-select = ["E501", "PLR0915", "T20", "PGH004", "RUF008", "RUF009", "RUF100"] +lint.extend-select = ["E501", "T20", "PGH004", "RUF008", "RUF009", "RUF100"] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external # so RUF100 leaves their directives alone: the strict gate (ruff-strict.toml) and upstream @@ -23,9 +23,4 @@ exclude = ["litellm/types/*", "litellm/__init__.py", "litellm/proxy/example_conf "litellm/llms/azure_ai/embed/__init__.py" = ["F401"] "litellm/llms/azure_ai/rerank/__init__.py" = ["F401"] "litellm/llms/bedrock/chat/__init__.py" = ["F401"] -"litellm/proxy/utils.py" = ["F401", "PLR0915"] -"litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py" = ["PLR0915"] -"litellm/proxy/guardrails/guardrail_hooks/guardrail_benchmarks/test_eval.py" = ["PLR0915"] -"litellm/responses/streaming_iterator.py" = ["PLR0915"] -"litellm/files/main.py" = ["PLR0915"] -"litellm/llms/litellm_proxy/skills/sandbox_executor.py" = ["PLR0915"] +"litellm/proxy/utils.py" = ["F401"] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 2e4b7f9ae74..362f4986c62 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1060,7 +1060,7 @@ def test_initialize_pass_through_endpoints_with_cost_per_request(): @pytest.mark.asyncio -async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): # noqa: PLR0915 +async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): """ Test that pass_through_request (parent method) correctly includes proxy_server_request in kwargs passed to the success handler. From be4fa702e7b7a84f0f3a7bbb5621766b62fb6555 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:59:21 -0700 Subject: [PATCH 2/3] ci(lint): ratcheted type-discipline gate (mutable collections, casts, guards, kwargs, suppressions) (#30500) * ci(lint): enforce type-discipline budget for casts and type guards Add a ratcheted gate that blocks net-new typing.cast() usage and bans TypeGuard/TypeIs outright, layered on the existing ruff-strict budget setup. - ruff-strict.toml: ban cast/TypeGuard/TypeIs (typing + typing_extensions) via flake8-tidy-imports banned-api (TID251) for a coarse import-level freeze. - ruff-strict-budget.json: bump TID251 baseline 2404 -> 2662 to absorb the ~258 pre-existing usages now matched by the new banned-api entries. - scripts/check_type_discipline.py: AST checker adding LIT006 (cast call sites, suppress with `# cast-ok: `) and LIT007 (TypeGuard/TypeIs annotations, suppress with `# guard-ok: `) for per-call-site granularity. - scripts/type_discipline_gate.py: baseline+slack gate with delta-vs-base, mirroring ruff_strict_gate.py. - type-discipline-budget.json: LIT006 baseline 1013 (slack 10), LIT007 0/0. - test-linting.yml: run the gate in CI against the PR base SHA. * ci(lint): enforce suppression-reason budgets and guard budgets against loosening - wire the **kwargs ban (LIT008) into the vendored type-discipline checker so it matches the budget that already referenced it - freeze LIT003/LIT004 (noqa / type-ignore without codes or reason) and LIT005 (*-ok suppression without a reason) at slack 0 so any net-new unexplained suppression trips the type-discipline gate - add scripts/budget_ratchet_check.py and a separate, non-gating budget-ratchet CI job that turns red when any *-budget.json ceiling is raised, a rule is dropped, or a budget file is deleted * ci(lint): ban mutable collections in annotations and all mutable construction Expand LIT001 from coarse builtins at interfaces to any mutable collection in any annotation (builtins, typing aliases, collections concretes, mutable ABCs) across signatures, class attributes, locals, and globals. Add LIT009 to flag mutable-collection construction (literals, comprehensions, constructors) so the unannotated seed-then-mutate pattern is caught too. Enumerate any-ok in LIT005 so its reason requirement holds even when only the stdlib checker runs. Budget LIT001 (21452) and LIT009 (25222) with slack 10 to ratchet down. * ci(lint): recommend pydantic at boundaries and add functional-refactor guidance Drop the msgspec mention from the cast banned-api messages so the recommended validation path matches the codebase's primary pattern (pydantic). Add a note to CLAUDE.md that lint / type-discipline failures should be resolved by refactoring to functional, immutable patterns rather than reaching for mutable structures or `# mutable-ok`. * style: make CLAUDE.md more concise * chore: update CLAUDE.md guidelines * ci(lint): renumber mutable construction LIT009 -> LIT002 next to LIT001 Group the mutable-collection family together: LIT001 (mutable collection in any annotation) and the construction rule now sit adjacent at LIT001/LIT002. The freed LIT009 slot is taken by the sibling Any gate (check_any_discipline.py, #30379), which moves its Any-typed-value rule LIT002 -> LIT009 in lockstep so the shared LIT namespace stays contiguous with no holes. Budget, gate docstring, and the checker's own docstring/messages are updated to match. * fix: numbering in CLAUDE.md * test(lint): test type-discipline checker, scope LIT007 to return types Add regression tests for check_type_discipline.py (every LIT rule, its suppression, and the comment scanner) and for budget_ratchet_check.py. Confine LIT007 to function return annotations, the only place TypeGuard/TypeIs are valid, so a runtime name that merely reads those identifiers is no longer flagged. Switch scan_comments to io.StringIO(source).readline, the standard readline that returns '' at EOF, dropping the iter(...).__next__ idiom. * fix(lint): best-effort worktree teardown so cleanup can't mask the real error base_counts ran `git worktree remove` through the raising `_run` in its finally, so a failed `git worktree add` (or a failure in the body) was masked by a second SystemExit from the cleanup. Tear the worktree down best-effort, like the sibling rmtree, so the original error propagates. * fix(lint): ratchet fails loudly on an unresolvable base; drop dead checker state Verify the merge-base ref resolves to a commit before trusting a missing-file result from git show, so an invalid or empty BASE_SHA now turns the budget-ratchet guard red instead of skipping every budget and passing vacuously Also drop the unused Comments.by_line field and the phantom --changed-only usage line from check_type_discipline's docstring, and cover the ref handling with tests * fix(lint): degrade malformed source to LIT000 instead of crashing the checker tokenize.generate_tokens raises IndentationError (a SyntaxError subclass) on a dedent mismatch, which escaped scan_comments' tokenize.TokenError handler and crashed the whole checker run, zeroing the gate for that invocation. Catch SyntaxError too so the file falls through to ast.parse and is reported as LIT000, matching the checker's graceful-degradation contract. Also add the trailing newline ruff-strict.toml lacked * perf(lint): skip the base worktree scan when no rule is over its ceiling cmd_check created a git worktree and re-scanned the base tree on every run, but a rule can only breach when its head count is already over baseline + slack; when none are, the base comparison cannot change the verdict. Short-circuit to OK in that case, which is every green PR, roughly halving the gate's work. Extract over_ceiling and cover it (and evaluate's drift-safety) with tests * fix(lint): exempt .dict()/.list()/.set() method calls from LIT002 _construction_kind matched dict/list/set as constructors via func.attr too, flagging common method calls like pydantic's model.dict() as mutable construction; 200 such false positives existed in litellm. Recognize dict/list/set construction only when unqualified while keeping the collections concretes (deque/defaultdict/...) matchable as attributes, since those are rarely method names. Ratchet the LIT002 baseline down 25222 -> 25022 to reflect the removed false positives * chore(lint): bump basedpyright ceilings to absorb staging base drift The basedpyright gate added in #30379 is a total-count check against basedpyright-code-budget.json and the linting workflow runs only on pull_request, so pushes to litellm_internal_staging never re-baseline it. Merging staging into this branch surfaced that drift: seven reportAny/reportUnknown* rules sit 10-149 errors above their committed ceiling even though this PR changes no files under litellm/, the only path basedpyright scans (pyrightconfig include is litellm). The new baselines match the counts CI measured on the merge commit, with the existing per-rule slack preserved * fix(lint): ratchet guard watches every budget file, not just two DEFAULT_BUDGETS only listed ruff-strict-budget.json and type-discipline-budget.json, so mypy-code-budget.json and basedpyright-code-budget.json were unguarded and their ceilings could rise with no signal, which is exactly the failure mode this guard exists to prevent. The gap became concrete when this PR bumped basedpyright-code-budget.json to absorb staging drift. All four budgets are now watched, so the budget-ratchet job surfaces that basedpyright bump for human review the same way it surfaces the TID251 raise. A regression test pins that every *-budget.json on disk is in DEFAULT_BUDGETS, failing loudly if a future budget escapes the ratchet * fix: add a lot more slack * fix(lint): restore LIT003 frozen slack to 0 The blanket slack bump set LIT003 (bare # noqa without codes or a reason) to a slack of 50, which contradicts the documented zero-tolerance invariant: the gate docstring and the PR description table both freeze LIT003/LIT004/LIT005 at slack 0 so any net-new unexplained suppression trips the gate. Slack 50 would let 50 new bare noqas through silently. The actual LIT003 count is 397, well under the 516 baseline, so restoring slack to 0 keeps the gate green while putting the freeze back. LIT004/LIT005/LIT007 were already correct at 0 * fix(lint): restore documented slack 10 for the buffered LIT rules The slack bump left LIT001/LIT002/LIT006/LIT008 at 2000/2500/100/100, 10-250x the "/ 10" the PR description table and the gate docstring document. That buffer was never needed: the gate already blames a rule only when its count exceeds the ceiling and grew vs the merge-base, so the violations the staging merge added in litellm/ sit in both head and base and are never charged to this PR. With slack back at the documented 10 the gate stays green, and the ceiling is tight again (LIT006 no longer waves through 99 net-new cast() calls). Baselines are unchanged; only the slack returns to its documented value * fix(lint): ratchet LIT003 baseline down to its actual count The LIT003 baseline was 516 while the current bare-noqa count is 397, leaving ~119 units of headroom that undercut the documented zero-tolerance freeze: the gate docstring claims any net-new bare noqa trips the gate, but with cap 516 a PR could add over a hundred first. Drop the baseline to the measured 397 so the freeze is exact (cap = 397 + slack 0), the same hard-zero-at-the-boundary shape LIT005 and LIT007 already use and pass in CI. PR table row updated to 397 / 0 * fix: increase slack * fix: increase slack * docs(lint): align gate docstring with buffered LIT003/LIT004 slack The budget now gives LIT003/LIT004 nonzero slack, so the gate's prose no longer claims they are frozen at slack 0; LIT005 remains the reasonless- suppression freeze and LIT007 the hard zero. --- .github/workflows/test-linting.yml | 33 ++ CLAUDE.md | 2 + ruff-strict-budget.json | 82 +-- ruff-strict.toml | 14 +- scripts/budget_ratchet_check.py | 157 ++++++ scripts/check_type_discipline.py | 476 ++++++++++++++++++ scripts/type_discipline_gate.py | 198 ++++++++ .../test_litellm/test_budget_ratchet_check.py | 97 ++++ .../test_check_type_discipline.py | 199 ++++++++ .../test_litellm/test_type_discipline_gate.py | 40 ++ type-discipline-budget.json | 34 ++ 11 files changed, 1290 insertions(+), 42 deletions(-) create mode 100644 scripts/budget_ratchet_check.py create mode 100644 scripts/check_type_discipline.py create mode 100644 scripts/type_discipline_gate.py create mode 100644 tests/test_litellm/test_budget_ratchet_check.py create mode 100644 tests/test_litellm/test_check_type_discipline.py create mode 100644 tests/test_litellm/test_type_discipline_gate.py create mode 100644 type-discipline-budget.json diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 2e967f3ed3f..d06b9a16e6d 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -77,6 +77,12 @@ jobs: run: | uv run --no-sync python scripts/ruff_strict_gate.py --base "$BASE_SHA" + - name: Check type-discipline budget (mutable collections / casts / type guards / kwargs / unexplained suppressions, delta vs base) + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} + run: | + uv run --no-sync python scripts/type_discipline_gate.py --base "$BASE_SHA" + - name: Print OpenAI version run: | uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')" @@ -100,6 +106,33 @@ jobs: run: | uv run --no-sync python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1) + # Intentionally NON-GATING. This job turns red when a *-budget.json ceiling is + # raised (or a rule/budget is dropped) so a loosening is obvious in review, but it + # must be kept OUT of the branch-protection required-checks list so a justified + # bump can still be merged by a human who has seen and accepted the red. + budget-ratchet: + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + contents: read + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Ratchet check (budgets may only decrease; non-gating) + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} + run: | + python scripts/budget_ratchet_check.py --base "$BASE_SHA" + any-discipline: # Separate job: the first run cold-builds litellm's type cache (~2 min, ~3 GB), # so keep it off the main lint job's time budget. Subsequent runs reuse the diff --git a/CLAUDE.md b/CLAUDE.md index 48dc3d81d94..a81ee1f3b91 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -42,6 +42,8 @@ If you're trying to create a new function that relies on untyped stuff, instead The Any-discipline gate (`make lint-any`, also a CI job) fails when a line you changed under `litellm/` holds a value typed `Any`, including the `X | Any`. Ideally `# any-ok: ` is never used; treat it as a last resort for a genuine typed/untyped boundary that Pydantic truly can't model +If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason + Ask to commit and push your work when you're done (or if you're confident that your code is good and works, just do it) When you must use real LLM models to, for example, write e2e tests, write a QA runbook, etc., make sure to use the latest models (doesn't have to be smartest, can also be a modern small, fast one. No strong preference for smart vs fast here, just use something modern) as of the year and month of the current date. Do a web search as necessary to figure that out diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index bb02ec01569..62ebdb559fc 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,27 +1,27 @@ { "ANN001": { "baseline": 2865, - "slack": 10 + "slack": 50 }, "ANN002": { "baseline": 64, - "slack": 3 + "slack": 5 }, "ANN003": { "baseline": 759, - "slack": 10 + "slack": 30 }, "ANN201": { "baseline": 1944, - "slack": 10 + "slack": 50 }, "ANN202": { "baseline": 858, - "slack": 10 + "slack": 30 }, "ANN204": { "baseline": 658, - "slack": 10 + "slack": 20 }, "ANN205": { "baseline": 117, @@ -33,7 +33,7 @@ }, "ANN401": { "baseline": 1886, - "slack": 10 + "slack": 50 }, "ASYNC230": { "baseline": 11, @@ -45,15 +45,15 @@ }, "B006": { "baseline": 180, - "slack": 3 + "slack": 10 }, "B008": { "baseline": 490, - "slack": 10 + "slack": 15 }, "B009": { "baseline": 79, - "slack": 10 + "slack": 5 }, "B010": { "baseline": 187, @@ -81,7 +81,7 @@ }, "BLE001": { "baseline": 2854, - "slack": 10 + "slack": 50 }, "C401": { "baseline": 8, @@ -109,7 +109,7 @@ }, "C901": { "baseline": 301, - "slack": 3 + "slack": 15 }, "D419": { "baseline": 6, @@ -125,7 +125,7 @@ }, "DTZ005": { "baseline": 229, - "slack": 10 + "slack": 15 }, "DTZ006": { "baseline": 10, @@ -165,7 +165,7 @@ }, "I001": { "baseline": 258, - "slack": 10 + "slack": 15 }, "LOG015": { "baseline": 5, @@ -189,11 +189,11 @@ }, "PERF403": { "baseline": 69, - "slack": 10 + "slack": 5 }, "PIE790": { "baseline": 263, - "slack": 10 + "slack": 15 }, "PIE800": { "baseline": 1, @@ -233,7 +233,7 @@ }, "PLR0913": { "baseline": 1813, - "slack": 3 + "slack": 50 }, "PLR1704": { "baseline": 3, @@ -245,7 +245,7 @@ }, "PLR1714": { "baseline": 252, - "slack": 10 + "slack": 15 }, "PLR1730": { "baseline": 7, @@ -265,11 +265,11 @@ }, "PLW0602": { "baseline": 215, - "slack": 10 + "slack": 15 }, "PLW0603": { "baseline": 183, - "slack": 3 + "slack": 10 }, "PLW1508": { "baseline": 188, @@ -301,15 +301,15 @@ }, "RET504": { "baseline": 709, - "slack": 10 + "slack": 20 }, "RUF010": { "baseline": 844, - "slack": 10 + "slack": 30 }, "RUF012": { "baseline": 158, - "slack": 3 + "slack": 10 }, "RUF015": { "baseline": 8, @@ -321,7 +321,7 @@ }, "RUF022": { "baseline": 80, - "slack": 10 + "slack": 5 }, "RUF023": { "baseline": 2, @@ -337,15 +337,15 @@ }, "RUF059": { "baseline": 69, - "slack": 10 + "slack": 5 }, "RUF100": { "baseline": 465, - "slack": 10 + "slack": 15 }, "S110": { "baseline": 222, - "slack": 10 + "slack": 15 }, "S112": { "baseline": 21, @@ -353,11 +353,11 @@ }, "SIM101": { "baseline": 58, - "slack": 10 + "slack": 5 }, "SIM102": { "baseline": 311, - "slack": 10 + "slack": 15 }, "SIM103": { "baseline": 119, @@ -412,20 +412,20 @@ "slack": 3 }, "TID251": { - "baseline": 2405, - "slack": 10 + "baseline": 2664, + "slack": 50 }, "TRY002": { "baseline": 528, - "slack": 10 + "slack": 20 }, "TRY004": { "baseline": 93, - "slack": 10 + "slack": 5 }, "TRY201": { "baseline": 409, - "slack": 10 + "slack": 15 }, "TRY203": { "baseline": 113, @@ -433,15 +433,15 @@ }, "TRY300": { "baseline": 853, - "slack": 10 + "slack": 30 }, "UP006": { "baseline": 12941, - "slack": 10 + "slack": 100 }, "UP007": { "baseline": 2520, - "slack": 10 + "slack": 50 }, "UP008": { "baseline": 2, @@ -469,7 +469,7 @@ }, "UP032": { "baseline": 609, - "slack": 10 + "slack": 20 }, "UP034": { "baseline": 1, @@ -477,7 +477,7 @@ }, "UP035": { "baseline": 2250, - "slack": 10 + "slack": 50 }, "UP036": { "baseline": 1, @@ -485,10 +485,10 @@ }, "UP037": { "baseline": 100, - "slack": 10 + "slack": 5 }, "UP045": { "baseline": 18417, - "slack": 10 + "slack": 100 } } diff --git a/ruff-strict.toml b/ruff-strict.toml index 8d517615244..1caa3567872 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -18,4 +18,16 @@ max-args = 5 "typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; create a Mapping alias with concrete value types if truly dynamic." "typing.Set".msg = "frozenset[X] or AbstractSet[X]." "typing.MutableSequence".msg = "Sequence[X]." -"typing.MutableMapping".msg = "See typing.Dict." \ No newline at end of file +"typing.MutableMapping".msg = "See typing.Dict." +# Unchecked casts: cast() lies to the type checker with no runtime guarantee. +# Validate into a concrete frozen type at the boundary (pydantic) instead. +# Per-call-site coverage lives in check_type_discipline.py (LIT006); this freezes +# new cast imports. Suppress (with a reason) via `# noqa: TID251 # `. +"typing.cast".msg = "No unchecked casts: validate into a frozen dataclass/NamedTuple/ReadOnly TypedDict at the boundary (pydantic)." +"typing_extensions.cast".msg = "Same as typing.cast." +# Unverified narrowing predicates: the checker never validates the guard body, so a +# wrong guard silently corrupts types. Banned outright (there are none today). +"typing.TypeGuard".msg = "Unverified narrowing. Parse into a concrete type, or use isinstance for a runtime-checked narrowing." +"typing_extensions.TypeGuard".msg = "Same as typing.TypeGuard." +"typing.TypeIs".msg = "Unverified narrowing (the body is trusted). Parse into a concrete type instead." +"typing_extensions.TypeIs".msg = "Same as typing.TypeIs." diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py new file mode 100644 index 00000000000..c4b2c3ee655 --- /dev/null +++ b/scripts/budget_ratchet_check.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Non-gating ratchet guard: budget ceilings may only fall, never rise. + +Every `*-budget.json` file (ruff-strict, type-discipline, mypy-code, basedpyright-code) is a +one-way ratchet: each rule's ceiling is `baseline + slack`, and the whole point is +to drive that number DOWN over time. This check compares every budget file against +its own content at the merge-base with the target branch and fails (exits 1, red) if: + + * a rule's ceiling went up, + * a rule was dropped from a budget (its ceiling effectively became infinite), or + * an entire budget file was deleted. + +New rules and lowered/equal ceilings are fine. + +This is deliberately NOT a gating check. It should turn the run red so that a +loosening is impossible to miss in review, but it must stay OUT of the +branch-protection required-checks list: a justified bump (e.g. banning a new API, +which mechanically raises a baseline) can then still be merged by a human who has +seen the red and accepted it. + +Usage: + python scripts/budget_ratchet_check.py [--base REF] [budget.json ...] + +Stdlib only. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path +from typing import NamedTuple + +REPO_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_BASE = "origin/litellm_internal_staging" +DEFAULT_BUDGETS: tuple[str, ...] = ( + "ruff-strict-budget.json", + "type-discipline-budget.json", + "mypy-code-budget.json", + "basedpyright-code-budget.json", +) + + +class Regression(NamedTuple): + budget: str + rule: str + detail: str + + +def _run(cmd: list[str]) -> subprocess.CompletedProcess[str]: + return subprocess.run(cmd, cwd=REPO_ROOT, capture_output=True, text=True) + + +def _merge_base(base: str) -> str: + """The common ancestor of `base` and HEAD, so unrelated base drift is ignored.""" + proc = _run(["git", "merge-base", base, "HEAD"]) + return proc.stdout.strip() or base + + +def _load_head(rel: str) -> dict | None: + path = REPO_ROOT / rel + if not path.exists(): + return None + return json.loads(path.read_text()) + + +def _ref_is_commit(ref: str) -> bool: + return _run(["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"]).returncode == 0 + + +def _load_base(rel: str, ref: str) -> dict | None: + """Budget content at `ref`, or None when the file did not exist there. + + `ref` is verified as a real commit by the caller, so a non-zero `git show` here means + the path was absent at that commit, not that the ref itself is unresolvable. + """ + proc = _run(["git", "show", f"{ref}:{rel}"]) + if proc.returncode != 0: + return None + return json.loads(proc.stdout) + + +def _caps(budget: dict) -> dict[str, int]: + """Map each rule to its ceiling (baseline + slack); skip malformed specs.""" + caps: dict[str, int] = {} + for rule, spec in budget.items(): + if isinstance(spec, dict): + caps[rule] = int(spec.get("baseline", 0)) + int(spec.get("slack", 0)) + return caps + + +def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]: + if base is None: + return [] # new budget file: nothing to ratchet against yet + if head is None: + return [Regression(rel, "*", "budget file was deleted (every ceiling removed)")] + + base_caps = _caps(base) + head_caps = _caps(head) + out: list[Regression] = [] + for rule, base_cap in sorted(base_caps.items()): + if rule not in head_caps: + out.append(Regression(rel, rule, f"rule dropped (ceiling {base_cap} -> removed)")) + elif head_caps[rule] > base_cap: + out.append(Regression(rel, rule, f"ceiling raised {base_cap} -> {head_caps[rule]}")) + return out + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("budgets", nargs="*", help="budget files to check") + args = parser.parse_args() + budgets = args.budgets or list(DEFAULT_BUDGETS) + + ref = _merge_base(args.base) + if not _ref_is_commit(ref): + print( + f"FAIL: base ref {ref!r} does not resolve to a commit, so the ratchet has nothing " + f"to compare against; refusing to pass vacuously (check the --base / BASE_SHA value)", + file=sys.stderr, + ) + return 1 + + regressions: list[Regression] = [] + checked: list[str] = [] + for rel in budgets: + base = _load_base(rel, ref) + head = _load_head(rel) + if base is None and head is None: + continue + if base is None: + print(f"skip {rel}: new file (no base at {args.base} to ratchet against)") + continue + checked.append(rel) + regressions.extend(regressions_for(rel, base, head)) + + if regressions: + print(f"FAIL: budget ceiling(s) loosened vs base {args.base} (merge-base {ref[:12]}):") + for reg in regressions: + print(f" {reg.budget} {reg.rule}: {reg.detail}") + print( + "Budgets are one-way ratchets and may only go down or stay flat. This " + "check is non-gating: if the increase is justified (e.g. a newly banned " + "API), a human can merge over the red after acknowledging it." + ) + return 1 + + suffix = f" ({', '.join(checked)})" if checked else "" + print(f"OK: no budget ceiling increased vs base {args.base}{suffix}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py new file mode 100644 index 00000000000..d83a1a7512f --- /dev/null +++ b/scripts/check_type_discipline.py @@ -0,0 +1,476 @@ +#!/usr/bin/env python3 +"""Type-discipline checker: the rules ruff can't enforce. + +Rules +----- +LIT001 Mutable collection in a type annotation, anywhere it appears: function + parameters, return types, class attributes, locals, and module globals. + Covers the builtins (dict/list/set, bare or parameterized), their typing + aliases (Dict/List/...), the collections concretes (deque/defaultdict/...), + and the mutable ABCs (MutableMapping/MutableSequence/MutableSet). A mutable + collection lets whoever holds it grow or rewrite it after the fact; annotate + a read-only view instead (Mapping/Sequence/AbstractSet/tuple[X, ...]/ + frozenset[X], or a frozen dataclass / NamedTuple / ReadOnly TypedDict) and + build it functionally (comprehension / map, not append-in-a-loop). + Suppress with `# mutable-ok: ` on the offending line. +LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehension, or + a call to a mutable constructor (list/dict/set/deque/defaultdict/Counter/...). + Catches the unannotated seed-then-mutate pattern LIT001 cannot see (`acc = []`). + Build the value in one shot and freeze it: a `tuple`/`frozenset` wrapping a + generator (`tuple(f(x) for x in xs)`), a tuple literal, or a frozen dataclass / + NamedTuple / ReadOnly TypedDict. Generator expressions and `tuple`/`frozenset` + calls are not construction and pass. Annotation-internal lists (`Callable[[int], + str]`) are exempt. Suppress with `# mutable-ok: `. +LIT003 noqa suppression without rule codes or without a reason. + Required shape: `# noqa: TID251 # ` +LIT004 type/pyright/mypy ignore without bracketed codes or without a reason. + Required shape: `# pyright: ignore[reportArgumentType] # ` +LIT005 A `# mutable-ok` / `# cast-ok` / `# guard-ok` / `# kwargs-ok` / `# any-ok` + suppression without a reason. (`any-ok` belongs to check_any_discipline.py; + it is enumerated here so the reason requirement holds even when only this + stdlib checker runs.) +LIT006 `cast(...)` call. typing.cast is an unchecked assertion (the moral equivalent + of TypeScript's `as`); it lies to the type checker with zero runtime guarantee. + Validate into a concrete frozen type at the boundary instead. + Suppress with `# cast-ok: ` on the call's first line. +LIT007 `TypeGuard[...]` / `TypeIs[...]` annotation. The narrowing predicate's body is + never verified by the checker, so a wrong guard silently corrupts types. + Prefer parsing into a concrete type. Suppress with `# guard-ok: `. +LIT008 `**kwargs` parameter. The keyword contract is erased and everything it carries + is effectively Any. ruff can force it to be typed (ANN003) but can't ban the + syntax. Declare explicit keyword params, or accept one frozen payload. `*args`, + by contrast, is fine when typed (it's just a tuple). Suppress: `# kwargs-ok: `. + +LIT000 and LIT009 are the sibling Any gate's (check_any_discipline.py, #30379): a mypy +build/read failure and an Any-typed value. They share this LIT namespace but are emitted +by that checker, not this one. + +Usage +----- + python check_type_discipline.py litellm/ tests/ + +Exit code 1 if any violation is found. Stdlib only. +""" + +from __future__ import annotations + +import ast +import io +import re +import sys +import tokenize +from dataclasses import dataclass +from pathlib import Path +from collections.abc import Iterable, Iterator, Sequence +from typing import NamedTuple + +# Mutable collection types, banned in *every* annotation. Name-based, so `dict`, +# `typing.Dict`, `collections.deque`, and `collections.abc.MutableMapping` all match +# however they were imported. The read-only interfaces (Mapping, Sequence, the +# immutable AbstractSet / `abc.Set`, Collection) and the immutable concretes (tuple, +# frozenset) are the escape hatch and are deliberately absent -- as is the bare name +# `Set`, which collides with the read-only `collections.abc.Set`. +MUTABLE_COLLECTIONS = frozenset(( + "dict", "list", "set", + "Dict", "List", "DefaultDict", "OrderedDict", "Counter", "Deque", "ChainMap", + "deque", "defaultdict", + "MutableMapping", "MutableSequence", "MutableSet", +)) + +# Callables whose result is a fresh *mutable* collection (LIT002). `tuple` and +# `frozenset` are deliberately absent -- they are the wrappers you reach for, and +# a generator expression fed to them is the blessed one-shot build. +MUTABLE_CONSTRUCTORS = frozenset(( + "dict", "list", "set", + "deque", "defaultdict", "OrderedDict", "Counter", "ChainMap", +)) +# A *qualified* call (`x.deque()`) counts as construction only for names that are rarely +# method names; `dict`/`list`/`set` are dropped here because `.dict()` / `.set()` / `.list()` +# are common methods (e.g. pydantic's `model.dict()`), not collection construction. A +# qualified `collections.deque(...)` still counts. +QUALIFIED_CONSTRUCTORS = MUTABLE_CONSTRUCTORS - frozenset(("dict", "list", "set")) +UNSAFE_GUARDS = frozenset(("TypeGuard", "TypeIs")) +MIN_REASON_LEN = 3 + +NOQA_RE = re.compile( + r"#\s*noqa" + r"(?P:\s*(?P[A-Z]+[0-9]+(?:\s*,\s*[A-Z]+[0-9]+)*))?" + r"(?P.*)", + re.IGNORECASE, +) +IGNORE_RE = re.compile( + r"#\s*(?:type|pyright|mypy):\s*ignore(?P\[[^\]]*\])?(?P.*)" +) +MUTABLE_OK_RE = re.compile(r"#\s*mutable-ok(?::\s*(?P.*))?") +CAST_OK_RE = re.compile(r"#\s*cast-ok(?::\s*(?P.*))?") +GUARD_OK_RE = re.compile(r"#\s*guard-ok(?::\s*(?P.*))?") +KWARGS_OK_RE = re.compile(r"#\s*kwargs-ok(?::\s*(?P.*))?") +ANY_OK_RE = re.compile(r"#\s*any-ok(?::\s*(?P.*))?") + +# Suppression tokens that must each carry a reason (LIT005). `any-ok` is owned by +# check_any_discipline.py but listed here so the reason requirement is enforced even +# when only this stdlib checker runs. +OK_SUPPRESSIONS: tuple[tuple[str, re.Pattern[str]], ...] = ( + ("mutable-ok", MUTABLE_OK_RE), + ("cast-ok", CAST_OK_RE), + ("guard-ok", GUARD_OK_RE), + ("kwargs-ok", KWARGS_OK_RE), + ("any-ok", ANY_OK_RE), +) + + +class Violation(NamedTuple): + path: Path + line: int + code: str + message: str + + def render(self) -> str: + return f"{self.path}:{self.line}: {self.code} {self.message}" + + +@dataclass(frozen=True, slots=True) +class Comments: + """The lines carrying each valid `*-ok` suppression.""" + + mutable_ok_lines: frozenset[int] + cast_ok_lines: frozenset[int] + guard_ok_lines: frozenset[int] + kwargs_ok_lines: frozenset[int] + + +# --------------------------------------------------------------------------- # +# Comment scanning (LIT003 / LIT004 / LIT005) +# --------------------------------------------------------------------------- # + + +def _reason_of(rest: str) -> str: + return rest.strip().lstrip("#-").strip() + + +def _valid_ok(regex: re.Pattern[str], text: str) -> bool: + """True iff `text` carries this suppression with a reason of usable length.""" + m = regex.search(text) + return bool(m) and len((m.group("reason") or "").strip()) >= MIN_REASON_LEN + + +def _comment_violations(path: Path, line_no: int, text: str) -> Iterator[Violation]: + """Pure: all LIT003/004/005 findings for one comment.""" + for token, regex in OK_SUPPRESSIONS: + m = regex.search(text) + if m and len((m.group("reason") or "").strip()) < MIN_REASON_LEN: + yield Violation(path, line_no, "LIT005", f"{token} requires a reason: `# {token}: `") + + m = NOQA_RE.search(text) + if m: + if not m.group("codes"): + yield Violation(path, line_no, "LIT003", "noqa requires rule codes: `# noqa: XXX123 # `") + elif len(_reason_of(m.group("rest"))) < MIN_REASON_LEN: + yield Violation(path, line_no, "LIT003", "noqa requires a reason: `# noqa: XXX123 # `") + + m = IGNORE_RE.search(text) + if m: + codes = m.group("codes") + if not codes or codes == "[]": + yield Violation(path, line_no, "LIT004", + "ignore requires codes: `# pyright: ignore[ruleName] # `") + elif len(_reason_of(m.group("rest"))) < MIN_REASON_LEN: + yield Violation(path, line_no, "LIT004", + "ignore requires a reason: `# pyright: ignore[ruleName] # `") + + +def scan_comments(path: Path, source: str) -> tuple[Comments, tuple[Violation, ...]]: + try: + tokens = tokenize.generate_tokens(io.StringIO(source).readline) + comment_toks = tuple((t.start[0], t.string) for t in tokens if t.type == tokenize.COMMENT) + except (tokenize.TokenError, SyntaxError): + # tokenize raises TokenError (EOF mid-construct) or a SyntaxError subclass + # (IndentationError / TabError) on malformed source; defer to ast.parse below, + # which re-raises and is reported as LIT000 rather than crashing the run. + return Comments(frozenset(), frozenset(), frozenset(), frozenset()), () + + def _lines_with(regex: re.Pattern[str]) -> frozenset[int]: + return frozenset(line for line, text in comment_toks if _valid_ok(regex, text)) + + return ( + Comments( + mutable_ok_lines=_lines_with(MUTABLE_OK_RE), + cast_ok_lines=_lines_with(CAST_OK_RE), + guard_ok_lines=_lines_with(GUARD_OK_RE), + kwargs_ok_lines=_lines_with(KWARGS_OK_RE), + ), + tuple(v for line, text in comment_toks for v in _comment_violations(path, line, text)), + ) + + +# --------------------------------------------------------------------------- # + + +def mutable_names_in(annotation: ast.expr) -> Iterator[str]: + """Yield mutable-collection names anywhere inside an annotation expression. + + Matches bare names (`dict`, `MutableMapping`) and dotted access (`typing.Dict`, + `collections.deque`, `collections.abc.MutableMapping`), descends through nesting + (`Mapping[str, list[int]]`, `tuple[set[int], ...]`) and string forward references. + """ + for node in ast.walk(annotation): + if isinstance(node, ast.Name) and node.id in MUTABLE_COLLECTIONS: + yield node.id + elif isinstance(node, ast.Attribute) and node.attr in MUTABLE_COLLECTIONS: + yield node.attr + elif isinstance(node, ast.Constant): + value: object = node.value # forward references arrive as string constants + if isinstance(value, str): + try: + inner = ast.parse(value, mode="eval").body + except SyntaxError: + continue + yield from mutable_names_in(inner) + + +def _mutable_ann(path: Path, line: int, name: str, where: str) -> Violation: + return Violation( + path, line, "LIT001", + f"mutable `{name}` in {where}: a mutable collection can be grown or rewritten " + f"by whoever holds it. Annotate a read-only view -- Mapping[...], Sequence[...], " + f"AbstractSet[...], tuple[X, ...], frozenset[X], or a frozen dataclass / " + f"NamedTuple / ReadOnly TypedDict -- and build it functionally, not by " + f"append-in-a-loop (suppress: `# mutable-ok: `)", + ) + + +def _annotation_violations( + path: Path, annotation: ast.expr | None, line: int, where: str, ok_lines: frozenset[int] +) -> Iterator[Violation]: + if annotation is None or line in ok_lines: + return + yield from (_mutable_ann(path, line, name, where) for name in mutable_names_in(annotation)) + + +def _function_violations( + path: Path, node: ast.FunctionDef | ast.AsyncFunctionDef, comments: Comments +) -> Iterator[Violation]: + mutable_ok = comments.mutable_ok_lines + args = node.args + for arg in (*args.posonlyargs, *args.args, *args.kwonlyargs): + yield from _annotation_violations( + path, arg.annotation, arg.lineno, f"parameter `{arg.arg}` of `{node.name}`", mutable_ok + ) + + # *args is allowed when typed (it's just a tuple); ruff ANN002 forces the + # annotation, so here we only add the LIT001 mutable-collection check on the element type. + if args.vararg is not None: + yield from _annotation_violations( + path, args.vararg.annotation, args.vararg.lineno, f"`*args` of `{node.name}`", mutable_ok + ) + + # **kwargs is banned outright (LIT008): it erases the keyword contract and forces + # Any-typing on everything it carries. ruff can require it be typed (ANN003) but + # cannot ban the syntax, so this rule does. + if args.kwarg is not None and args.kwarg.lineno not in comments.kwargs_ok_lines: + yield Violation( + path, args.kwarg.lineno, "LIT008", + f"`**{args.kwarg.arg}` is banned: it erases the keyword contract and forces " + f"Any-typing; declare explicit keyword parameters, or accept one frozen payload " + f"(frozen dataclass / NamedTuple / ReadOnly TypedDict) " + f"(suppress: `# kwargs-ok: `)", + ) + + if node.returns is not None: + yield from _annotation_violations( + path, node.returns, node.returns.lineno, f"return type of `{node.name}`", mutable_ok + ) + + +def iter_annotation_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]: + # Every annotation is in scope: signatures (params / *args / return) plus every + # `x: T` -- class attribute, local, or module global. The latter three are all + # ast.AnnAssign, so one walk covers them; only the signature annotations (which + # are not AnnAssign) need the dedicated helper. + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + yield from _function_violations(path, node, comments) + elif isinstance(node, ast.AnnAssign): + target = node.target.id if isinstance(node.target, ast.Name) else "" + yield from _annotation_violations( + path, node.annotation, node.lineno, + f"the type of `{target}`", comments.mutable_ok_lines, + ) + + +# --------------------------------------------------------------------------- # +# Unchecked casts (LIT006) and unverified narrowing predicates (LIT007) +# --------------------------------------------------------------------------- # + + +def _is_cast_call(node: ast.Call) -> bool: + """`cast(...)` or `typing.cast(...)`, however the name was imported/aliased. + + Name-based like MUTABLE_COLLECTIONS: a stray method called `.cast()` is a rare + false positive, suppressible with `# cast-ok: `. + """ + func = node.func + return (isinstance(func, ast.Name) and func.id == "cast") or ( + isinstance(func, ast.Attribute) and func.attr == "cast" + ) + + +def iter_cast_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]: + for node in ast.walk(tree): + if isinstance(node, ast.Call) and _is_cast_call(node) and node.lineno not in comments.cast_ok_lines: + yield Violation( + path, node.lineno, "LIT006", + "cast() is an unchecked assertion (the type checker takes it on faith); " + "validate into a frozen dataclass/NamedTuple/ReadOnly TypedDict at the " + "boundary instead (suppress: `# cast-ok: `)", + ) + + +def iter_guard_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]: + # TypeGuard/TypeIs are legal only as a function's return annotation (`-> TypeGuard[int]`), + # so the walk is confined to `node.returns`; a runtime name that merely happens to read + # `TypeGuard` is not a narrowing predicate. ruff bans the import; this flags the use. + for node in ast.walk(tree): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) or node.returns is None: + continue + for sub in ast.walk(node.returns): + name = ( + sub.id if isinstance(sub, ast.Name) + else sub.attr if isinstance(sub, ast.Attribute) + else None + ) + if name in UNSAFE_GUARDS and sub.lineno not in comments.guard_ok_lines: + yield Violation( + path, sub.lineno, "LIT007", + f"`{name}` narrowing predicate: the checker never verifies the body, so a " + f"wrong guard silently corrupts types; parse into a concrete type instead " + f"(suppress: `# guard-ok: `)", + ) + + +# --------------------------------------------------------------------------- # +# Mutable-collection construction (LIT002) +# --------------------------------------------------------------------------- # + + +def _annotations_of(node: ast.AST) -> tuple[ast.expr | None, ...]: + """The annotation expressions a node carries (signatures and `x: T`).""" + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + a = node.args + params = (*a.posonlyargs, *a.args, *a.kwonlyargs, a.vararg, a.kwarg) + return (*(p.annotation for p in params if p is not None), node.returns) + if isinstance(node, ast.AnnAssign): + return (node.annotation,) + return () + + +def _annotation_node_ids(tree: ast.AST) -> frozenset[int]: + """ids() of every node living inside an annotation. + + A list display inside an annotation (`Callable[[int], str]`) is type syntax, + not construction, so the LIT002 walk must skip those subtrees. + """ + return frozenset( + id(sub) + for node in ast.walk(tree) + for ann in _annotations_of(node) + if ann is not None + for sub in ast.walk(ann) + ) + + +def _construction_kind(node: ast.expr) -> str | None: + """Human label if `node` builds a mutable collection, else None.""" + if isinstance(node, ast.List): + return "list literal" + if isinstance(node, ast.ListComp): + return "list comprehension" + if isinstance(node, ast.Set): + return "set literal" + if isinstance(node, ast.SetComp): + return "set comprehension" + if isinstance(node, ast.Dict): + return "dict literal" + if isinstance(node, ast.DictComp): + return "dict comprehension" + if isinstance(node, ast.Call): + func = node.func + if isinstance(func, ast.Name) and func.id in MUTABLE_CONSTRUCTORS: + return f"`{func.id}()` constructor" + if isinstance(func, ast.Attribute) and func.attr in QUALIFIED_CONSTRUCTORS: + return f"`{func.attr}()` constructor" + return None + + +def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]: + in_annotation = _annotation_node_ids(tree) + for node in ast.walk(tree): + if not isinstance(node, ast.expr) or id(node) in in_annotation: + continue + kind = _construction_kind(node) + if kind is None or node.lineno in comments.mutable_ok_lines: + continue + yield Violation( + path, node.lineno, "LIT002", + f"mutable {kind}: this builds a collection that can be grown or rewritten. " + f"Build it in one shot and freeze it -- a tuple/frozenset wrapping a generator " + f"(`tuple(f(x) for x in xs)`), a tuple literal, or a frozen dataclass / NamedTuple " + f"/ ReadOnly TypedDict (suppress: `# mutable-ok: `)", + ) + + +# --------------------------------------------------------------------------- # +# Driver +# --------------------------------------------------------------------------- # + + +def check_file(path: Path) -> tuple[Violation, ...]: + try: + source = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError) as exc: + return (Violation(path, 0, "LIT000", f"could not read file: {exc}"),) + + comments, violations = scan_comments(path, source) + + try: + tree = ast.parse(source, filename=str(path)) + except SyntaxError as exc: + return (*violations, Violation(path, exc.lineno or 0, "LIT000", f"syntax error: {exc.msg}")) + + return ( + *violations, + *iter_annotation_violations(path, tree, comments), + *iter_cast_violations(path, tree, comments), + *iter_guard_violations(path, tree, comments), + *iter_construction_violations(path, tree, comments), + ) + + +def collect_paths(raw: Iterable[str]) -> Iterator[Path]: + for item in raw: + p = Path(item) + if p.is_dir(): + yield from sorted(p.rglob("*.py")) + elif p.suffix == ".py": + yield p + + +def main(argv: Sequence[str]) -> int: + paths = tuple(a for a in argv if not a.startswith("-")) + if not paths: + print("usage: check_type_discipline.py ...", file=sys.stderr) + return 2 + + violations = sorted(v for path in collect_paths(paths) for v in check_file(path)) + for v in violations: + print(v.render()) + + if violations: + print(f"\n{len(violations)} violation(s).", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) + \ No newline at end of file diff --git a/scripts/type_discipline_gate.py b/scripts/type_discipline_gate.py new file mode 100644 index 00000000000..c111486e56a --- /dev/null +++ b/scripts/type_discipline_gate.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +"""Total-count gate for the LIT* rules in scripts/check_type_discipline.py. + +Sibling of scripts/ruff_strict_gate.py. Each rule listed in +type-discipline-budget.json has a hard ceiling (baseline + slack). The gate counts +each rule across the whole `litellm` tree and fails when a rule is both over its +ceiling and higher than the base it merges into, so a change is blamed for the +violations it adds, never for drift that already exists in the base. + +Rules not present in the budget are ignored, but today every rule the checker +emits is gated: LIT001 (mutable collection in any annotation), LIT002 +(mutable-collection construction), LIT003/LIT004 (noqa / ignore without codes or +reason), LIT006 (cast), and LIT008 (`**kwargs`) carry slack-buffered ceilings to +ratchet down; LIT005 (`*-ok` suppression without a reason) is frozen at slack 0 +so any net-new reasonless suppression trips the gate; and LIT007 (TypeGuard/TypeIs) +is a hard zero. Re-baseline with `--update` to ratchet a ceiling down. +""" + +import argparse +import json +import re +import shutil +import subprocess +import sys +import tempfile +from collections import Counter +from pathlib import Path +from typing import NamedTuple + +REPO_ROOT = Path(__file__).resolve().parent.parent +CHECKER = REPO_ROOT / "scripts" / "check_type_discipline.py" +BUDGET_PATH = REPO_ROOT / "type-discipline-budget.json" +TARGET = "litellm" +DEFAULT_BASE = "origin/litellm_internal_staging" + +_HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") +_LINE = re.compile(r"^(?P.+?):(?P\d+): (?PLIT\d+) ") + + +class Violation(NamedTuple): + file: str + line: int + code: str + + +class Breach(NamedTuple): + rule: str + total: int + cap: int + added: int + + +def _run(cmd: list, cwd: Path = REPO_ROOT) -> str: + proc = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) + if proc.returncode not in (0, 1): + sys.stderr.write(proc.stderr) + raise SystemExit(f"{cmd[0]} exited {proc.returncode}") + return proc.stdout + + +def _check(root: Path, checker: Path) -> list: + # Resolve root first: on macOS tempfile dirs (/var/...) resolve to /private/var/..., + # and the checker prints already-resolved absolute paths, so relative_to would fail. + root = root.resolve() + out = _run([sys.executable, str(checker), str(root / TARGET)], cwd=root) + found = [] + for line in out.splitlines(): + m = _LINE.match(line) + if m is None: + continue + name = Path(m.group("file")) + full = name if name.is_absolute() else root / name + rel = full.resolve().relative_to(root).as_posix() + found.append(Violation(rel, int(m.group("line")), m.group("code"))) + return found + + +def head_violations() -> list: + return _check(REPO_ROOT, CHECKER) + + +def count_by_rule(violations: list) -> dict: + return dict(Counter(v.code for v in violations)) + + +def base_counts(ref: str) -> dict: + parent = Path(tempfile.mkdtemp(prefix="lit_base_")) + worktree = parent / "wt" + try: + _run(["git", "worktree", "add", "--detach", str(worktree), ref]) + # Measure the base with the *current* rule logic, not whatever shipped at base. + (worktree / "scripts").mkdir(parents=True, exist_ok=True) + checker = worktree / "scripts" / "check_type_discipline.py" + shutil.copy(CHECKER, checker) + return count_by_rule(_check(worktree, checker)) + finally: + # Best-effort teardown: cleanup must never raise, or it masks the real error when + # the body (or the `worktree add` itself) failed. rmtree is already best-effort. + subprocess.run( + ["git", "worktree", "remove", "--force", str(worktree)], + cwd=REPO_ROOT, capture_output=True, text=True, + ) + shutil.rmtree(parent, ignore_errors=True) + + +def over_ceiling(head: dict, budget: dict) -> frozenset: + """Rules whose head count already exceeds baseline + slack. + + A rule can only breach when it is over its ceiling, so when none are the base + comparison cannot change the verdict and the base worktree scan can be skipped. + """ + return frozenset( + rule for rule, spec in budget.items() + if head.get(rule, 0) > spec["baseline"] + spec["slack"] + ) + + +def evaluate(head: dict, base: dict, budget: dict) -> list: + breaches = [] + for rule, spec in budget.items(): + cap = spec["baseline"] + spec["slack"] + total = head.get(rule, 0) + if total > cap and total > base.get(rule, 0): + breaches.append(Breach(rule, total, cap, total - base.get(rule, 0))) + return sorted(breaches) + + +def parse_changed_lines(diff_text: str) -> dict: + changed: dict = {} + path = None + for line in diff_text.splitlines(): + if line.startswith("+++ b/"): + path = line[6:] + elif path and (match := _HUNK.match(line)): + start = int(match.group(1)) + count = int(match.group(2)) if match.group(2) is not None else 1 + changed.setdefault(path, set()).update(range(start, start + count)) + return changed + + +def introduced(violations: list, changed: dict) -> list: + return [v for v in violations if v.line in changed.get(v.file, set())] + + +def cmd_check(base: str) -> None: + budget = json.loads(BUDGET_PATH.read_text()) + head = head_violations() + head_counts = count_by_rule(head) + if not over_ceiling(head_counts, budget): + print(f"OK: every LIT rule is within its codebase ceiling (base {base})") + return + base_point = _run(["git", "merge-base", base, "HEAD"]).strip() or base + breaches = evaluate(head_counts, base_counts(base_point), budget) + if not breaches: + print(f"OK: every LIT rule is within its codebase ceiling (base {base})") + return + new = introduced( + head, + parse_changed_lines( + _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET]) + ), + ) + print(f"FAIL: LIT-rule totals exceed their ceiling (base {base}):") + for breach in breaches: + print( + f" {breach.rule}: total {breach.total} over cap {breach.cap} (this change added {breach.added})" + ) + for violation in sorted(v for v in new if v.code == breach.rule): + print(f" {violation.file}:{violation.line}") + print( + "Remove the new violations, give each a reason (`# noqa: XXX # `, " + "`# pyright: ignore[rule] # `, `# mutable-ok: `, " + "`# cast-ok: `, `# guard-ok: `, `# kwargs-ok: `), or " + "remove an equal number elsewhere; the ceiling is baseline + slack in " + "type-discipline-budget.json." + ) + raise SystemExit(1) + + +def cmd_update() -> None: + budget = json.loads(BUDGET_PATH.read_text()) + head = count_by_rule(head_violations()) + for rule in budget: + budget[rule]["baseline"] = head.get(rule, 0) + BUDGET_PATH.write_text(json.dumps(budget, indent=2, sort_keys=True) + "\n") + print("Re-captured per-rule baselines from the current tree") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--update", action="store_true") + args = parser.parse_args() + cmd_update() if args.update else cmd_check(args.base) + + +if __name__ == "__main__": + main() diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py new file mode 100644 index 00000000000..a7ae9deaf87 --- /dev/null +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -0,0 +1,97 @@ +"""Tests for scripts/budget_ratchet_check.py. + +The guard's whole contract is "ceilings may only fall": a raised ceiling, a dropped +rule, or a deleted file is a regression, while a lowered/equal ceiling, a brand-new +rule, or a brand-new budget file is fine. Each branch is pinned here. +""" + +import importlib.util +import subprocess +import sys +from pathlib import Path + +_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py" +_spec = importlib.util.spec_from_file_location("budget_ratchet_check", _MODULE_PATH) +ratchet = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(ratchet) + + +def _spec_of(baseline, slack): + return {"baseline": baseline, "slack": slack} + + +def test_caps_sum_baseline_and_slack_and_skip_malformed(): + caps = ratchet._caps({"LIT006": _spec_of(1013, 10), "junk": 5}) + assert caps == {"LIT006": 1023} # malformed (non-dict) spec ignored + + +def test_raised_ceiling_is_a_regression(): + base = {"LIT006": _spec_of(1013, 10)} + head = {"LIT006": _spec_of(1013, 11)} # cap 1023 -> 1024 + regs = ratchet.regressions_for("b.json", base, head) + assert [r.rule for r in regs] == ["LIT006"] + assert "1023 -> 1024" in regs[0].detail + + +def test_lowered_or_equal_ceiling_is_clean(): + base = {"LIT006": _spec_of(1013, 10)} + assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1000, 10)}) == [] + assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1013, 10)}) == [] + # slack traded for baseline at the same ceiling is fine + assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1023, 0)}) == [] + + +def test_dropped_rule_is_a_regression(): + regs = ratchet.regressions_for("b.json", {"LIT007": _spec_of(0, 0)}, {}) + assert [r.rule for r in regs] == ["LIT007"] + assert "dropped" in regs[0].detail + + +def test_new_rule_in_head_is_clean(): + assert ratchet.regressions_for("b.json", {}, {"LIT009": _spec_of(5, 0)}) == [] + + +def test_deleted_budget_file_is_a_regression(): + regs = ratchet.regressions_for("b.json", {"LIT006": _spec_of(1, 0)}, None) + assert [r.rule for r in regs] == ["*"] + assert "deleted" in regs[0].detail + + +def test_new_budget_file_has_nothing_to_ratchet(): + assert ratchet.regressions_for("b.json", None, {"LIT006": _spec_of(1, 0)}) == [] + + +def test_default_budgets_watch_every_budget_file_in_the_repo(): + # This job is the repo's only ceiling-raise alarm, so every *-budget.json on disk must be + # watched; a budget left out of DEFAULT_BUDGETS (e.g. basedpyright-code-budget.json) can be + # loosened with no signal. Equality also catches a phantom entry that no longer exists. + repo_root = _MODULE_PATH.parents[1] + on_disk = frozenset(p.name for p in repo_root.glob("*budget*.json")) + assert on_disk == frozenset(ratchet.DEFAULT_BUDGETS) + + +# --------------------------------------------------------------------------- # +# Base-ref resolution: a bad ref must fail loudly, never pass vacuously +# --------------------------------------------------------------------------- # + + +def test_ref_is_commit_distinguishes_real_from_bogus(): + assert ratchet._ref_is_commit("HEAD") is True + assert ratchet._ref_is_commit("definitely-not-a-real-ref-zzz") is False + + +def test_load_base_reads_a_present_file_and_none_for_an_absent_one(): + # A real budget file exists at HEAD; a made-up path is absent at the same (valid) ref. + assert ratchet._load_base("type-discipline-budget.json", "HEAD") is not None + assert ratchet._load_base("scripts/no-such-budget-xyz.json", "HEAD") is None + + +def test_unresolvable_base_ref_exits_nonzero_instead_of_skipping(): + proc = subprocess.run( + [sys.executable, str(_MODULE_PATH), "--base", "definitely-not-a-real-ref-zzz"], + cwd=_MODULE_PATH.parents[1], + capture_output=True, + text=True, + ) + assert proc.returncode == 1 + assert "does not resolve to a commit" in proc.stderr diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py new file mode 100644 index 00000000000..436904b017c --- /dev/null +++ b/tests/test_litellm/test_check_type_discipline.py @@ -0,0 +1,199 @@ +"""Tests for scripts/check_type_discipline.py. + +Each rule is exercised on a snippet that violates it and on one that does not, so a +mutation that drops a rule, inverts a suppression, or breaks the comment scanner makes +a test fail. The comment-scanner cases are the regression for the readline path: if +`scan_comments` ever stops tokenizing comments, the LIT003/LIT005 assertions go red. +""" + +import importlib.util +import json +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_MODULE_PATH = _REPO_ROOT / "scripts" / "check_type_discipline.py" +_spec = importlib.util.spec_from_file_location("check_type_discipline", _MODULE_PATH) +checker = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = checker # let the frozen dataclass resolve its own module +_spec.loader.exec_module(checker) + + +def _codes(tmp_path, source): + f = tmp_path / "snippet.py" + f.write_text(source, encoding="utf-8") + return [v.code for v in checker.check_file(f)] + + +# --------------------------------------------------------------------------- # +# Comment scanning (the readline path) — LIT003 / LIT004 / LIT005 +# --------------------------------------------------------------------------- # + + +def test_scan_comments_tokenizes_every_comment(): + # Direct regression for scan_comments: a bare noqa (LIT003) only surfaces if the comment + # was tokenized, and the valid cast-ok suppression line must be captured. A crash in the + # readline path would leave both empty. + source = "x = 1 # noqa\ny = 2 # cast-ok: validated upstream by the caller\n" + comments, violations = checker.scan_comments(Path("snippet.py"), source) + assert [v.code for v in violations] == ["LIT003"] + assert comments.cast_ok_lines == frozenset({2}) + + +def test_scan_comments_does_not_crash_on_malformed_source(): + # A dedent mismatch makes tokenize raise IndentationError (a SyntaxError subclass); + # scan_comments must swallow it, not propagate and crash the whole run. + comments, violations = checker.scan_comments(Path("x.py"), "if True:\n a = 1\n b = 2\n") + assert violations == () + assert comments.cast_ok_lines == frozenset() + + +def test_malformed_source_degrades_to_lit000(tmp_path): + # The checker's contract is "bad file -> LIT000, never crash". An untokenizable file + # falls through scan_comments to ast.parse, which is reported as a single LIT000. + assert _codes(tmp_path, "if True:\n a = 1\n b = 2\n") == ["LIT000"] + + +def test_noqa_without_codes_is_flagged(tmp_path): + assert "LIT003" in _codes(tmp_path, "x = 1 # noqa\n") + + +def test_noqa_with_codes_and_reason_is_clean(tmp_path): + assert "LIT003" not in _codes(tmp_path, "x = 1 # noqa: TID251 # legacy import, removed in #123\n") + + +def test_ignore_without_reason_is_flagged(tmp_path): + assert "LIT004" in _codes(tmp_path, "x = 1 # type: ignore[arg-type]\n") + + +def test_ignore_with_codes_and_reason_is_clean(tmp_path): + assert "LIT004" not in _codes(tmp_path, "x = 1 # pyright: ignore[reportArgumentType] # upstream stub is wrong\n") + + +def test_ok_suppression_without_reason_is_flagged(tmp_path): + codes = _codes(tmp_path, "y = [] # mutable-ok\n") + assert "LIT005" in codes # reasonless suppression + assert "LIT002" in codes # and it does not suppress, so the construction still trips + + +# --------------------------------------------------------------------------- # +# Mutable annotations (LIT001) and construction (LIT002) +# --------------------------------------------------------------------------- # + + +def test_mutable_annotation_is_flagged(tmp_path): + assert "LIT001" in _codes(tmp_path, "x: dict[str, int]\n") + + +def test_typing_alias_and_forward_ref_annotations_are_flagged(tmp_path): + assert "LIT001" in _codes(tmp_path, "from typing import List\nx: List[int]\n") + assert "LIT001" in _codes(tmp_path, 'x: "dict[str, int]"\n') + + +def test_readonly_annotations_are_clean(tmp_path): + for ann in ("Mapping[str, int]", "Sequence[int]", "tuple[int, ...]", "frozenset[int]"): + assert "LIT001" not in _codes(tmp_path, f"from typing import Mapping, Sequence\nx: {ann}\n") + + +def test_mutable_construction_is_flagged(tmp_path): + assert "LIT002" in _codes(tmp_path, "y = []\n") + assert "LIT002" in _codes(tmp_path, "z = dict(a=1)\n") + + +def test_construction_inside_annotation_is_exempt(tmp_path): + # `Callable[[int], str]` carries a list display that is type syntax, not construction. + assert "LIT002" not in _codes( + tmp_path, "from typing import Callable\ndef f(cb: Callable[[int], str]) -> None:\n return None\n" + ) + + +def test_generator_and_tuple_are_not_construction(tmp_path): + assert "LIT002" not in _codes(tmp_path, "g = tuple(i for i in range(3))\n") + assert "LIT002" not in _codes(tmp_path, "t = (1, 2, 3)\n") + + +def test_dict_list_set_method_calls_are_not_construction(tmp_path): + # `.dict()` / `.list()` / `.set()` are common method names (e.g. pydantic model.dict()), + # not collection construction; only the unqualified builtins count. + assert "LIT002" not in _codes(tmp_path, "d = model.dict()\n") + assert "LIT002" not in _codes(tmp_path, "s = obj.set()\n") + assert "LIT002" in _codes(tmp_path, "d = dict(a=1)\n") # unqualified still counts + + +def test_qualified_collections_constructors_still_count(tmp_path): + # collections concretes are rarely method names, so a qualified call still flags. + assert "LIT002" in _codes(tmp_path, "import collections\nq = collections.deque()\n") + assert "LIT002" in _codes(tmp_path, "import collections\nm = collections.defaultdict(list)\n") + + +def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): + codes = _codes(tmp_path, "x: dict[str, int] = {} # mutable-ok: in-place buffer mutated hot path\n") + assert "LIT001" not in codes + assert "LIT002" not in codes + + +# --------------------------------------------------------------------------- # +# Casts (LIT006) +# --------------------------------------------------------------------------- # + + +def test_cast_call_is_flagged(tmp_path): + assert "LIT006" in _codes(tmp_path, "from typing import cast\ny = cast(int, object())\n") + + +def test_cast_ok_with_reason_suppresses(tmp_path): + assert "LIT006" not in _codes( + tmp_path, "from typing import cast\ny = cast(int, object()) # cast-ok: validated by schema above\n" + ) + + +# --------------------------------------------------------------------------- # +# Narrowing predicates (LIT007) — must fire only in return annotations +# --------------------------------------------------------------------------- # + + +def test_guard_in_return_annotation_is_flagged(tmp_path): + src = "from typing import TypeGuard\ndef is_int(v: object) -> TypeGuard[int]:\n return isinstance(v, int)\n" + assert "LIT007" in _codes(tmp_path, src) + + +def test_guard_name_outside_annotation_is_not_flagged(tmp_path): + # A runtime name or attribute that merely reads `TypeGuard`/`TypeIs` is not a predicate. + assert "LIT007" not in _codes(tmp_path, "TypeGuard = 1\nx = TypeGuard + 1\n") + assert "LIT007" not in _codes(tmp_path, "import obj\n_ = obj.TypeIs\n") + + +def test_guard_ok_with_reason_suppresses(tmp_path): + src = ( + "from typing import TypeGuard\n" + "def is_int(v: object) -> TypeGuard[int]: # guard-ok: predicate proven by the assert below\n" + " assert isinstance(v, int)\n" + " return True\n" + ) + assert "LIT007" not in _codes(tmp_path, src) + + +# --------------------------------------------------------------------------- # +# **kwargs (LIT008) — typed *args stays clean +# --------------------------------------------------------------------------- # + + +def test_kwargs_parameter_is_flagged(tmp_path): + assert "LIT008" in _codes(tmp_path, "def f(**kwargs) -> None:\n return None\n") + + +def test_typed_args_is_clean_but_kwargs_ok_suppresses(tmp_path): + assert "LIT008" not in _codes(tmp_path, "def f(*args: int) -> None:\n return None\n") + assert "LIT008" not in _codes( + tmp_path, "def f(**kwargs: int) -> None: # kwargs-ok: passthrough to a third-party sink\n return None\n" + ) + + +# --------------------------------------------------------------------------- # +# Budget integrity: every emittable LIT rule (bar the LIT000 read/parse error) is gated +# --------------------------------------------------------------------------- # + + +def test_budget_covers_exactly_the_checker_rules(): + budget = json.loads((_REPO_ROOT / "type-discipline-budget.json").read_text()) + assert set(budget) == {f"LIT00{n}" for n in range(1, 9)} diff --git a/tests/test_litellm/test_type_discipline_gate.py b/tests/test_litellm/test_type_discipline_gate.py new file mode 100644 index 00000000000..d7d827685a6 --- /dev/null +++ b/tests/test_litellm/test_type_discipline_gate.py @@ -0,0 +1,40 @@ +"""Tests for scripts/type_discipline_gate.py. + +The gate's correctness lives in two pure functions: `over_ceiling` (which decides +whether the expensive base worktree scan is even needed) and `evaluate` (the +drift-safe breach check). Both are pinned here. +""" + +import importlib.util +from pathlib import Path + +_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "type_discipline_gate.py" +_spec = importlib.util.spec_from_file_location("type_discipline_gate", _MODULE_PATH) +gate = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(gate) + + +def _budget(baseline, slack): + return {"LIT006": {"baseline": baseline, "slack": slack}} + + +def test_over_ceiling_flags_only_counts_above_baseline_plus_slack(): + budget = _budget(10, 2) # cap 12 + assert gate.over_ceiling({"LIT006": 12}, budget) == frozenset() # at cap + assert gate.over_ceiling({"LIT006": 13}, budget) == frozenset({"LIT006"}) # over cap + assert gate.over_ceiling({}, budget) == frozenset() # missing rule counts as zero + + +def test_over_ceiling_is_independent_across_rules(): + budget = {"LIT001": {"baseline": 5, "slack": 0}, "LIT006": {"baseline": 10, "slack": 0}} + assert gate.over_ceiling({"LIT001": 6, "LIT006": 10}, budget) == frozenset({"LIT001"}) + + +def test_evaluate_blames_only_a_rule_over_cap_and_over_base(): + budget = _budget(10, 0) # cap 10 + # over cap and grown vs base -> breach + assert [b.rule for b in gate.evaluate({"LIT006": 12}, {"LIT006": 9}, budget)] == ["LIT006"] + # over cap but flat vs base (pre-existing drift) -> not blamed + assert gate.evaluate({"LIT006": 12}, {"LIT006": 12}, budget) == [] + # within cap -> not blamed regardless of base + assert gate.evaluate({"LIT006": 10}, {"LIT006": 0}, budget) == [] diff --git a/type-discipline-budget.json b/type-discipline-budget.json new file mode 100644 index 00000000000..a6588ac89aa --- /dev/null +++ b/type-discipline-budget.json @@ -0,0 +1,34 @@ +{ + "LIT001": { + "baseline": 21452, + "slack": 2000 + }, + "LIT002": { + "baseline": 25022, + "slack": 2500 + }, + "LIT003": { + "baseline": 397, + "slack": 25 + }, + "LIT004": { + "baseline": 2515, + "slack": 50 + }, + "LIT005": { + "baseline": 0, + "slack": 0 + }, + "LIT006": { + "baseline": 1013, + "slack": 100 + }, + "LIT007": { + "baseline": 0, + "slack": 0 + }, + "LIT008": { + "baseline": 914, + "slack": 90 + } +} From cd26f7d77af73308d90270320c0b66d7be8a7850 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 16 Jun 2026 17:23:37 -0700 Subject: [PATCH 3/3] feat(proxy): add verification_uri_complete to CLI SSO device flow (#30571) * feat(proxy): add verification_uri_complete to CLI SSO device flow Add an opt-in verification_uri_complete to POST /sso/cli/start. The URL is the existing /sso/key/generate?source=litellm-cli&key= browser-start URL with an added user_code query param. The code is carried through the OAuth flow via the same state channel that already carries login_id, and the post-SSO verify page pre-fills the user_code input (HTML-escaped) so same-host clients confirm rather than transcribe. The manual flow is unchanged and remains the default: when no user_code is present the verify page renders the empty input byte-for-byte as before, and submission still hashes and compare_digest-checks both the user_code and the browser_complete_token. Pre-filling is a UX shortcut, not an auth bypass. Resolves LIT-3693 * fix(proxy): validate CLI SSO user_code and clarify pre-filled verify page Address Greptile review on the verification_uri_complete flow. Guard the user_code query param with the canonical server-issued format ([A-HJ-NP-Z2-9]{4}-[A-HJ-NP-Z2-9]{4}) before it is threaded into the OAuth state, so an actor who knows a login_id cannot bloat the size-limited state with an arbitrary value; a non-conforming code falls back to the manual flow. Make the verify-page instruction conditional so the pre-filled page reads "Confirm the verification code below" instead of pointing at a terminal that, in the daemon use case, does not exist. * fix(proxy): modern union syntax for new CLI SSO params and regen dashboard types Use str | None instead of Optional[str] on the CLI SSO signatures touched by this PR so the ruff strict-rule budget (UP045) stays under its ceiling, and regenerate ui/litellm-dashboard/src/lib/http/schema.d.ts so the dashboard API types pick up the new optional user_code query param on /sso/key/generate. * fix(proxy): gate CLI SSO verification_uri_complete behind operator opt-in (default off) Gate verification_uri_complete behind a new general_settings flag allow_cli_sso_verification_uri_complete, default false. When off, /sso/cli/start does not return verification_uri_complete and /sso/key/generate ignores the user_code query param, so the default deployment keeps the existing manual flow. Same-host clients, where the device that starts the flow and the browser run on the same machine, opt in explicitly. The submitted code is still hashed and compare_digest-checked and browser_complete_token is still required. Documents the flag on ConfigGeneralSettings and regenerates the dashboard API types. --- litellm/proxy/_types.py | 4 + litellm/proxy/management_endpoints/ui_sso.py | 109 +++++- .../proxy/management_endpoints/test_ui_sso.py | 326 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 4 files changed, 432 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c71127a4c3a..765e90bc896 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2152,6 +2152,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): master_key: Optional[str] = Field( None, description="require a key for all calls to proxy" ) + allow_cli_sso_verification_uri_complete: bool | None = Field( + None, + description="opt-in to RFC 8628 verification_uri_complete for the CLI SSO device flow, pre-filling the user_code in the browser. Off by default; intended for same-host clients where the device that starts the flow and the browser run on the same machine", + ) database_url: Optional[str] = Field( None, description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 2bf12880a75..91a5c109acf 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -145,6 +145,9 @@ _CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS = 60 _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30 _CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" _CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$") +_CLI_SSO_USER_CODE_RE = re.compile( + rf"^[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}-[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}$" +) _CLI_SSO_SCALAR_TYPES = (str, int, float, bool) _CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$") _CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset( @@ -182,6 +185,45 @@ def _is_valid_cli_sso_login_id(login_id: Optional[str]) -> bool: return isinstance(login_id, str) and bool(_CLI_SSO_LOGIN_ID_RE.fullmatch(login_id)) +def _is_valid_cli_sso_user_code(user_code: str | None) -> bool: + return isinstance(user_code, str) and bool( + _CLI_SSO_USER_CODE_RE.fullmatch(user_code) + ) + + +def _cli_sso_verification_uri_complete_enabled() -> bool: + from litellm.proxy.proxy_server import general_settings + + return bool( + general_settings.get( # any-ok: operator opt-in read from the untyped general_settings dict + "allow_cli_sso_verification_uri_complete", False + ) + ) + + +def _cli_sso_start_response_body( + *, + login_id: str, + poll_secret: str, + user_code: str, + verification_uri_complete: str | None, +) -> dict[str, str | int]: + if verification_uri_complete is None: + return { + "login_id": login_id, + "poll_secret": poll_secret, + "user_code": user_code, + "expires_in": CLI_SSO_SESSION_TTL_SECONDS, + } + return { + "login_id": login_id, + "poll_secret": poll_secret, + "user_code": user_code, + "verification_uri_complete": verification_uri_complete, + "expires_in": CLI_SSO_SESSION_TTL_SECONDS, + } + + def _get_cli_sso_start_rate_limit_cache_key( request: Request, use_x_forwarded_for: Optional[bool] = False ) -> str: @@ -478,10 +520,20 @@ def _cli_poll_attribution_metadata_from_session( def _render_cli_sso_verification_page( - verify_url: str, browser_complete_token: str + verify_url: str, + browser_complete_token: str, + prefill_user_code: str | None = None, ) -> str: escaped_verify_url = escape(verify_url, quote=True) escaped_browser_complete_token = escape(browser_complete_token, quote=True) + user_code_value_attr = ( + f' value="{escape(prefill_user_code, quote=True)}"' if prefill_user_code else "" + ) + instructions = ( + "Confirm the verification code below to finish this login." + if prefill_user_code + else "Enter the verification code shown in your terminal to finish this login." + ) return f""" @@ -535,11 +587,11 @@ def _render_cli_sso_verification_page(

Complete CLI Login

-

Enter the verification code shown in your terminal to finish this login.

+

{instructions}

- +
@@ -573,12 +625,29 @@ async def cli_sso_start(request: Request): } _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) - return { - "login_id": login_id, - "poll_secret": poll_secret, - "user_code": user_code, - "expires_in": CLI_SSO_SESSION_TTL_SECONDS, - } + verification_uri_complete: str | None = ( + ( + get_custom_url( + request_base_url=str(request.base_url), route="sso/key/generate" + ) + + "?" + + urlencode( + { + "source": LITELLM_CLI_SOURCE_IDENTIFIER, + "key": login_id, + "user_code": user_code, + } + ) + ) + if _cli_sso_verification_uri_complete_enabled() + else None + ) + return _cli_sso_start_response_body( + login_id=login_id, + poll_secret=poll_secret, + user_code=user_code, + verification_uri_complete=verification_uri_complete, + ) @router.post( @@ -829,6 +898,7 @@ async def google_login( key: Optional[str] = None, existing_key: Optional[str] = None, return_to: Optional[str] = None, + user_code: str | None = None, ): """ Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env @@ -897,6 +967,7 @@ async def google_login( cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( source=source, key=key, + user_code=(user_code if _cli_sso_verification_uri_complete_enabled() else None), ) # check if user defined a custom auth sso sign in handler, if yes, use it @@ -1921,14 +1992,16 @@ async def auth_callback(request: Request, state: Optional[str] = None): ) if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): - # State format: {PREFIX}:{login_id} - state_parts = state.split(":", 1) + # State format: {PREFIX}:{login_id}[:{user_code}] + state_parts = state.split(":", 2) key_id = state_parts[1] if len(state_parts) > 1 else None + prefill_user_code = state_parts[2] if len(state_parts) > 2 else None verbose_proxy_logger.info("CLI SSO callback detected") return await cli_sso_callback( request=request, key=key_id, + prefill_user_code=prefill_user_code, result=result, received_response=received_response, ) @@ -2008,6 +2081,7 @@ async def _complete_cli_sso_callback_session( prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + prefill_user_code: str | None = None, ): from fastapi.responses import HTMLResponse @@ -2071,6 +2145,7 @@ async def _complete_cli_sso_callback_session( content=_render_cli_sso_verification_page( verify_url=verify_url, browser_complete_token=browser_complete_token, + prefill_user_code=prefill_user_code, ), status_code=200, ) @@ -2081,6 +2156,7 @@ async def cli_sso_callback( key: Optional[str] = None, result: Optional[Union[OpenID, dict]] = None, received_response: Optional[dict] = None, + prefill_user_code: str | None = None, ): """CLI SSO callback - stores session info for JWT generation on polling""" verbose_proxy_logger.info("CLI SSO callback") @@ -2137,6 +2213,7 @@ async def cli_sso_callback( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + prefill_user_code=prefill_user_code, ) except ProxyException: raise @@ -3053,21 +3130,27 @@ class SSOAuthenticationHandler: @staticmethod def _get_cli_state( - source: Optional[str], key: Optional[str], existing_key: Optional[str] = None + source: str | None, + key: str | None, + existing_key: str | None = None, + user_code: str | None = None, ) -> Optional[str]: """ Checks the request 'source' if a cli state token was passed in This is used to authenticate through the CLI login flow. - The state parameter format is: {PREFIX}:{login_id} + The state parameter format is: {PREFIX}:{login_id}[:{user_code}] - The state parameter is used to pass data through the OAuth flow without changing the callback URL + - user_code is appended only for the opt-in verification_uri_complete flow so the verify page can pre-fill it """ from litellm.constants import ( LITELLM_CLI_SESSION_TOKEN_PREFIX, ) if source == LITELLM_CLI_SOURCE_IDENTIFIER and key: + if _is_valid_cli_sso_user_code(user_code): + return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}:{user_code}" return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" else: return None diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 2efec3e0b34..acca357e641 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2120,6 +2120,292 @@ class TestCLIKeyRegenerationFlow: assert exc_info.value.status_code == 429 mock_cache.set_cache.assert_not_called() + @pytest.mark.asyncio + async def test_cli_sso_start_returns_verification_uri_complete_when_enabled(self): + """Test CLI SSO start returns a verification_uri_complete that round-trips the user_code only when the operator opts in""" + from urllib.parse import parse_qs, urlparse + + from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER + from litellm.proxy.management_endpoints.ui_sso import cli_sso_start + + mock_request = MagicMock(spec=Request) + mock_request.client = SimpleNamespace(host="127.0.0.1") + mock_request.headers = {} + mock_request.base_url = "https://proxy.example.com/" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 1 + + with ( + patch.dict( + os.environ, + {"PROXY_BASE_URL": "https://proxy.example.com", "SERVER_ROOT_PATH": ""}, + ), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_cli_sso_verification_uri_complete": True}, + ), + ): + result = await cli_sso_start(request=mock_request) + + verification_uri_complete = result["verification_uri_complete"] + parsed = urlparse(verification_uri_complete) + query = parse_qs(parsed.query) + + assert parsed.path.endswith("/sso/key/generate") + assert query["source"] == [LITELLM_CLI_SOURCE_IDENTIFIER] + assert query["key"] == [result["login_id"]] + assert query["user_code"] == [result["user_code"]] + + @pytest.mark.asyncio + async def test_cli_sso_start_omits_verification_uri_complete_by_default(self): + """Test CLI SSO start does NOT advertise verification_uri_complete unless the operator enables it (default off)""" + from litellm.proxy.management_endpoints.ui_sso import cli_sso_start + + mock_request = MagicMock(spec=Request) + mock_request.client = SimpleNamespace(host="127.0.0.1") + mock_request.headers = {} + mock_request.base_url = "https://proxy.example.com/" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 1 + + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + result = await cli_sso_start(request=mock_request) + + assert "verification_uri_complete" not in result + assert result["user_code"] + assert result["login_id"].startswith("cli-") + + def test_cli_sso_verification_uri_complete_enabled_reads_general_settings(self): + """Test the operator opt-in flag is read from general_settings and defaults off""" + from litellm.proxy.management_endpoints.ui_sso import ( + _cli_sso_verification_uri_complete_enabled, + ) + + with patch("litellm.proxy.proxy_server.general_settings", {}): + assert _cli_sso_verification_uri_complete_enabled() is False + with patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_cli_sso_verification_uri_complete": True}, + ): + assert _cli_sso_verification_uri_complete_enabled() is True + + @pytest.mark.asyncio + async def test_google_login_only_threads_user_code_when_enabled(self): + """Test google_login forwards user_code into the OAuth state only when the operator opt-in is on, dropping it otherwise""" + from litellm.proxy.management_endpoints.ui_sso import google_login + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.example.com/" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = {"poll_secret_hash": "h"} + + async def drive(enabled: bool): + with ( + patch.dict(os.environ, {}, clear=True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch( + "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", + None, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_cli_sso_verification_uri_complete": enabled}, + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.show_missing_vars_in_env", + return_value=None, + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.SSOAuthenticationHandler.get_redirect_url_for_sso", + return_value="https://proxy.example.com/sso/callback", + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.SSOAuthenticationHandler._get_cli_state", + return_value=None, + ) as mock_get_cli_state, + ): + try: + await google_login( + request=mock_request, + source="litellm-cli", + key="cli-validsessionkey123456", + user_code="WXYZ-2345", + ) + except Exception: + pass + return mock_get_cli_state.call_args.kwargs["user_code"] + + assert await drive(enabled=True) == "WXYZ-2345" + assert await drive(enabled=False) is None + + def test_get_cli_state_appends_user_code_for_prefill(self): + """Test the OAuth state carries the user_code only for the opt-in prefill flow""" + from litellm.constants import ( + LITELLM_CLI_SESSION_TOKEN_PREFIX, + LITELLM_CLI_SOURCE_IDENTIFIER, + ) + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + manual_state = SSOAuthenticationHandler._get_cli_state( + source=LITELLM_CLI_SOURCE_IDENTIFIER, key="cli-abc123" + ) + prefill_state = SSOAuthenticationHandler._get_cli_state( + source=LITELLM_CLI_SOURCE_IDENTIFIER, + key="cli-abc123", + user_code="WXYZ-2345", + ) + + assert manual_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123" + assert ( + prefill_state == f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123:WXYZ-2345" + ) + assert ( + SSOAuthenticationHandler._get_cli_state( + source="not-cli", key="cli-abc123", user_code="WXYZ-2345" + ) + is None + ) + + def test_get_cli_state_drops_malformed_user_code(self): + """Test a user_code that is not a server-issued code is dropped before reaching the size-limited OAuth state""" + from litellm.constants import ( + LITELLM_CLI_SESSION_TOKEN_PREFIX, + LITELLM_CLI_SOURCE_IDENTIFIER, + ) + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + manual_only = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-abc123" + for bad_user_code in ("A" * 4096, "not-a-code", "WXYZ2345", "WXYZ-234", ""): + assert ( + SSOAuthenticationHandler._get_cli_state( + source=LITELLM_CLI_SOURCE_IDENTIFIER, + key="cli-abc123", + user_code=bad_user_code, + ) + == manual_only + ) + + def test_is_valid_cli_sso_user_code_matches_generated_format(self): + """Test the user_code validator accepts a freshly generated code and rejects malformed input""" + from litellm.proxy.management_endpoints.ui_sso import ( + _generate_cli_sso_user_code, + _is_valid_cli_sso_user_code, + ) + + assert _is_valid_cli_sso_user_code(_generate_cli_sso_user_code()) + assert _is_valid_cli_sso_user_code("WXYZ-2345") + assert not _is_valid_cli_sso_user_code("WXYZ-2340") # 0 is not in the alphabet + assert not _is_valid_cli_sso_user_code("wxyz-2345") + assert not _is_valid_cli_sso_user_code("WXYZ2345") + assert not _is_valid_cli_sso_user_code("A" * 64) + assert not _is_valid_cli_sso_user_code(None) + + def test_cli_state_round_trips_user_code_to_callback_parser(self): + """Test the callback's state parser recovers login_id and user_code from the state _get_cli_state builds""" + from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + state = SSOAuthenticationHandler._get_cli_state( + source=LITELLM_CLI_SOURCE_IDENTIFIER, + key="cli-abc123", + user_code="WXYZ-2345", + ) + + state_parts = state.split(":", 2) + key_id = state_parts[1] if len(state_parts) > 1 else None + prefill_user_code = state_parts[2] if len(state_parts) > 2 else None + + assert key_id == "cli-abc123" + assert prefill_user_code == "WXYZ-2345" + + def test_render_cli_sso_verification_page_prefills_user_code(self): + """Test the verify page pre-fills the user_code input (HTML-escaped) when provided""" + from litellm.proxy.management_endpoints.ui_sso import ( + _render_cli_sso_verification_page, + ) + + html = _render_cli_sso_verification_page( + verify_url="https://proxy.example.com/sso/cli/complete/cli-abc123", + browser_complete_token="browser-token", + prefill_user_code='WXYZ-2345">