diff --git a/AGENTS.md b/AGENTS.md index d830802..32395a6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,11 +4,14 @@ ## 1. 项目上下文 此仓库包含用于 **LiteLLM** 的自定义回调脚本,核心功能是限制模型用量(特别是火山引擎/VolcEngine)。 -- **核心逻辑**:通过 `quota_interceptor` 拦截请求,并通过 `DailyLimitHandler` 记录用量到 Redis。 +- **核心逻辑**:通过 `quota_interceptor` 拦截请求,并通过 `DailyLimitHandler` 记录用量到 Redis。后台 `active_monitor_loop` 守护线程定期从火山引擎 API 同步配额到 Redis。 +- **架构**:所有逻辑统一在 `volc_limit_sync.py` 单文件中。LiteLLM 通过 `config.yaml` 中 `callbacks: ["volc_limit_sync.proxy_handler_instance"]` 加载回调,模块加载时自动启动后台同步线程。 - **关键文件**: - - `volc_limit_sync.py`: 主同步逻辑,从火山引擎 API 拉取配额并更新 Redis。 - - `daily_limit.py`: 轻量级回调,用于 LiteLLM 运行时检查。 - - `config.yaml`: 包含模型定义的 LiteLLM 配置文件。 + - `volc_limit_sync.py`: **唯一核心文件**,包含同步守护线程(`active_monitor_loop`)、LiteLLM 回调(`DailyLimitHandler` + `quota_interceptor`)、限额计算、模型屏蔽/恢复等全部逻辑。 + - `config.yaml`: LiteLLM 配置文件,包含模型定义和回调配置。 + - `docker-compose.yml`: 容器编排,仅挂载 `config.yaml` 和 `volc_limit_sync.py`。 + - `debug_volc_sync.py` / `debug_volc_sync_manual.py`: 调试脚本,用于手动验证 API 和 Redis 状态。 + - `daily_limit.py.deprecated`: 已废弃,功能已全部合并到 `volc_limit_sync.py`,仅保留作历史参考。 ## 2. 环境与运行 (Build & Run) @@ -40,7 +43,7 @@ pip install litellm redis requests ### 3.1 基础规范 - **Python 版本**:兼容 Python 3.8+。 - **命名规范**: - - 变量/函数:`snake_case` (e.g., `clean_model_name`, `get_daily_limit`) + - 变量/函数:`snake_case` (e.g., `clean_model_name`, `get_effective_limit`) - 常量:`UPPER_CASE` (e.g., `LOG_LEVEL`, `REDIS_HOST`) - 类名:`PascalCase` (e.g., `DailyLimitHandler`) - **类型提示 (Typing)**: diff --git a/daily_limit.py b/daily_limit.py.deprecated similarity index 100% rename from daily_limit.py rename to daily_limit.py.deprecated diff --git a/docker-compose.yml b/docker-compose.yml index 965126e..90e0893 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -8,9 +8,7 @@ services: volumes: - ./config.yaml:/app/config.yaml # 【关键】挂载自定义脚本到容器根目录 - - ./daily_limit.py:/app/daily_limit.py - ./volc_limit_sync.py:/app/volc_limit_sync.py - - ./debug_volc_sync_manual.py:/app/debug_volc_sync_manual.py command: - "--config=/app/config.yaml" ports: diff --git a/volc_limit_sync.py b/volc_limit_sync.py index 08c9506..3b64189 100644 --- a/volc_limit_sync.py +++ b/volc_limit_sync.py @@ -1025,14 +1025,36 @@ def active_monitor_loop(): beijing_time_ts = now + (8 * 3600) today_str = time.strftime("%Y-%m-%d", time.gmtime(beijing_time_ts)) - # Idle Check - usage_key_global = f"daily_usage:{today_str}:global_total" - current_usage_global = int(redis_client.get(usage_key_global) or 0) - last_usage_global = sum(last_usage_map.values()) if last_usage_map else 0 - delta = current_usage_global - last_usage_global + # Idle Check - Only exit idle if a "low quota" model has usage + LOW_QUOTA_THRESHOLD = 800000 # 判定为“低额度”的阈值 + model_usages = {} + any_low_quota_usage = False + + for m in MONITORED_MODELS: + m_usage = 0 + target_hashes = [p[1] for p in MONITORED_MODEL_PAIRS if p[0] == m] + for h in target_hashes: + val = redis_client.get(f"daily_usage:{today_str}:{m}:{h}") + if val: m_usage += float(val) + val_no_hash = redis_client.get(f"daily_usage:{today_str}:{m}") + if val_no_hash: m_usage += float(val_no_hash) + model_usages[m] = m_usage + + # Calculate delta for this specific model + m_delta = m_usage - last_usage_map.get(m, 0.0) + + # Check if this model is considered "low quota" + rem_key = f"limit:dynamic:volcengine:remaining:{m}" + rem_val = redis_client.get(rem_key) + # 未知额度视为充足,避免首次启动时所有模型被误判为"低额度" + remaining = float(rem_val) if rem_val is not None else float(LOW_QUOTA_THRESHOLD) + + if m_delta > 0 and remaining < LOW_QUOTA_THRESHOLD: + any_low_quota_usage = True + time_diff = now - last_sync_time - if delta <= 0: + if not any_low_quota_usage: if idle_start_time == 0: idle_start_time = now else: @@ -1085,10 +1107,8 @@ def active_monitor_loop(): burst_reason = "" for model in MONITORED_MODELS: - # 1. Get Current Usage - curr_usage_key = f"limit:dynamic:volcengine:usage:{model}" - curr_usage_val = redis_client.get(curr_usage_key) - curr_usage = float(curr_usage_val) if curr_usage_val else 0.0 + # 1. Get Current Usage (calculated above) + curr_usage = model_usages.get(model, 0.0) # 2. Get Last Synced Usage (Snapshot) last_usage = last_usage_map.get(model, 0.0) @@ -1325,15 +1345,8 @@ def active_monitor_loop(): global_min_remaining = temp_min_remaining # Update snapshots for next iteration delta calculation - current_usage_global_snapshot = 0 - for m in MONITORED_MODELS: - uk = f"limit:dynamic:volcengine:usage:{m}" - uv = redis_client.get(uk) - u_val = float(uv) if uv else 0.0 + for m, u_val in model_usages.items(): last_usage_map[m] = u_val - current_usage_global_snapshot += u_val - - last_usage_map["global"] = current_usage_global_snapshot last_sync_time = now if global_min_remaining > 800000: