初始化空白分支文件
This commit is contained in:
@@ -0,0 +1,23 @@
|
|||||||
|
# 如果你的服务器在中国,可能需要设置代理
|
||||||
|
# HTTPS_PROXY=http://192.168.1.x:7890
|
||||||
|
# HTTP_PROXY=http://192.168.1.x:7890
|
||||||
|
|
||||||
|
LITELLM_LOG=INFO
|
||||||
|
#LITELLM_LOG=DEBUG
|
||||||
|
|
||||||
|
DAILY_LIMIT_LOG_LEVEL=DEBUG
|
||||||
|
|
||||||
|
# Master Key 配置
|
||||||
|
LITELLM_MASTER_KEY=sk-7v9m3CvOR2r9rDr3yhTg4WiHjWPfTeH6p5qwEIjwE0lpfOKK
|
||||||
|
|
||||||
|
# 你的 NewAPI 配置
|
||||||
|
NEWAPI_BASE_URL=http://192.168.1.4:23002/v1
|
||||||
|
NEWAPI_API_KEY=sk-xFRYNClzkWXIzC9tJPbOI5LMk66uZfqrR6FYH9SWVw1dKKCN
|
||||||
|
|
||||||
|
# 利用 LiteLLM 的默认回退机制,直接把 NewAPI 伪装成标准的 OpenAI 接口
|
||||||
|
OPENAI_API_BASE=http://192.168.1.4:23002/v1
|
||||||
|
OPENAI_API_KEY=sk-xFRYNClzkWXIzC9tJPbOI5LMk66uZfqrR6FYH9SWVw1dKKCN
|
||||||
|
|
||||||
|
#火山平台抖音
|
||||||
|
HUOSHAN_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
|
||||||
|
HUOSHAN_API_KEY=aacf11dd-2cdd-4d9a-a819-3e3a0310847b
|
||||||
+304
@@ -0,0 +1,304 @@
|
|||||||
|
##yaml锚点方便配置。手动设定价格,方便ui中查看
|
||||||
|
x-definitions: &token_pricing
|
||||||
|
input_cost_per_token: 1
|
||||||
|
output_cost_per_token: 1
|
||||||
|
|
||||||
|
model_list:
|
||||||
|
|
||||||
|
#=======================火山平台,限时免费
|
||||||
|
|
||||||
|
#日常使用ds v3.2 ,强制think
|
||||||
|
- model_name: deepseek-v3.2
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/deepseek-v3-2-251201
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
# 说明中建议设置30分钟(1800秒)以上,避免深度思考过程中断连接
|
||||||
|
timeout: 1800
|
||||||
|
# 关键配置2:启用思考模式
|
||||||
|
# extra_body 中的参数会被直接合并到请求的 JSON body 中
|
||||||
|
extra_body:
|
||||||
|
thinking:
|
||||||
|
type: "enabled"
|
||||||
|
|
||||||
|
- model_name: deepseek-v3.2-nothink
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/deepseek-v3-2-251201
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
extra_body:
|
||||||
|
thinking:
|
||||||
|
type: "disabled"
|
||||||
|
|
||||||
|
|
||||||
|
- model_name: doubao-seed-1-8
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/doubao-seed-1-8-251228
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
|
||||||
|
- model_name: doubao-seed-1-8-nothink
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/doubao-seed-1-8-251228
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
extra_body:
|
||||||
|
reasoning_effort: "minimal"
|
||||||
|
|
||||||
|
- model_name: glm-4.7
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/glm-4-7-251222
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
|
||||||
|
- model_name: glm-4.7-nothink
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/glm-4-7-251222
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
extra_body:
|
||||||
|
thinking:
|
||||||
|
type: "disabled"
|
||||||
|
|
||||||
|
#=======================fanyi
|
||||||
|
|
||||||
|
- model_name: test-newapi-dbao
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/doubao-seed-1-8-251228
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: test-ggllmm
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/glm-4-7-251222
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
extra_body:
|
||||||
|
thinking:
|
||||||
|
type: "disabled"
|
||||||
|
|
||||||
|
- model_name: test-dbao
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/doubao-seed-1-8-251228
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
extra_body:
|
||||||
|
reasoning_effort: "minimal"
|
||||||
|
|
||||||
|
- model_name: fanyi
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/doubao-seed-1-8-251228
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
extra_body:
|
||||||
|
reasoning_effort: "minimal"
|
||||||
|
|
||||||
|
|
||||||
|
- model_name: fanyi
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/glm-4-7-251222
|
||||||
|
api_base: os.environ/HUOSHAN_BASE_URL
|
||||||
|
api_key: os.environ/HUOSHAN_API_KEY
|
||||||
|
extra_body:
|
||||||
|
thinking:
|
||||||
|
type: "disabled"
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
#=======================Claude系列
|
||||||
|
|
||||||
|
#=======================Claude-Opus-v4.5系列
|
||||||
|
- model_name: Claude-Opus-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-opus-4-5-thinking
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Opus-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-opus-4-5
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Opus-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-4.5-opus
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Opus-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-opus-4-5-20251101
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Opus-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/cursor2-claude-4.1-opus
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
#=======================Claude-Sonnet-v4.5系列
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-sonnet-4-5-20250929
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-sonnet-4-5
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-4.5-sonnet-official-reverse
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-4.5-sonnet
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-4.5-sonnet-docs
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-4.5-sonnet-2cc
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/claude-sonnet-4-5-thinking
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: Claude-Sonnet-v4.5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/cursor2-claude-4.5-sonnet
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
#=======================gpt-5系列
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/cursor2-gpt-5
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/cursor2-gpt-5-codex
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/gpt-5.1
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/gpt-5.1-codex
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/gpt-5.1-codex-2cx
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/gpt-5.1-codex-max-2cx
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/gpt-5.2
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
- model_name: gpt-5
|
||||||
|
litellm_params:
|
||||||
|
<<: *token_pricing
|
||||||
|
model: openai/gpt-5.2-2cx
|
||||||
|
api_base: os.environ/NEWAPI_BASE_URL
|
||||||
|
api_key: os.environ/NEWAPI_API_KEY
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
litellm_settings:
|
||||||
|
num_retries: 2 # retry call 3 times on each model_name (e.g. zephyr-beta)
|
||||||
|
request_timeout: 120 # raise Timeout error if call takes longer than 10s. Sets litellm.request_timeout
|
||||||
|
#fallbacks: [{"zephyr-beta": ["gpt-4o"]}] # fallback to gpt-4o if call fails num_retries
|
||||||
|
#context_window_fallbacks: [{"zephyr-beta": ["gpt-3.5-turbo-16k"]}, {"gpt-4o": ["gpt-3.5-turbo-16k"]}] # fallback to gpt-3.5-turbo-16k if context window error
|
||||||
|
cache: False # 开启缓存
|
||||||
|
# cache_params:
|
||||||
|
# type: "redis" # 指定使用 Redis
|
||||||
|
# supported_call_types: ["completion", "embedding", "acompletion"] # 支持的类型
|
||||||
|
allowed_fails: 3 # cooldown model if it fails > 1 call in a minute.
|
||||||
|
# 关键配置:注册自定义回调,实现单上游模型每日token限制的关键
|
||||||
|
callbacks: ["daily_limit.proxy_handler_instance"]
|
||||||
|
|
||||||
|
router_settings: # router_settings are optional
|
||||||
|
routing_strategy: simple-shuffle # Literal["simple-shuffle", "least-busy", "usage-based-routing","latency-based-routing"], default="simple-shuffle"
|
||||||
|
#model_group_alias: {"gpt-4": "gpt-4o"} # all requests with `gpt-4` will be routed to models with `gpt-4o`
|
||||||
|
num_retries: 3
|
||||||
|
timeout: 300 # 30 seconds
|
||||||
|
#cooldown_time: 60
|
||||||
|
|
||||||
|
general_settings:
|
||||||
|
store_model_in_db: true
|
||||||
|
store_prompts_in_spend_logs: true # <--- 这是最关键的一行
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
+348
@@ -0,0 +1,348 @@
|
|||||||
|
import os
|
||||||
|
import time
|
||||||
|
import hashlib
|
||||||
|
import litellm
|
||||||
|
from litellm.integrations.custom_logger import CustomLogger
|
||||||
|
from litellm.exceptions import ServiceUnavailableError, ContextWindowExceededError
|
||||||
|
import redis
|
||||||
|
import traceback
|
||||||
|
|
||||||
|
# ==========================================
|
||||||
|
# 🛠️ 调试配置
|
||||||
|
# ==========================================
|
||||||
|
# 日志级别配置
|
||||||
|
# 可选值: "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 v19 (Specific Quotas & Docker Track)...")
|
||||||
|
|
||||||
|
# ==========================================
|
||||||
|
# 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:38430740": 1500000,
|
||||||
|
"doubao-seed-1-8-251228:38430740": 14000, # 测试限额
|
||||||
|
"glm-4-7-251222:38430740": 1500000,
|
||||||
|
"doubao-seed-1-8-251228:9fbb42b8": 1500000, # 测试重名
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
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):
|
||||||
|
if not api_base:
|
||||||
|
return "default"
|
||||||
|
|
||||||
|
# === 新增标准化代码 ===
|
||||||
|
# 1. 去除首尾空格 (strip)
|
||||||
|
# 2. 去除末尾的斜杠 (rstrip('/'))
|
||||||
|
# 这样 '.../v3/' 和 '.../v3' 会生成相同的 Hash
|
||||||
|
normalized_base = api_base.strip().rstrip('/')
|
||||||
|
# ====================
|
||||||
|
|
||||||
|
h = hashlib.md5(api_base.encode()).hexdigest()[:8]
|
||||||
|
if DEBUG_MODE:
|
||||||
|
# Debugging hash mismatch
|
||||||
|
print(f"🔍 [Hash Debug] Base='{api_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="DEBUG",
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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,77 @@
|
|||||||
|
services:
|
||||||
|
litellm:
|
||||||
|
image: ghcr.io/berriai/litellm:main-latest
|
||||||
|
container_name: litellm
|
||||||
|
restart: unless-stopped
|
||||||
|
volumes:
|
||||||
|
- ./config.yaml:/app/config.yaml
|
||||||
|
# 【关键】挂载自定义脚本到容器根目录
|
||||||
|
- ./daily_limit.py:/app/daily_limit.py
|
||||||
|
command:
|
||||||
|
- "--config=/app/config.yaml"
|
||||||
|
ports:
|
||||||
|
- "23003:4000"
|
||||||
|
environment:
|
||||||
|
# ⚠️ 必须设置这个环境变量,否则 Python 找不到你的脚本
|
||||||
|
PYTHONPATH: "/app"
|
||||||
|
DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
|
||||||
|
STORE_MODEL_IN_DB: "True"
|
||||||
|
# --- Modified: 添加 Redis 配置 ---
|
||||||
|
REDIS_HOST: "redis" # 对应下方 redis 服务的名字
|
||||||
|
REDIS_PORT: "6379"
|
||||||
|
# -------------------------------
|
||||||
|
env_file:
|
||||||
|
- .env
|
||||||
|
depends_on:
|
||||||
|
- db
|
||||||
|
- redis # --- New: 依赖 redis 启动 ---
|
||||||
|
healthcheck:
|
||||||
|
test:
|
||||||
|
- CMD-SHELL
|
||||||
|
- python3 -c "import urllib.request; urllib.request.urlopen('http://localhost:4000/health/liveliness')"
|
||||||
|
interval: 300s
|
||||||
|
timeout: 10s
|
||||||
|
retries: 3
|
||||||
|
start_period: 40s
|
||||||
|
|
||||||
|
# --- New: 添加 Redis 服务 ---
|
||||||
|
redis:
|
||||||
|
image: redis:alpine # 使用 alpine 版本,体积非常小
|
||||||
|
container_name: litellm_redis
|
||||||
|
restart: unless-stopped
|
||||||
|
# --maxmemory 512mb : 限制 Redis 最多使用 512MB 内存
|
||||||
|
# --maxmemory-policy allkeys-lru : 内存满了之后,自动删除最久未使用的缓存(无论是否过期)
|
||||||
|
# --save "" : 如果你不需要持久化缓存到硬盘(重启后清空),可以把 save 设为空,减少磁盘IO
|
||||||
|
# command: redis-server --save 60 1 --maxmemory 128mb --maxmemory-policy allkeys-lru --loglevel warning
|
||||||
|
command: redis-server --maxmemory 32mb --maxmemory-policy allkeys-lru --save 300 1 --loglevel warning
|
||||||
|
# 配置详解:
|
||||||
|
# --maxmemory 32mb: 严格限制内存使用 32MB
|
||||||
|
# --maxmemory-policy allkeys-lru: 内存满时挤掉旧数据(安全兜底)
|
||||||
|
# --save 300 1: 每300秒(5分钟)如果至少有1个key变化,就保存快照到磁盘 (降低IO)
|
||||||
|
volumes:
|
||||||
|
# 映射数据目录,确保重启容器后,今天的计数还在
|
||||||
|
- redis_data:/data
|
||||||
|
# ---------------------------
|
||||||
|
|
||||||
|
db:
|
||||||
|
image: postgres:16
|
||||||
|
restart: unless-stopped
|
||||||
|
container_name: litellm_db
|
||||||
|
environment:
|
||||||
|
POSTGRES_DB: litellm
|
||||||
|
POSTGRES_USER: llmproxy
|
||||||
|
POSTGRES_PASSWORD: dbpassword9090
|
||||||
|
volumes:
|
||||||
|
- postgres_data:/var/lib/postgresql/data
|
||||||
|
healthcheck:
|
||||||
|
test: ["CMD-SHELL", "pg_isready -d litellm -U llmproxy"]
|
||||||
|
interval: 5s
|
||||||
|
timeout: 5s
|
||||||
|
retries: 10
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
postgres_data:
|
||||||
|
driver: local
|
||||||
|
name: litellm_postgres_data
|
||||||
|
redis_data: # 如果开启了 redis 持久化,这里也要打开
|
||||||
|
driver: local
|
||||||
@@ -0,0 +1,153 @@
|
|||||||
|
import os
|
||||||
|
import yaml
|
||||||
|
import hashlib
|
||||||
|
import re
|
||||||
|
from typing import Dict, Any
|
||||||
|
|
||||||
|
# ==========================================
|
||||||
|
# 🛠️ 配置
|
||||||
|
# ==========================================
|
||||||
|
CONFIG_FILE = "config.yaml"
|
||||||
|
ENV_FILE = ".env"
|
||||||
|
|
||||||
|
|
||||||
|
def load_env(env_path: str) -> Dict[str, str]:
|
||||||
|
"""读取 .env 文件并返回字典"""
|
||||||
|
env_vars = {}
|
||||||
|
if not os.path.exists(env_path):
|
||||||
|
print(f"⚠️ Warning: {env_path} not found. Env var substitution might fail.")
|
||||||
|
return env_vars
|
||||||
|
|
||||||
|
with open(env_path, "r", encoding="utf-8") as f:
|
||||||
|
for line in f:
|
||||||
|
line = line.strip()
|
||||||
|
if not line or line.startswith("#"):
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
key, value = line.split("=", 1)
|
||||||
|
env_vars[key.strip()] = value.strip().strip('"').strip("'")
|
||||||
|
except ValueError:
|
||||||
|
continue
|
||||||
|
return env_vars
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_env_var(value: str, env_vars: Dict[str, str]) -> str:
|
||||||
|
"""
|
||||||
|
处理环境变量替换,支持以下两种格式:
|
||||||
|
1. ${VAR_NAME} - 标准格式
|
||||||
|
2. os.environ/VAR_NAME - LiteLLM 特有格式
|
||||||
|
"""
|
||||||
|
if not isinstance(value, str):
|
||||||
|
return value
|
||||||
|
|
||||||
|
# 1. 处理 os.environ/VAR_NAME 格式
|
||||||
|
if value.startswith("os.environ/"):
|
||||||
|
var_name = value.replace("os.environ/", "")
|
||||||
|
return env_vars.get(var_name, os.getenv(var_name, value))
|
||||||
|
|
||||||
|
# 2. 处理 ${VAR_NAME} 或 $VAR_NAME 格式
|
||||||
|
pattern = re.compile(r"\$\{?([A-Z0-9_]+)\}?")
|
||||||
|
|
||||||
|
def replace_match(match):
|
||||||
|
var_name = match.group(1)
|
||||||
|
# 优先从 .env 读取,其次从系统环境变量读取
|
||||||
|
return env_vars.get(var_name, os.getenv(var_name, match.group(0)))
|
||||||
|
|
||||||
|
return pattern.sub(replace_match, value)
|
||||||
|
|
||||||
|
|
||||||
|
def clean_model_name(name):
|
||||||
|
if not name:
|
||||||
|
return ""
|
||||||
|
return name.replace("openai/", "").replace("azure/", "").replace("anthropic/", "")
|
||||||
|
|
||||||
|
|
||||||
|
def generate_api_base_hash(api_base):
|
||||||
|
if not api_base:
|
||||||
|
return "default"
|
||||||
|
# Logic matches daily_limit.py v18
|
||||||
|
|
||||||
|
# === 新增标准化代码 ===
|
||||||
|
# 1. 去除首尾空格 (strip)
|
||||||
|
# 2. 去除末尾的斜杠 (rstrip('/'))
|
||||||
|
# 这样 '.../v3/' 和 '.../v3' 会生成相同的 Hash
|
||||||
|
normalized_base = api_base.strip().rstrip('/')
|
||||||
|
# ====================
|
||||||
|
|
||||||
|
|
||||||
|
return hashlib.md5(api_base.encode()).hexdigest()[:8]
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
print("🔍 Generating Daily Limit Keys from config.yaml...\n")
|
||||||
|
|
||||||
|
# 1. Load Env
|
||||||
|
env_vars = load_env(ENV_FILE)
|
||||||
|
|
||||||
|
# 2. Load Config
|
||||||
|
if not os.path.exists(CONFIG_FILE):
|
||||||
|
print(f"❌ Error: {CONFIG_FILE} not found!")
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
with open(CONFIG_FILE, "r", encoding="utf-8") as f:
|
||||||
|
config = yaml.safe_load(f)
|
||||||
|
except Exception as e:
|
||||||
|
print(f"❌ Error parsing yaml: {e}")
|
||||||
|
return
|
||||||
|
|
||||||
|
model_list = config.get("model_list", [])
|
||||||
|
if not model_list:
|
||||||
|
print("⚠️ No models found in config.yaml")
|
||||||
|
return
|
||||||
|
|
||||||
|
results = []
|
||||||
|
seen_keys = set()
|
||||||
|
|
||||||
|
# Table Header
|
||||||
|
# Model Alias | Upstream Model | API Base (Env/URL) | Generated Key
|
||||||
|
header = f"{'Model Name (Alias)':<25} | {'Upstream Model':<30} | {'API Base (Env/URL)':<40} | {'Generated Key'}"
|
||||||
|
print(header)
|
||||||
|
print("-" * len(header))
|
||||||
|
|
||||||
|
for model in model_list:
|
||||||
|
litellm_params = model.get("litellm_params", {})
|
||||||
|
|
||||||
|
# Get raw values
|
||||||
|
model_alias = model.get("model_name", "unknown")
|
||||||
|
upstream_model = litellm_params.get("model", "")
|
||||||
|
api_base_raw = litellm_params.get("api_base", "")
|
||||||
|
|
||||||
|
# Resolve env vars for hash generation
|
||||||
|
api_base_resolved = resolve_env_var(api_base_raw, env_vars)
|
||||||
|
|
||||||
|
# Generate components
|
||||||
|
clean_upstream = clean_model_name(upstream_model)
|
||||||
|
base_hash = generate_api_base_hash(api_base_resolved)
|
||||||
|
|
||||||
|
# Generate final key
|
||||||
|
final_key = f"{clean_upstream}:{base_hash}"
|
||||||
|
|
||||||
|
# Deduplicate
|
||||||
|
if final_key in seen_keys:
|
||||||
|
continue
|
||||||
|
seen_keys.add(final_key)
|
||||||
|
|
||||||
|
# Truncate api_base_raw if it's too long for display
|
||||||
|
display_api_base = api_base_raw
|
||||||
|
if len(display_api_base) > 37:
|
||||||
|
display_api_base = display_api_base[:34] + "..."
|
||||||
|
|
||||||
|
print(f"{model_alias:<25} | {clean_upstream:<30} | {display_api_base:<40} | {final_key}")
|
||||||
|
results.append(f'"{final_key}": 1500000')
|
||||||
|
|
||||||
|
print("\n📋 [Copy & Paste into daily_limit.py MODEL_QUOTAS]:")
|
||||||
|
print("-" * 50)
|
||||||
|
print("MODEL_QUOTAS = {")
|
||||||
|
print(",\n".join(f" {item}" for item in results))
|
||||||
|
print("}")
|
||||||
|
print("-" * 50)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user