diff --git a/volc_limit_sync.py b/volc_limit_sync.py index 265ecb0..f9dbebc 100644 --- a/volc_limit_sync.py +++ b/volc_limit_sync.py @@ -17,7 +17,11 @@ from urllib.parse import urlparse try: import litellm from litellm.integrations.custom_logger import CustomLogger - from litellm.exceptions import ServiceUnavailableError, ContextWindowExceededError + from litellm.exceptions import ( + ServiceUnavailableError, + ContextWindowExceededError, + RateLimitError, + ) LITELLM_AVAILABLE = True except ImportError: @@ -33,6 +37,9 @@ except ImportError: class ContextWindowExceededError(Exception): pass + class RateLimitError(Exception): + pass + # ========================================== # 🛠️ Configuration & Constants @@ -626,6 +633,55 @@ def startup_audit(): log(f"❌ Startup audit failed: {e}", level="ERROR") +def _check_request_block(model_name): + """ + 请求前闸门:对外模型名 → 判定是否超限。 + 只拦截 MONITORED_MODEL_PAIRS 内的火山监控条目;newapi 等非监控渠道天然放行。 + 返回 (clean_model, usage, limit) 表示超限,None 表示放行。任何异常 fail-open。 + """ + if not redis_client or not LITELLM_AVAILABLE: + return None + try: + from litellm.proxy.proxy_server import llm_router + + if not llm_router or not llm_router.model_list: + return None + + for entry in llm_router.model_list: + if entry.get("model_name") != model_name: + continue + # 已被限额踢出的条目跳过(路由层已匹配不到) + if "_daily_limit_backup" in entry: + continue + + params = entry.get("litellm_params", {}) + raw_model = params.get("model", "") + api_base = params.get("api_base", "") + + clean = clean_model_name(raw_model) + h = generate_api_base_hash(api_base) + + # 只拦已纳入监控的火山模型(模型 + 端点粒度) + if MONITORED_MODEL_PAIRS and (clean, h) not in MONITORED_MODEL_PAIRS: + continue + + effective_limit = get_effective_limit(clean, api_base) + if effective_limit is None: + continue + + usage_key = get_usage_key(clean, api_base) + val = redis_client.get(usage_key) + usage = float(val or 0) + + if usage >= float(effective_limit): + return (clean, usage, effective_limit) + + return None + except Exception as e: + log(f"❌ Pre-block check failed: {e}", level="ERROR") + return None + + def quota_interceptor(kwargs, user_api_key_dict=None): try: # 1. 尝试恢复(定期) @@ -646,6 +702,32 @@ class DailyLimitHandler(CustomLogger): litellm.callbacks.append(self) dprint("✅ DailyLimitHandler registered") + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + """ + 请求前限额闸门:在 route_request 之前拒绝超限请求(429), + 不进入 router → 不触发 num_retries / fallback 放大。 + """ + if call_type not in ("completion", "acompletion", "text_completion"): + return None + if not data or not data.get("model"): + return None + + blocked = _check_request_block(data["model"]) + if not blocked: + return None + + clean, usage, limit = blocked + log( + f"🚫 [PRE-BLOCK] {data['model']} daily limit reached ({int(usage):,}/{int(limit):,})", + level="WARNING", + ) + raise RateLimitError( + message=f"Daily limit reached for model '{data['model']}' ({int(usage):,}/{int(limit):,})", + model=data["model"], + llm_provider="volcengine", + request_data=data, + ) + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self._record_usage(kwargs, response_obj)