From 2d98a5c45ddfd2e3f96c751c2449389c8c67893e Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Mon, 29 Dec 2025 07:07:02 +0900 Subject: [PATCH] fix: lint --- .../proxy/hooks/proxy_track_cost_callback.py | 32 ++++--- litellm/proxy/proxy_server.py | 86 +++++++++---------- litellm/proxy/utils.py | 66 +++++++------- 3 files changed, 93 insertions(+), 91 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b6d661c454a..daf55a12590 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -111,11 +111,11 @@ class _ProxyDBLogger(CustomLogger): ) _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["status"] = "failure" - _metadata["error_information"] = ( - StandardLoggingPayloadSetup.get_error_information( - original_exception=original_exception, - traceback_str=traceback_str, - ) + _metadata[ + "error_information" + ] = StandardLoggingPayloadSetup.get_error_information( + original_exception=original_exception, + traceback_str=traceback_str, ) existing_metadata: dict = request_data.get("metadata", None) or {} @@ -123,25 +123,31 @@ class _ProxyDBLogger(CustomLogger): if "litellm_params" not in request_data: request_data["litellm_params"] = {} - + existing_litellm_params = request_data.get("litellm_params", {}) existing_litellm_metadata = existing_litellm_params.get("metadata", {}) or {} - + # Preserve tags from existing metadata if existing_litellm_metadata.get("tags"): existing_metadata["tags"] = existing_litellm_metadata.get("tags") - + request_data["litellm_params"]["proxy_server_request"] = ( - request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} + request_data.get("proxy_server_request") + or existing_litellm_params.get("proxy_server_request") + or {} ) request_data["litellm_params"]["metadata"] = existing_metadata - + # Preserve model name and custom_llm_provider if "model" not in request_data: - request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "") + request_data["model"] = existing_litellm_params.get( + "model" + ) or request_data.get("model", "") if "custom_llm_provider" not in request_data: - request_data["custom_llm_provider"] = existing_litellm_params.get("custom_llm_provider") or request_data.get("custom_llm_provider", "") - + request_data["custom_llm_provider"] = existing_litellm_params.get( + "custom_llm_provider" + ) or request_data.get("custom_llm_provider", "") + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, response_cost=0.0, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4d812498fe7..75fac7d37a9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -533,9 +533,9 @@ except ImportError: server_root_path = os.getenv("SERVER_ROOT_PATH", "") _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional["EnterpriseLicenseData"] = ( - _license_check.airgapped_license_data -) +premium_user_data: Optional[ + "EnterpriseLicenseData" +] = _license_check.airgapped_license_data global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -1083,9 +1083,7 @@ try: # In non-root Docker, we restructure in /var/lib/litellm/ui. try: _restructure_ui_html_files(ui_path) - verbose_proxy_logger.info( - f"Restructured UI directory: {ui_path}" - ) + verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}") except PermissionError as e: verbose_proxy_logger.exception( f"Permission error while restructuring UI directory {ui_path}: {e}" @@ -1171,9 +1169,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional["ClientSession"] = ( - None # Global shared session for connection reuse -) +shared_aiohttp_session: Optional[ + "ClientSession" +] = None # Global shared session for connection reuse user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1181,9 +1179,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[RedisCache] = ( - None # redis cache used for tracking spend, tpm/rpm limits -) +redis_usage_cache: Optional[ + RedisCache +] = None # redis cache used for tracking spend, tpm/rpm limits polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None @@ -1324,11 +1322,7 @@ def cost_tracking(): if prisma_client is None: return - from litellm.proxy.utils import ProxyUpdateSpend - - store_proxy_response = ( - ProxyUpdateSpend.should_store_proxy_response_in_spend_logs() - ) + store_proxy_response = ProxyUpdateSpend.should_store_proxy_response_in_spend_logs() disable_spend = ProxyUpdateSpend.disable_spend_updates() if store_proxy_response is None and not disable_spend: @@ -1340,9 +1334,7 @@ def cost_tracking(): proxy_db_logger = _ProxyDBLogger() litellm.logging_callback_manager.add_litellm_callback(proxy_db_logger) - litellm.logging_callback_manager.add_litellm_async_success_callback( - proxy_db_logger - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(proxy_db_logger) async def update_cache( # noqa: PLR0915 @@ -1539,9 +1531,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[LiteLLM_TeamTable] = ( - await user_api_key_cache.async_get_cache(key=_id) - ) + existing_spend_obj: Optional[ + LiteLLM_TeamTable + ] = await user_api_key_cache.async_get_cache(key=_id) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -1893,7 +1885,6 @@ class ProxyConfig: "environment_variables" in config_to_save and config_to_save["environment_variables"] ): - # decrypt the environment_variables - in case a caller function has already encrypted the environment_variables decrypted_env_vars = self._decrypt_and_set_db_env_variables( environment_variables=config_to_save["environment_variables"], @@ -2811,21 +2802,21 @@ class ProxyConfig: verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}") if _alerting_callbacks is None: return - + # Ensure proxy_logging_obj.alerting is set for all alerting types _alerting_value = general_settings.get("alerting", None) - verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}") + verbose_proxy_logger.debug( + f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}" + ) proxy_logging_obj.update_values( alerting=_alerting_value, alerting_threshold=general_settings.get("alerting_threshold", 600), alert_types=general_settings.get("alert_types", None), - alert_to_webhook_url=general_settings.get( - "alert_to_webhook_url", None - ), + alert_to_webhook_url=general_settings.get("alert_to_webhook_url", None), alerting_args=general_settings.get("alerting_args", None), redis_cache=redis_usage_cache, ) - + for _alert in _alerting_callbacks: if _alert == "slack": # [OLD] v0 implementation - already handled by update_values above @@ -3296,7 +3287,7 @@ class ProxyConfig: proxy_logging_obj: ProxyLogging """ _general_settings = config_data.get("general_settings", {}) - + if _general_settings is not None and "alerting" in _general_settings: if ( general_settings is not None @@ -3311,7 +3302,8 @@ class ProxyConfig: _merged_alerting = list(_yaml_alerting.union(_db_alerting)) # Preserve order: YAML values first, then DB values _merged_alerting = list(general_settings["alerting"]) + [ - item for item in _general_settings["alerting"] + item + for item in _general_settings["alerting"] if item not in general_settings["alerting"] ] verbose_proxy_logger.debug( @@ -3622,7 +3614,6 @@ class ProxyConfig: await self._init_vector_stores_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="vector_store_indexes"): - await self._init_vector_store_indexes_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="mcp"): @@ -3821,10 +3812,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[Guardrail] = ( - await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client - ) + guardrails_in_db: List[ + Guardrail + ] = await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -4151,9 +4142,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ["AZURE_API_VERSION"] = ( - api_version # set this for azure - litellm can read this from the env - ) + os.environ[ + "AZURE_API_VERSION" + ] = api_version # set this for azure - litellm can read this from the env if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -4671,10 +4662,14 @@ class ProxyStartupEvent: replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info("Responses cost check job scheduled successfully") + verbose_proxy_logger.info( + "Responses cost check job scheduled successfully" + ) except Exception as e: - verbose_proxy_logger.error(f"Failed to setup responses cost checking: {e}") + verbose_proxy_logger.error( + f"Failed to setup responses cost checking: {e}" + ) verbose_proxy_logger.debug( "Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..." ) @@ -5954,7 +5949,6 @@ async def realtime_websocket_endpoint( ), user_api_key_dict=Depends(user_api_key_auth_websocket), ): - await websocket.accept() # Only use explicit parameters, not all query params @@ -9551,9 +9545,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[idx].field_description = ( - sub_field_info.description - ) + nested_fields[ + idx + ].field_description = sub_field_info.description idx += 1 _stored_in_db = None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2cdbae4df6a..e33f0bac234 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -151,25 +151,25 @@ def _get_email_logger_class(): """ Determine which email logger class to use based on environment variables. Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback) - + Returns: The email logger class to use, or None if BaseEmailLogger is not available """ if BaseEmailLogger is None: return None - + # Check for SendGrid API key if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"): return SendGridEmailLogger - + # Check for Resend API key if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"): return ResendEmailLogger - + # Check for SMTP configuration if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"): return SMTPEmailLogger - + # Fallback to BaseEmailLogger (though it won't actually send emails) return BaseEmailLogger @@ -452,7 +452,6 @@ class ProxyLogging: litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # type: ignore for callback in litellm.callbacks: if isinstance(callback, str): - callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore cast(_custom_logger_compatible_callbacks_literal, callback), internal_usage_cache=self.internal_usage_cache.dual_cache, @@ -1038,7 +1037,6 @@ class ProxyLogging: data.pop("prompt_id", None) if custom_logger and prompt_spec is not None: - ( model, messages, @@ -1288,7 +1286,6 @@ class ProxyLogging: call_type=call_type, ) else: - guardrail_task = callback.async_moderation_hook( data=data, user_api_key_dict=user_api_key_auth_dict, # type: ignore @@ -1337,7 +1334,7 @@ class ProxyLogging: if self.alerting is None: # do nothing if alerting is not switched on return - + if "slack" in self.alerting: await self.slack_alerting_instance.budget_alerts( type=type, @@ -1548,7 +1545,10 @@ class ProxyLogging: traceback_str=traceback_str, ) # If callback returned an HTTPException, use it (first one wins) - if isinstance(hook_result, HTTPException) and transformed_exception is None: + if ( + isinstance(hook_result, HTTPException) + and transformed_exception is None + ): transformed_exception = hook_result except HTTPException as e: # If callback raised an HTTPException, use it (first one wins) @@ -1849,7 +1849,6 @@ class ProxyLogging: current_response = response for callback in litellm.callbacks: - _callback: Optional[CustomLogger] = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -3568,11 +3567,13 @@ class ProxyUpdateSpend: ) # Atomically read and remove logs to process (protected by lock) async with prisma_client._spend_log_transactions_lock: - logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] + logs_to_process = prisma_client.spend_log_transactions[ + :MAX_LOGS_PER_INTERVAL + ] # Remove the logs we're about to process - prisma_client.spend_log_transactions = ( - prisma_client.spend_log_transactions[len(logs_to_process):] - ) + prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ + len(logs_to_process) : + ] start_time = time.time() try: for i in range(n_retry_times + 1): @@ -3685,9 +3686,7 @@ async def update_spend( # noqa: PLR0915 # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - verbose_proxy_logger.debug( - "Spend Logs transactions: {}".format(queue_size) - ) + verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size)) # Process spend log transactions when called directly. # This keeps backwards compatibility with the old behavior. @@ -3709,19 +3708,19 @@ async def update_spend_logs_job( ): """ Job to process spend_log_transactions queue. - + This job is triggered based on queue size rather than time. Processes spend log transactions when the queue reaches a threshold. """ n_retry_times = 3 - + # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - + if queue_size == 0: return - + await ProxyUpdateSpend.update_spend_logs( n_retry_times=n_retry_times, prisma_client=prisma_client, @@ -3738,7 +3737,7 @@ async def _monitor_spend_logs_queue( """ Background task that monitors the spend_log_transactions queue size and triggers processing when the threshold is reached. - + Args: prisma_client: Prisma client instance db_writer_client: Optional HTTP handler for external spend logs endpoint @@ -3748,23 +3747,23 @@ async def _monitor_spend_logs_queue( SPEND_LOG_QUEUE_POLL_INTERVAL, SPEND_LOG_QUEUE_SIZE_THRESHOLD, ) - + threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL max_backoff = 30.0 # Maximum backoff interval in seconds backoff_multiplier = 1.5 # Exponential backoff multiplier current_interval = base_interval - + verbose_proxy_logger.info( f"Starting spend logs queue monitor (threshold: {threshold}, poll_interval: {base_interval}s)" ) - + while True: try: # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - + if queue_size > 0: if queue_size >= threshold: verbose_proxy_logger.debug( @@ -3777,8 +3776,10 @@ async def _monitor_spend_logs_queue( f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff" ) # Exponential backoff when below threshold but still processing - current_interval = min(current_interval * backoff_multiplier, max_backoff) - + current_interval = min( + current_interval * backoff_multiplier, max_backoff + ) + await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, @@ -3786,8 +3787,10 @@ async def _monitor_spend_logs_queue( ) else: # Exponential backoff when no logs to process - current_interval = min(current_interval * backoff_multiplier, max_backoff) - + current_interval = min( + current_interval * backoff_multiplier, max_backoff + ) + await asyncio.sleep(current_interval) except Exception as e: verbose_proxy_logger.error( @@ -3798,7 +3801,6 @@ async def _monitor_spend_logs_queue( await asyncio.sleep(current_interval) - def _raise_failed_update_spend_exception( e: Exception, start_time: float, proxy_logging_obj: ProxyLogging ):