fix: 限额改 v4 冻结快照，余额红线防扣费
v2 的 stable_limit = usage + balance - BUFFER 依赖余额 API 随消耗实时递减， 实测滞后（08-12 快照 1,558,307 在 usage 涨 45 万期间不动）导致池膨胀、限额 虚高，放行到扣费区（当日超余额 13.8 万已扣费）。 扣费口径（用户确认）：是否扣费只取决于当下免费额度，超出即扣费，采集量只 决定明日发放。v4：stable_limit = balance - BUFFER 冻结快照，本地 usage 实时 撞线即拦，等价于本地实时扣减余额；effective = min(2M 配额, stable_limit)。 新增 docs/volc-quota-notes.md 额度机制速查；merge-auto-consume.md 补 §4.5。
This commit is contained in:
vibe-code
@@ -0,0 +1,382 @@
|
||||
import os
|
||||
import time
|
||||
import hashlib
|
||||
from urllib.parse import urlparse # <--- 新增导入
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.exceptions import ServiceUnavailableError, ContextWindowExceededError
|
||||
import redis
|
||||
import traceback
|
||||
|
||||
# ==========================================
|
||||
# ⚠️ 入口说明
|
||||
# ==========================================
|
||||
# 当前 config.yaml 的 callbacks 使用的是 volc_limit_sync.proxy_handler_instance，
|
||||
# 不是本文件。除非你把 callbacks 切回 daily_limit.py，否则这里的 MODEL_QUOTAS
|
||||
# 修改不会影响当前 Docker 容器的监控与限额行为。
|
||||
|
||||
# ==========================================
|
||||
# 🛠️ 调试配置
|
||||
# ==========================================
|
||||
# 日志级别配置
|
||||
# 可选值: "DEBUG", "INFO", "WARNING", "ERROR"
|
||||
# 修改此变量以控制日志详细程度，或使用环境变量 DAILY_LIMIT_LOG_LEVEL
|
||||
LOG_LEVEL = os.getenv("DAILY_LIMIT_LOG_LEVEL", "INFO").upper()
|
||||
|
||||
|
||||
def get_log_level_value(level_name):
|
||||
levels = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40}
|
||||
return levels.get(level_name, 20)
|
||||
|
||||
|
||||
CURRENT_LEVEL_VAL = get_log_level_value(LOG_LEVEL)
|
||||
DEBUG_MODE = CURRENT_LEVEL_VAL <= 10
|
||||
|
||||
|
||||
def log(msg, level="INFO"):
|
||||
msg_val = get_log_level_value(level)
|
||||
if msg_val >= CURRENT_LEVEL_VAL:
|
||||
icon = {"DEBUG": "🔍", "INFO": "ℹ️", "WARNING": "⚠️", "ERROR": "❌"}.get(level, "📝")
|
||||
print(f"{icon} [DailyLimit] {msg}", flush=True)
|
||||
|
||||
|
||||
def dprint(msg):
|
||||
log(msg, level="INFO")
|
||||
|
||||
|
||||
dprint("Loading daily_limit.py module v20 (legacy entry; current config uses volc_limit_sync.py)...")
|
||||
|
||||
# ==========================================
|
||||
# Redis 与 限额配置
|
||||
# ==========================================
|
||||
REDIS_HOST = os.getenv("REDIS_HOST", "litellm_redis")
|
||||
REDIS_PORT = os.getenv("REDIS_PORT", "6379")
|
||||
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", None)
|
||||
|
||||
MODEL_QUOTAS = {
|
||||
"deepseek-v3-2-251201:ba0383d3": 1600000,
|
||||
"doubao-seed-1-8-251228:ba0383d3": 1800000,
|
||||
"glm-4-7-251222:ba0383d3": 1800000,
|
||||
"doubao-seed-2-0-pro-260215:ba0383d3": 1900000,
|
||||
"doubao-seed-2-0-lite-260215:ba0383d3": 1900000,
|
||||
}
|
||||
|
||||
|
||||
try:
|
||||
redis_host = os.getenv("REDIS_HOST", "redis")
|
||||
redis_port = int(os.getenv("REDIS_PORT", 6379))
|
||||
redis_client = redis.Redis(host=redis_host, port=redis_port, db=0, decode_responses=True)
|
||||
log(f"✅ Redis connected successfully to {redis_host}:{redis_port}", level="INFO")
|
||||
except Exception as e:
|
||||
log(f"❌ Redis connection failed: {str(e)}", level="ERROR")
|
||||
redis_client = None
|
||||
|
||||
|
||||
def clean_model_name(name):
|
||||
if not name:
|
||||
return ""
|
||||
return name.replace("openai/", "").replace("azure/", "").replace("anthropic/", "")
|
||||
|
||||
|
||||
def generate_api_base_hash(api_base):
|
||||
"""
|
||||
生成 api_base 的 Hash，强制忽略 URL 的路径部分。
|
||||
例如:
|
||||
'https://api.com/v1/' -> 'https://api.com'
|
||||
'https://api.com/v3' -> 'https://api.com'
|
||||
两者生成的 Hash 将完全一致。
|
||||
"""
|
||||
if not api_base:
|
||||
return "default"
|
||||
|
||||
try:
|
||||
# 1. 预处理：去空格
|
||||
s = api_base.strip()
|
||||
|
||||
# 2. 解析 URL
|
||||
parsed = urlparse(s)
|
||||
|
||||
# 3. 提取 Scheme 和 Netloc (域名+端口)
|
||||
# 如果解析成功且包含 scheme 和 netloc，则重组
|
||||
if parsed.scheme and parsed.netloc:
|
||||
clean_base = f"{parsed.scheme}://{parsed.netloc}".lower()
|
||||
else:
|
||||
# 兜底：如果不是标准 URL，或者解析失败，尝试简单的去斜杠和小写化
|
||||
# 这种情况比较少见，通常 litellm 会校验 URL 格式
|
||||
clean_base = s.rstrip('/').lower()
|
||||
|
||||
except Exception:
|
||||
# 极端情况下的兜底
|
||||
clean_base = api_base.strip().rstrip('/').lower()
|
||||
|
||||
# 4. 计算 Hash
|
||||
h = hashlib.md5(clean_base.encode()).hexdigest()[:8]
|
||||
|
||||
if DEBUG_MODE:
|
||||
# Debugging hash mismatch
|
||||
# 打印原始值和清洗后的值，方便确认
|
||||
print(f"🔍 [Hash Debug] Base='{api_base}' -> Clean='{clean_base}' -> Hash='{h}'", flush=True)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
# ==========================================
|
||||
# ⏰ 时间逻辑：北京时间 9:00 (UTC 01:00) 刷新
|
||||
# ==========================================
|
||||
def get_redis_key(clean_name, api_base=None):
|
||||
now = time.time()
|
||||
# UTC 01:00 视为新的一天，所以减去 1小时 (3600秒)
|
||||
adjusted_time = now - 3600
|
||||
date_str = time.strftime("%Y-%m-%d", time.gmtime(adjusted_time))
|
||||
|
||||
suffix = ""
|
||||
if api_base:
|
||||
suffix = f":{generate_api_base_hash(api_base)}"
|
||||
|
||||
return f"daily_usage:{date_str}:{clean_name}{suffix}"
|
||||
|
||||
|
||||
# ==========================================
|
||||
# 🎯 限额查找逻辑 (支持精确匹配)
|
||||
# ==========================================
|
||||
def get_daily_limit(clean_model, api_base):
|
||||
# 1. 尝试精确匹配 (Model + Hash)
|
||||
# 用户可以在 MODEL_QUOTAS 中配置 "model_name:hash": limit
|
||||
current_hash = generate_api_base_hash(api_base)
|
||||
specific_key = f"{clean_model}:{current_hash}"
|
||||
|
||||
limit = MODEL_QUOTAS.get(specific_key)
|
||||
if limit:
|
||||
return limit
|
||||
|
||||
# 2. 尝试精确模型名匹配
|
||||
limit = MODEL_QUOTAS.get(clean_model)
|
||||
if limit:
|
||||
return limit
|
||||
|
||||
# 3. 尝试模糊模型名匹配 (向下兼容)
|
||||
for quota_model, limit_val in MODEL_QUOTAS.items():
|
||||
# 排除包含冒号的特定规则，防止误匹配
|
||||
if ":" not in quota_model and quota_model in clean_model:
|
||||
return limit_val
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ==========================================
|
||||
# 👻 核心逻辑：软更名踢出 & 自动恢复
|
||||
# ==========================================
|
||||
def soft_evict_from_router(target_upstream_model, target_api_base=None):
|
||||
"""
|
||||
不删除节点，而是修改节点的 model_name，让 Router 找不到它。
|
||||
同时备份原始信息以便恢复。
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
if not llm_router or not llm_router.model_list:
|
||||
return
|
||||
|
||||
evicted = False
|
||||
target_upstream_clean = clean_model_name(target_upstream_model)
|
||||
|
||||
# Iterate over the actual list in memory
|
||||
for i, model_entry in enumerate(llm_router.model_list):
|
||||
params = model_entry.get("litellm_params", {})
|
||||
upstream_model = params.get("model", "")
|
||||
current_api_base = params.get("api_base", "")
|
||||
|
||||
# 匹配上游模型名称
|
||||
if target_upstream_clean not in clean_model_name(upstream_model):
|
||||
continue
|
||||
|
||||
# Strict api_base matching if provided
|
||||
# ⚠️ 注意：这里比较的是原始字符串，为了确保精准匹配当前运行的实例
|
||||
if target_api_base and current_api_base != target_api_base:
|
||||
continue
|
||||
|
||||
current_alias = model_entry.get("model_name")
|
||||
|
||||
# 如果它还没有被标记为 blocked
|
||||
if current_alias and "blocked" not in current_alias:
|
||||
log(
|
||||
f"👻 [ROUTER HIDE] Renaming alias '{current_alias}' -> '{current_alias}-blocked' for {target_upstream_model}",
|
||||
level="INFO",
|
||||
)
|
||||
|
||||
# Backup original state
|
||||
model_entry["_daily_limit_backup"] = {"model_name": current_alias, "api_base": current_api_base}
|
||||
|
||||
# 1. 修改别名 - Modify IN PLACE
|
||||
llm_router.model_list[i]["model_name"] = f"{current_alias}-blocked"
|
||||
|
||||
# 2. 破坏连接
|
||||
llm_router.model_list[i]["litellm_params"]["api_base"] = "http://0.0.0.0"
|
||||
|
||||
evicted = True
|
||||
|
||||
if evicted:
|
||||
# Re-initialize router to reflect changes immediately if possible
|
||||
# But changing model_list in place is usually enough for next lookup
|
||||
dprint("✅ Node effectively removed from routing pool.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Router eviction failed: {e}", flush=True)
|
||||
|
||||
|
||||
last_restore_check = 0
|
||||
RESTORE_CHECK_INTERVAL = 60 # Check every 60 seconds
|
||||
|
||||
|
||||
def restore_blocked_models():
|
||||
"""
|
||||
检查是否有 blocked 的模型配额已重置，如果是则恢复。
|
||||
"""
|
||||
global last_restore_check
|
||||
now = time.time()
|
||||
if now - last_restore_check < RESTORE_CHECK_INTERVAL:
|
||||
return
|
||||
|
||||
last_restore_check = now
|
||||
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
if not llm_router or not llm_router.model_list:
|
||||
return
|
||||
|
||||
restored_count = 0
|
||||
for i, model_entry in enumerate(llm_router.model_list):
|
||||
backup = model_entry.get("_daily_limit_backup")
|
||||
if not backup:
|
||||
continue
|
||||
|
||||
# 获取必要信息检查配额
|
||||
params = model_entry.get("litellm_params", {})
|
||||
upstream_model = params.get("model", "")
|
||||
clean_model = clean_model_name(upstream_model)
|
||||
original_api_base = backup.get("api_base")
|
||||
|
||||
# Check quota
|
||||
daily_limit = get_daily_limit(clean_model, original_api_base)
|
||||
|
||||
should_restore = False
|
||||
if not daily_limit:
|
||||
should_restore = True
|
||||
else:
|
||||
key = get_redis_key(clean_model, original_api_base)
|
||||
current_usage = redis_client.get(key)
|
||||
current_val = float(current_usage) if current_usage else 0.0
|
||||
|
||||
if current_val < float(daily_limit):
|
||||
should_restore = True
|
||||
|
||||
if should_restore:
|
||||
dprint(f"♻️ [RESTORE] Restoring {backup['model_name']} (Quota OK)")
|
||||
|
||||
# Restore IN PLACE
|
||||
llm_router.model_list[i]["model_name"] = backup["model_name"]
|
||||
llm_router.model_list[i]["litellm_params"]["api_base"] = original_api_base
|
||||
del llm_router.model_list[i]["_daily_limit_backup"]
|
||||
restored_count += 1
|
||||
|
||||
if restored_count > 0:
|
||||
dprint(f"✅ Restored {restored_count} models to routing pool.")
|
||||
|
||||
except Exception as e:
|
||||
log(f"❌ Restore check failed: {e}", level="ERROR")
|
||||
|
||||
|
||||
# ==========================================
|
||||
# Input Callback
|
||||
# ==========================================
|
||||
def quota_interceptor(kwargs, user_api_key_dict=None):
|
||||
"""
|
||||
Standard LiteLLM input callback signature:
|
||||
def input_callback(kwargs, user_api_key_dict):
|
||||
"""
|
||||
try:
|
||||
restore_blocked_models()
|
||||
# check_and_hide(kwargs, source="input_callback") # Removed to avoid rigid blocking
|
||||
except (ServiceUnavailableError, ContextWindowExceededError):
|
||||
raise
|
||||
except Exception:
|
||||
# Don't let other errors block the request
|
||||
pass
|
||||
|
||||
|
||||
# ==========================================
|
||||
# Handler & Logging
|
||||
# ==========================================
|
||||
class DailyLimitHandler(CustomLogger):
|
||||
def __init__(self):
|
||||
# Register the callback properly
|
||||
if quota_interceptor not in litellm.input_callback:
|
||||
litellm.input_callback.append(quota_interceptor)
|
||||
dprint("✅ quota_interceptor registered")
|
||||
if self not in litellm.callbacks:
|
||||
litellm.callbacks.append(self)
|
||||
dprint("✅ DailyLimitHandler registered")
|
||||
|
||||
# 统计逻辑
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self._record_usage(kwargs, response_obj)
|
||||
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self._record_usage(kwargs, response_obj)
|
||||
|
||||
def _record_usage(self, kwargs, response_obj):
|
||||
if not redis_client:
|
||||
return
|
||||
try:
|
||||
litellm_params = kwargs.get("litellm_params") or {}
|
||||
raw_model = litellm_params.get("model") or kwargs.get("model")
|
||||
clean_model = clean_model_name(raw_model)
|
||||
api_base = litellm_params.get("api_base")
|
||||
|
||||
usage = None
|
||||
if hasattr(response_obj, "usage"):
|
||||
usage = response_obj.usage
|
||||
elif isinstance(response_obj, dict):
|
||||
usage = response_obj.get("usage")
|
||||
if hasattr(usage, "model_dump"):
|
||||
usage = usage.model_dump()
|
||||
elif hasattr(usage, "dict"):
|
||||
usage = usage.dict()
|
||||
|
||||
if not usage:
|
||||
return
|
||||
|
||||
total_tokens = usage.get("total_tokens", 0)
|
||||
if clean_model and total_tokens > 0:
|
||||
key = get_redis_key(clean_model, api_base)
|
||||
current_total = redis_client.incrby(key, int(total_tokens))
|
||||
redis_client.expire(key, 172800)
|
||||
|
||||
# Get Limit for display
|
||||
daily_limit = get_daily_limit(clean_model, api_base)
|
||||
display_limit = f"/{int(daily_limit):,}" if daily_limit else ""
|
||||
|
||||
# Show usage if DEBUG level is enabled
|
||||
key_display = key.split(":", 2)[-1]
|
||||
log(
|
||||
f"📊 [Usage] {key_display}: +{total_tokens} (Today: {current_total:,}{display_limit})",
|
||||
level="INFO",
|
||||
)
|
||||
|
||||
# Enforcement Logic
|
||||
if daily_limit and current_total >= float(daily_limit):
|
||||
log(f"🚫 [BLOCK] {clean_model} Limit: {current_total}/{daily_limit}", level="WARNING")
|
||||
soft_evict_from_router(raw_model, api_base)
|
||||
else:
|
||||
if DEBUG_MODE:
|
||||
if not clean_model:
|
||||
log("⚠️ [Usage Skip] No clean model name found", level="DEBUG")
|
||||
if total_tokens <= 0:
|
||||
log(f"⚠️ [Usage Skip] Total tokens is {total_tokens}", level="DEBUG")
|
||||
|
||||
except Exception as e:
|
||||
log(f"❌ _record_usage failed: {e}\n{traceback.format_exc()}", level="ERROR")
|
||||
|
||||
|
||||
proxy_handler_instance = DailyLimitHandler()
|
||||
@@ -0,0 +1,34 @@
|
||||
import json
|
||||
import re
|
||||
|
||||
transcript_path = r"C:\Users\Kimber Xu\.gemini\antigravity-ide\brain\ef230fd3-39a1-47f8-aa16-4969fb505243\.system_generated\logs\transcript.jsonl"
|
||||
file_path = r"d:\PyCharm_projects\litellm-config\volc_limit_sync.py"
|
||||
|
||||
lines_dict = {}
|
||||
|
||||
with open(transcript_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
try:
|
||||
step = json.loads(line)
|
||||
if step.get("type") == "SYSTEM_RESPONSE":
|
||||
content = step.get("content", "")
|
||||
if "File Path: `file:///d:/PyCharm_projects/litellm-config/volc_limit_sync.py`" in content:
|
||||
# Extract lines
|
||||
for match in re.finditer(r"^(\d+):\s(.*)$", content, re.MULTILINE):
|
||||
line_num = int(match.group(1))
|
||||
line_text = match.group(2)
|
||||
lines_dict[line_num] = line_text
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
if lines_dict:
|
||||
max_line = max(lines_dict.keys())
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
for i in range(1, max_line + 1):
|
||||
if i in lines_dict:
|
||||
f.write(lines_dict[i] + "\n")
|
||||
else:
|
||||
f.write("\n")
|
||||
print(f"Recovered {max_line} lines to {file_path}")
|
||||
else:
|
||||
print("Failed to find lines in transcript.")
|
||||
Reference in New Issue
Block a user