mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
Initial commit
This commit is contained in:
3
src/services/__init__.py
Normal file
3
src/services/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
服务层模块
|
||||
"""
|
||||
13
src/services/auth/__init__.py
Normal file
13
src/services/auth/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
认证服务模块
|
||||
|
||||
包含认证服务、JWT 黑名单等功能。
|
||||
"""
|
||||
|
||||
from src.services.auth.jwt_blacklist import JWTBlacklistService
|
||||
from src.services.auth.service import AuthService
|
||||
|
||||
__all__ = [
|
||||
"AuthService",
|
||||
"JWTBlacklistService",
|
||||
]
|
||||
191
src/services/auth/jwt_blacklist.py
Normal file
191
src/services/auth/jwt_blacklist.py
Normal file
@@ -0,0 +1,191 @@
|
||||
"""
|
||||
JWT Token 黑名单服务
|
||||
|
||||
使用 Redis 存储被撤销的 JWT Token,防止已登出或被撤销的 Token 继续使用
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
# 安全策略配置:当 Redis 不可用时的行为
|
||||
# True = fail-closed(安全优先,拒绝访问)
|
||||
# False = fail-open(可用性优先,允许访问)
|
||||
BLACKLIST_FAIL_CLOSED = os.getenv("JWT_BLACKLIST_FAIL_CLOSED", "true").lower() == "true"
|
||||
|
||||
|
||||
class JWTBlacklistService:
|
||||
"""JWT Token 黑名单服务"""
|
||||
|
||||
# Redis key 前缀
|
||||
BLACKLIST_PREFIX = "jwt:blacklist:"
|
||||
|
||||
@staticmethod
|
||||
def _get_token_hash(token: str) -> str:
|
||||
"""
|
||||
获取 Token 的哈希值(用于 Redis key)
|
||||
|
||||
使用 SHA256 哈希避免直接存储完整 Token
|
||||
"""
|
||||
return hashlib.sha256(token.encode()).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
async def add_to_blacklist(token: str, exp_timestamp: int, reason: str = "logout") -> bool:
|
||||
"""
|
||||
将 Token 添加到黑名单
|
||||
|
||||
Args:
|
||||
token: JWT token 字符串
|
||||
exp_timestamp: Token 的过期时间戳(Unix timestamp)
|
||||
reason: 添加到黑名单的原因(logout, revoked, security)
|
||||
|
||||
Returns:
|
||||
是否成功添加到黑名单
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning("Redis 不可用,无法将 Token 添加到黑名单(降级模式)")
|
||||
return False
|
||||
|
||||
try:
|
||||
token_hash = JWTBlacklistService._get_token_hash(token)
|
||||
redis_key = f"{JWTBlacklistService.BLACKLIST_PREFIX}{token_hash}"
|
||||
|
||||
# 计算 TTL(Token 过期前的剩余时间)
|
||||
now = datetime.now(timezone.utc).timestamp()
|
||||
ttl_seconds = max(int(exp_timestamp - now), 0)
|
||||
|
||||
if ttl_seconds <= 0:
|
||||
# Token 已经过期,不需要加入黑名单
|
||||
logger.debug(f"Token 已过期,无需加入黑名单: {token[:10]}...")
|
||||
return True
|
||||
|
||||
# 存储到 Redis,设置 TTL 为 Token 过期时间
|
||||
# 值存储为原因字符串
|
||||
await redis_client.setex(redis_key, ttl_seconds, reason)
|
||||
|
||||
logger.info(f"Token 已加入黑名单: {token[:10]}... (原因: {reason}, TTL: {ttl_seconds}s)")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"添加 Token 到黑名单失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def is_blacklisted(token: str) -> bool:
|
||||
"""
|
||||
检查 Token 是否在黑名单中
|
||||
|
||||
Args:
|
||||
token: JWT token 字符串
|
||||
|
||||
Returns:
|
||||
Token 是否在黑名单中
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
# Redis 不可用时,根据安全策略决定行为
|
||||
if BLACKLIST_FAIL_CLOSED:
|
||||
logger.warning("Redis 不可用,采用 fail-closed 策略拒绝访问(可通过 JWT_BLACKLIST_FAIL_CLOSED=false 改变)")
|
||||
return True # 返回 True 表示在黑名单中,拒绝访问
|
||||
else:
|
||||
logger.debug("Redis 不可用,采用 fail-open 策略允许访问")
|
||||
return False
|
||||
|
||||
try:
|
||||
token_hash = JWTBlacklistService._get_token_hash(token)
|
||||
redis_key = f"{JWTBlacklistService.BLACKLIST_PREFIX}{token_hash}"
|
||||
|
||||
# 检查 key 是否存在
|
||||
exists = await redis_client.exists(redis_key)
|
||||
|
||||
if exists:
|
||||
# 获取黑名单原因(可选)
|
||||
reason = await redis_client.get(redis_key)
|
||||
logger.warning(f"检测到黑名单 Token: {token[:10]}... (原因: {reason})")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"检查 Token 黑名单状态失败: {e}")
|
||||
# 发生错误时,根据安全策略决定行为
|
||||
if BLACKLIST_FAIL_CLOSED:
|
||||
logger.warning("黑名单检查失败,采用 fail-closed 策略拒绝访问")
|
||||
return True # 安全优先,拒绝访问
|
||||
else:
|
||||
logger.warning("黑名单检查失败,采用 fail-open 策略允许访问")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def remove_from_blacklist(token: str) -> bool:
|
||||
"""
|
||||
从黑名单中移除 Token(用于测试或特殊情况)
|
||||
|
||||
Args:
|
||||
token: JWT token 字符串
|
||||
|
||||
Returns:
|
||||
是否成功移除
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning("Redis 不可用,无法从黑名单中移除 Token")
|
||||
return False
|
||||
|
||||
try:
|
||||
token_hash = JWTBlacklistService._get_token_hash(token)
|
||||
redis_key = f"{JWTBlacklistService.BLACKLIST_PREFIX}{token_hash}"
|
||||
|
||||
deleted = await redis_client.delete(redis_key)
|
||||
|
||||
if deleted:
|
||||
logger.info(f"Token 已从黑名单移除: {token[:10]}...")
|
||||
else:
|
||||
logger.debug(f"Token 不在黑名单中: {token[:10]}...")
|
||||
|
||||
return bool(deleted)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"从黑名单移除 Token 失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def get_blacklist_stats() -> dict:
|
||||
"""
|
||||
获取黑名单统计信息
|
||||
|
||||
Returns:
|
||||
包含统计信息的字典
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
return {"available": False, "total_blacklisted": 0, "error": "Redis 不可用"}
|
||||
|
||||
try:
|
||||
# 扫描黑名单 key
|
||||
pattern = f"{JWTBlacklistService.BLACKLIST_PREFIX}*"
|
||||
cursor = 0
|
||||
total = 0
|
||||
|
||||
while True:
|
||||
cursor, keys = await redis_client.scan(cursor=cursor, match=pattern, count=100)
|
||||
total += len(keys)
|
||||
|
||||
if cursor == 0:
|
||||
break
|
||||
|
||||
return {"available": True, "total_blacklisted": total}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取黑名单统计失败: {e}")
|
||||
return {"available": False, "total_blacklisted": 0, "error": str(e)}
|
||||
282
src/services/auth/service.py
Normal file
282
src/services/auth/service.py
Normal file
@@ -0,0 +1,282 @@
|
||||
"""
|
||||
认证服务
|
||||
"""
|
||||
|
||||
import os
|
||||
import secrets
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import jwt
|
||||
from fastapi import HTTPException, status
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
from src.config import config
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, User, UserRole
|
||||
from src.services.auth.jwt_blacklist import JWTBlacklistService
|
||||
from src.services.cache.user_cache import UserCacheService
|
||||
from src.services.user.apikey import ApiKeyService
|
||||
|
||||
|
||||
# JWT配置从config读取
|
||||
if not config.jwt_secret_key:
|
||||
# 如果没有配置,生成一个随机密钥并警告
|
||||
if config.environment == "production":
|
||||
raise ValueError("JWT_SECRET_KEY must be set in production environment!")
|
||||
config.jwt_secret_key = secrets.token_urlsafe(32)
|
||||
logger.warning(f"JWT_SECRET_KEY未在环境变量中找到,已生成随机密钥用于开发: {config.jwt_secret_key[:10]}...")
|
||||
logger.warning("生产环境请设置JWT_SECRET_KEY环境变量!")
|
||||
|
||||
JWT_SECRET_KEY = config.jwt_secret_key
|
||||
JWT_ALGORITHM = config.jwt_algorithm
|
||||
JWT_EXPIRATION_HOURS = config.jwt_expiration_hours
|
||||
# Refresh token 有效期设为7天
|
||||
REFRESH_TOKEN_EXPIRATION_DAYS = 7
|
||||
|
||||
|
||||
class AuthService:
|
||||
"""认证服务"""
|
||||
|
||||
@staticmethod
|
||||
def create_access_token(data: dict) -> str:
|
||||
"""创建JWT访问令牌"""
|
||||
to_encode = data.copy()
|
||||
expire = datetime.now(timezone.utc) + timedelta(hours=JWT_EXPIRATION_HOURS)
|
||||
to_encode.update({"exp": expire, "type": "access"})
|
||||
encoded_jwt = jwt.encode(to_encode, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
|
||||
return encoded_jwt
|
||||
|
||||
@staticmethod
|
||||
def create_refresh_token(data: dict) -> str:
|
||||
"""创建JWT刷新令牌"""
|
||||
to_encode = data.copy()
|
||||
expire = datetime.now(timezone.utc) + timedelta(days=REFRESH_TOKEN_EXPIRATION_DAYS)
|
||||
to_encode.update({"exp": expire, "type": "refresh"})
|
||||
encoded_jwt = jwt.encode(to_encode, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
|
||||
return encoded_jwt
|
||||
|
||||
@staticmethod
|
||||
async def verify_token(token: str, token_type: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""验证JWT令牌
|
||||
|
||||
Args:
|
||||
token: JWT token字符串
|
||||
token_type: 期望的token类型 ('access' 或 'refresh'),None表示不验证类型
|
||||
"""
|
||||
try:
|
||||
payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
|
||||
|
||||
# 验证token类型(如果指定)
|
||||
if token_type:
|
||||
actual_type = payload.get("type")
|
||||
if actual_type != token_type:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=f"Token类型错误: 期望 {token_type}, 实际 {actual_type}",
|
||||
)
|
||||
|
||||
# 检查 Token 是否在黑名单中
|
||||
is_blacklisted = await JWTBlacklistService.is_blacklisted(token)
|
||||
if is_blacklisted:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Token已被撤销"
|
||||
)
|
||||
|
||||
return payload
|
||||
except jwt.ExpiredSignatureError:
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Token已过期")
|
||||
except jwt.InvalidTokenError:
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的Token")
|
||||
|
||||
@staticmethod
|
||||
async def authenticate_user(db: Session, email: str, password: str) -> Optional[User]:
|
||||
"""用户登录认证"""
|
||||
# 使用缓存查询用户
|
||||
user = await UserCacheService.get_user_by_email(db, email)
|
||||
|
||||
if not user:
|
||||
logger.warning(f"登录失败 - 用户不存在: {email}")
|
||||
return None
|
||||
|
||||
if not user.verify_password(password):
|
||||
logger.warning(f"登录失败 - 密码错误: {email}")
|
||||
return None
|
||||
|
||||
if not user.is_active:
|
||||
logger.warning(f"登录失败 - 用户已禁用: {email}")
|
||||
return None
|
||||
|
||||
# 更新最后登录时间
|
||||
# 需要重新从数据库获取以便更新(缓存的对象是分离的)
|
||||
db_user = db.query(User).filter(User.id == user.id).first()
|
||||
if db_user:
|
||||
db_user.last_login_at = datetime.now(timezone.utc)
|
||||
db.commit() # 立即提交事务,释放数据库锁
|
||||
# 清除缓存,因为用户信息已更新
|
||||
await UserCacheService.invalidate_user_cache(user.id, user.email)
|
||||
|
||||
logger.info(f"用户登录成功: {email} (ID: {user.id})")
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
def authenticate_api_key(db: Session, api_key: str) -> Optional[tuple[User, ApiKey]]:
|
||||
"""API密钥认证"""
|
||||
# 对API密钥进行哈希查找,预加载 user 关系以支持后续访问限制检查
|
||||
key_hash = ApiKey.hash_key(api_key)
|
||||
key_record = (
|
||||
db.query(ApiKey)
|
||||
.options(joinedload(ApiKey.user))
|
||||
.filter(ApiKey.key_hash == key_hash)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not key_record:
|
||||
# 只记录认证失败事件,不记录任何 key 信息以防止信息泄露
|
||||
logger.warning("API认证失败 - 密钥不存在或无效")
|
||||
return None
|
||||
|
||||
if not key_record.is_active:
|
||||
logger.warning("API认证失败 - 密钥已禁用")
|
||||
return None
|
||||
|
||||
# 检查过期时间
|
||||
if key_record.expires_at:
|
||||
# 确保 expires_at 是 aware datetime
|
||||
expires_at = key_record.expires_at
|
||||
if expires_at.tzinfo is None:
|
||||
# 如果没有时区信息,假定为 UTC
|
||||
expires_at = expires_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
if expires_at < datetime.now(timezone.utc):
|
||||
logger.warning("API认证失败 - 密钥已过期")
|
||||
return None
|
||||
|
||||
# 检查余额限制(仅独立Key)
|
||||
is_balance_ok, remaining = ApiKeyService.check_balance(key_record)
|
||||
if not is_balance_ok:
|
||||
# 获取剩余余额用于日志
|
||||
remaining_balance = ApiKeyService.get_remaining_balance(key_record)
|
||||
logger.warning(f"API认证失败 - 余额不足 "
|
||||
f"(已用: ${key_record.balance_used_usd:.4f}, 剩余: ${remaining_balance:.4f})")
|
||||
return None
|
||||
|
||||
# 获取用户
|
||||
user = key_record.user
|
||||
if not user.is_active:
|
||||
logger.warning(f"API认证失败 - 用户已禁用: {user.email}")
|
||||
return None
|
||||
|
||||
# 更新最后使用时间
|
||||
key_record.last_used_at = datetime.now(timezone.utc)
|
||||
db.commit() # 立即提交事务,释放数据库锁,避免阻塞后续请求
|
||||
|
||||
logger.debug(f"API认证成功: 用户 {user.email} (Key: {api_key[:10]}...)")
|
||||
return user, key_record
|
||||
|
||||
@staticmethod
|
||||
def check_user_quota(user: User, estimated_cost: float = 0) -> bool:
|
||||
"""检查用户配额"""
|
||||
if user.role == UserRole.ADMIN:
|
||||
return True # 管理员无限制
|
||||
|
||||
# NULL 表示无限制
|
||||
if user.quota_usd is None:
|
||||
return True
|
||||
|
||||
# 检查美元配额
|
||||
if user.used_usd + estimated_cost > user.quota_usd:
|
||||
logger.warning(f"用户配额不足: {user.email} (已用: ${user.used_usd:.2f}, 配额: ${user.quota_usd:.2f})")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def check_permission(user: User, required_role: UserRole = UserRole.USER) -> bool:
|
||||
"""检查用户权限"""
|
||||
if user.role == UserRole.ADMIN:
|
||||
return True
|
||||
|
||||
if user.role.value >= required_role.value:
|
||||
return True
|
||||
|
||||
logger.warning(f"权限不足: 用户 {user.email} 角色 {user.role.value} < 需要 {required_role.value}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def logout(token: str) -> bool:
|
||||
"""
|
||||
用户登出,将 Token 加入黑名单
|
||||
|
||||
Args:
|
||||
token: JWT token字符串
|
||||
|
||||
Returns:
|
||||
是否成功登出
|
||||
"""
|
||||
try:
|
||||
# 解码 Token 获取过期时间(不验证黑名单)
|
||||
payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
|
||||
exp_timestamp = payload.get("exp")
|
||||
|
||||
if not exp_timestamp:
|
||||
logger.warning("Token 缺少过期时间,无法加入黑名单")
|
||||
return False
|
||||
|
||||
# 将 Token 加入黑名单
|
||||
success = await JWTBlacklistService.add_to_blacklist(
|
||||
token=token, exp_timestamp=exp_timestamp, reason="logout"
|
||||
)
|
||||
|
||||
if success:
|
||||
user_id = payload.get("sub")
|
||||
logger.info(f"用户登出成功: user_id={user_id}")
|
||||
|
||||
return success
|
||||
|
||||
except jwt.InvalidTokenError as e:
|
||||
logger.warning(f"登出失败 - 无效的 Token: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"登出失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def revoke_token(token: str, reason: str = "revoked") -> bool:
|
||||
"""
|
||||
撤销 Token(管理员操作)
|
||||
|
||||
Args:
|
||||
token: JWT token字符串
|
||||
reason: 撤销原因
|
||||
|
||||
Returns:
|
||||
是否成功撤销
|
||||
"""
|
||||
try:
|
||||
# 解码 Token 获取过期时间
|
||||
payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
|
||||
exp_timestamp = payload.get("exp")
|
||||
|
||||
if not exp_timestamp:
|
||||
logger.warning("Token 缺少过期时间,无法撤销")
|
||||
return False
|
||||
|
||||
# 将 Token 加入黑名单
|
||||
success = await JWTBlacklistService.add_to_blacklist(
|
||||
token=token, exp_timestamp=exp_timestamp, reason=reason
|
||||
)
|
||||
|
||||
if success:
|
||||
user_id = payload.get("sub")
|
||||
logger.warning(f"Token 已被撤销: user_id={user_id}, reason={reason}")
|
||||
|
||||
return success
|
||||
|
||||
except jwt.InvalidTokenError as e:
|
||||
logger.warning(f"撤销失败 - 无效的 Token: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"撤销 Token 失败: {e}")
|
||||
return False
|
||||
19
src/services/cache/__init__.py
vendored
Normal file
19
src/services/cache/__init__.py
vendored
Normal file
@@ -0,0 +1,19 @@
|
||||
"""
|
||||
缓存服务模块
|
||||
|
||||
包含缓存后端、缓存亲和性、缓存同步等功能。
|
||||
|
||||
注意:由于循环依赖问题,部分类需要直接从子模块导入:
|
||||
from src.services.cache.affinity_manager import CacheAffinityManager
|
||||
from src.services.cache.aware_scheduler import CacheAwareScheduler
|
||||
"""
|
||||
|
||||
# 只导出不会导致循环依赖的基础类
|
||||
from src.services.cache.backend import BaseCacheBackend, LocalCache, RedisCache, get_cache_backend
|
||||
|
||||
__all__ = [
|
||||
"BaseCacheBackend",
|
||||
"LocalCache",
|
||||
"RedisCache",
|
||||
"get_cache_backend",
|
||||
]
|
||||
668
src/services/cache/affinity_manager.py
vendored
Normal file
668
src/services/cache/affinity_manager.py
vendored
Normal file
@@ -0,0 +1,668 @@
|
||||
"""
|
||||
缓存亲和性管理器 (Cache Affinity Manager) - 支持 Redis 或内存存储
|
||||
|
||||
职责:
|
||||
1. 跟踪请求API Key的Provider+Key缓存状态
|
||||
2. 管理缓存有效期
|
||||
3. 提供缓存统计和分析
|
||||
4. 自动失效不支持缓存的Provider
|
||||
|
||||
设计原理:
|
||||
- 每个API Key使用某个Provider的Key后,在缓存TTL期内,应该继续使用同一个Key
|
||||
- 这样可以最大化利用提供商的Prompt Caching机制
|
||||
- 当Key故障时,自动失效该Key的缓存亲和性
|
||||
- 当Provider关闭缓存支持时,自动失效所有相关亲和性
|
||||
|
||||
注意:
|
||||
- affinity_key 参数通常为请求使用的 API Key ID(api_key_id)
|
||||
- 这样可以支持"独立余额Key"场景,每个Key有自己的缓存亲和性
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, Dict, List, NamedTuple, Optional, Tuple
|
||||
|
||||
from src.config.constants import CacheTTL
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
|
||||
class CacheAffinity(NamedTuple):
|
||||
"""缓存亲和性信息"""
|
||||
|
||||
provider_id: str
|
||||
endpoint_id: str
|
||||
key_id: str
|
||||
api_format: str # API格式 (claude/openai)
|
||||
model_name: str # 模型名称
|
||||
created_at: float # 创建时间戳
|
||||
expire_at: float # 过期时间戳
|
||||
request_count: int # 使用次数
|
||||
|
||||
|
||||
class CacheAffinityManager:
|
||||
"""
|
||||
缓存亲和性管理器(支持 Redis 或内存存储)
|
||||
|
||||
存储结构:
|
||||
----------------------
|
||||
Key格式: cache_affinity:{affinity_key}:{api_format}:{model_name}
|
||||
- affinity_key: 通常为请求使用的 API Key ID(支持独立余额Key场景)
|
||||
- api_format: API格式 (claude/openai)
|
||||
- model_name: 模型名称(区分不同模型的缓存亲和性)
|
||||
Value格式: JSON/Dict
|
||||
{
|
||||
"provider_id": "xxx",
|
||||
"endpoint_id": "yyy",
|
||||
"key_id": "zzz",
|
||||
"model_name": "claude-3-5-sonnet-20241022",
|
||||
"created_at": 1234567890.123,
|
||||
"expire_at": 1234567890.123,
|
||||
"request_count": 5
|
||||
}
|
||||
TTL: 自动过期
|
||||
|
||||
设计改进:
|
||||
- 每个API Key可以对多个API格式和模型分别维护缓存亲和性
|
||||
- 不同模型请求使用独立的缓存亲和性,避免模型切换导致的缓存失效
|
||||
- 某个端点故障切换不会影响其他端点的亲和性
|
||||
- 更精确的缓存命中率统计
|
||||
- 支持"独立余额Key"场景,每个Key有独立的缓存亲和性
|
||||
"""
|
||||
|
||||
# 默认缓存TTL(秒)- 使用统一常量
|
||||
DEFAULT_CACHE_TTL = CacheTTL.CACHE_AFFINITY
|
||||
|
||||
def __init__(self, redis_client=None, default_ttl: int = DEFAULT_CACHE_TTL):
|
||||
"""
|
||||
初始化缓存亲和性管理器
|
||||
|
||||
Args:
|
||||
redis_client: Redis客户端(可选)
|
||||
default_ttl: 默认缓存TTL(秒)
|
||||
"""
|
||||
self.redis = redis_client
|
||||
self.default_ttl = default_ttl
|
||||
self._memory_store: Dict[str, Dict[str, Any]] = {}
|
||||
self._memory_lock: Optional[asyncio.Lock] = None
|
||||
|
||||
# L1 缓存(即使使用 Redis 也启用,减少网络往返)
|
||||
self._l1_cache_ttl = int(os.getenv("CACHE_AFFINITY_L1_TTL", str(CacheTTL.L1_LOCAL)))
|
||||
self._l1_cache: Dict[str, Tuple[float, Dict[str, Any]]] = {}
|
||||
self._l1_lock = asyncio.Lock()
|
||||
self._l1_max_size = int(os.getenv("CACHE_AFFINITY_L1_MAX_SIZE", "1000")) # 最大缓存条目数
|
||||
self._l1_last_cleanup = time.time()
|
||||
|
||||
# 请求级别锁,避免同一用户+端点同时更新造成抖动
|
||||
self._request_locks: Dict[str, asyncio.Lock] = {}
|
||||
|
||||
# 统计信息
|
||||
self._stats = {
|
||||
"total_affinities": 0,
|
||||
"cache_hits": 0,
|
||||
"cache_misses": 0,
|
||||
"cache_invalidations": 0,
|
||||
"provider_switches": 0,
|
||||
"key_switches": 0,
|
||||
}
|
||||
|
||||
if self.redis:
|
||||
logger.debug("CacheAffinityManager: 使用Redis存储")
|
||||
else:
|
||||
logger.debug("CacheAffinityManager: Redis不可用,回退到内存存储(仅适用于单实例/开发环境)")
|
||||
|
||||
def _is_memory_backend(self) -> bool:
|
||||
"""是否处于内存模式"""
|
||||
return self.redis is None
|
||||
|
||||
def _get_memory_lock(self) -> asyncio.Lock:
|
||||
"""懒初始化内存锁"""
|
||||
if self._memory_lock is None:
|
||||
self._memory_lock = asyncio.Lock()
|
||||
return self._memory_lock
|
||||
|
||||
def _get_cache_key(self, affinity_key: str, api_format: str, model_name: str) -> str:
|
||||
"""
|
||||
生成Redis Key
|
||||
|
||||
Args:
|
||||
affinity_key: 亲和性标识符(通常为API Key ID)
|
||||
api_format: API格式 (claude/openai)
|
||||
model_name: 模型名称
|
||||
|
||||
Returns:
|
||||
格式化的缓存键: cache_affinity:{affinity_key}:{api_format}:{model_name}
|
||||
"""
|
||||
return f"cache_affinity:{affinity_key}:{api_format}:{model_name}"
|
||||
|
||||
async def _get_l1_entry(self, cache_key: str) -> Optional[Dict[str, Any]]:
|
||||
async with self._l1_lock:
|
||||
record = self._l1_cache.get(cache_key)
|
||||
if not record:
|
||||
return None
|
||||
expire_at, payload = record
|
||||
if time.time() > expire_at:
|
||||
self._l1_cache.pop(cache_key, None)
|
||||
return None
|
||||
return dict(payload)
|
||||
|
||||
async def _set_l1_entry(self, cache_key: str, payload: Optional[Dict[str, Any]]):
|
||||
async with self._l1_lock:
|
||||
if not payload:
|
||||
self._l1_cache.pop(cache_key, None)
|
||||
return
|
||||
expire_at = time.time() + max(1, self._l1_cache_ttl)
|
||||
self._l1_cache[cache_key] = (expire_at, dict(payload))
|
||||
|
||||
# 定期清理过期条目(每 60 秒最多一次)
|
||||
current_time = time.time()
|
||||
if current_time - self._l1_last_cleanup > 60:
|
||||
self._cleanup_l1_cache_unlocked(current_time)
|
||||
self._l1_last_cleanup = current_time
|
||||
|
||||
def _cleanup_l1_cache_unlocked(self, current_time: float) -> int:
|
||||
"""清理过期的 L1 缓存条目(需要在持有锁的情况下调用)
|
||||
|
||||
Returns:
|
||||
清理的条目数量
|
||||
"""
|
||||
expired_keys = [
|
||||
key for key, (expire_at, _) in self._l1_cache.items()
|
||||
if current_time > expire_at
|
||||
]
|
||||
for key in expired_keys:
|
||||
self._l1_cache.pop(key, None)
|
||||
|
||||
# 如果缓存仍然过大,按过期时间排序移除最旧的条目
|
||||
if len(self._l1_cache) > self._l1_max_size:
|
||||
sorted_items = sorted(
|
||||
self._l1_cache.items(),
|
||||
key=lambda x: x[1][0] # 按 expire_at 排序
|
||||
)
|
||||
# 移除最旧的 20% 条目
|
||||
remove_count = len(self._l1_cache) - int(self._l1_max_size * 0.8)
|
||||
for key, _ in sorted_items[:remove_count]:
|
||||
self._l1_cache.pop(key, None)
|
||||
expired_keys.extend([k for k, _ in sorted_items[:remove_count]])
|
||||
|
||||
if expired_keys:
|
||||
logger.debug(f"L1 缓存清理: 移除 {len(expired_keys)} 个条目,当前 {len(self._l1_cache)} 个")
|
||||
|
||||
return len(expired_keys)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _acquire_request_lock(self, cache_key: str):
|
||||
lock = self._request_locks.get(cache_key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
self._request_locks[cache_key] = lock
|
||||
await lock.acquire()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
lock.release()
|
||||
|
||||
async def _load_affinity_dict(self, cache_key: str) -> Optional[Dict[str, Any]]:
|
||||
"""读取缓存亲和性字典"""
|
||||
# 先尝试L1缓存
|
||||
l1_value = await self._get_l1_entry(cache_key)
|
||||
if l1_value is not None:
|
||||
return l1_value
|
||||
|
||||
if not self._is_memory_backend():
|
||||
data = await self.redis.get(cache_key)
|
||||
if not data:
|
||||
return None
|
||||
value = json.loads(data)
|
||||
await self._set_l1_entry(cache_key, value)
|
||||
return value
|
||||
|
||||
lock = self._get_memory_lock()
|
||||
async with lock:
|
||||
record = self._memory_store.get(cache_key)
|
||||
if record:
|
||||
await self._set_l1_entry(cache_key, record)
|
||||
return dict(record) if record else None
|
||||
|
||||
async def _save_affinity_dict(
|
||||
self, cache_key: str, ttl: int, affinity_dict: Dict[str, Any]
|
||||
) -> None:
|
||||
"""存储缓存亲和性字典"""
|
||||
if not self._is_memory_backend():
|
||||
await self.redis.setex(cache_key, ttl, json.dumps(affinity_dict))
|
||||
await self._set_l1_entry(cache_key, affinity_dict)
|
||||
return
|
||||
|
||||
lock = self._get_memory_lock()
|
||||
async with lock:
|
||||
self._memory_store[cache_key] = dict(affinity_dict)
|
||||
await self._set_l1_entry(cache_key, affinity_dict)
|
||||
|
||||
async def _delete_affinity_key(self, cache_key: str) -> None:
|
||||
"""删除缓存亲和性"""
|
||||
if not self._is_memory_backend():
|
||||
await self.redis.delete(cache_key)
|
||||
else:
|
||||
lock = self._get_memory_lock()
|
||||
async with lock:
|
||||
self._memory_store.pop(cache_key, None)
|
||||
|
||||
await self._set_l1_entry(cache_key, None)
|
||||
|
||||
async def _snapshot_memory_items(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""复制内存存储内容(仅内存模式使用)"""
|
||||
lock = self._get_memory_lock()
|
||||
async with lock:
|
||||
return {k: dict(v) for k, v in self._memory_store.items()}
|
||||
|
||||
async def get_affinity(
|
||||
self, affinity_key: str, api_format: str, model_name: str
|
||||
) -> Optional[CacheAffinity]:
|
||||
"""
|
||||
获取指定亲和性标识符对特定API格式和模型的缓存亲和性
|
||||
|
||||
Args:
|
||||
affinity_key: 亲和性标识符(通常为API Key ID)
|
||||
api_format: API格式 (claude/openai)
|
||||
model_name: 模型名称
|
||||
|
||||
Returns:
|
||||
CacheAffinity对象,如果不存在或已过期则返回None
|
||||
"""
|
||||
try:
|
||||
cache_key = self._get_cache_key(affinity_key, api_format, model_name)
|
||||
async with self._acquire_request_lock(cache_key):
|
||||
affinity_dict = await self._load_affinity_dict(cache_key)
|
||||
|
||||
if not affinity_dict:
|
||||
self._stats["cache_misses"] += 1
|
||||
return None
|
||||
|
||||
# 检查是否过期(双重检查,防止TTL未及时清理)
|
||||
current_time = time.time()
|
||||
if current_time > affinity_dict["expire_at"]:
|
||||
await self._delete_affinity_key(cache_key)
|
||||
self._stats["cache_misses"] += 1
|
||||
return None
|
||||
|
||||
self._stats["cache_hits"] += 1
|
||||
|
||||
return CacheAffinity(
|
||||
provider_id=affinity_dict["provider_id"],
|
||||
endpoint_id=affinity_dict["endpoint_id"],
|
||||
key_id=affinity_dict["key_id"],
|
||||
api_format=affinity_dict.get("api_format", api_format),
|
||||
model_name=affinity_dict.get("model_name", model_name),
|
||||
created_at=affinity_dict["created_at"],
|
||||
expire_at=affinity_dict["expire_at"],
|
||||
request_count=affinity_dict["request_count"],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"获取缓存亲和性失败: {e}")
|
||||
self._stats["cache_misses"] += 1
|
||||
return None
|
||||
|
||||
async def set_affinity(
|
||||
self,
|
||||
affinity_key: str,
|
||||
provider_id: str,
|
||||
endpoint_id: str,
|
||||
key_id: str,
|
||||
api_format: str,
|
||||
model_name: str,
|
||||
supports_caching: bool = True,
|
||||
ttl: Optional[int] = None,
|
||||
) -> None:
|
||||
"""
|
||||
设置指定亲和性标识符对特定API格式和模型的缓存亲和性
|
||||
|
||||
Args:
|
||||
affinity_key: 亲和性标识符(通常为API Key ID)
|
||||
provider_id: Provider ID
|
||||
endpoint_id: Endpoint ID
|
||||
key_id: Key ID
|
||||
api_format: API格式 (claude/openai)
|
||||
model_name: 模型名称
|
||||
supports_caching: 该Provider是否支持缓存
|
||||
ttl: 缓存有效期(秒),如果不提供则使用默认值
|
||||
|
||||
注意:每次调用都会刷新过期时间(滑动窗口机制),以保持对同一个Provider/Endpoint/Key的亲和性
|
||||
"""
|
||||
if not supports_caching:
|
||||
# 不支持缓存的Provider不记录亲和性
|
||||
logger.debug(f"Provider {provider_id[:8]}... 不支持缓存,跳过亲和性记录")
|
||||
return
|
||||
|
||||
ttl = ttl or self.default_ttl
|
||||
current_time = time.time()
|
||||
expire_at = current_time + ttl # 每次都刷新过期时间
|
||||
cache_key = self._get_cache_key(affinity_key, api_format, model_name)
|
||||
|
||||
try:
|
||||
async with self._acquire_request_lock(cache_key):
|
||||
existing_dict = await self._load_affinity_dict(cache_key)
|
||||
existing_affinity: Optional[CacheAffinity] = None
|
||||
if existing_dict and current_time <= existing_dict.get("expire_at", 0):
|
||||
existing_affinity = CacheAffinity(
|
||||
provider_id=existing_dict["provider_id"],
|
||||
endpoint_id=existing_dict["endpoint_id"],
|
||||
key_id=existing_dict["key_id"],
|
||||
api_format=existing_dict.get("api_format", api_format),
|
||||
model_name=existing_dict.get("model_name", model_name),
|
||||
created_at=existing_dict["created_at"],
|
||||
expire_at=existing_dict["expire_at"],
|
||||
request_count=existing_dict.get("request_count", 0),
|
||||
)
|
||||
|
||||
if existing_affinity:
|
||||
created_at = existing_affinity.created_at
|
||||
request_count = existing_affinity.request_count + 1
|
||||
|
||||
# 检查是否切换了 Provider/Endpoint/Key
|
||||
if (
|
||||
existing_affinity.provider_id != provider_id
|
||||
or existing_affinity.endpoint_id != endpoint_id
|
||||
or existing_affinity.key_id != key_id
|
||||
):
|
||||
self._stats["key_switches"] += 1
|
||||
logger.debug(f"Key {affinity_key[:8]}... 在 {api_format} 格式下切换后端: "
|
||||
f"[{existing_affinity.provider_id[:8]}.../{existing_affinity.endpoint_id[:8]}.../"
|
||||
f"{existing_affinity.key_id[:8]}...] → "
|
||||
f"[{provider_id[:8]}.../{endpoint_id[:8]}.../{key_id[:8]}...], 重置计数器")
|
||||
created_at = current_time
|
||||
request_count = 1
|
||||
else:
|
||||
logger.debug(f"刷新缓存亲和性: key={affinity_key[:8]}..., api_format={api_format}, "
|
||||
f"provider={provider_id[:8]}..., endpoint={endpoint_id[:8]}..., "
|
||||
f"provider_key={key_id[:8]}..., ttl+={ttl}s")
|
||||
else:
|
||||
created_at = current_time
|
||||
request_count = 1
|
||||
self._stats["total_affinities"] += 1
|
||||
|
||||
affinity_dict = {
|
||||
"provider_id": provider_id,
|
||||
"endpoint_id": endpoint_id,
|
||||
"key_id": key_id,
|
||||
"api_format": api_format,
|
||||
"model_name": model_name,
|
||||
"created_at": created_at,
|
||||
"expire_at": expire_at,
|
||||
"request_count": request_count,
|
||||
}
|
||||
|
||||
await self._save_affinity_dict(cache_key, ttl, affinity_dict)
|
||||
|
||||
logger.debug(f"设置缓存亲和性: key={affinity_key[:8]}..., api_format={api_format}, "
|
||||
f"model={model_name}, provider={provider_id[:8]}..., endpoint={endpoint_id[:8]}..., "
|
||||
f"provider_key={key_id[:8]}..., ttl={ttl}s")
|
||||
except Exception as e:
|
||||
logger.exception(f"设置缓存亲和性失败: {e}")
|
||||
|
||||
async def invalidate_affinity(
|
||||
self,
|
||||
affinity_key: str,
|
||||
api_format: str,
|
||||
model_name: str,
|
||||
key_id: Optional[str] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
endpoint_id: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
失效指定亲和性标识符对特定API格式和模型的缓存亲和性
|
||||
|
||||
Args:
|
||||
affinity_key: 亲和性标识符(通常为API Key ID)
|
||||
api_format: API格式 (claude/openai)
|
||||
model_name: 模型名称
|
||||
key_id: Provider Key ID(可选,如果提供则只在Key匹配时失效)
|
||||
provider_id: Provider ID(可选,如果提供则只在Provider匹配时失效)
|
||||
endpoint_id: Endpoint ID(可选,如果提供则只在Endpoint匹配时失效)
|
||||
"""
|
||||
existing_affinity = await self.get_affinity(affinity_key, api_format, model_name)
|
||||
|
||||
if not existing_affinity:
|
||||
return
|
||||
|
||||
# 检查是否匹配过滤条件
|
||||
should_invalidate = True
|
||||
|
||||
if key_id and existing_affinity.key_id != key_id:
|
||||
should_invalidate = False
|
||||
|
||||
if provider_id and existing_affinity.provider_id != provider_id:
|
||||
should_invalidate = False
|
||||
|
||||
if endpoint_id and existing_affinity.endpoint_id != endpoint_id:
|
||||
should_invalidate = False
|
||||
|
||||
if not should_invalidate:
|
||||
logger.debug(f"跳过失效: affinity_key={affinity_key[:8]}..., api_format={api_format}, "
|
||||
f"model={model_name}, 过滤条件不匹配 (key={key_id}, provider={provider_id}, endpoint={endpoint_id})")
|
||||
return
|
||||
|
||||
try:
|
||||
cache_key = self._get_cache_key(affinity_key, api_format, model_name)
|
||||
async with self._acquire_request_lock(cache_key):
|
||||
await self._delete_affinity_key(cache_key)
|
||||
|
||||
self._stats["cache_invalidations"] += 1
|
||||
|
||||
logger.debug(f"失效缓存亲和性: affinity_key={affinity_key[:8]}..., api_format={api_format}, "
|
||||
f"model={model_name}, provider={existing_affinity.provider_id[:8]}..., "
|
||||
f"endpoint={existing_affinity.endpoint_id[:8]}..., "
|
||||
f"provider_key={existing_affinity.key_id[:8]}...")
|
||||
except Exception as e:
|
||||
logger.exception(f"删除缓存亲和性失败: {e}")
|
||||
|
||||
async def invalidate_all_for_provider(self, provider_id: str) -> int:
|
||||
"""
|
||||
失效所有与指定Provider相关的缓存亲和性
|
||||
|
||||
用途:当Provider关闭缓存支持时调用
|
||||
|
||||
Args:
|
||||
provider_id: Provider ID
|
||||
|
||||
Returns:
|
||||
失效的亲和性数量
|
||||
"""
|
||||
try:
|
||||
invalidated_count = 0
|
||||
|
||||
if not self._is_memory_backend():
|
||||
pattern = "cache_affinity:*"
|
||||
keys = await self.redis.keys(pattern)
|
||||
else:
|
||||
keys = list((await self._snapshot_memory_items()).keys())
|
||||
|
||||
for key in keys:
|
||||
affinity_dict = await self._load_affinity_dict(key)
|
||||
if not affinity_dict:
|
||||
continue
|
||||
|
||||
if affinity_dict.get("provider_id") == provider_id:
|
||||
await self._delete_affinity_key(key)
|
||||
invalidated_count += 1
|
||||
self._stats["cache_invalidations"] += 1
|
||||
|
||||
if invalidated_count > 0:
|
||||
logger.debug(f"批量失效Provider缓存亲和性: provider={provider_id[:8]}..., "
|
||||
f"失效数量={invalidated_count}")
|
||||
|
||||
return invalidated_count
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"批量失效Provider缓存亲和性失败: {e}")
|
||||
return 0
|
||||
|
||||
async def clear_all(self) -> int:
|
||||
"""
|
||||
清除所有缓存亲和性(管理功能)
|
||||
|
||||
Returns:
|
||||
清除的数量
|
||||
"""
|
||||
try:
|
||||
if not self._is_memory_backend():
|
||||
keys = await self.redis.keys("cache_affinity:*")
|
||||
if keys:
|
||||
await self.redis.delete(*keys)
|
||||
logger.debug(f"清除所有Redis缓存亲和性: {len(keys)} 个")
|
||||
return len(keys)
|
||||
return 0
|
||||
|
||||
lock = self._get_memory_lock()
|
||||
async with lock:
|
||||
count = len(self._memory_store)
|
||||
self._memory_store.clear()
|
||||
if count:
|
||||
logger.debug(f"清除所有内存缓存亲和性: {count} 个")
|
||||
return count
|
||||
except Exception as e:
|
||||
logger.exception(f"清除缓存亲和性失败: {e}")
|
||||
return 0
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
cache_hit_rate = 0.0
|
||||
total_requests = self._stats["cache_hits"] + self._stats["cache_misses"]
|
||||
|
||||
if total_requests > 0:
|
||||
cache_hit_rate = self._stats["cache_hits"] / total_requests
|
||||
|
||||
storage_type = "redis" if not self._is_memory_backend() else "memory"
|
||||
|
||||
return {
|
||||
"storage_type": storage_type,
|
||||
"total_affinities": self._stats["total_affinities"],
|
||||
"cache_hits": self._stats["cache_hits"],
|
||||
"cache_misses": self._stats["cache_misses"],
|
||||
"cache_hit_rate": cache_hit_rate,
|
||||
"cache_invalidations": self._stats["cache_invalidations"],
|
||||
"provider_switches": self._stats["provider_switches"],
|
||||
"key_switches": self._stats["key_switches"],
|
||||
"config": {
|
||||
"default_ttl": self.default_ttl,
|
||||
},
|
||||
}
|
||||
|
||||
async def list_affinities(self) -> List[Dict[str, Any]]:
|
||||
"""获取所有缓存亲和性列表
|
||||
|
||||
返回的每条记录包含:
|
||||
- affinity_key: 亲和性标识符(通常是 API Key ID)
|
||||
- provider_id, endpoint_id, key_id: Provider 相关信息
|
||||
- api_format, model_name: API 格式和模型名称
|
||||
- created_at, expire_at, request_count: 缓存元数据
|
||||
"""
|
||||
results: List[Dict[str, Any]] = []
|
||||
|
||||
try:
|
||||
pattern = "cache_affinity:*"
|
||||
cursor = 0
|
||||
|
||||
if not self._is_memory_backend():
|
||||
while True:
|
||||
cursor, keys = await self.redis.scan(cursor=cursor, match=pattern, count=200)
|
||||
|
||||
if keys:
|
||||
values = await self.redis.mget(*keys)
|
||||
for cache_key, data in zip(keys, values):
|
||||
if not data:
|
||||
continue
|
||||
|
||||
try:
|
||||
affinity = json.loads(data)
|
||||
# 解析 cache_affinity:{affinity_key}:{api_format}:{model_name}
|
||||
parts = cache_key.split(":")
|
||||
affinity_key_value = parts[1] if len(parts) > 1 else cache_key
|
||||
api_format = (
|
||||
parts[2]
|
||||
if len(parts) > 2
|
||||
else affinity.get("api_format", "unknown")
|
||||
)
|
||||
model_name = (
|
||||
parts[3]
|
||||
if len(parts) > 3
|
||||
else affinity.get("model_name", "unknown")
|
||||
)
|
||||
|
||||
affinity["affinity_key"] = affinity_key_value
|
||||
if "api_format" not in affinity:
|
||||
affinity["api_format"] = api_format
|
||||
if "model_name" not in affinity:
|
||||
affinity["model_name"] = model_name
|
||||
results.append(affinity)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.exception(f"解析缓存亲和性记录失败: {cache_key} - {e}")
|
||||
|
||||
if cursor == 0:
|
||||
break
|
||||
else:
|
||||
snapshot = await self._snapshot_memory_items()
|
||||
expired_keys: List[str] = []
|
||||
current_time = time.time()
|
||||
|
||||
for cache_key, affinity in snapshot.items():
|
||||
if current_time > affinity["expire_at"]:
|
||||
expired_keys.append(cache_key)
|
||||
continue
|
||||
|
||||
# 解析 cache_affinity:{affinity_key}:{api_format}:{model_name}
|
||||
parts = cache_key.split(":")
|
||||
affinity_key_value = parts[1] if len(parts) > 1 else cache_key
|
||||
api_format = (
|
||||
parts[2] if len(parts) > 2 else affinity.get("api_format", "unknown")
|
||||
)
|
||||
model_name = (
|
||||
parts[3] if len(parts) > 3 else affinity.get("model_name", "unknown")
|
||||
)
|
||||
|
||||
affinity_with_key = dict(affinity)
|
||||
affinity_with_key["affinity_key"] = affinity_key_value
|
||||
if "api_format" not in affinity_with_key:
|
||||
affinity_with_key["api_format"] = api_format
|
||||
if "model_name" not in affinity_with_key:
|
||||
affinity_with_key["model_name"] = model_name
|
||||
results.append(affinity_with_key)
|
||||
|
||||
# 清理过期的键
|
||||
if expired_keys:
|
||||
async with self._get_memory_lock():
|
||||
for key in expired_keys:
|
||||
self._memory_store.pop(key, None)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"获取缓存亲和性列表失败: {e}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# 全局单例
|
||||
_affinity_manager: Optional[CacheAffinityManager] = None
|
||||
|
||||
|
||||
async def get_affinity_manager(redis_client=None) -> CacheAffinityManager:
|
||||
"""
|
||||
获取全局CacheAffinityManager实例(若Redis不可用则降级为内存模式)
|
||||
|
||||
Args:
|
||||
redis_client: Redis客户端(可选)
|
||||
|
||||
Returns:
|
||||
CacheAffinityManager实例
|
||||
"""
|
||||
global _affinity_manager
|
||||
|
||||
if _affinity_manager is None:
|
||||
_affinity_manager = CacheAffinityManager(redis_client)
|
||||
elif redis_client and _affinity_manager.redis is None:
|
||||
# 当最初使用内存后 Redis 可用时,升级为 Redis 存储
|
||||
_affinity_manager = CacheAffinityManager(redis_client)
|
||||
|
||||
return _affinity_manager
|
||||
1316
src/services/cache/aware_scheduler.py
vendored
Normal file
1316
src/services/cache/aware_scheduler.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
330
src/services/cache/backend.py
vendored
Normal file
330
src/services/cache/backend.py
vendored
Normal file
@@ -0,0 +1,330 @@
|
||||
"""
|
||||
缓存后端抽象层
|
||||
|
||||
提供统一的缓存接口,支持多种后端实现:
|
||||
1. LocalCache: 内存缓存(单实例,线程安全)
|
||||
2. RedisCache: Redis 缓存(分布式)
|
||||
|
||||
使用场景:
|
||||
- ModelMappingResolver: 模型映射与别名解析缓存
|
||||
- ModelMapper: 模型映射缓存
|
||||
- 其他需要缓存的服务
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from src.core.logger import logger
|
||||
|
||||
from src.clients.redis_client import get_redis_client_sync
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
class BaseCacheBackend(ABC):
|
||||
"""缓存后端抽象基类"""
|
||||
|
||||
@abstractmethod
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
"""获取缓存值"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def set(self, key: str, value: Any, ttl: int = 300) -> None:
|
||||
"""设置缓存值"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, key: str) -> None:
|
||||
"""删除缓存值"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def clear(self, pattern: Optional[str] = None) -> None:
|
||||
"""清空缓存(支持模式匹配)"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""检查键是否存在"""
|
||||
pass
|
||||
|
||||
|
||||
class LocalCache(BaseCacheBackend):
|
||||
"""本地内存缓存后端(LRU + TTL,线程安全)"""
|
||||
|
||||
def __init__(self, max_size: int = 1000, default_ttl: int = 300):
|
||||
"""
|
||||
初始化本地缓存
|
||||
|
||||
Args:
|
||||
max_size: 最大缓存条目数
|
||||
default_ttl: 默认过期时间(秒)
|
||||
"""
|
||||
self._cache: OrderedDict = OrderedDict()
|
||||
self._expiry: Dict[str, float] = {}
|
||||
self._max_size = max_size
|
||||
self._default_ttl = default_ttl
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
"""获取缓存值(线程安全)"""
|
||||
async with self._lock:
|
||||
if key not in self._cache:
|
||||
return None
|
||||
|
||||
# 检查过期
|
||||
if key in self._expiry and time.time() > self._expiry[key]:
|
||||
# 过期,删除
|
||||
del self._cache[key]
|
||||
del self._expiry[key]
|
||||
return None
|
||||
|
||||
# 更新访问顺序(LRU)
|
||||
self._cache.move_to_end(key)
|
||||
return self._cache[key]
|
||||
|
||||
async def set(self, key: str, value: Any, ttl: int = None) -> None:
|
||||
"""设置缓存值(线程安全)"""
|
||||
async with self._lock:
|
||||
if ttl is None:
|
||||
ttl = self._default_ttl
|
||||
|
||||
# 如果键已存在,更新访问顺序
|
||||
if key in self._cache:
|
||||
self._cache.move_to_end(key)
|
||||
|
||||
self._cache[key] = value
|
||||
self._expiry[key] = time.time() + ttl
|
||||
|
||||
# 检查容量限制,淘汰最旧项
|
||||
if len(self._cache) > self._max_size:
|
||||
oldest_key = next(iter(self._cache))
|
||||
del self._cache[oldest_key]
|
||||
if oldest_key in self._expiry:
|
||||
del self._expiry[oldest_key]
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
"""删除缓存值(线程安全)"""
|
||||
async with self._lock:
|
||||
if key in self._cache:
|
||||
del self._cache[key]
|
||||
if key in self._expiry:
|
||||
del self._expiry[key]
|
||||
|
||||
async def clear(self, pattern: Optional[str] = None) -> None:
|
||||
"""清空缓存(线程安全)"""
|
||||
async with self._lock:
|
||||
if pattern is None:
|
||||
# 清空所有
|
||||
self._cache.clear()
|
||||
self._expiry.clear()
|
||||
else:
|
||||
# 模式匹配删除(简单实现:支持前缀匹配)
|
||||
prefix = pattern.rstrip("*")
|
||||
keys_to_delete = [k for k in self._cache.keys() if k.startswith(prefix)]
|
||||
for key in keys_to_delete:
|
||||
del self._cache[key]
|
||||
if key in self._expiry:
|
||||
del self._expiry[key]
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""检查键是否存在(线程安全)"""
|
||||
async with self._lock:
|
||||
if key not in self._cache:
|
||||
return False
|
||||
|
||||
# 检查过期
|
||||
if key in self._expiry and time.time() > self._expiry[key]:
|
||||
del self._cache[key]
|
||||
del self._expiry[key]
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""获取缓存统计信息"""
|
||||
return {
|
||||
"backend": "local",
|
||||
"size": len(self._cache),
|
||||
"max_size": self._max_size,
|
||||
"default_ttl": self._default_ttl,
|
||||
}
|
||||
|
||||
|
||||
class RedisCache(BaseCacheBackend):
|
||||
"""Redis 缓存后端(分布式)"""
|
||||
|
||||
def __init__(
|
||||
self, redis_client: aioredis.Redis, key_prefix: str = "cache", default_ttl: int = 300
|
||||
):
|
||||
"""
|
||||
初始化 Redis 缓存
|
||||
|
||||
Args:
|
||||
redis_client: Redis 客户端实例
|
||||
key_prefix: 缓存键前缀
|
||||
default_ttl: 默认过期时间(秒)
|
||||
"""
|
||||
self._redis = redis_client
|
||||
self._key_prefix = key_prefix
|
||||
self._default_ttl = default_ttl
|
||||
|
||||
def _make_key(self, key: str) -> str:
|
||||
"""构造完整的 Redis 键"""
|
||||
return f"{self._key_prefix}:{key}"
|
||||
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
"""获取缓存值"""
|
||||
try:
|
||||
redis_key = self._make_key(key)
|
||||
value = await self._redis.get(redis_key)
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
# 尝试 JSON 反序列化
|
||||
try:
|
||||
return json.loads(value)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# 如果不是 JSON,直接返回字符串
|
||||
return value
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 获取缓存失败: {key}, 错误: {e}")
|
||||
return None
|
||||
|
||||
async def set(self, key: str, value: Any, ttl: int = None) -> None:
|
||||
"""设置缓存值"""
|
||||
if ttl is None:
|
||||
ttl = self._default_ttl
|
||||
|
||||
try:
|
||||
redis_key = self._make_key(key)
|
||||
|
||||
# 序列化值
|
||||
if isinstance(value, (dict, list, tuple)):
|
||||
serialized = json.dumps(value)
|
||||
elif isinstance(value, (int, float, bool)):
|
||||
serialized = json.dumps(value)
|
||||
else:
|
||||
serialized = str(value)
|
||||
|
||||
await self._redis.setex(redis_key, ttl, serialized)
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 设置缓存失败: {key}, 错误: {e}")
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
"""删除缓存值"""
|
||||
try:
|
||||
redis_key = self._make_key(key)
|
||||
await self._redis.delete(redis_key)
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 删除缓存失败: {key}, 错误: {e}")
|
||||
|
||||
async def clear(self, pattern: Optional[str] = None) -> None:
|
||||
"""清空缓存"""
|
||||
try:
|
||||
if pattern is None:
|
||||
# 清空所有带前缀的键
|
||||
pattern = "*"
|
||||
|
||||
redis_pattern = self._make_key(pattern)
|
||||
cursor = 0
|
||||
deleted_count = 0
|
||||
|
||||
while True:
|
||||
cursor, keys = await self._redis.scan(cursor, match=redis_pattern, count=100)
|
||||
if keys:
|
||||
await self._redis.delete(*keys)
|
||||
deleted_count += len(keys)
|
||||
if cursor == 0:
|
||||
break
|
||||
|
||||
logger.info(f"[RedisCache] 清空缓存: {redis_pattern}, 删除 {deleted_count} 个键")
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 清空缓存失败: {pattern}, 错误: {e}")
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""检查键是否存在"""
|
||||
try:
|
||||
redis_key = self._make_key(key)
|
||||
return await self._redis.exists(redis_key) > 0
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 检查键存在失败: {key}, 错误: {e}")
|
||||
return False
|
||||
|
||||
async def publish_invalidation(self, channel: str, key: str) -> None:
|
||||
"""发布缓存失效消息(用于分布式同步)"""
|
||||
try:
|
||||
message = json.dumps({"key": key, "timestamp": time.time()})
|
||||
await self._redis.publish(channel, message)
|
||||
logger.debug(f"[RedisCache] 发布缓存失效: {channel} -> {key}")
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 发布缓存失效失败: {channel}, {key}, 错误: {e}")
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""获取缓存统计信息"""
|
||||
return {
|
||||
"backend": "redis",
|
||||
"key_prefix": self._key_prefix,
|
||||
"default_ttl": self._default_ttl,
|
||||
}
|
||||
|
||||
|
||||
# 缓存后端工厂
|
||||
_cache_backends: Dict[str, BaseCacheBackend] = {}
|
||||
|
||||
|
||||
async def get_cache_backend(
|
||||
name: str, backend_type: str = "auto", max_size: int = 1000, ttl: int = 300
|
||||
) -> BaseCacheBackend:
|
||||
"""
|
||||
获取缓存后端实例
|
||||
|
||||
Args:
|
||||
name: 缓存名称(用于区分不同的缓存实例)
|
||||
backend_type: 后端类型 (auto/local/redis)
|
||||
max_size: LocalCache 的最大容量
|
||||
ttl: 默认过期时间(秒)
|
||||
|
||||
Returns:
|
||||
BaseCacheBackend 实例
|
||||
"""
|
||||
cache_key = f"{name}:{backend_type}"
|
||||
|
||||
if cache_key in _cache_backends:
|
||||
return _cache_backends[cache_key]
|
||||
|
||||
# 根据类型创建缓存后端
|
||||
if backend_type == "redis":
|
||||
# 尝试使用 Redis
|
||||
redis_client = get_redis_client_sync()
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning(f"[CacheBackend] Redis 未初始化,{name} 降级为本地缓存")
|
||||
backend = LocalCache(max_size=max_size, default_ttl=ttl)
|
||||
else:
|
||||
backend = RedisCache(redis_client=redis_client, key_prefix=name, default_ttl=ttl)
|
||||
logger.info(f"[CacheBackend] {name} 使用 Redis 缓存")
|
||||
|
||||
elif backend_type == "local":
|
||||
# 强制使用本地缓存
|
||||
backend = LocalCache(max_size=max_size, default_ttl=ttl)
|
||||
logger.info(f"[CacheBackend] {name} 使用本地缓存")
|
||||
|
||||
else: # auto
|
||||
# 自动选择:优先 Redis,降级到 Local
|
||||
redis_client = get_redis_client_sync()
|
||||
|
||||
if redis_client is not None:
|
||||
backend = RedisCache(redis_client=redis_client, key_prefix=name, default_ttl=ttl)
|
||||
logger.debug(f"[CacheBackend] {name} 自动选择 Redis 缓存")
|
||||
else:
|
||||
backend = LocalCache(max_size=max_size, default_ttl=ttl)
|
||||
logger.debug(f"[CacheBackend] {name} 自动选择本地缓存(Redis 不可用)")
|
||||
|
||||
_cache_backends[cache_key] = backend
|
||||
return backend
|
||||
125
src/services/cache/invalidation.py
vendored
Normal file
125
src/services/cache/invalidation.py
vendored
Normal file
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
缓存失效服务
|
||||
|
||||
统一管理各种缓存的失效逻辑,支持:
|
||||
1. GlobalModel 变更时失效相关缓存
|
||||
2. ModelMapping 变更时失效别名/降级缓存
|
||||
3. Model 变更时失效模型映射缓存
|
||||
4. 支持同步和异步缓存后端
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Optional
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
class CacheInvalidationService:
|
||||
"""
|
||||
缓存失效服务
|
||||
|
||||
提供统一的缓存失效接口,当数据库模型变更时自动清理相关缓存
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""初始化缓存失效服务"""
|
||||
self._mapping_resolver = None
|
||||
self._model_mappers = [] # 可能有多个 ModelMapperMiddleware 实例
|
||||
|
||||
def set_mapping_resolver(self, mapping_resolver):
|
||||
"""设置模型映射解析器实例"""
|
||||
self._mapping_resolver = mapping_resolver
|
||||
logger.debug(f"[CacheInvalidation] 模型映射解析器已注册 (实例: {id(mapping_resolver)})")
|
||||
|
||||
def register_model_mapper(self, model_mapper):
|
||||
"""注册 ModelMapper 实例"""
|
||||
if model_mapper not in self._model_mappers:
|
||||
self._model_mappers.append(model_mapper)
|
||||
logger.debug(f"[CacheInvalidation] ModelMapper 已注册 (实例: {id(model_mapper)},总数: {len(self._model_mappers)})")
|
||||
|
||||
def on_global_model_changed(self, model_name: str):
|
||||
"""
|
||||
GlobalModel 变更时的缓存失效
|
||||
|
||||
Args:
|
||||
model_name: 变更的 GlobalModel.name
|
||||
"""
|
||||
logger.info(f"[CacheInvalidation] GlobalModel 变更: {model_name}")
|
||||
|
||||
# 异步失效模型解析器中的缓存
|
||||
if self._mapping_resolver:
|
||||
asyncio.create_task(self._mapping_resolver.invalidate_global_model_cache())
|
||||
|
||||
# 失效所有 ModelMapper 中与此模型相关的缓存
|
||||
for mapper in self._model_mappers:
|
||||
# 清空所有缓存(因为不知道哪些 provider 使用了这个模型)
|
||||
mapper.clear_cache()
|
||||
logger.debug(f"[CacheInvalidation] 已清空 ModelMapper 缓存")
|
||||
|
||||
def on_model_mapping_changed(self, source_model: str, provider_id: Optional[str] = None):
|
||||
"""
|
||||
ModelMapping 变更时的缓存失效
|
||||
|
||||
Args:
|
||||
source_model: 变更的源模型名
|
||||
provider_id: 相关 Provider(None 表示全局)
|
||||
"""
|
||||
logger.info(f"[CacheInvalidation] ModelMapping 变更: {source_model} (provider={provider_id})")
|
||||
|
||||
if self._mapping_resolver:
|
||||
asyncio.create_task(
|
||||
self._mapping_resolver.invalidate_mapping_cache(source_model, provider_id)
|
||||
)
|
||||
|
||||
for mapper in self._model_mappers:
|
||||
if provider_id:
|
||||
mapper.refresh_cache(provider_id)
|
||||
else:
|
||||
mapper.clear_cache()
|
||||
|
||||
def on_model_changed(self, provider_id: str, global_model_id: str):
|
||||
"""
|
||||
Model 变更时的缓存失效
|
||||
|
||||
Args:
|
||||
provider_id: Provider ID
|
||||
global_model_id: GlobalModel ID
|
||||
"""
|
||||
logger.info(f"[CacheInvalidation] Model 变更: provider={provider_id[:8]}..., "
|
||||
f"global_model={global_model_id[:8]}...")
|
||||
|
||||
# 失效 ModelMapper 中特定 Provider 的缓存
|
||||
for mapper in self._model_mappers:
|
||||
mapper.refresh_cache(provider_id)
|
||||
|
||||
def clear_all_caches(self):
|
||||
"""清空所有缓存"""
|
||||
logger.info("[CacheInvalidation] 清空所有缓存")
|
||||
|
||||
if self._mapping_resolver:
|
||||
asyncio.create_task(self._mapping_resolver.clear_cache())
|
||||
|
||||
for mapper in self._model_mappers:
|
||||
mapper.clear_cache()
|
||||
|
||||
|
||||
# 全局单例
|
||||
_cache_invalidation_service: Optional[CacheInvalidationService] = None
|
||||
|
||||
|
||||
def get_cache_invalidation_service() -> CacheInvalidationService:
|
||||
"""
|
||||
获取全局缓存失效服务实例
|
||||
|
||||
Returns:
|
||||
CacheInvalidationService 实例
|
||||
"""
|
||||
global _cache_invalidation_service
|
||||
|
||||
if _cache_invalidation_service is None:
|
||||
_cache_invalidation_service = CacheInvalidationService()
|
||||
logger.debug("[CacheInvalidation] 初始化缓存失效服务")
|
||||
|
||||
return _cache_invalidation_service
|
||||
325
src/services/cache/model_cache.py
vendored
Normal file
325
src/services/cache/model_cache.py
vendored
Normal file
@@ -0,0 +1,325 @@
|
||||
"""
|
||||
Model 映射缓存服务 - 减少模型映射和别名查询
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.constants import CacheTTL
|
||||
from src.core.cache_service import CacheService
|
||||
from src.core.logger import logger
|
||||
from src.models.database import GlobalModel, Model, ModelMapping
|
||||
|
||||
|
||||
|
||||
class ModelCacheService:
|
||||
"""Model 映射缓存服务"""
|
||||
|
||||
# 缓存 TTL(秒)- 使用统一常量
|
||||
CACHE_TTL = CacheTTL.MODEL
|
||||
|
||||
@staticmethod
|
||||
async def get_model_by_id(db: Session, model_id: str) -> Optional[Model]:
|
||||
"""
|
||||
获取 Model(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
model_id: Model ID
|
||||
|
||||
Returns:
|
||||
Model 对象或 None
|
||||
"""
|
||||
cache_key = f"model:id:{model_id}"
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"Model 缓存命中: {model_id}")
|
||||
return ModelCacheService._dict_to_model(cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
model = db.query(Model).filter(Model.id == model_id).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if model:
|
||||
model_dict = ModelCacheService._model_to_dict(model)
|
||||
await CacheService.set(cache_key, model_dict, ttl_seconds=ModelCacheService.CACHE_TTL)
|
||||
logger.debug(f"Model 已缓存: {model_id}")
|
||||
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
async def get_global_model_by_id(db: Session, global_model_id: str) -> Optional[GlobalModel]:
|
||||
"""
|
||||
获取 GlobalModel(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
global_model_id: GlobalModel ID
|
||||
|
||||
Returns:
|
||||
GlobalModel 对象或 None
|
||||
"""
|
||||
cache_key = f"global_model:id:{global_model_id}"
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"GlobalModel 缓存命中: {global_model_id}")
|
||||
return ModelCacheService._dict_to_global_model(cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
global_model = db.query(GlobalModel).filter(GlobalModel.id == global_model_id).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if global_model:
|
||||
global_model_dict = ModelCacheService._global_model_to_dict(global_model)
|
||||
await CacheService.set(
|
||||
cache_key, global_model_dict, ttl_seconds=ModelCacheService.CACHE_TTL
|
||||
)
|
||||
logger.debug(f"GlobalModel 已缓存: {global_model_id}")
|
||||
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
async def get_model_by_provider_and_global_model(
|
||||
db: Session, provider_id: str, global_model_id: str
|
||||
) -> Optional[Model]:
|
||||
"""
|
||||
通过 Provider ID 和 GlobalModel ID 获取 Model(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
provider_id: Provider ID
|
||||
global_model_id: GlobalModel ID
|
||||
|
||||
Returns:
|
||||
Model 对象或 None
|
||||
"""
|
||||
cache_key = f"model:provider_global:{provider_id}:{global_model_id}"
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"Model 缓存命中(provider+global): {provider_id[:8]}...+{global_model_id[:8]}...")
|
||||
return ModelCacheService._dict_to_model(cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
model = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_id,
|
||||
Model.global_model_id == global_model_id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
# 3. 写入缓存
|
||||
if model:
|
||||
model_dict = ModelCacheService._model_to_dict(model)
|
||||
await CacheService.set(cache_key, model_dict, ttl_seconds=ModelCacheService.CACHE_TTL)
|
||||
logger.debug(f"Model 已缓存(provider+global): {provider_id[:8]}...+{global_model_id[:8]}...")
|
||||
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
async def get_global_model_by_name(db: Session, name: str) -> Optional[GlobalModel]:
|
||||
"""
|
||||
通过名称获取 GlobalModel(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
name: GlobalModel 名称
|
||||
|
||||
Returns:
|
||||
GlobalModel 对象或 None
|
||||
"""
|
||||
cache_key = f"global_model:name:{name}"
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"GlobalModel 缓存命中(名称): {name}")
|
||||
return ModelCacheService._dict_to_global_model(cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
global_model = db.query(GlobalModel).filter(GlobalModel.name == name).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if global_model:
|
||||
global_model_dict = ModelCacheService._global_model_to_dict(global_model)
|
||||
await CacheService.set(
|
||||
cache_key, global_model_dict, ttl_seconds=ModelCacheService.CACHE_TTL
|
||||
)
|
||||
logger.debug(f"GlobalModel 已缓存(名称): {name}")
|
||||
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
async def resolve_alias(
|
||||
db: Session, source_model: str, provider_id: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
解析模型别名(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
source_model: 源模型名称或别名
|
||||
provider_id: Provider ID(可选,用于 Provider 特定别名)
|
||||
|
||||
Returns:
|
||||
目标 GlobalModel ID 或 None
|
||||
"""
|
||||
# 构造缓存键
|
||||
if provider_id:
|
||||
cache_key = f"alias:provider:{provider_id}:{source_model}"
|
||||
else:
|
||||
cache_key = f"alias:global:{source_model}"
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_result = await CacheService.get(cache_key)
|
||||
if cached_result:
|
||||
logger.debug(f"别名缓存命中: {source_model} (provider: {provider_id or 'global'})")
|
||||
return cached_result
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
query = db.query(ModelMapping).filter(ModelMapping.source_model == source_model)
|
||||
|
||||
if provider_id:
|
||||
# Provider 特定别名优先
|
||||
query = query.filter(ModelMapping.provider_id == provider_id)
|
||||
else:
|
||||
# 全局别名
|
||||
query = query.filter(ModelMapping.provider_id.is_(None))
|
||||
|
||||
mapping = query.first()
|
||||
|
||||
# 3. 写入缓存
|
||||
target_global_model_id = mapping.target_global_model_id if mapping else None
|
||||
await CacheService.set(
|
||||
cache_key, target_global_model_id, ttl_seconds=ModelCacheService.CACHE_TTL
|
||||
)
|
||||
|
||||
if mapping:
|
||||
logger.debug(f"别名已缓存: {source_model} → {target_global_model_id}")
|
||||
|
||||
return target_global_model_id
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_model_cache(
|
||||
model_id: str, provider_id: Optional[str] = None, global_model_id: Optional[str] = None
|
||||
):
|
||||
"""清除 Model 缓存
|
||||
|
||||
Args:
|
||||
model_id: Model ID
|
||||
provider_id: Provider ID(用于清除 provider_global 缓存)
|
||||
global_model_id: GlobalModel ID(用于清除 provider_global 缓存)
|
||||
"""
|
||||
# 清除 model:id 缓存
|
||||
await CacheService.delete(f"model:id:{model_id}")
|
||||
|
||||
# 清除 provider_global 缓存(如果提供了必要参数)
|
||||
if provider_id and global_model_id:
|
||||
await CacheService.delete(f"model:provider_global:{provider_id}:{global_model_id}")
|
||||
logger.debug(f"Model 缓存已清除: {model_id}, provider_global:{provider_id[:8]}...:{global_model_id[:8]}...")
|
||||
else:
|
||||
logger.debug(f"Model 缓存已清除: {model_id}")
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_global_model_cache(global_model_id: str, name: Optional[str] = None):
|
||||
"""清除 GlobalModel 缓存"""
|
||||
await CacheService.delete(f"global_model:id:{global_model_id}")
|
||||
if name:
|
||||
await CacheService.delete(f"global_model:name:{name}")
|
||||
logger.debug(f"GlobalModel 缓存已清除: {global_model_id}")
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_alias_cache(source_model: str, provider_id: Optional[str] = None):
|
||||
"""清除别名缓存"""
|
||||
if provider_id:
|
||||
cache_key = f"alias:provider:{provider_id}:{source_model}"
|
||||
else:
|
||||
cache_key = f"alias:global:{source_model}"
|
||||
|
||||
await CacheService.delete(cache_key)
|
||||
logger.debug(f"别名缓存已清除: {source_model}")
|
||||
|
||||
@staticmethod
|
||||
def _model_to_dict(model: Model) -> dict:
|
||||
"""将 Model 对象转换为字典"""
|
||||
return {
|
||||
"id": model.id,
|
||||
"provider_id": model.provider_id,
|
||||
"global_model_id": model.global_model_id,
|
||||
"provider_model_name": model.provider_model_name,
|
||||
"is_active": model.is_active,
|
||||
"is_available": model.is_available if hasattr(model, "is_available") else True,
|
||||
"price_per_request": (
|
||||
float(model.price_per_request) if model.price_per_request else None
|
||||
),
|
||||
"tiered_pricing": model.tiered_pricing,
|
||||
"supports_vision": model.supports_vision,
|
||||
"supports_function_calling": model.supports_function_calling,
|
||||
"supports_streaming": model.supports_streaming,
|
||||
"supports_extended_thinking": model.supports_extended_thinking,
|
||||
"config": model.config,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _dict_to_model(model_dict: dict) -> Model:
|
||||
"""从字典重建 Model 对象"""
|
||||
model = Model(
|
||||
id=model_dict["id"],
|
||||
provider_id=model_dict["provider_id"],
|
||||
global_model_id=model_dict["global_model_id"],
|
||||
provider_model_name=model_dict["provider_model_name"],
|
||||
is_active=model_dict["is_active"],
|
||||
is_available=model_dict.get("is_available", True),
|
||||
price_per_request=model_dict.get("price_per_request"),
|
||||
tiered_pricing=model_dict.get("tiered_pricing"),
|
||||
supports_vision=model_dict.get("supports_vision"),
|
||||
supports_function_calling=model_dict.get("supports_function_calling"),
|
||||
supports_streaming=model_dict.get("supports_streaming"),
|
||||
supports_extended_thinking=model_dict.get("supports_extended_thinking"),
|
||||
config=model_dict.get("config"),
|
||||
)
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def _global_model_to_dict(global_model: GlobalModel) -> dict:
|
||||
"""将 GlobalModel 对象转换为字典"""
|
||||
return {
|
||||
"id": global_model.id,
|
||||
"name": global_model.name,
|
||||
"display_name": global_model.display_name,
|
||||
"family": global_model.family,
|
||||
"group_id": global_model.group_id,
|
||||
"supports_vision": global_model.supports_vision,
|
||||
"supports_thinking": global_model.supports_thinking,
|
||||
"context_window": global_model.context_window,
|
||||
"max_output_tokens": global_model.max_output_tokens,
|
||||
"is_active": global_model.is_active,
|
||||
"description": global_model.description,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _dict_to_global_model(global_model_dict: dict) -> GlobalModel:
|
||||
"""从字典重建 GlobalModel 对象"""
|
||||
global_model = GlobalModel(
|
||||
id=global_model_dict["id"],
|
||||
name=global_model_dict["name"],
|
||||
display_name=global_model_dict.get("display_name"),
|
||||
family=global_model_dict.get("family"),
|
||||
group_id=global_model_dict.get("group_id"),
|
||||
supports_vision=global_model_dict.get("supports_vision", False),
|
||||
supports_thinking=global_model_dict.get("supports_thinking", False),
|
||||
context_window=global_model_dict.get("context_window"),
|
||||
max_output_tokens=global_model_dict.get("max_output_tokens"),
|
||||
is_active=global_model_dict.get("is_active", True),
|
||||
description=global_model_dict.get("description"),
|
||||
)
|
||||
return global_model
|
||||
254
src/services/cache/provider_cache.py
vendored
Normal file
254
src/services/cache/provider_cache.py
vendored
Normal file
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
Provider 配置缓存服务 - 减少 Provider/Endpoint/APIKey 查询
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.constants import CacheTTL
|
||||
from src.core.cache_service import CacheKeys, CacheService
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
|
||||
|
||||
|
||||
class ProviderCacheService:
|
||||
"""Provider 配置缓存服务"""
|
||||
|
||||
# 缓存 TTL(秒)- 使用统一常量
|
||||
CACHE_TTL = CacheTTL.PROVIDER
|
||||
|
||||
@staticmethod
|
||||
async def get_provider_by_id(db: Session, provider_id: str) -> Optional[Provider]:
|
||||
"""
|
||||
获取 Provider(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
provider_id: Provider ID
|
||||
|
||||
Returns:
|
||||
Provider 对象或 None
|
||||
"""
|
||||
cache_key = CacheKeys.provider_by_id(provider_id)
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"Provider 缓存命中: {provider_id}")
|
||||
return ProviderCacheService._dict_to_provider(cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if provider:
|
||||
provider_dict = ProviderCacheService._provider_to_dict(provider)
|
||||
await CacheService.set(
|
||||
cache_key, provider_dict, ttl_seconds=ProviderCacheService.CACHE_TTL
|
||||
)
|
||||
logger.debug(f"Provider 已缓存: {provider_id}")
|
||||
|
||||
return provider
|
||||
|
||||
@staticmethod
|
||||
async def get_endpoint_by_id(db: Session, endpoint_id: str) -> Optional[ProviderEndpoint]:
|
||||
"""
|
||||
获取 Endpoint(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
endpoint_id: Endpoint ID
|
||||
|
||||
Returns:
|
||||
ProviderEndpoint 对象或 None
|
||||
"""
|
||||
cache_key = CacheKeys.endpoint_by_id(endpoint_id)
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"Endpoint 缓存命中: {endpoint_id}")
|
||||
return ProviderCacheService._dict_to_endpoint(cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
endpoint = db.query(ProviderEndpoint).filter(ProviderEndpoint.id == endpoint_id).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if endpoint:
|
||||
endpoint_dict = ProviderCacheService._endpoint_to_dict(endpoint)
|
||||
await CacheService.set(
|
||||
cache_key, endpoint_dict, ttl_seconds=ProviderCacheService.CACHE_TTL
|
||||
)
|
||||
logger.debug(f"Endpoint 已缓存: {endpoint_id}")
|
||||
|
||||
return endpoint
|
||||
|
||||
@staticmethod
|
||||
async def get_api_key_by_id(db: Session, api_key_id: str) -> Optional[ProviderAPIKey]:
|
||||
"""
|
||||
获取 API Key(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
api_key_id: API Key ID
|
||||
|
||||
Returns:
|
||||
ProviderAPIKey 对象或 None
|
||||
"""
|
||||
cache_key = CacheKeys.api_key_by_id(api_key_id)
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"API Key 缓存命中: {api_key_id}")
|
||||
return ProviderCacheService._dict_to_api_key(cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
api_key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == api_key_id).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if api_key:
|
||||
api_key_dict = ProviderCacheService._api_key_to_dict(api_key)
|
||||
await CacheService.set(
|
||||
cache_key, api_key_dict, ttl_seconds=ProviderCacheService.CACHE_TTL
|
||||
)
|
||||
logger.debug(f"API Key 已缓存: {api_key_id}")
|
||||
|
||||
return api_key
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_provider_cache(provider_id: str):
|
||||
"""
|
||||
清除 Provider 缓存
|
||||
|
||||
Args:
|
||||
provider_id: Provider ID
|
||||
"""
|
||||
await CacheService.delete(CacheKeys.provider_by_id(provider_id))
|
||||
logger.debug(f"Provider 缓存已清除: {provider_id}")
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_endpoint_cache(endpoint_id: str):
|
||||
"""
|
||||
清除 Endpoint 缓存
|
||||
|
||||
Args:
|
||||
endpoint_id: Endpoint ID
|
||||
"""
|
||||
await CacheService.delete(CacheKeys.endpoint_by_id(endpoint_id))
|
||||
logger.debug(f"Endpoint 缓存已清除: {endpoint_id}")
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_api_key_cache(api_key_id: str):
|
||||
"""
|
||||
清除 API Key 缓存
|
||||
|
||||
Args:
|
||||
api_key_id: API Key ID
|
||||
"""
|
||||
await CacheService.delete(CacheKeys.api_key_by_id(api_key_id))
|
||||
logger.debug(f"API Key 缓存已清除: {api_key_id}")
|
||||
|
||||
@staticmethod
|
||||
def _provider_to_dict(provider: Provider) -> dict:
|
||||
"""将 Provider 对象转换为字典(用于缓存)"""
|
||||
return {
|
||||
"id": provider.id,
|
||||
"name": provider.name,
|
||||
"api_format": provider.api_format,
|
||||
"base_url": provider.base_url,
|
||||
"is_active": provider.is_active,
|
||||
"priority": provider.priority,
|
||||
"rpm_limit": provider.rpm_limit,
|
||||
"rpm_used": provider.rpm_used,
|
||||
"rpm_reset_at": provider.rpm_reset_at.isoformat() if provider.rpm_reset_at else None,
|
||||
"config": provider.config,
|
||||
"description": provider.description,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _dict_to_provider(provider_dict: dict) -> Provider:
|
||||
"""从字典重建 Provider 对象(分离的对象,不在 Session 中)"""
|
||||
from datetime import datetime
|
||||
|
||||
provider = Provider(
|
||||
id=provider_dict["id"],
|
||||
name=provider_dict["name"],
|
||||
api_format=provider_dict["api_format"],
|
||||
base_url=provider_dict.get("base_url"),
|
||||
is_active=provider_dict["is_active"],
|
||||
priority=provider_dict.get("priority", 0),
|
||||
rpm_limit=provider_dict.get("rpm_limit"),
|
||||
rpm_used=provider_dict.get("rpm_used", 0),
|
||||
config=provider_dict.get("config"),
|
||||
description=provider_dict.get("description"),
|
||||
)
|
||||
|
||||
if provider_dict.get("rpm_reset_at"):
|
||||
provider.rpm_reset_at = datetime.fromisoformat(provider_dict["rpm_reset_at"])
|
||||
|
||||
return provider
|
||||
|
||||
@staticmethod
|
||||
def _endpoint_to_dict(endpoint: ProviderEndpoint) -> dict:
|
||||
"""将 Endpoint 对象转换为字典"""
|
||||
return {
|
||||
"id": endpoint.id,
|
||||
"provider_id": endpoint.provider_id,
|
||||
"name": endpoint.name,
|
||||
"base_url": endpoint.base_url,
|
||||
"is_active": endpoint.is_active,
|
||||
"priority": endpoint.priority,
|
||||
"weight": endpoint.weight,
|
||||
"custom_path": endpoint.custom_path,
|
||||
"config": endpoint.config,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _dict_to_endpoint(endpoint_dict: dict) -> ProviderEndpoint:
|
||||
"""从字典重建 Endpoint 对象"""
|
||||
endpoint = ProviderEndpoint(
|
||||
id=endpoint_dict["id"],
|
||||
provider_id=endpoint_dict["provider_id"],
|
||||
name=endpoint_dict["name"],
|
||||
base_url=endpoint_dict["base_url"],
|
||||
is_active=endpoint_dict["is_active"],
|
||||
priority=endpoint_dict.get("priority", 0),
|
||||
weight=endpoint_dict.get("weight", 1.0),
|
||||
custom_path=endpoint_dict.get("custom_path"),
|
||||
config=endpoint_dict.get("config"),
|
||||
)
|
||||
return endpoint
|
||||
|
||||
@staticmethod
|
||||
def _api_key_to_dict(api_key: ProviderAPIKey) -> dict:
|
||||
"""将 API Key 对象转换为字典"""
|
||||
return {
|
||||
"id": api_key.id,
|
||||
"endpoint_id": api_key.endpoint_id,
|
||||
"key_value": api_key.key_value,
|
||||
"is_active": api_key.is_active,
|
||||
"max_rpm": api_key.max_rpm,
|
||||
"current_rpm": api_key.current_rpm,
|
||||
"health_score": api_key.health_score,
|
||||
"circuit_breaker_state": api_key.circuit_breaker_state,
|
||||
"adaptive_concurrency_limit": api_key.adaptive_concurrency_limit,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _dict_to_api_key(api_key_dict: dict) -> ProviderAPIKey:
|
||||
"""从字典重建 API Key 对象"""
|
||||
api_key = ProviderAPIKey(
|
||||
id=api_key_dict["id"],
|
||||
endpoint_id=api_key_dict["endpoint_id"],
|
||||
key_value=api_key_dict["key_value"],
|
||||
is_active=api_key_dict["is_active"],
|
||||
max_rpm=api_key_dict.get("max_rpm"),
|
||||
current_rpm=api_key_dict.get("current_rpm", 0),
|
||||
health_score=api_key_dict.get("health_score", 1.0),
|
||||
circuit_breaker_state=api_key_dict.get("circuit_breaker_state"),
|
||||
adaptive_concurrency_limit=api_key_dict.get("adaptive_concurrency_limit"),
|
||||
)
|
||||
return api_key
|
||||
209
src/services/cache/sync.py
vendored
Normal file
209
src/services/cache/sync.py
vendored
Normal file
@@ -0,0 +1,209 @@
|
||||
"""
|
||||
缓存同步服务(Redis Pub/Sub)
|
||||
|
||||
提供分布式缓存失效同步功能,用于多实例部署场景。
|
||||
当一个实例修改数据并失效本地缓存时,通过 Redis pub/sub 通知其他实例同步失效。
|
||||
|
||||
使用场景:
|
||||
1. 多实例部署时,确保所有实例的缓存一致性
|
||||
2. GlobalModel/ModelMapping 变更时,同步失效所有实例的缓存
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Callable, Dict, Optional
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from src.core.logger import logger
|
||||
|
||||
from src.clients.redis_client import get_redis_client_sync
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
class CacheSyncService:
|
||||
"""
|
||||
缓存同步服务
|
||||
|
||||
通过 Redis pub/sub 实现分布式缓存失效同步
|
||||
"""
|
||||
|
||||
# Redis 频道名称
|
||||
CHANNEL_GLOBAL_MODEL = "cache:invalidate:global_model"
|
||||
CHANNEL_MODEL_MAPPING = "cache:invalidate:model_mapping"
|
||||
CHANNEL_MODEL = "cache:invalidate:model"
|
||||
CHANNEL_CLEAR_ALL = "cache:invalidate:clear_all"
|
||||
|
||||
def __init__(self, redis_client: aioredis.Redis):
|
||||
"""
|
||||
初始化缓存同步服务
|
||||
|
||||
Args:
|
||||
redis_client: Redis 客户端实例
|
||||
"""
|
||||
self._redis = redis_client
|
||||
self._pubsub: Optional[aioredis.client.PubSub] = None
|
||||
self._listener_task: Optional[asyncio.Task] = None
|
||||
self._handlers: Dict[str, Callable] = {}
|
||||
self._running = False
|
||||
|
||||
async def start(self):
|
||||
"""启动缓存同步服务(订阅 Redis 频道)"""
|
||||
if self._running:
|
||||
logger.warning("[CacheSync] 服务已在运行")
|
||||
return
|
||||
|
||||
try:
|
||||
self._pubsub = self._redis.pubsub()
|
||||
|
||||
# 订阅所有缓存失效频道
|
||||
await self._pubsub.subscribe(
|
||||
self.CHANNEL_GLOBAL_MODEL,
|
||||
self.CHANNEL_MODEL_MAPPING,
|
||||
self.CHANNEL_MODEL,
|
||||
self.CHANNEL_CLEAR_ALL,
|
||||
)
|
||||
|
||||
# 启动监听任务
|
||||
self._listener_task = asyncio.create_task(self._listen())
|
||||
self._running = True
|
||||
|
||||
logger.info("[CacheSync] 缓存同步服务已启动,订阅频道: "
|
||||
f"{self.CHANNEL_GLOBAL_MODEL}, {self.CHANNEL_MODEL_MAPPING}, "
|
||||
f"{self.CHANNEL_MODEL}, {self.CHANNEL_CLEAR_ALL}")
|
||||
except Exception as e:
|
||||
logger.error(f"[CacheSync] 启动失败: {e}")
|
||||
raise
|
||||
|
||||
async def stop(self):
|
||||
"""停止缓存同步服务"""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
|
||||
# 取消监听任务
|
||||
if self._listener_task:
|
||||
self._listener_task.cancel()
|
||||
try:
|
||||
await self._listener_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 取消订阅
|
||||
if self._pubsub:
|
||||
await self._pubsub.unsubscribe()
|
||||
await self._pubsub.close()
|
||||
|
||||
logger.info("[CacheSync] 缓存同步服务已停止")
|
||||
|
||||
def register_handler(self, channel: str, handler: Callable):
|
||||
"""
|
||||
注册缓存失效处理器
|
||||
|
||||
Args:
|
||||
channel: Redis 频道名称
|
||||
handler: 处理函数(接收消息数据作为参数)
|
||||
"""
|
||||
self._handlers[channel] = handler
|
||||
logger.debug(f"[CacheSync] 注册处理器: {channel}")
|
||||
|
||||
async def _listen(self):
|
||||
"""监听 Redis pub/sub 消息"""
|
||||
logger.info("[CacheSync] 开始监听缓存失效消息")
|
||||
|
||||
try:
|
||||
async for message in self._pubsub.listen():
|
||||
if message["type"] == "message":
|
||||
channel = message["channel"]
|
||||
data = message["data"]
|
||||
|
||||
# 解析消息
|
||||
try:
|
||||
payload = json.loads(data)
|
||||
logger.debug(f"[CacheSync] 收到消息: {channel} -> {payload}")
|
||||
|
||||
# 调用注册的处理器
|
||||
if channel in self._handlers:
|
||||
handler = self._handlers[channel]
|
||||
await handler(payload)
|
||||
else:
|
||||
logger.warning(f"[CacheSync] 未找到处理器: {channel}")
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"[CacheSync] 消息解析失败: {data}, 错误: {e}")
|
||||
except Exception as e:
|
||||
logger.error(f"[CacheSync] 处理消息失败: {channel}, 错误: {e}")
|
||||
except asyncio.CancelledError:
|
||||
logger.info("[CacheSync] 监听任务已取消")
|
||||
except Exception as e:
|
||||
logger.error(f"[CacheSync] 监听失败: {e}")
|
||||
|
||||
async def publish_global_model_changed(self, model_name: str):
|
||||
"""发布 GlobalModel 变更通知"""
|
||||
await self._publish(self.CHANNEL_GLOBAL_MODEL, {"model_name": model_name})
|
||||
|
||||
async def publish_model_mapping_changed(
|
||||
self, source_model: str, provider_id: Optional[str] = None
|
||||
):
|
||||
"""发布 ModelMapping 变更通知"""
|
||||
await self._publish(
|
||||
self.CHANNEL_MODEL_MAPPING, {"source_model": source_model, "provider_id": provider_id}
|
||||
)
|
||||
|
||||
async def publish_model_changed(self, provider_id: str, global_model_id: str):
|
||||
"""发布 Model 变更通知"""
|
||||
await self._publish(
|
||||
self.CHANNEL_MODEL, {"provider_id": provider_id, "global_model_id": global_model_id}
|
||||
)
|
||||
|
||||
async def publish_clear_all(self):
|
||||
"""发布清空所有缓存通知"""
|
||||
await self._publish(self.CHANNEL_CLEAR_ALL, {})
|
||||
|
||||
async def _publish(self, channel: str, data: dict):
|
||||
"""发布消息到 Redis 频道"""
|
||||
try:
|
||||
message = json.dumps(data)
|
||||
await self._redis.publish(channel, message)
|
||||
logger.debug(f"[CacheSync] 发布消息: {channel} -> {data}")
|
||||
except Exception as e:
|
||||
logger.error(f"[CacheSync] 发布消息失败: {channel}, 错误: {e}")
|
||||
|
||||
|
||||
# 全局单例
|
||||
_cache_sync_service: Optional[CacheSyncService] = None
|
||||
|
||||
|
||||
async def get_cache_sync_service(redis_client: aioredis.Redis = None) -> Optional[CacheSyncService]:
|
||||
"""
|
||||
获取缓存同步服务实例
|
||||
|
||||
Args:
|
||||
redis_client: Redis 客户端实例(首次调用时需要提供)
|
||||
|
||||
Returns:
|
||||
CacheSyncService 实例,如果 Redis 不可用返回 None
|
||||
"""
|
||||
global _cache_sync_service
|
||||
|
||||
if _cache_sync_service is None:
|
||||
if redis_client is None:
|
||||
# 尝试获取全局 Redis 客户端
|
||||
redis_client = get_redis_client_sync()
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning("[CacheSync] Redis 不可用,分布式缓存同步已禁用")
|
||||
return None
|
||||
|
||||
_cache_sync_service = CacheSyncService(redis_client)
|
||||
logger.info("[CacheSync] 缓存同步服务已初始化")
|
||||
|
||||
return _cache_sync_service
|
||||
|
||||
|
||||
async def close_cache_sync_service():
|
||||
"""关闭缓存同步服务"""
|
||||
global _cache_sync_service
|
||||
|
||||
if _cache_sync_service:
|
||||
await _cache_sync_service.stop()
|
||||
_cache_sync_service = None
|
||||
155
src/services/cache/user_cache.py
vendored
Normal file
155
src/services/cache/user_cache.py
vendored
Normal file
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
用户缓存服务 - 减少数据库查询
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.constants import CacheTTL
|
||||
from src.core.cache_service import CacheKeys, CacheService
|
||||
from src.core.logger import logger
|
||||
from src.models.database import User
|
||||
|
||||
|
||||
|
||||
class UserCacheService:
|
||||
"""用户缓存服务"""
|
||||
|
||||
# 缓存 TTL(秒)- 使用统一常量
|
||||
CACHE_TTL = CacheTTL.USER
|
||||
|
||||
@staticmethod
|
||||
async def get_user_by_id(db: Session, user_id: str) -> Optional[User]:
|
||||
"""
|
||||
获取用户(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID
|
||||
|
||||
Returns:
|
||||
User 对象或 None
|
||||
"""
|
||||
cache_key = CacheKeys.user_by_id(user_id)
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"用户缓存命中: {user_id}")
|
||||
# 从缓存数据重建 User 对象
|
||||
return UserCacheService._dict_to_user(db, cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if user:
|
||||
user_dict = UserCacheService._user_to_dict(user)
|
||||
await CacheService.set(cache_key, user_dict, ttl_seconds=UserCacheService.CACHE_TTL)
|
||||
logger.debug(f"用户已缓存: {user_id}")
|
||||
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
async def get_user_by_email(db: Session, email: str) -> Optional[User]:
|
||||
"""
|
||||
通过邮箱获取用户(带缓存)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
email: 用户邮箱
|
||||
|
||||
Returns:
|
||||
User 对象或 None
|
||||
"""
|
||||
cache_key = CacheKeys.user_by_email(email)
|
||||
|
||||
# 1. 尝试从缓存获取
|
||||
cached_data = await CacheService.get(cache_key)
|
||||
if cached_data:
|
||||
logger.debug(f"用户缓存命中(邮箱): {email}")
|
||||
return UserCacheService._dict_to_user(db, cached_data)
|
||||
|
||||
# 2. 缓存未命中,查询数据库
|
||||
user = db.query(User).filter(User.email == email).first()
|
||||
|
||||
# 3. 写入缓存
|
||||
if user:
|
||||
user_dict = UserCacheService._user_to_dict(user)
|
||||
await CacheService.set(cache_key, user_dict, ttl_seconds=UserCacheService.CACHE_TTL)
|
||||
logger.debug(f"用户已缓存(邮箱): {email}")
|
||||
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_user_cache(user_id: str, email: Optional[str] = None):
|
||||
"""
|
||||
清除用户缓存
|
||||
|
||||
Args:
|
||||
user_id: 用户ID
|
||||
email: 用户邮箱(可选)
|
||||
"""
|
||||
# 删除 ID 缓存
|
||||
await CacheService.delete(CacheKeys.user_by_id(user_id))
|
||||
|
||||
# 删除邮箱缓存
|
||||
if email:
|
||||
await CacheService.delete(CacheKeys.user_by_email(email))
|
||||
|
||||
logger.debug(f"用户缓存已清除: {user_id}")
|
||||
|
||||
@staticmethod
|
||||
def _user_to_dict(user: User) -> dict:
|
||||
"""将 User 对象转换为字典(用于缓存)"""
|
||||
return {
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"username": user.username,
|
||||
"role": user.role.value if user.role else None,
|
||||
"is_active": user.is_active,
|
||||
"quota_usd": float(user.quota_usd) if user.quota_usd is not None else None,
|
||||
"used_usd": float(user.used_usd),
|
||||
"created_at": user.created_at.isoformat() if user.created_at else None,
|
||||
"last_login_at": user.last_login_at.isoformat() if user.last_login_at else None,
|
||||
"model_capability_settings": user.model_capability_settings,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _dict_to_user(db: Session, user_dict: dict) -> User:
|
||||
"""
|
||||
从字典重建 User 对象
|
||||
|
||||
注意:这是一个"分离"的对象,不在 Session 中
|
||||
如果需要修改,需要使用 db.merge() 或重新查询
|
||||
"""
|
||||
from datetime import datetime
|
||||
|
||||
from src.models.database import UserRole
|
||||
|
||||
user = User(
|
||||
id=user_dict["id"],
|
||||
email=user_dict["email"],
|
||||
username=user_dict["username"],
|
||||
is_active=user_dict["is_active"],
|
||||
used_usd=user_dict["used_usd"],
|
||||
)
|
||||
|
||||
# 设置可选字段
|
||||
if user_dict.get("role"):
|
||||
user.role = UserRole(user_dict["role"])
|
||||
|
||||
if user_dict.get("quota_usd") is not None:
|
||||
user.quota_usd = user_dict["quota_usd"]
|
||||
|
||||
if user_dict.get("created_at"):
|
||||
user.created_at = datetime.fromisoformat(user_dict["created_at"])
|
||||
|
||||
if user_dict.get("last_login_at"):
|
||||
user.last_login_at = datetime.fromisoformat(user_dict["last_login_at"])
|
||||
|
||||
if user_dict.get("model_capability_settings") is not None:
|
||||
user.model_capability_settings = user_dict["model_capability_settings"]
|
||||
|
||||
return user
|
||||
7
src/services/capability/__init__.py
Normal file
7
src/services/capability/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
能力服务模块
|
||||
"""
|
||||
|
||||
from .resolver import CapabilityResolver
|
||||
|
||||
__all__ = ["CapabilityResolver"]
|
||||
174
src/services/capability/resolver.py
Normal file
174
src/services/capability/resolver.py
Normal file
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
能力需求解析器
|
||||
|
||||
负责从各种来源解析请求的能力需求:
|
||||
1. 用户模型级配置 (User.model_capability_settings)
|
||||
2. 用户 API Key 强制配置 (ApiKey.force_capabilities)
|
||||
3. 请求头 X-Require-Capability(显式声明)
|
||||
4. Adapter 的 detect_capability_requirements(如 Claude 的 anthropic-beta)
|
||||
5. 显式传入 (用于重试升级)
|
||||
"""
|
||||
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
from src.core.key_capabilities import (
|
||||
CAPABILITY_DEFINITIONS,
|
||||
CapabilityConfigMode,
|
||||
get_user_configurable_capabilities,
|
||||
)
|
||||
from src.core.logger import logger
|
||||
|
||||
# Adapter 检测器类型:接受 headers 和可选的 request_body,返回能力需求字典
|
||||
AdapterDetectorType = Callable[[Dict[str, str], Optional[Dict[str, Any]]], Dict[str, bool]]
|
||||
|
||||
|
||||
class CapabilityResolver:
|
||||
"""能力需求解析器"""
|
||||
|
||||
@staticmethod
|
||||
def resolve_requirements(
|
||||
user: Optional[Any] = None,
|
||||
user_api_key: Optional[Any] = None,
|
||||
model_name: Optional[str] = None,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
explicit_requirements: Optional[Dict[str, bool]] = None,
|
||||
adapter_detector: Optional[AdapterDetectorType] = None,
|
||||
) -> Dict[str, bool]:
|
||||
"""
|
||||
解析请求的能力需求
|
||||
|
||||
来源优先级(后者覆盖前者):
|
||||
1. 用户模型级配置 (User.model_capability_settings)
|
||||
2. 用户 API Key 强制配置 (ApiKey.force_capabilities)
|
||||
3. 请求头 X-Require-Capability(显式声明)
|
||||
4. Adapter 的 detect_capability_requirements(如 Claude 的 anthropic-beta)
|
||||
5. 显式传入的 explicit_requirements(用于重试升级)
|
||||
|
||||
Args:
|
||||
user: User 对象
|
||||
user_api_key: 用户 ApiKey 对象
|
||||
model_name: 模型名称(用于查找模型级配置)
|
||||
request_headers: 请求头
|
||||
request_body: 请求体(可选,部分 Adapter 可能需要)
|
||||
explicit_requirements: 显式传入的需求(重试时使用)
|
||||
adapter_detector: Adapter 的能力检测方法
|
||||
|
||||
Returns:
|
||||
能力需求字典,如 {"cache_1h": True, "context_1m": False}
|
||||
"""
|
||||
requirements: Dict[str, bool] = {}
|
||||
|
||||
# 1. 从用户模型级配置获取(仅用户可配置型能力)
|
||||
if user and model_name:
|
||||
model_settings = getattr(user, "model_capability_settings", None) or {}
|
||||
model_caps = model_settings.get(model_name, {})
|
||||
if model_caps:
|
||||
for cap_name, cap_value in model_caps.items():
|
||||
cap_def = CAPABILITY_DEFINITIONS.get(cap_name)
|
||||
if cap_def and cap_def.config_mode == CapabilityConfigMode.USER_CONFIGURABLE:
|
||||
requirements[cap_name] = bool(cap_value)
|
||||
logger.debug(
|
||||
f"[CapabilityResolver] 从用户模型配置获取 {cap_name}={cap_value} "
|
||||
f"(model={model_name})"
|
||||
)
|
||||
|
||||
# 2. 从用户 API Key 强制配置获取(覆盖模型级配置)
|
||||
if user_api_key:
|
||||
force_caps = getattr(user_api_key, "force_capabilities", None) or {}
|
||||
if force_caps:
|
||||
for cap_name, cap_value in force_caps.items():
|
||||
cap_def = CAPABILITY_DEFINITIONS.get(cap_name)
|
||||
if cap_def and cap_def.config_mode == CapabilityConfigMode.USER_CONFIGURABLE:
|
||||
requirements[cap_name] = bool(cap_value)
|
||||
logger.debug(
|
||||
f"[CapabilityResolver] 从 API Key 强制配置获取 {cap_name}={cap_value}"
|
||||
)
|
||||
|
||||
# 3. 从请求头 X-Require-Capability 获取(显式声明)
|
||||
if request_headers:
|
||||
header_caps = request_headers.get("X-Require-Capability", "")
|
||||
if header_caps:
|
||||
for cap in header_caps.split(","):
|
||||
cap = cap.strip()
|
||||
if not cap:
|
||||
continue
|
||||
if cap.startswith("-"):
|
||||
# -cache_1h 表示不需要
|
||||
cap_name = cap[1:]
|
||||
requirements[cap_name] = False
|
||||
else:
|
||||
requirements[cap] = True
|
||||
logger.debug(
|
||||
f"[CapabilityResolver] 从请求头获取 {cap_name if cap.startswith('-') else cap}"
|
||||
)
|
||||
|
||||
# 4. 从 Adapter 的 detect_capability_requirements 获取
|
||||
if adapter_detector and request_headers:
|
||||
detected = adapter_detector(request_headers, request_body)
|
||||
for cap_name, cap_value in detected.items():
|
||||
# 只有尚未设置的能力才从 Adapter 检测
|
||||
if cap_name not in requirements:
|
||||
requirements[cap_name] = cap_value
|
||||
logger.debug(
|
||||
f"[CapabilityResolver] 从 Adapter 检测到 {cap_name}={cap_value}"
|
||||
)
|
||||
|
||||
# 5. 显式覆盖(重试时使用)
|
||||
if explicit_requirements:
|
||||
for cap_name, cap_value in explicit_requirements.items():
|
||||
requirements[cap_name] = cap_value
|
||||
logger.debug(f"[CapabilityResolver] 显式覆盖 {cap_name}={cap_value}")
|
||||
|
||||
return requirements
|
||||
|
||||
@staticmethod
|
||||
def get_default_requirements_for_model(
|
||||
user: Optional[Any] = None,
|
||||
model_name: Optional[str] = None,
|
||||
) -> Dict[str, bool]:
|
||||
"""
|
||||
获取用户对特定模型的默认能力需求
|
||||
|
||||
仅返回用户可配置型能力的配置。
|
||||
|
||||
Args:
|
||||
user: User 对象
|
||||
model_name: 模型名称
|
||||
|
||||
Returns:
|
||||
能力需求字典
|
||||
"""
|
||||
requirements: Dict[str, bool] = {}
|
||||
|
||||
if not user or not model_name:
|
||||
return requirements
|
||||
|
||||
model_settings = getattr(user, "model_capability_settings", None) or {}
|
||||
model_caps = model_settings.get(model_name, {})
|
||||
|
||||
for cap_def in get_user_configurable_capabilities():
|
||||
if cap_def.name in model_caps:
|
||||
requirements[cap_def.name] = bool(model_caps[cap_def.name])
|
||||
|
||||
return requirements
|
||||
|
||||
@staticmethod
|
||||
def merge_requirements(
|
||||
base: Optional[Dict[str, bool]],
|
||||
override: Optional[Dict[str, bool]],
|
||||
) -> Dict[str, bool]:
|
||||
"""
|
||||
合并两个能力需求字典
|
||||
|
||||
Args:
|
||||
base: 基础需求
|
||||
override: 覆盖需求
|
||||
|
||||
Returns:
|
||||
合并后的需求
|
||||
"""
|
||||
result = dict(base or {})
|
||||
if override:
|
||||
result.update(override)
|
||||
return result
|
||||
14
src/services/health/__init__.py
Normal file
14
src/services/health/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
健康监控服务模块
|
||||
|
||||
包含健康监控相关功能:
|
||||
- health_monitor: 健康度监控单例
|
||||
- HealthMonitor: 健康监控类
|
||||
"""
|
||||
|
||||
from .monitor import HealthMonitor, health_monitor
|
||||
|
||||
__all__ = [
|
||||
"health_monitor",
|
||||
"HealthMonitor",
|
||||
]
|
||||
452
src/services/health/endpoint.py
Normal file
452
src/services/health/endpoint.py
Normal file
@@ -0,0 +1,452 @@
|
||||
"""
|
||||
端点健康状态服务
|
||||
|
||||
提供统一的端点健康监控功能,支持:
|
||||
1. 按 API 格式聚合的健康状态
|
||||
2. 基于时间窗口的状态追踪
|
||||
3. 管理员和普通用户的差异化视图
|
||||
4. Redis 缓存优化
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import case, func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint, RequestCandidate
|
||||
|
||||
|
||||
# 缓存配置
|
||||
CACHE_TTL_SECONDS = 30 # 缓存 30 秒
|
||||
CACHE_KEY_PREFIX = "health:endpoint:"
|
||||
|
||||
|
||||
def _get_redis_client():
|
||||
"""获取 Redis 客户端,失败返回 None"""
|
||||
try:
|
||||
from src.clients.redis_client import redis_client
|
||||
return redis_client
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class EndpointHealthService:
|
||||
"""端点健康状态服务"""
|
||||
|
||||
@staticmethod
|
||||
def get_endpoint_health_by_format(
|
||||
db: Session,
|
||||
lookback_hours: int = 6,
|
||||
include_admin_fields: bool = False,
|
||||
use_cache: bool = True,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
获取按 API 格式聚合的端点健康状态
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
lookback_hours: 回溯小时数
|
||||
include_admin_fields: 是否包含管理员字段(provider_count, key_count等)
|
||||
use_cache: 是否使用缓存(仅对普通用户视图有效)
|
||||
|
||||
Returns:
|
||||
按 API 格式聚合的健康状态列表
|
||||
"""
|
||||
# 尝试从缓存获取
|
||||
cache_key = f"{CACHE_KEY_PREFIX}format:{lookback_hours}:{include_admin_fields}"
|
||||
if use_cache:
|
||||
cached = EndpointHealthService._get_from_cache(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
# 查询所有活跃的端点(一次性获取所有需要的数据)
|
||||
endpoints = (
|
||||
db.query(ProviderEndpoint).join(Provider).filter(Provider.is_active.is_(True)).all()
|
||||
)
|
||||
|
||||
# 收集所有 endpoint_ids
|
||||
all_endpoint_ids = [ep.id for ep in endpoints]
|
||||
|
||||
# 批量查询所有密钥
|
||||
all_keys = (
|
||||
db.query(ProviderAPIKey)
|
||||
.filter(ProviderAPIKey.endpoint_id.in_(all_endpoint_ids))
|
||||
.all()
|
||||
) if all_endpoint_ids else []
|
||||
|
||||
# 按 endpoint_id 分组密钥
|
||||
keys_by_endpoint: Dict[str, List[ProviderAPIKey]] = defaultdict(list)
|
||||
for key in all_keys:
|
||||
keys_by_endpoint[key.endpoint_id].append(key)
|
||||
|
||||
# 按 API 格式聚合
|
||||
format_stats = defaultdict(
|
||||
lambda: {
|
||||
"total_endpoints": 0,
|
||||
"total_keys": 0,
|
||||
"active_keys": 0,
|
||||
"health_scores": [],
|
||||
"endpoint_ids": [],
|
||||
"provider_ids": set(),
|
||||
"key_ids": [],
|
||||
}
|
||||
)
|
||||
|
||||
for ep in endpoints:
|
||||
api_format = ep.api_format if ep.api_format else "UNKNOWN"
|
||||
|
||||
# 统计端点数
|
||||
format_stats[api_format]["total_endpoints"] += 1
|
||||
format_stats[api_format]["endpoint_ids"].append(ep.id)
|
||||
format_stats[api_format]["provider_ids"].add(ep.provider_id)
|
||||
|
||||
# 从预加载的密钥中获取
|
||||
keys = keys_by_endpoint.get(ep.id, [])
|
||||
format_stats[api_format]["total_keys"] += len(keys)
|
||||
|
||||
# 统计活跃密钥和健康度
|
||||
if ep.is_active:
|
||||
for key in keys:
|
||||
format_stats[api_format]["key_ids"].append(key.id)
|
||||
if key.is_active and not key.circuit_breaker_open:
|
||||
format_stats[api_format]["active_keys"] += 1
|
||||
health_score = key.health_score if key.health_score is not None else 1.0
|
||||
format_stats[api_format]["health_scores"].append(health_score)
|
||||
|
||||
# 批量生成所有格式的时间线数据
|
||||
all_key_ids = []
|
||||
format_key_mapping: Dict[str, List[str]] = {}
|
||||
for api_format, stats in format_stats.items():
|
||||
key_ids = stats["key_ids"]
|
||||
format_key_mapping[api_format] = key_ids
|
||||
all_key_ids.extend(key_ids)
|
||||
|
||||
# 一次性查询所有时间线数据
|
||||
timeline_data_map = EndpointHealthService._generate_timeline_batch(
|
||||
db, format_key_mapping, now, lookback_hours
|
||||
)
|
||||
|
||||
# 生成结果
|
||||
result = []
|
||||
|
||||
for api_format, stats in format_stats.items():
|
||||
timeline_data = timeline_data_map.get(api_format, {
|
||||
"timeline": ["unknown"] * 100,
|
||||
"time_range_start": None,
|
||||
"time_range_end": None,
|
||||
})
|
||||
timeline = timeline_data["timeline"]
|
||||
time_range_start = timeline_data.get("time_range_start")
|
||||
time_range_end = timeline_data.get("time_range_end")
|
||||
|
||||
# 基于时间线计算实际健康度
|
||||
if timeline:
|
||||
healthy_count = sum(1 for status in timeline if status == "healthy")
|
||||
warning_count = sum(1 for status in timeline if status == "warning")
|
||||
unhealthy_count = sum(1 for status in timeline if status == "unhealthy")
|
||||
known_count = healthy_count + warning_count + unhealthy_count
|
||||
|
||||
if known_count > 0:
|
||||
avg_health = (healthy_count * 1.0 + warning_count * 0.8) / known_count
|
||||
else:
|
||||
if stats["health_scores"]:
|
||||
avg_health = sum(stats["health_scores"]) / len(stats["health_scores"])
|
||||
elif stats["total_keys"] == 0:
|
||||
avg_health = 0.0
|
||||
else:
|
||||
avg_health = 0.1
|
||||
else:
|
||||
avg_health = 0.0
|
||||
|
||||
item = {
|
||||
"api_format": api_format,
|
||||
"display_name": EndpointHealthService._format_display_name(api_format),
|
||||
"health_score": avg_health,
|
||||
"timeline": timeline,
|
||||
"time_range_start": time_range_start.isoformat() if time_range_start else None,
|
||||
"time_range_end": time_range_end.isoformat() if time_range_end else None,
|
||||
}
|
||||
|
||||
if include_admin_fields:
|
||||
item.update(
|
||||
{
|
||||
"total_endpoints": stats["total_endpoints"],
|
||||
"total_keys": stats["total_keys"],
|
||||
"active_keys": stats["active_keys"],
|
||||
"provider_count": len(stats["provider_ids"]),
|
||||
}
|
||||
)
|
||||
|
||||
result.append(item)
|
||||
|
||||
result.sort(key=lambda x: x["health_score"], reverse=True)
|
||||
|
||||
# 写入缓存
|
||||
if use_cache:
|
||||
EndpointHealthService._set_to_cache(cache_key, result)
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _generate_timeline_batch(
|
||||
db: Session,
|
||||
format_key_mapping: Dict[str, List[str]],
|
||||
now: datetime,
|
||||
lookback_hours: int,
|
||||
segments: int = 100,
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
批量生成多个 API 格式的时间线数据(基于 RequestCandidate 表)
|
||||
|
||||
使用 RequestCandidate 表可以:
|
||||
1. 记录所有尝试(包括 fallback 中失败的尝试)
|
||||
2. 准确反映每个 Provider/Key 的真实健康状态
|
||||
3. 失败的请求会显示为红色节点
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
format_key_mapping: API格式 -> key_ids 的映射
|
||||
now: 当前时间
|
||||
lookback_hours: 回溯小时数
|
||||
segments: 时间段数量
|
||||
|
||||
Returns:
|
||||
API格式 -> 时间线数据的映射
|
||||
"""
|
||||
# 收集所有 key_ids
|
||||
all_key_ids = []
|
||||
for key_ids in format_key_mapping.values():
|
||||
all_key_ids.extend(key_ids)
|
||||
|
||||
if not all_key_ids:
|
||||
return {
|
||||
api_format: {
|
||||
"timeline": ["unknown"] * 100,
|
||||
"time_range_start": None,
|
||||
"time_range_end": None,
|
||||
}
|
||||
for api_format in format_key_mapping.keys()
|
||||
}
|
||||
|
||||
# 计算时间范围
|
||||
interval_minutes = (lookback_hours * 60) // segments
|
||||
start_time = now - timedelta(hours=lookback_hours)
|
||||
|
||||
# 使用 RequestCandidate 表查询所有尝试记录
|
||||
# 只统计最终状态:success, failed, skipped
|
||||
final_statuses = ["success", "failed", "skipped"]
|
||||
|
||||
segment_expr = func.floor(
|
||||
func.extract('epoch', RequestCandidate.created_at - start_time) / (interval_minutes * 60)
|
||||
).label('segment_idx')
|
||||
|
||||
candidate_stats = (
|
||||
db.query(
|
||||
RequestCandidate.key_id,
|
||||
segment_expr,
|
||||
func.count(RequestCandidate.id).label('total_count'),
|
||||
func.sum(
|
||||
case(
|
||||
(RequestCandidate.status == "success", 1),
|
||||
else_=0
|
||||
)
|
||||
).label('success_count'),
|
||||
func.sum(
|
||||
case(
|
||||
(RequestCandidate.status == "failed", 1),
|
||||
else_=0
|
||||
)
|
||||
).label('failed_count'),
|
||||
func.min(RequestCandidate.created_at).label('min_time'),
|
||||
func.max(RequestCandidate.created_at).label('max_time'),
|
||||
)
|
||||
.filter(
|
||||
RequestCandidate.key_id.in_(all_key_ids),
|
||||
RequestCandidate.created_at >= start_time,
|
||||
RequestCandidate.created_at <= now,
|
||||
RequestCandidate.status.in_(final_statuses),
|
||||
)
|
||||
.group_by(RequestCandidate.key_id, segment_expr)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 构建 key_id -> api_format 的反向映射
|
||||
key_to_format: Dict[str, str] = {}
|
||||
for api_format, key_ids in format_key_mapping.items():
|
||||
for key_id in key_ids:
|
||||
key_to_format[key_id] = api_format
|
||||
|
||||
# 按 api_format 和 segment 聚合数据
|
||||
format_segment_data: Dict[str, Dict[int, Dict]] = defaultdict(lambda: defaultdict(lambda: {
|
||||
"total": 0,
|
||||
"success": 0,
|
||||
"failed": 0,
|
||||
"min_time": None,
|
||||
"max_time": None,
|
||||
}))
|
||||
|
||||
for row in candidate_stats:
|
||||
key_id = row.key_id
|
||||
segment_idx = int(row.segment_idx) if row.segment_idx is not None else 0
|
||||
api_format = key_to_format.get(key_id)
|
||||
|
||||
if api_format and 0 <= segment_idx < segments:
|
||||
seg_data = format_segment_data[api_format][segment_idx]
|
||||
seg_data["total"] += row.total_count or 0
|
||||
seg_data["success"] += row.success_count or 0
|
||||
seg_data["failed"] += row.failed_count or 0
|
||||
|
||||
if row.min_time:
|
||||
if seg_data["min_time"] is None or row.min_time < seg_data["min_time"]:
|
||||
seg_data["min_time"] = row.min_time
|
||||
if row.max_time:
|
||||
if seg_data["max_time"] is None or row.max_time > seg_data["max_time"]:
|
||||
seg_data["max_time"] = row.max_time
|
||||
|
||||
# 生成各格式的时间线
|
||||
result: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
for api_format in format_key_mapping.keys():
|
||||
timeline = []
|
||||
earliest_time = None
|
||||
latest_time = None
|
||||
|
||||
segment_data = format_segment_data.get(api_format, {})
|
||||
|
||||
for i in range(segments):
|
||||
seg = segment_data.get(i)
|
||||
if not seg or seg["total"] == 0:
|
||||
timeline.append("unknown")
|
||||
else:
|
||||
# 更新时间范围
|
||||
if seg["min_time"]:
|
||||
if earliest_time is None or seg["min_time"] < earliest_time:
|
||||
earliest_time = seg["min_time"]
|
||||
if seg["max_time"]:
|
||||
if latest_time is None or seg["max_time"] > latest_time:
|
||||
latest_time = seg["max_time"]
|
||||
|
||||
# 计算成功率 = success / (success + failed)
|
||||
# skipped 不算失败,不影响成功率
|
||||
actual_completed = seg["success"] + seg["failed"]
|
||||
if actual_completed > 0:
|
||||
success_rate = seg["success"] / actual_completed
|
||||
else:
|
||||
# 只有 skipped,视为健康
|
||||
success_rate = 1.0
|
||||
|
||||
if success_rate >= 0.95:
|
||||
timeline.append("healthy")
|
||||
elif success_rate >= 0.7:
|
||||
timeline.append("warning")
|
||||
else:
|
||||
timeline.append("unhealthy")
|
||||
|
||||
result[api_format] = {
|
||||
"timeline": timeline,
|
||||
"time_range_start": earliest_time,
|
||||
"time_range_end": latest_time if latest_time else now,
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _generate_timeline_from_usage(
|
||||
db: Session,
|
||||
endpoint_ids: List[str],
|
||||
now: datetime,
|
||||
lookback_hours: int,
|
||||
segments: int = 100,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
从真实使用记录生成时间线数据(兼容旧接口,使用批量查询优化)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
endpoint_ids: 端点ID列表
|
||||
now: 当前时间
|
||||
lookback_hours: 回溯小时数
|
||||
segments: 时间段数量
|
||||
|
||||
Returns:
|
||||
包含时间线和时间范围的字典
|
||||
"""
|
||||
if not endpoint_ids:
|
||||
return {
|
||||
"timeline": ["unknown"] * 100,
|
||||
"time_range_start": None,
|
||||
"time_range_end": None,
|
||||
}
|
||||
|
||||
# 先查询该 API 格式下的所有密钥
|
||||
key_ids = [
|
||||
k.id
|
||||
for k in db.query(ProviderAPIKey.id)
|
||||
.filter(ProviderAPIKey.endpoint_id.in_(endpoint_ids))
|
||||
.all()
|
||||
]
|
||||
|
||||
if not key_ids:
|
||||
return {
|
||||
"timeline": ["unknown"] * 100,
|
||||
"time_range_start": None,
|
||||
"time_range_end": None,
|
||||
}
|
||||
|
||||
# 使用批量查询
|
||||
format_key_mapping = {"_single": key_ids}
|
||||
result = EndpointHealthService._generate_timeline_batch(
|
||||
db, format_key_mapping, now, lookback_hours, segments
|
||||
)
|
||||
|
||||
return result.get("_single", {
|
||||
"timeline": ["unknown"] * 100,
|
||||
"time_range_start": None,
|
||||
"time_range_end": None,
|
||||
})
|
||||
|
||||
@staticmethod
|
||||
def _format_display_name(api_format: str) -> str:
|
||||
"""格式化 API 格式的显示名称"""
|
||||
format_names = {
|
||||
"CLAUDE": "Claude API",
|
||||
"CLAUDE_CLI": "Claude CLI",
|
||||
"CLAUDE_COMPATIBLE": "Claude 兼容",
|
||||
"OPENAI": "OpenAI API",
|
||||
"OPENAI_CLI": "OpenAI CLI",
|
||||
"OPENAI_COMPATIBLE": "OpenAI 兼容",
|
||||
}
|
||||
return format_names.get(api_format, api_format)
|
||||
|
||||
@staticmethod
|
||||
def _get_from_cache(key: str) -> Optional[List[Dict[str, Any]]]:
|
||||
"""从 Redis 缓存获取数据"""
|
||||
redis_client = _get_redis_client()
|
||||
if not redis_client:
|
||||
return None
|
||||
|
||||
try:
|
||||
data = redis_client.get(key)
|
||||
if data:
|
||||
return json.loads(data)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to get from cache: {e}")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _set_to_cache(key: str, data: List[Dict[str, Any]]) -> None:
|
||||
"""写入 Redis 缓存"""
|
||||
redis_client = _get_redis_client()
|
||||
if not redis_client:
|
||||
return
|
||||
|
||||
try:
|
||||
redis_client.setex(key, CACHE_TTL_SECONDS, json.dumps(data, default=str))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to set cache: {e}")
|
||||
641
src/services/health/monitor.py
Normal file
641
src/services/health/monitor.py
Normal file
@@ -0,0 +1,641 @@
|
||||
"""
|
||||
健康监控器 - Endpoint 和 Key 的健康度追踪
|
||||
|
||||
功能:
|
||||
1. 基于滑动窗口的错误率计算
|
||||
2. 三态熔断器:关闭 -> 打开 -> 半开 -> 关闭
|
||||
3. 半开状态允许少量请求验证服务恢复
|
||||
4. 提供健康度查询和管理 API
|
||||
"""
|
||||
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from sqlalchemy import case, func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.constants import CircuitBreakerDefaults
|
||||
from src.core.batch_committer import get_batch_committer
|
||||
from src.core.logger import logger
|
||||
from src.core.metrics import health_open_circuits
|
||||
from src.models.database import ProviderAPIKey, ProviderEndpoint
|
||||
|
||||
|
||||
class CircuitState:
|
||||
"""熔断器状态"""
|
||||
|
||||
CLOSED = "closed" # 关闭(正常)
|
||||
OPEN = "open" # 打开(熔断)
|
||||
HALF_OPEN = "half_open" # 半开(验证恢复)
|
||||
|
||||
|
||||
class HealthMonitor:
|
||||
"""健康监控器(滑动窗口 + 半开状态模式)"""
|
||||
|
||||
# === 滑动窗口配置 ===
|
||||
WINDOW_SIZE = int(os.getenv("HEALTH_WINDOW_SIZE", str(CircuitBreakerDefaults.WINDOW_SIZE)))
|
||||
WINDOW_SECONDS = int(
|
||||
os.getenv("HEALTH_WINDOW_SECONDS", str(CircuitBreakerDefaults.WINDOW_SECONDS))
|
||||
)
|
||||
MIN_REQUESTS = int(
|
||||
os.getenv("HEALTH_MIN_REQUESTS", str(CircuitBreakerDefaults.MIN_REQUESTS_FOR_DECISION))
|
||||
)
|
||||
ERROR_RATE_THRESHOLD = float(
|
||||
os.getenv("HEALTH_ERROR_RATE_THRESHOLD", str(CircuitBreakerDefaults.ERROR_RATE_THRESHOLD))
|
||||
)
|
||||
|
||||
# === 半开状态配置 ===
|
||||
HALF_OPEN_DURATION = int(
|
||||
os.getenv(
|
||||
"HEALTH_HALF_OPEN_DURATION", str(CircuitBreakerDefaults.HALF_OPEN_DURATION_SECONDS)
|
||||
)
|
||||
)
|
||||
HALF_OPEN_SUCCESS_THRESHOLD = int(
|
||||
os.getenv(
|
||||
"HEALTH_HALF_OPEN_SUCCESS", str(CircuitBreakerDefaults.HALF_OPEN_SUCCESS_THRESHOLD)
|
||||
)
|
||||
)
|
||||
HALF_OPEN_FAILURE_THRESHOLD = int(
|
||||
os.getenv(
|
||||
"HEALTH_HALF_OPEN_FAILURE", str(CircuitBreakerDefaults.HALF_OPEN_FAILURE_THRESHOLD)
|
||||
)
|
||||
)
|
||||
|
||||
# === 恢复配置 ===
|
||||
INITIAL_RECOVERY_SECONDS = int(
|
||||
os.getenv(
|
||||
"HEALTH_INITIAL_RECOVERY_SECONDS", str(CircuitBreakerDefaults.INITIAL_RECOVERY_SECONDS)
|
||||
)
|
||||
)
|
||||
RECOVERY_BACKOFF = int(
|
||||
os.getenv(
|
||||
"HEALTH_RECOVERY_BACKOFF", str(CircuitBreakerDefaults.RECOVERY_BACKOFF_MULTIPLIER)
|
||||
)
|
||||
)
|
||||
MAX_RECOVERY_SECONDS = int(
|
||||
os.getenv("HEALTH_MAX_RECOVERY_SECONDS", str(CircuitBreakerDefaults.MAX_RECOVERY_SECONDS))
|
||||
)
|
||||
|
||||
# === 兼容旧参数(用于健康度展示)===
|
||||
SUCCESS_INCREMENT = float(
|
||||
os.getenv("HEALTH_SUCCESS_INCREMENT", str(CircuitBreakerDefaults.SUCCESS_INCREMENT))
|
||||
)
|
||||
FAILURE_DECREMENT = float(
|
||||
os.getenv("HEALTH_FAILURE_DECREMENT", str(CircuitBreakerDefaults.FAILURE_DECREMENT))
|
||||
)
|
||||
PROBE_RECOVERY_SCORE = float(
|
||||
os.getenv("HEALTH_PROBE_RECOVERY_SCORE", str(CircuitBreakerDefaults.PROBE_RECOVERY_SCORE))
|
||||
)
|
||||
|
||||
# === 其他配置 ===
|
||||
ALLOW_AUTO_RECOVER = os.getenv("HEALTH_AUTO_RECOVER_ENABLED", "true").lower() == "true"
|
||||
CIRCUIT_HISTORY_LIMIT = int(os.getenv("HEALTH_CIRCUIT_HISTORY_LIMIT", "200"))
|
||||
|
||||
# 进程级别状态缓存
|
||||
_circuit_history: List[Dict[str, Any]] = []
|
||||
_open_circuit_keys: int = 0
|
||||
|
||||
# ==================== 核心方法 ====================
|
||||
|
||||
@classmethod
|
||||
def record_success(
|
||||
cls,
|
||||
db: Session,
|
||||
key_id: Optional[str] = None,
|
||||
response_time_ms: Optional[int] = None,
|
||||
) -> None:
|
||||
"""记录成功请求"""
|
||||
try:
|
||||
if not key_id:
|
||||
return
|
||||
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
return
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
now_ts = now.timestamp()
|
||||
|
||||
# 1. 更新滑动窗口
|
||||
cls._add_to_window(key, now_ts, success=True)
|
||||
|
||||
# 2. 更新健康度(用于展示)
|
||||
new_score = min(float(key.health_score or 0) + cls.SUCCESS_INCREMENT, 1.0)
|
||||
key.health_score = new_score # type: ignore[assignment]
|
||||
|
||||
# 3. 更新统计
|
||||
key.consecutive_failures = 0 # type: ignore[assignment]
|
||||
key.last_failure_at = None # type: ignore[assignment]
|
||||
key.success_count = int(key.success_count or 0) + 1 # type: ignore[assignment]
|
||||
key.request_count = int(key.request_count or 0) + 1 # type: ignore[assignment]
|
||||
if response_time_ms:
|
||||
key.total_response_time_ms = int(key.total_response_time_ms or 0) + response_time_ms # type: ignore[assignment]
|
||||
|
||||
# 4. 处理熔断器状态
|
||||
state = cls._get_circuit_state(key, now)
|
||||
|
||||
if state == CircuitState.HALF_OPEN:
|
||||
# 半开状态:记录成功
|
||||
key.half_open_successes = int(key.half_open_successes or 0) + 1 # type: ignore[assignment]
|
||||
|
||||
if int(key.half_open_successes or 0) >= cls.HALF_OPEN_SUCCESS_THRESHOLD:
|
||||
# 达到成功阈值,关闭熔断器
|
||||
cls._close_circuit(key, now, reason="半开状态验证成功")
|
||||
|
||||
elif state == CircuitState.OPEN:
|
||||
# 打开状态下的成功(探测成功),进入半开状态
|
||||
cls._enter_half_open(key, now)
|
||||
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"记录成功请求失败: {e}")
|
||||
db.rollback()
|
||||
|
||||
@classmethod
|
||||
def record_failure(
|
||||
cls,
|
||||
db: Session,
|
||||
key_id: Optional[str] = None,
|
||||
error_type: Optional[str] = None,
|
||||
) -> None:
|
||||
"""记录失败请求"""
|
||||
try:
|
||||
if not key_id:
|
||||
return
|
||||
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
return
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
now_ts = now.timestamp()
|
||||
|
||||
# 1. 更新滑动窗口
|
||||
cls._add_to_window(key, now_ts, success=False)
|
||||
|
||||
# 2. 更新健康度(用于展示)
|
||||
new_score = max(float(key.health_score or 1) - cls.FAILURE_DECREMENT, 0.0)
|
||||
key.health_score = new_score # type: ignore[assignment]
|
||||
|
||||
# 3. 更新统计
|
||||
key.consecutive_failures = int(key.consecutive_failures or 0) + 1 # type: ignore[assignment]
|
||||
key.last_failure_at = now # type: ignore[assignment]
|
||||
key.error_count = int(key.error_count or 0) + 1 # type: ignore[assignment]
|
||||
key.request_count = int(key.request_count or 0) + 1 # type: ignore[assignment]
|
||||
|
||||
# 4. 处理熔断器状态
|
||||
state = cls._get_circuit_state(key, now)
|
||||
|
||||
if state == CircuitState.HALF_OPEN:
|
||||
# 半开状态:记录失败
|
||||
key.half_open_failures = int(key.half_open_failures or 0) + 1 # type: ignore[assignment]
|
||||
|
||||
if int(key.half_open_failures or 0) >= cls.HALF_OPEN_FAILURE_THRESHOLD:
|
||||
# 达到失败阈值,重新打开熔断器
|
||||
cls._open_circuit(key, now, reason="半开状态验证失败")
|
||||
|
||||
elif state == CircuitState.CLOSED:
|
||||
# 关闭状态:检查是否需要打开熔断器
|
||||
error_rate = cls._calculate_error_rate(key, now_ts)
|
||||
window = key.request_results_window or []
|
||||
|
||||
if len(window) >= cls.MIN_REQUESTS and error_rate >= cls.ERROR_RATE_THRESHOLD:
|
||||
cls._open_circuit(
|
||||
key, now, reason=f"错误率 {error_rate:.0%} 超过阈值 {cls.ERROR_RATE_THRESHOLD:.0%}"
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[WARN] Key 健康度下降: {key_id[:8]}... -> {new_score:.2f} "
|
||||
f"(连续失败 {key.consecutive_failures} 次, error_type={error_type})"
|
||||
)
|
||||
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"记录失败请求失败: {e}")
|
||||
db.rollback()
|
||||
|
||||
# ==================== 滑动窗口方法 ====================
|
||||
|
||||
@classmethod
|
||||
def _add_to_window(cls, key: ProviderAPIKey, now_ts: float, success: bool) -> None:
|
||||
"""添加请求结果到滑动窗口"""
|
||||
window: List[Dict[str, Any]] = key.request_results_window or []
|
||||
|
||||
# 添加新记录
|
||||
window.append({"ts": now_ts, "ok": success})
|
||||
|
||||
# 清理过期记录
|
||||
cutoff_ts = now_ts - cls.WINDOW_SECONDS
|
||||
window = [r for r in window if r["ts"] > cutoff_ts]
|
||||
|
||||
# 限制窗口大小
|
||||
if len(window) > cls.WINDOW_SIZE:
|
||||
window = window[-cls.WINDOW_SIZE :]
|
||||
|
||||
key.request_results_window = window # type: ignore[assignment]
|
||||
|
||||
@classmethod
|
||||
def _calculate_error_rate(cls, key: ProviderAPIKey, now_ts: float) -> float:
|
||||
"""计算滑动窗口内的错误率"""
|
||||
window: List[Dict[str, Any]] = key.request_results_window or []
|
||||
if not window:
|
||||
return 0.0
|
||||
|
||||
# 过滤过期记录
|
||||
cutoff_ts = now_ts - cls.WINDOW_SECONDS
|
||||
valid_records = [r for r in window if r["ts"] > cutoff_ts]
|
||||
|
||||
if not valid_records:
|
||||
return 0.0
|
||||
|
||||
failures = sum(1 for r in valid_records if not r["ok"])
|
||||
return failures / len(valid_records)
|
||||
|
||||
# ==================== 熔断器状态方法 ====================
|
||||
|
||||
@classmethod
|
||||
def _get_circuit_state(cls, key: ProviderAPIKey, now: datetime) -> str:
|
||||
"""获取当前熔断器状态"""
|
||||
if not key.circuit_breaker_open:
|
||||
return CircuitState.CLOSED
|
||||
|
||||
# 检查是否在半开状态
|
||||
if key.half_open_until and now < key.half_open_until:
|
||||
return CircuitState.HALF_OPEN
|
||||
|
||||
# 检查是否到了探测时间(进入半开)
|
||||
if key.next_probe_at and now >= key.next_probe_at:
|
||||
return CircuitState.HALF_OPEN
|
||||
|
||||
return CircuitState.OPEN
|
||||
|
||||
@classmethod
|
||||
def _open_circuit(cls, key: ProviderAPIKey, now: datetime, reason: str) -> None:
|
||||
"""打开熔断器"""
|
||||
was_open = key.circuit_breaker_open
|
||||
|
||||
key.circuit_breaker_open = True # type: ignore[assignment]
|
||||
key.circuit_breaker_open_at = now # type: ignore[assignment]
|
||||
key.half_open_until = None # type: ignore[assignment]
|
||||
key.half_open_successes = 0 # type: ignore[assignment]
|
||||
key.half_open_failures = 0 # type: ignore[assignment]
|
||||
|
||||
# 计算下次探测时间(进入半开状态的时间)
|
||||
consecutive = int(key.consecutive_failures or 0)
|
||||
recovery_seconds = cls._calculate_recovery_seconds(consecutive)
|
||||
key.next_probe_at = now + timedelta(seconds=recovery_seconds) # type: ignore[assignment]
|
||||
|
||||
if not was_open:
|
||||
cls._open_circuit_keys += 1
|
||||
health_open_circuits.set(cls._open_circuit_keys)
|
||||
|
||||
logger.warning(
|
||||
f"[OPEN] Key 熔断器打开: {key.id[:8]}... | 原因: {reason} | "
|
||||
f"{recovery_seconds}秒后进入半开状态"
|
||||
)
|
||||
|
||||
cls._push_circuit_event(
|
||||
{
|
||||
"event": "opened",
|
||||
"key_id": key.id,
|
||||
"reason": reason,
|
||||
"recovery_seconds": recovery_seconds,
|
||||
"timestamp": now.isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _enter_half_open(cls, key: ProviderAPIKey, now: datetime) -> None:
|
||||
"""进入半开状态"""
|
||||
key.half_open_until = now + timedelta(seconds=cls.HALF_OPEN_DURATION) # type: ignore[assignment]
|
||||
key.half_open_successes = 0 # type: ignore[assignment]
|
||||
key.half_open_failures = 0 # type: ignore[assignment]
|
||||
|
||||
logger.info(
|
||||
f"[HALF-OPEN] Key 进入半开状态: {key.id[:8]}... | "
|
||||
f"需要 {cls.HALF_OPEN_SUCCESS_THRESHOLD} 次成功关闭熔断器"
|
||||
)
|
||||
|
||||
cls._push_circuit_event(
|
||||
{
|
||||
"event": "half_open",
|
||||
"key_id": key.id,
|
||||
"timestamp": now.isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _close_circuit(cls, key: ProviderAPIKey, now: datetime, reason: str) -> None:
|
||||
"""关闭熔断器"""
|
||||
key.circuit_breaker_open = False # type: ignore[assignment]
|
||||
key.circuit_breaker_open_at = None # type: ignore[assignment]
|
||||
key.next_probe_at = None # type: ignore[assignment]
|
||||
key.half_open_until = None # type: ignore[assignment]
|
||||
key.half_open_successes = 0 # type: ignore[assignment]
|
||||
key.half_open_failures = 0 # type: ignore[assignment]
|
||||
|
||||
# 快速恢复健康度
|
||||
key.health_score = max(float(key.health_score or 0), cls.PROBE_RECOVERY_SCORE) # type: ignore[assignment]
|
||||
|
||||
cls._open_circuit_keys = max(0, cls._open_circuit_keys - 1)
|
||||
health_open_circuits.set(cls._open_circuit_keys)
|
||||
|
||||
logger.info(f"[CLOSED] Key 熔断器关闭: {key.id[:8]}... | 原因: {reason}")
|
||||
|
||||
cls._push_circuit_event(
|
||||
{
|
||||
"event": "closed",
|
||||
"key_id": key.id,
|
||||
"reason": reason,
|
||||
"timestamp": now.isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _calculate_recovery_seconds(cls, consecutive_failures: int) -> int:
|
||||
"""计算恢复等待时间(指数退避)"""
|
||||
# 指数退避:30s -> 60s -> 120s -> 240s -> 300s(上限)
|
||||
exponent = min(consecutive_failures // 5, 4) # 每5次失败增加一级
|
||||
seconds = cls.INITIAL_RECOVERY_SECONDS * (cls.RECOVERY_BACKOFF**exponent)
|
||||
return min(int(seconds), cls.MAX_RECOVERY_SECONDS)
|
||||
|
||||
# ==================== 状态查询方法 ====================
|
||||
|
||||
@classmethod
|
||||
def is_circuit_breaker_closed(cls, resource: ProviderAPIKey) -> bool:
|
||||
"""检查熔断器是否允许请求通过"""
|
||||
if not resource.circuit_breaker_open:
|
||||
return True
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
state = cls._get_circuit_state(resource, now)
|
||||
|
||||
# 半开状态允许请求通过
|
||||
if state == CircuitState.HALF_OPEN:
|
||||
return True
|
||||
|
||||
# 检查是否到了探测时间
|
||||
if resource.next_probe_at and now >= resource.next_probe_at:
|
||||
# 自动进入半开状态
|
||||
cls._enter_half_open(resource, now)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def get_circuit_breaker_status(
|
||||
cls, resource: ProviderAPIKey
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
"""获取熔断器详细状态"""
|
||||
if not resource.circuit_breaker_open:
|
||||
return True, None
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
state = cls._get_circuit_state(resource, now)
|
||||
|
||||
if state == CircuitState.HALF_OPEN:
|
||||
successes = int(resource.half_open_successes or 0)
|
||||
return True, f"半开状态({successes}/{cls.HALF_OPEN_SUCCESS_THRESHOLD}成功)"
|
||||
|
||||
if resource.next_probe_at:
|
||||
if now >= resource.next_probe_at:
|
||||
return True, None
|
||||
|
||||
remaining = resource.next_probe_at - now
|
||||
remaining_seconds = int(remaining.total_seconds())
|
||||
if remaining_seconds >= 60:
|
||||
time_str = f"{remaining_seconds // 60}min{remaining_seconds % 60}s"
|
||||
else:
|
||||
time_str = f"{remaining_seconds}s"
|
||||
return False, f"熔断中({time_str}后半开)"
|
||||
|
||||
return False, "熔断中"
|
||||
|
||||
@classmethod
|
||||
def get_key_health(cls, db: Session, key_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""获取 Key 健康状态"""
|
||||
try:
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
return None
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
now_ts = now.timestamp()
|
||||
|
||||
# 计算当前错误率
|
||||
error_rate = cls._calculate_error_rate(key, now_ts)
|
||||
window = key.request_results_window or []
|
||||
valid_window = [r for r in window if r["ts"] > now_ts - cls.WINDOW_SECONDS]
|
||||
|
||||
avg_response_time_ms = (
|
||||
int(key.total_response_time_ms or 0) / int(key.success_count or 1)
|
||||
if key.success_count
|
||||
else 0
|
||||
)
|
||||
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"health_score": float(key.health_score or 1.0),
|
||||
"error_rate": error_rate,
|
||||
"window_size": len(valid_window),
|
||||
"consecutive_failures": int(key.consecutive_failures or 0),
|
||||
"last_failure_at": key.last_failure_at.isoformat() if key.last_failure_at else None,
|
||||
"is_active": key.is_active,
|
||||
"statistics": {
|
||||
"request_count": int(key.request_count or 0),
|
||||
"success_count": int(key.success_count or 0),
|
||||
"error_count": int(key.error_count or 0),
|
||||
"success_rate": (
|
||||
int(key.success_count or 0) / int(key.request_count or 1)
|
||||
if key.request_count
|
||||
else 0.0
|
||||
),
|
||||
"avg_response_time_ms": round(avg_response_time_ms, 2),
|
||||
},
|
||||
"circuit_breaker": {
|
||||
"state": cls._get_circuit_state(key, now),
|
||||
"open": key.circuit_breaker_open,
|
||||
"open_at": (
|
||||
key.circuit_breaker_open_at.isoformat()
|
||||
if key.circuit_breaker_open_at
|
||||
else None
|
||||
),
|
||||
"next_probe_at": (
|
||||
key.next_probe_at.isoformat() if key.next_probe_at else None
|
||||
),
|
||||
"half_open_until": (
|
||||
key.half_open_until.isoformat() if key.half_open_until else None
|
||||
),
|
||||
"half_open_successes": int(key.half_open_successes or 0),
|
||||
"half_open_failures": int(key.half_open_failures or 0),
|
||||
},
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取 Key 健康状态失败: {e}")
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_endpoint_health(cls, db: Session, endpoint_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""获取 Endpoint 健康状态"""
|
||||
try:
|
||||
endpoint = (
|
||||
db.query(ProviderEndpoint).filter(ProviderEndpoint.id == endpoint_id).first()
|
||||
)
|
||||
if not endpoint:
|
||||
return None
|
||||
|
||||
return {
|
||||
"endpoint_id": endpoint.id,
|
||||
"health_score": float(endpoint.health_score or 1.0),
|
||||
"consecutive_failures": int(endpoint.consecutive_failures or 0),
|
||||
"last_failure_at": (
|
||||
endpoint.last_failure_at.isoformat() if endpoint.last_failure_at else None
|
||||
),
|
||||
"is_active": endpoint.is_active,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取 Endpoint 健康状态失败: {e}")
|
||||
return None
|
||||
|
||||
# ==================== 管理方法 ====================
|
||||
|
||||
@classmethod
|
||||
def reset_health(cls, db: Session, key_id: Optional[str] = None) -> bool:
|
||||
"""重置健康度"""
|
||||
try:
|
||||
if key_id:
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if key:
|
||||
key.health_score = 1.0 # type: ignore[assignment]
|
||||
key.consecutive_failures = 0 # type: ignore[assignment]
|
||||
key.last_failure_at = None # type: ignore[assignment]
|
||||
key.request_results_window = [] # type: ignore[assignment]
|
||||
key.circuit_breaker_open = False # type: ignore[assignment]
|
||||
key.circuit_breaker_open_at = None # type: ignore[assignment]
|
||||
key.next_probe_at = None # type: ignore[assignment]
|
||||
key.half_open_until = None # type: ignore[assignment]
|
||||
key.half_open_successes = 0 # type: ignore[assignment]
|
||||
key.half_open_failures = 0 # type: ignore[assignment]
|
||||
logger.info(f"[RESET] 重置 Key 健康度: {key_id}")
|
||||
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"重置健康度失败: {e}")
|
||||
db.rollback()
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def manually_enable(cls, db: Session, key_id: Optional[str] = None) -> bool:
|
||||
"""手动启用 Key"""
|
||||
try:
|
||||
if key_id:
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if key and not key.is_active:
|
||||
key.is_active = True # type: ignore[assignment]
|
||||
key.consecutive_failures = 0 # type: ignore[assignment]
|
||||
logger.info(f"[OK] 手动启用 Key: {key_id}")
|
||||
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"手动启用失败: {e}")
|
||||
db.rollback()
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def get_all_health_status(cls, db: Session) -> Dict[str, Any]:
|
||||
"""获取所有健康状态摘要"""
|
||||
try:
|
||||
endpoint_stats = db.query(
|
||||
func.count(ProviderEndpoint.id).label("total"),
|
||||
func.sum(case((ProviderEndpoint.is_active == True, 1), else_=0)).label("active"),
|
||||
func.sum(case((ProviderEndpoint.health_score < 0.5, 1), else_=0)).label(
|
||||
"unhealthy"
|
||||
),
|
||||
).first()
|
||||
|
||||
key_stats = db.query(
|
||||
func.count(ProviderAPIKey.id).label("total"),
|
||||
func.sum(case((ProviderAPIKey.is_active == True, 1), else_=0)).label("active"),
|
||||
func.sum(case((ProviderAPIKey.health_score < 0.5, 1), else_=0)).label("unhealthy"),
|
||||
func.sum(case((ProviderAPIKey.circuit_breaker_open == True, 1), else_=0)).label(
|
||||
"circuit_open"
|
||||
),
|
||||
).first()
|
||||
|
||||
return {
|
||||
"endpoints": {
|
||||
"total": endpoint_stats.total or 0 if endpoint_stats else 0,
|
||||
"active": int(endpoint_stats.active or 0) if endpoint_stats else 0,
|
||||
"unhealthy": int(endpoint_stats.unhealthy or 0) if endpoint_stats else 0,
|
||||
},
|
||||
"keys": {
|
||||
"total": key_stats.total or 0 if key_stats else 0,
|
||||
"active": int(key_stats.active or 0) if key_stats else 0,
|
||||
"unhealthy": int(key_stats.unhealthy or 0) if key_stats else 0,
|
||||
"circuit_open": int(key_stats.circuit_open or 0) if key_stats else 0,
|
||||
},
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取健康状态摘要失败: {e}")
|
||||
return {
|
||||
"endpoints": {"total": 0, "active": 0, "unhealthy": 0},
|
||||
"keys": {"total": 0, "active": 0, "unhealthy": 0, "circuit_open": 0},
|
||||
}
|
||||
|
||||
# ==================== 历史记录方法 ====================
|
||||
|
||||
@classmethod
|
||||
def _push_circuit_event(cls, event: Dict[str, Any]) -> None:
|
||||
cls._circuit_history.append(event)
|
||||
if len(cls._circuit_history) > cls.CIRCUIT_HISTORY_LIMIT:
|
||||
cls._circuit_history.pop(0)
|
||||
|
||||
@classmethod
|
||||
def get_circuit_history(cls, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
return cls._circuit_history[-limit:]
|
||||
|
||||
# ==================== 兼容旧方法 ====================
|
||||
|
||||
@classmethod
|
||||
def is_eligible_for_probe(
|
||||
cls,
|
||||
db: Session,
|
||||
endpoint_id: Optional[str] = None,
|
||||
key_id: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""检查是否有资格进行探测(兼容旧接口)"""
|
||||
if not cls.ALLOW_AUTO_RECOVER:
|
||||
return False
|
||||
|
||||
if endpoint_id:
|
||||
return False # Endpoint 不支持探测
|
||||
|
||||
if key_id:
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if key and key.circuit_breaker_open:
|
||||
now = datetime.now(timezone.utc)
|
||||
state = cls._get_circuit_state(key, now)
|
||||
return state == CircuitState.HALF_OPEN
|
||||
|
||||
return False
|
||||
|
||||
|
||||
# 全局健康监控器实例
|
||||
health_monitor = HealthMonitor()
|
||||
health_open_circuits.set(0)
|
||||
19
src/services/model/__init__.py
Normal file
19
src/services/model/__init__.py
Normal file
@@ -0,0 +1,19 @@
|
||||
"""
|
||||
模型服务模块
|
||||
|
||||
包含模型管理、模型映射、成本计算等功能。
|
||||
"""
|
||||
|
||||
from src.services.model.cost import ModelCostService
|
||||
from src.services.model.global_model import GlobalModelService
|
||||
from src.services.model.mapper import ModelMapperMiddleware
|
||||
from src.services.model.mapping_resolver import ModelMappingResolver
|
||||
from src.services.model.service import ModelService
|
||||
|
||||
__all__ = [
|
||||
"ModelService",
|
||||
"GlobalModelService",
|
||||
"ModelMapperMiddleware",
|
||||
"ModelMappingResolver",
|
||||
"ModelCostService",
|
||||
]
|
||||
946
src/services/model/cost.py
Normal file
946
src/services/model/cost.py
Normal file
@@ -0,0 +1,946 @@
|
||||
"""
|
||||
模型成本服务
|
||||
负责统一的价格解析、缓存以及成本计算逻辑。
|
||||
支持固定价格、按次计费和阶梯计费三种模式。
|
||||
|
||||
计费策略:
|
||||
- 不同 API format 可以有不同的计费逻辑
|
||||
- 通过 PricingStrategy 抽象,支持自定义总输入上下文计算、缓存 TTL 差异化等
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import GlobalModel, Model, ModelMapping, Provider
|
||||
|
||||
|
||||
ProviderRef = Union[str, Provider, None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class TieredPriceResult:
|
||||
"""阶梯计费价格查询结果"""
|
||||
input_price_per_1m: float
|
||||
output_price_per_1m: float
|
||||
cache_creation_price_per_1m: Optional[float] = None
|
||||
cache_read_price_per_1m: Optional[float] = None
|
||||
tier_index: int = 0 # 命中的阶梯索引
|
||||
|
||||
|
||||
@dataclass
|
||||
class CostBreakdown:
|
||||
"""成本明细"""
|
||||
input_cost: float
|
||||
output_cost: float
|
||||
cache_creation_cost: float
|
||||
cache_read_cost: float
|
||||
cache_cost: float
|
||||
request_cost: float
|
||||
total_cost: float
|
||||
|
||||
|
||||
class ModelCostService:
|
||||
"""集中负责模型价格与成本计算,避免在 mapper/usage 中重复实现。"""
|
||||
|
||||
_price_cache: Dict[str, Dict[str, float]] = {}
|
||||
_cache_price_cache: Dict[str, Dict[str, float]] = {}
|
||||
_tiered_pricing_cache: Dict[str, Optional[dict]] = {}
|
||||
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 阶梯计费相关方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def get_tier_for_tokens(
|
||||
tiered_pricing: dict,
|
||||
total_input_tokens: int
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
根据总输入 token 数确定价格阶梯。
|
||||
|
||||
Args:
|
||||
tiered_pricing: 阶梯计费配置 {"tiers": [...]}
|
||||
total_input_tokens: 总输入 token 数(input_tokens + cache_read_tokens)
|
||||
|
||||
Returns:
|
||||
匹配的阶梯配置,如果未找到返回 None
|
||||
"""
|
||||
if not tiered_pricing or "tiers" not in tiered_pricing:
|
||||
return None
|
||||
|
||||
tiers = tiered_pricing.get("tiers", [])
|
||||
if not tiers:
|
||||
return None
|
||||
|
||||
for tier in tiers:
|
||||
up_to = tier.get("up_to")
|
||||
if up_to is None or total_input_tokens <= up_to:
|
||||
return tier
|
||||
|
||||
# 如果所有阶梯都有上限且都超过了,返回最后一个阶梯
|
||||
return tiers[-1] if tiers else None
|
||||
|
||||
@staticmethod
|
||||
def get_cache_read_price_for_ttl(
|
||||
tier: dict,
|
||||
cache_ttl_minutes: Optional[int] = None
|
||||
) -> Optional[float]:
|
||||
"""
|
||||
根据缓存 TTL 获取缓存读取价格。
|
||||
|
||||
Args:
|
||||
tier: 当前阶梯配置
|
||||
cache_ttl_minutes: 缓存时长(分钟),如果为 None 使用默认价格
|
||||
|
||||
Returns:
|
||||
缓存读取价格
|
||||
"""
|
||||
# 首先检查是否有 TTL 差异化定价
|
||||
ttl_pricing = tier.get("cache_ttl_pricing")
|
||||
if ttl_pricing and cache_ttl_minutes is not None:
|
||||
# 找到匹配或最接近的 TTL 价格
|
||||
matched_price = None
|
||||
for ttl_config in ttl_pricing:
|
||||
ttl_limit = ttl_config.get("ttl_minutes", 0)
|
||||
if cache_ttl_minutes <= ttl_limit:
|
||||
matched_price = ttl_config.get("cache_read_price_per_1m")
|
||||
break
|
||||
if matched_price is not None:
|
||||
return matched_price
|
||||
# 如果超过所有配置的 TTL,使用最后一个
|
||||
if ttl_pricing:
|
||||
return ttl_pricing[-1].get("cache_read_price_per_1m")
|
||||
|
||||
# 使用默认的缓存读取价格
|
||||
return tier.get("cache_read_price_per_1m")
|
||||
|
||||
async def get_tiered_pricing_async(
|
||||
self, provider: ProviderRef, model: str
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
异步获取模型的阶梯计费配置。
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 模型名称
|
||||
|
||||
Returns:
|
||||
阶梯计费配置,如果未配置返回 None
|
||||
"""
|
||||
result = await self.get_tiered_pricing_with_source_async(provider, model)
|
||||
return result.get("pricing") if result else None
|
||||
|
||||
async def get_tiered_pricing_with_source_async(
|
||||
self, provider: ProviderRef, model: str
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
异步获取模型的阶梯计费配置及来源信息。
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 模型名称
|
||||
|
||||
Returns:
|
||||
包含 pricing 和 source 的字典:
|
||||
- pricing: 阶梯计费配置
|
||||
- source: 'provider' 或 'global'
|
||||
"""
|
||||
provider_name = self._provider_name(provider)
|
||||
cache_key = f"{provider_name}:{model}:tiered_with_source"
|
||||
|
||||
if cache_key in self._tiered_pricing_cache:
|
||||
return self._tiered_pricing_cache[cache_key]
|
||||
|
||||
provider_obj = self._resolve_provider(provider)
|
||||
result = None
|
||||
|
||||
if provider_obj:
|
||||
from src.services.model.mapping_resolver import resolve_model_to_global_name
|
||||
|
||||
global_model_name = await resolve_model_to_global_name(
|
||||
self.db, model, provider_obj.id
|
||||
)
|
||||
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.name == global_model_name,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if global_model:
|
||||
model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_obj.id,
|
||||
Model.global_model_id == global_model.id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if model_obj:
|
||||
# 判断定价来源
|
||||
if model_obj.tiered_pricing is not None:
|
||||
result = {
|
||||
"pricing": model_obj.tiered_pricing,
|
||||
"source": "provider"
|
||||
}
|
||||
elif global_model.default_tiered_pricing is not None:
|
||||
result = {
|
||||
"pricing": global_model.default_tiered_pricing,
|
||||
"source": "global"
|
||||
}
|
||||
|
||||
self._tiered_pricing_cache[cache_key] = result
|
||||
return result
|
||||
|
||||
def get_tiered_pricing(self, provider: ProviderRef, model: str) -> Optional[dict]:
|
||||
"""同步获取模型的阶梯计费配置。"""
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
return loop.run_until_complete(self.get_tiered_pricing_async(provider, model))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 公共方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_model_price_async(self, provider: ProviderRef, model: str) -> Tuple[float, float]:
|
||||
"""
|
||||
异步版本: 返回给定 provider/model 的 (input_price, output_price)。
|
||||
|
||||
注意:如果模型配置了阶梯计费,此方法返回第一个阶梯的价格作为默认值。
|
||||
实际计费时应使用 compute_cost_with_tiered_pricing 方法。
|
||||
|
||||
计费逻辑(基于 mapping_type):
|
||||
1. 查找 ModelMapping(如果存在)
|
||||
2. 如果 mapping_type='alias':使用目标 GlobalModel 的价格
|
||||
3. 如果 mapping_type='mapping':尝试使用 source_model 对应的 GlobalModel 价格
|
||||
- 如果 source_model 对应的 GlobalModel 存在且有 Model 实现,使用那个价格
|
||||
- 否则回退到目标 GlobalModel 的价格
|
||||
4. 如果没有找到任何 ModelMapping,尝试直接匹配 GlobalModel.name
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 用户请求的模型名(可能是 GlobalModel.name 或别名)
|
||||
|
||||
Returns:
|
||||
(input_price, output_price) 元组
|
||||
"""
|
||||
provider_name = self._provider_name(provider)
|
||||
cache_key = f"{provider_name}:{model}"
|
||||
|
||||
if cache_key in self._price_cache:
|
||||
prices = self._price_cache[cache_key]
|
||||
return prices["input"], prices["output"]
|
||||
|
||||
provider_obj = self._resolve_provider(provider)
|
||||
input_price = None
|
||||
output_price = None
|
||||
|
||||
if provider_obj:
|
||||
# 步骤 1: 查找 ModelMapping 以确定 mapping_type
|
||||
from src.models.database import ModelMapping
|
||||
|
||||
mapping = None
|
||||
# 先查 Provider 特定映射
|
||||
mapping = (
|
||||
self.db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == model,
|
||||
ModelMapping.provider_id == provider_obj.id,
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
# 再查全局映射
|
||||
if not mapping:
|
||||
mapping = (
|
||||
self.db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == model,
|
||||
ModelMapping.provider_id.is_(None),
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if mapping:
|
||||
# 有映射,根据 mapping_type 决定计费模型
|
||||
if mapping.mapping_type == "mapping":
|
||||
# mapping 模式:尝试使用 source_model 对应的 GlobalModel 价格
|
||||
source_global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.name == model,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if source_global_model:
|
||||
source_model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_obj.id,
|
||||
Model.global_model_id == source_global_model.id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if source_model_obj:
|
||||
# 检查是否有阶梯计费
|
||||
tiered = source_model_obj.get_effective_tiered_pricing()
|
||||
if tiered and tiered.get("tiers"):
|
||||
first_tier = tiered["tiers"][0]
|
||||
input_price = first_tier.get("input_price_per_1m", 0)
|
||||
output_price = first_tier.get("output_price_per_1m", 0)
|
||||
else:
|
||||
input_price = source_model_obj.get_effective_input_price()
|
||||
output_price = source_model_obj.get_effective_output_price()
|
||||
logger.debug(f"[mapping模式] 使用源模型价格: {model} "
|
||||
f"(输入: ${input_price}/M, 输出: ${output_price}/M)")
|
||||
|
||||
# alias 模式或 mapping 模式未找到源模型价格:使用目标 GlobalModel 价格
|
||||
if input_price is None:
|
||||
target_global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.id == mapping.target_global_model_id,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if target_global_model:
|
||||
target_model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_obj.id,
|
||||
Model.global_model_id == target_global_model.id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if target_model_obj:
|
||||
# 检查是否有阶梯计费
|
||||
tiered = target_model_obj.get_effective_tiered_pricing()
|
||||
if tiered and tiered.get("tiers"):
|
||||
first_tier = tiered["tiers"][0]
|
||||
input_price = first_tier.get("input_price_per_1m", 0)
|
||||
output_price = first_tier.get("output_price_per_1m", 0)
|
||||
else:
|
||||
input_price = target_model_obj.get_effective_input_price()
|
||||
output_price = target_model_obj.get_effective_output_price()
|
||||
mode_label = (
|
||||
"alias模式"
|
||||
if mapping.mapping_type == "alias"
|
||||
else "mapping模式(回退)"
|
||||
)
|
||||
logger.debug(f"[{mode_label}] 使用目标模型价格: {model} -> {target_global_model.name} "
|
||||
f"(输入: ${input_price}/M, 输出: ${output_price}/M)")
|
||||
else:
|
||||
# 没有映射,尝试直接匹配 GlobalModel.name
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.name == model,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if global_model:
|
||||
model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_obj.id,
|
||||
Model.global_model_id == global_model.id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if model_obj:
|
||||
# 检查是否有阶梯计费
|
||||
tiered = model_obj.get_effective_tiered_pricing()
|
||||
if tiered and tiered.get("tiers"):
|
||||
first_tier = tiered["tiers"][0]
|
||||
input_price = first_tier.get("input_price_per_1m", 0)
|
||||
output_price = first_tier.get("output_price_per_1m", 0)
|
||||
else:
|
||||
input_price = model_obj.get_effective_input_price()
|
||||
output_price = model_obj.get_effective_output_price()
|
||||
logger.debug(f"找到模型价格配置: {provider_name}/{model} "
|
||||
f"(输入: ${input_price}/M, 输出: ${output_price}/M)")
|
||||
|
||||
# 如果没有找到价格配置,使用 0.0 并记录警告
|
||||
if input_price is None:
|
||||
input_price = 0.0
|
||||
if output_price is None:
|
||||
output_price = 0.0
|
||||
|
||||
# 检查是否有按次计费配置(按次计费模型的 token 价格可以为 0)
|
||||
if input_price == 0.0 and output_price == 0.0:
|
||||
# 异步检查按次计费价格
|
||||
price_per_request = await self.get_request_price_async(provider, model)
|
||||
if price_per_request is None or price_per_request == 0.0:
|
||||
logger.warning(f"未找到模型价格配置: {provider_name}/{model},请在 GlobalModel 中配置价格")
|
||||
|
||||
self._price_cache[cache_key] = {"input": input_price, "output": output_price}
|
||||
return input_price, output_price
|
||||
|
||||
def get_model_price(self, provider: ProviderRef, model: str) -> Tuple[float, float]:
|
||||
"""
|
||||
返回给定 provider/model 的 (input_price, output_price)。
|
||||
|
||||
新架构逻辑:
|
||||
1. 使用 ModelMappingResolver 解析别名(如果是)
|
||||
2. 解析为 GlobalModel.name
|
||||
3. 查找该 Provider 的 Model 实现
|
||||
4. 获取价格配置
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 用户请求的模型名(可能是 GlobalModel.name 或别名)
|
||||
|
||||
Returns:
|
||||
(input_price, output_price) 元组
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# 在同步上下文中调用异步方法
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
return loop.run_until_complete(self.get_model_price_async(provider, model))
|
||||
|
||||
async def get_cache_prices_async(
|
||||
self, provider: ProviderRef, model: str, input_price: float
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
"""
|
||||
异步版本: 返回缓存创建/读取价格(每 1M tokens)。
|
||||
|
||||
新架构逻辑:
|
||||
1. 使用 ModelMappingResolver 解析别名(如果是)
|
||||
2. 解析为 GlobalModel.name
|
||||
3. 查找该 Provider 的 Model 实现
|
||||
4. 获取缓存价格配置
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 用户请求的模型名(可能是 GlobalModel.name 或别名)
|
||||
input_price: 基础输入价格(用于 Claude 模型的默认估算)
|
||||
|
||||
Returns:
|
||||
(cache_creation_price, cache_read_price) 元组
|
||||
"""
|
||||
provider_name = self._provider_name(provider)
|
||||
cache_key = f"{provider_name}:{model}"
|
||||
|
||||
if cache_key in self._cache_price_cache:
|
||||
prices = self._cache_price_cache[cache_key]
|
||||
return prices["creation"], prices["read"]
|
||||
|
||||
provider_obj = self._resolve_provider(provider)
|
||||
cache_creation_price = None
|
||||
cache_read_price = None
|
||||
|
||||
if provider_obj:
|
||||
# 步骤 1: 检查是否是别名
|
||||
from src.services.model.mapping_resolver import resolve_model_to_global_name
|
||||
|
||||
global_model_name = await resolve_model_to_global_name(self.db, model, provider_obj.id)
|
||||
|
||||
# 步骤 2: 查找 GlobalModel
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.name == global_model_name,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
# 步骤 3: 查找该 Provider 的 Model 实现
|
||||
if global_model:
|
||||
model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_obj.id,
|
||||
Model.global_model_id == global_model.id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if model_obj:
|
||||
# 检查是否有阶梯计费配置
|
||||
tiered = model_obj.get_effective_tiered_pricing()
|
||||
if tiered and tiered.get("tiers"):
|
||||
# 使用第一个阶梯的缓存价格作为默认值
|
||||
first_tier = tiered["tiers"][0]
|
||||
cache_creation_price = first_tier.get("cache_creation_price_per_1m")
|
||||
cache_read_price = first_tier.get("cache_read_price_per_1m")
|
||||
else:
|
||||
# 使用 get_effective_* 方法,会自动回退到 GlobalModel 的默认值
|
||||
cache_creation_price = model_obj.get_effective_cache_creation_price()
|
||||
cache_read_price = model_obj.get_effective_cache_read_price()
|
||||
|
||||
# 默认缓存价格估算(如果没有配置)- 基于输入价格计算
|
||||
if cache_creation_price is None or cache_read_price is None:
|
||||
if cache_creation_price is None:
|
||||
cache_creation_price = input_price * 1.25
|
||||
if cache_read_price is None:
|
||||
cache_read_price = input_price * 0.1
|
||||
|
||||
self._cache_price_cache[cache_key] = {
|
||||
"creation": cache_creation_price,
|
||||
"read": cache_read_price,
|
||||
}
|
||||
return cache_creation_price, cache_read_price
|
||||
|
||||
async def get_request_price_async(self, provider: ProviderRef, model: str) -> Optional[float]:
|
||||
"""
|
||||
异步版本: 返回按次计费价格(每次请求的固定费用)。
|
||||
|
||||
新架构逻辑:
|
||||
1. 使用 ModelMappingResolver 解析别名(如果是)
|
||||
2. 解析为 GlobalModel.name
|
||||
3. 查找该 Provider 的 Model 实现
|
||||
4. 获取按次计费价格配置
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 用户请求的模型名(可能是 GlobalModel.name 或别名)
|
||||
|
||||
Returns:
|
||||
按次计费价格,如果没有配置则返回 None
|
||||
"""
|
||||
provider_obj = self._resolve_provider(provider)
|
||||
price_per_request = None
|
||||
|
||||
if provider_obj:
|
||||
# 步骤 1: 检查是否是别名
|
||||
from src.services.model.mapping_resolver import resolve_model_to_global_name
|
||||
|
||||
global_model_name = await resolve_model_to_global_name(self.db, model, provider_obj.id)
|
||||
|
||||
# 步骤 2: 查找 GlobalModel
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.name == global_model_name,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
# 步骤 3: 查找该 Provider 的 Model 实现
|
||||
if global_model:
|
||||
model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_obj.id,
|
||||
Model.global_model_id == global_model.id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if model_obj:
|
||||
# 使用 get_effective_* 方法,会自动回退到 GlobalModel 的默认值
|
||||
price_per_request = model_obj.get_effective_price_per_request()
|
||||
|
||||
return price_per_request
|
||||
|
||||
def get_request_price(self, provider: ProviderRef, model: str) -> Optional[float]:
|
||||
"""
|
||||
返回按次计费价格(每次请求的固定费用)。
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 用户请求的模型名(可能是 GlobalModel.name 或别名)
|
||||
|
||||
Returns:
|
||||
按次计费价格,如果没有配置则返回 None
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# 在同步上下文中调用异步方法
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
return loop.run_until_complete(self.get_request_price_async(provider, model))
|
||||
|
||||
def get_cache_prices(
|
||||
self, provider: ProviderRef, model: str, input_price: float
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
"""
|
||||
返回缓存创建/读取价格(每 1M tokens)。
|
||||
|
||||
新架构逻辑:
|
||||
1. 使用 ModelMappingResolver 解析别名(如果是)
|
||||
2. 解析为 GlobalModel.name
|
||||
3. 查找该 Provider 的 Model 实现
|
||||
4. 获取缓存价格配置
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 用户请求的模型名(可能是 GlobalModel.name 或别名)
|
||||
input_price: 基础输入价格(用于 Claude 模型的默认估算)
|
||||
|
||||
Returns:
|
||||
(cache_creation_price, cache_read_price) 元组
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# 在同步上下文中调用异步方法
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
return loop.run_until_complete(self.get_cache_prices_async(provider, model, input_price))
|
||||
|
||||
def calculate_cost(
|
||||
self,
|
||||
provider: Provider,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
) -> Dict[str, float]:
|
||||
"""返回与旧 ModelMapper.calculate_cost 相同结构的费用信息。"""
|
||||
input_price, output_price = self.get_model_price(provider, model)
|
||||
input_cost, output_cost, _, _, _, _, total_cost = self.compute_cost(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
input_price_per_1m=input_price,
|
||||
output_price_per_1m=output_price,
|
||||
)
|
||||
return {
|
||||
"input_cost": round(input_cost, 6),
|
||||
"output_cost": round(output_cost, 6),
|
||||
"total_cost": round(total_cost, 6),
|
||||
"input_price_per_1m": input_price,
|
||||
"output_price_per_1m": output_price,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def compute_cost(
|
||||
*,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
input_price_per_1m: float,
|
||||
output_price_per_1m: float,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
cache_creation_price_per_1m: Optional[float] = None,
|
||||
cache_read_price_per_1m: Optional[float] = None,
|
||||
price_per_request: Optional[float] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float]:
|
||||
"""成本计算核心逻辑(固定价格模式),供 UsageService 等复用。
|
||||
|
||||
Returns:
|
||||
Tuple of (input_cost, output_cost, cache_creation_cost,
|
||||
cache_read_cost, cache_cost, request_cost, total_cost)
|
||||
"""
|
||||
input_cost = (input_tokens / 1_000_000) * input_price_per_1m
|
||||
output_cost = (output_tokens / 1_000_000) * output_price_per_1m
|
||||
|
||||
cache_creation_cost = 0.0
|
||||
cache_read_cost = 0.0
|
||||
if cache_creation_input_tokens > 0 and cache_creation_price_per_1m is not None:
|
||||
cache_creation_cost = (
|
||||
cache_creation_input_tokens / 1_000_000
|
||||
) * cache_creation_price_per_1m
|
||||
if cache_read_input_tokens > 0 and cache_read_price_per_1m is not None:
|
||||
cache_read_cost = (cache_read_input_tokens / 1_000_000) * cache_read_price_per_1m
|
||||
|
||||
cache_cost = cache_creation_cost + cache_read_cost
|
||||
|
||||
# 按次计费成本
|
||||
request_cost = price_per_request if price_per_request is not None else 0.0
|
||||
|
||||
total_cost = input_cost + output_cost + cache_cost + request_cost
|
||||
return (
|
||||
input_cost,
|
||||
output_cost,
|
||||
cache_creation_cost,
|
||||
cache_read_cost,
|
||||
cache_cost,
|
||||
request_cost,
|
||||
total_cost,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def compute_cost_with_tiered_pricing(
|
||||
*,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
tiered_pricing: Optional[dict] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
price_per_request: Optional[float] = None,
|
||||
# 回退价格(当没有阶梯配置时使用)
|
||||
fallback_input_price_per_1m: float = 0.0,
|
||||
fallback_output_price_per_1m: float = 0.0,
|
||||
fallback_cache_creation_price_per_1m: Optional[float] = None,
|
||||
fallback_cache_read_price_per_1m: Optional[float] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float, Optional[int]]:
|
||||
"""
|
||||
支持阶梯计费的成本计算核心逻辑。
|
||||
|
||||
阶梯判定:使用 input_tokens + cache_read_input_tokens(总输入上下文)
|
||||
|
||||
Args:
|
||||
input_tokens: 输入 token 数
|
||||
output_tokens: 输出 token 数
|
||||
cache_creation_input_tokens: 缓存创建 token 数
|
||||
cache_read_input_tokens: 缓存读取 token 数
|
||||
tiered_pricing: 阶梯计费配置
|
||||
cache_ttl_minutes: 缓存时长(分钟),用于 TTL 差异化定价
|
||||
price_per_request: 按次计费价格
|
||||
fallback_*: 回退价格配置
|
||||
|
||||
Returns:
|
||||
Tuple of (input_cost, output_cost, cache_creation_cost,
|
||||
cache_read_cost, cache_cost, request_cost, total_cost, tier_index)
|
||||
tier_index: 命中的阶梯索引(0-based),如果未使用阶梯计费则为 None
|
||||
"""
|
||||
# 计算总输入上下文(用于阶梯判定)
|
||||
total_input_context = input_tokens + cache_read_input_tokens
|
||||
|
||||
tier_index = None
|
||||
input_price_per_1m = fallback_input_price_per_1m
|
||||
output_price_per_1m = fallback_output_price_per_1m
|
||||
cache_creation_price_per_1m = fallback_cache_creation_price_per_1m
|
||||
cache_read_price_per_1m = fallback_cache_read_price_per_1m
|
||||
|
||||
# 如果有阶梯配置,查找匹配的阶梯
|
||||
if tiered_pricing and tiered_pricing.get("tiers"):
|
||||
tier = ModelCostService.get_tier_for_tokens(tiered_pricing, total_input_context)
|
||||
if tier:
|
||||
# 找到阶梯索引
|
||||
tier_index = tiered_pricing["tiers"].index(tier)
|
||||
|
||||
input_price_per_1m = tier.get("input_price_per_1m", fallback_input_price_per_1m)
|
||||
output_price_per_1m = tier.get("output_price_per_1m", fallback_output_price_per_1m)
|
||||
cache_creation_price_per_1m = tier.get(
|
||||
"cache_creation_price_per_1m", fallback_cache_creation_price_per_1m
|
||||
)
|
||||
|
||||
# 获取缓存读取价格(考虑 TTL 差异化)
|
||||
cache_read_price_per_1m = ModelCostService.get_cache_read_price_for_ttl(
|
||||
tier, cache_ttl_minutes
|
||||
)
|
||||
if cache_read_price_per_1m is None:
|
||||
cache_read_price_per_1m = fallback_cache_read_price_per_1m
|
||||
|
||||
logger.debug(
|
||||
f"[阶梯计费] 总输入上下文: {total_input_context}, "
|
||||
f"命中阶梯: {tier_index + 1}, "
|
||||
f"输入价格: ${input_price_per_1m}/M, "
|
||||
f"输出价格: ${output_price_per_1m}/M"
|
||||
)
|
||||
|
||||
# 计算成本
|
||||
input_cost = (input_tokens / 1_000_000) * input_price_per_1m
|
||||
output_cost = (output_tokens / 1_000_000) * output_price_per_1m
|
||||
|
||||
cache_creation_cost = 0.0
|
||||
cache_read_cost = 0.0
|
||||
if cache_creation_input_tokens > 0 and cache_creation_price_per_1m is not None:
|
||||
cache_creation_cost = (
|
||||
cache_creation_input_tokens / 1_000_000
|
||||
) * cache_creation_price_per_1m
|
||||
if cache_read_input_tokens > 0 and cache_read_price_per_1m is not None:
|
||||
cache_read_cost = (cache_read_input_tokens / 1_000_000) * cache_read_price_per_1m
|
||||
|
||||
cache_cost = cache_creation_cost + cache_read_cost
|
||||
|
||||
# 按次计费成本
|
||||
request_cost = price_per_request if price_per_request is not None else 0.0
|
||||
|
||||
total_cost = input_cost + output_cost + cache_cost + request_cost
|
||||
|
||||
return (
|
||||
input_cost,
|
||||
output_cost,
|
||||
cache_creation_cost,
|
||||
cache_read_cost,
|
||||
cache_cost,
|
||||
request_cost,
|
||||
total_cost,
|
||||
tier_index,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def clear_cache(cls):
|
||||
"""清理价格相关缓存。"""
|
||||
cls._price_cache.clear()
|
||||
cls._cache_price_cache.clear()
|
||||
cls._tiered_pricing_cache.clear()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部辅助
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _provider_name(self, provider: ProviderRef) -> str:
|
||||
if isinstance(provider, Provider):
|
||||
return provider.name
|
||||
return provider or "unknown"
|
||||
|
||||
def _resolve_provider(self, provider: ProviderRef) -> Optional[Provider]:
|
||||
if isinstance(provider, Provider):
|
||||
return provider
|
||||
if not provider or provider == "unknown":
|
||||
return None
|
||||
return self.db.query(Provider).filter(Provider.name == provider).first()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 基于策略模式的计费方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def compute_cost_with_strategy_async(
|
||||
self,
|
||||
provider: ProviderRef,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
api_format: Optional[str] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float, Optional[int]]:
|
||||
"""
|
||||
使用计费策略计算成本(异步版本)
|
||||
|
||||
根据 api_format 选择对应的 Adapter 计费逻辑,支持阶梯计费和 TTL 差异化。
|
||||
|
||||
Args:
|
||||
provider: Provider 对象或提供商名称
|
||||
model: 模型名称
|
||||
input_tokens: 输入 token 数
|
||||
output_tokens: 输出 token 数
|
||||
cache_creation_input_tokens: 缓存创建 token 数
|
||||
cache_read_input_tokens: 缓存读取 token 数
|
||||
api_format: API 格式(用于选择计费策略)
|
||||
cache_ttl_minutes: 缓存时长(分钟),用于 TTL 差异化定价
|
||||
|
||||
Returns:
|
||||
Tuple of (input_cost, output_cost, cache_creation_cost,
|
||||
cache_read_cost, cache_cost, request_cost, total_cost, tier_index)
|
||||
"""
|
||||
# 获取价格配置
|
||||
input_price, output_price = await self.get_model_price_async(provider, model)
|
||||
cache_creation_price, cache_read_price = await self.get_cache_prices_async(
|
||||
provider, model, input_price
|
||||
)
|
||||
request_price = await self.get_request_price_async(provider, model)
|
||||
tiered_pricing = await self.get_tiered_pricing_async(provider, model)
|
||||
|
||||
# 获取对应 API 格式的 Adapter 实例来计算成本
|
||||
# 优先检查 Chat Adapter,然后检查 CLI Adapter
|
||||
from src.api.handlers.base.chat_adapter_base import get_adapter_instance
|
||||
from src.api.handlers.base.cli_adapter_base import get_cli_adapter_instance
|
||||
|
||||
adapter = None
|
||||
if api_format:
|
||||
adapter = get_adapter_instance(api_format)
|
||||
if adapter is None:
|
||||
adapter = get_cli_adapter_instance(api_format)
|
||||
|
||||
if adapter:
|
||||
# 使用 Adapter 的计费方法
|
||||
result = adapter.compute_cost(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
input_price_per_1m=input_price,
|
||||
output_price_per_1m=output_price,
|
||||
cache_creation_price_per_1m=cache_creation_price,
|
||||
cache_read_price_per_1m=cache_read_price,
|
||||
price_per_request=request_price,
|
||||
tiered_pricing=tiered_pricing,
|
||||
cache_ttl_minutes=cache_ttl_minutes,
|
||||
)
|
||||
return (
|
||||
result["input_cost"],
|
||||
result["output_cost"],
|
||||
result["cache_creation_cost"],
|
||||
result["cache_read_cost"],
|
||||
result["cache_cost"],
|
||||
result["request_cost"],
|
||||
result["total_cost"],
|
||||
result["tier_index"],
|
||||
)
|
||||
else:
|
||||
# 回退到默认计算逻辑(无 Adapter 时使用静态方法)
|
||||
return self.compute_cost_with_tiered_pricing(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
tiered_pricing=tiered_pricing,
|
||||
cache_ttl_minutes=cache_ttl_minutes,
|
||||
price_per_request=request_price,
|
||||
fallback_input_price_per_1m=input_price,
|
||||
fallback_output_price_per_1m=output_price,
|
||||
fallback_cache_creation_price_per_1m=cache_creation_price,
|
||||
fallback_cache_read_price_per_1m=cache_read_price,
|
||||
)
|
||||
|
||||
def compute_cost_with_strategy(
|
||||
self,
|
||||
provider: ProviderRef,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
api_format: Optional[str] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float, Optional[int]]:
|
||||
"""
|
||||
使用计费策略计算成本(同步版本)
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
return loop.run_until_complete(
|
||||
self.compute_cost_with_strategy_async(
|
||||
provider=provider,
|
||||
model=model,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
api_format=api_format,
|
||||
cache_ttl_minutes=cache_ttl_minutes,
|
||||
)
|
||||
)
|
||||
299
src/services/model/global_model.py
Normal file
299
src/services/model/global_model.py
Normal file
@@ -0,0 +1,299 @@
|
||||
"""
|
||||
GlobalModel 服务层
|
||||
|
||||
提供 GlobalModel 的 CRUD 操作、查询和统计功能
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.models.database import GlobalModel, Model, ModelMapping
|
||||
from src.models.pydantic_models import GlobalModelUpdate
|
||||
|
||||
|
||||
|
||||
class GlobalModelService:
|
||||
"""GlobalModel 服务"""
|
||||
|
||||
@staticmethod
|
||||
def get_global_model(db: Session, global_model_id: str) -> GlobalModel:
|
||||
"""
|
||||
获取单个 GlobalModel
|
||||
|
||||
Args:
|
||||
global_model_id: GlobalModel 的 UUID 或 name
|
||||
"""
|
||||
# 先尝试通过 ID 查找
|
||||
global_model = db.query(GlobalModel).filter(GlobalModel.id == global_model_id).first()
|
||||
|
||||
# 如果没找到,尝试通过 name 查找
|
||||
if not global_model:
|
||||
global_model = db.query(GlobalModel).filter(GlobalModel.name == global_model_id).first()
|
||||
|
||||
if not global_model:
|
||||
raise NotFoundException(f"GlobalModel {global_model_id} not found")
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
def get_global_model_by_name(db: Session, name: str) -> Optional[GlobalModel]:
|
||||
"""通过名称获取 GlobalModel"""
|
||||
return db.query(GlobalModel).filter(GlobalModel.name == name).first()
|
||||
|
||||
@staticmethod
|
||||
def list_global_models(
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: Optional[bool] = None,
|
||||
search: Optional[str] = None,
|
||||
) -> List[GlobalModel]:
|
||||
"""列出 GlobalModel"""
|
||||
query = db.query(GlobalModel)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(GlobalModel.is_active == is_active)
|
||||
|
||||
if search:
|
||||
search_pattern = f"%{search}%"
|
||||
query = query.filter(
|
||||
(GlobalModel.name.ilike(search_pattern))
|
||||
| (GlobalModel.display_name.ilike(search_pattern))
|
||||
| (GlobalModel.description.ilike(search_pattern))
|
||||
)
|
||||
|
||||
# 按名称排序
|
||||
query = query.order_by(GlobalModel.name)
|
||||
|
||||
return query.offset(skip).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def create_global_model(
|
||||
db: Session,
|
||||
name: str,
|
||||
display_name: str,
|
||||
description: Optional[str] = None,
|
||||
official_url: Optional[str] = None,
|
||||
icon_url: Optional[str] = None,
|
||||
is_active: Optional[bool] = True,
|
||||
# 按次计费配置
|
||||
default_price_per_request: Optional[float] = None,
|
||||
# 阶梯计费配置(必填)
|
||||
default_tiered_pricing: dict = None,
|
||||
# 默认能力配置
|
||||
default_supports_vision: Optional[bool] = None,
|
||||
default_supports_function_calling: Optional[bool] = None,
|
||||
default_supports_streaming: Optional[bool] = None,
|
||||
default_supports_extended_thinking: Optional[bool] = None,
|
||||
# Key 能力配置
|
||||
supported_capabilities: Optional[List[str]] = None,
|
||||
) -> GlobalModel:
|
||||
"""创建 GlobalModel"""
|
||||
# 检查名称是否已存在
|
||||
existing = GlobalModelService.get_global_model_by_name(db, name)
|
||||
if existing:
|
||||
raise InvalidRequestException(f"GlobalModel with name '{name}' already exists")
|
||||
|
||||
global_model = GlobalModel(
|
||||
name=name,
|
||||
display_name=display_name,
|
||||
description=description,
|
||||
official_url=official_url,
|
||||
icon_url=icon_url,
|
||||
is_active=is_active,
|
||||
# 按次计费配置
|
||||
default_price_per_request=default_price_per_request,
|
||||
# 阶梯计费配置
|
||||
default_tiered_pricing=default_tiered_pricing,
|
||||
# 默认能力配置
|
||||
default_supports_vision=default_supports_vision,
|
||||
default_supports_function_calling=default_supports_function_calling,
|
||||
default_supports_streaming=default_supports_streaming,
|
||||
default_supports_extended_thinking=default_supports_extended_thinking,
|
||||
# Key 能力配置
|
||||
supported_capabilities=supported_capabilities,
|
||||
)
|
||||
|
||||
db.add(global_model)
|
||||
db.commit()
|
||||
db.refresh(global_model)
|
||||
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
def update_global_model(
|
||||
db: Session,
|
||||
global_model_id: str,
|
||||
update_data: GlobalModelUpdate,
|
||||
) -> GlobalModel:
|
||||
"""
|
||||
更新 GlobalModel
|
||||
|
||||
使用 exclude_unset=True 来区分"未提供字段"和"显式设置为 None":
|
||||
- 未提供的字段不会被更新
|
||||
- 显式设置为 None 的字段会被更新为 None(置空)
|
||||
"""
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
# 只更新显式设置的字段(包括显式设置为 None 的情况)
|
||||
data_dict = update_data.model_dump(exclude_unset=True)
|
||||
|
||||
# 处理阶梯计费配置:如果是 TieredPricingConfig 对象,转换为 dict
|
||||
if "default_tiered_pricing" in data_dict:
|
||||
tiered_pricing = data_dict["default_tiered_pricing"]
|
||||
if tiered_pricing is not None and hasattr(tiered_pricing, "model_dump"):
|
||||
data_dict["default_tiered_pricing"] = tiered_pricing.model_dump()
|
||||
|
||||
for field, value in data_dict.items():
|
||||
setattr(global_model, field, value)
|
||||
|
||||
db.commit()
|
||||
db.refresh(global_model)
|
||||
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
def delete_global_model(db: Session, global_model_id: str) -> None:
|
||||
"""
|
||||
删除 GlobalModel
|
||||
|
||||
默认行为: 级联删除所有关联的 Provider 模型实现
|
||||
"""
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
# 查找所有关联的 Model(使用 global_model.id,预加载 provider 关联)
|
||||
associated_models = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.provider))
|
||||
.filter(Model.global_model_id == global_model.id)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 级联删除所有关联的 Provider 模型实现
|
||||
if associated_models:
|
||||
logger.info(f"删除 GlobalModel {global_model.name} 的 {len(associated_models)} 个关联 Provider 模型")
|
||||
for model in associated_models:
|
||||
db.delete(model)
|
||||
|
||||
# 删除 GlobalModel
|
||||
db.delete(global_model)
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def get_global_model_stats(db: Session, global_model_id: str) -> Dict:
|
||||
"""获取 GlobalModel 统计信息"""
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
# 统计关联的 Model 数量(使用 global_model.id,预加载 provider 关联)
|
||||
models = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.provider))
|
||||
.filter(Model.global_model_id == global_model.id)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 统计支持的 Provider 数量
|
||||
provider_ids = set(model.provider_id for model in models)
|
||||
|
||||
# 从阶梯计费中提取价格范围
|
||||
input_prices = []
|
||||
output_prices = []
|
||||
for m in models:
|
||||
tiered = m.get_effective_tiered_pricing()
|
||||
if tiered and tiered.get("tiers"):
|
||||
first_tier = tiered["tiers"][0]
|
||||
if first_tier.get("input_price_per_1m") is not None:
|
||||
input_prices.append(first_tier["input_price_per_1m"])
|
||||
if first_tier.get("output_price_per_1m") is not None:
|
||||
output_prices.append(first_tier["output_price_per_1m"])
|
||||
|
||||
return {
|
||||
"global_model_id": global_model.id,
|
||||
"name": global_model.name,
|
||||
"total_models": len(models),
|
||||
"total_providers": len(provider_ids),
|
||||
"price_range": {
|
||||
"min_input": min(input_prices) if input_prices else None,
|
||||
"max_input": max(input_prices) if input_prices else None,
|
||||
"min_output": min(output_prices) if output_prices else None,
|
||||
"max_output": max(output_prices) if output_prices else None,
|
||||
},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def batch_assign_to_providers(
|
||||
db: Session,
|
||||
global_model_id: str,
|
||||
provider_ids: List[str],
|
||||
create_models: bool = False,
|
||||
) -> Dict:
|
||||
"""批量为多个 Provider 添加 GlobalModel 实现"""
|
||||
from .service import ModelService
|
||||
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
results = {
|
||||
"success": [],
|
||||
"errors": [],
|
||||
}
|
||||
|
||||
for provider_id in provider_ids:
|
||||
try:
|
||||
# 检查该 Provider 是否已有该 GlobalModel 的实现(使用 global_model.id)
|
||||
existing_model = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_id,
|
||||
Model.global_model_id == global_model.id,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if existing_model:
|
||||
results["errors"].append(
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"error": "Model already exists for this provider",
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if create_models:
|
||||
# 创建新的 Model(价格和能力设为 None,继承 GlobalModel 默认值)
|
||||
model = Model(
|
||||
provider_id=provider_id,
|
||||
global_model_id=global_model.id,
|
||||
provider_model_name=global_model.name, # 默认使用 GlobalModel name
|
||||
# 计费设为 None,使用 GlobalModel 默认值
|
||||
price_per_request=None,
|
||||
tiered_pricing=None,
|
||||
# 能力设为 None,使用 GlobalModel 默认值
|
||||
supports_vision=None,
|
||||
supports_function_calling=None,
|
||||
supports_streaming=None,
|
||||
supports_extended_thinking=None,
|
||||
is_active=True,
|
||||
)
|
||||
db.add(model)
|
||||
db.commit()
|
||||
|
||||
results["success"].append(
|
||||
{"provider_id": provider_id, "model_id": model.id, "created": True}
|
||||
)
|
||||
else:
|
||||
results["errors"].append(
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"error": "create_models=False, no existing model found",
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
results["errors"].append({"provider_id": provider_id, "error": str(e)})
|
||||
|
||||
db.commit()
|
||||
return results
|
||||
442
src/services/model/mapper.py
Normal file
442
src/services/model/mapper.py
Normal file
@@ -0,0 +1,442 @@
|
||||
"""
|
||||
模型映射中间件
|
||||
根据数据库中的配置,将用户请求的模型映射到提供商的实际模型
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.cache_utils import SyncLRUCache
|
||||
from src.core.logger import logger
|
||||
from src.models.claude import ClaudeMessagesRequest
|
||||
from src.models.database import GlobalModel, Model, ModelMapping, Provider, ProviderEndpoint
|
||||
from src.services.cache.model_cache import ModelCacheService
|
||||
from src.services.model.mapping_resolver import (
|
||||
get_model_mapping_resolver,
|
||||
resolve_model_to_global_name,
|
||||
)
|
||||
|
||||
|
||||
|
||||
class ModelMapperMiddleware:
|
||||
"""
|
||||
模型映射中间件
|
||||
负责将用户请求的模型名映射到提供商的实际模型名
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session, cache_max_size: int = 1000, cache_ttl: int = 300):
|
||||
"""
|
||||
初始化模型映射中间件
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
cache_max_size: 缓存最大容量(默认 1000)
|
||||
cache_ttl: 缓存过期时间(秒,默认 300)
|
||||
"""
|
||||
self.db = db
|
||||
self._cache = SyncLRUCache(max_size=cache_max_size, ttl=cache_ttl)
|
||||
|
||||
logger.debug(f"[ModelMapper] 初始化(max_size={cache_max_size}, ttl={cache_ttl}s)")
|
||||
|
||||
# 注册到缓存失效服务
|
||||
try:
|
||||
from src.services.cache.invalidation import get_cache_invalidation_service
|
||||
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.register_model_mapper(self)
|
||||
logger.debug("[ModelMapper] 已注册到缓存失效服务")
|
||||
except Exception as e:
|
||||
logger.warning(f"[ModelMapper] 注册缓存失效服务失败: {e}")
|
||||
|
||||
async def apply_mapping(
|
||||
self, request: ClaudeMessagesRequest, provider: Provider
|
||||
) -> ClaudeMessagesRequest:
|
||||
"""
|
||||
应用模型映射到请求
|
||||
|
||||
Args:
|
||||
request: 原始请求
|
||||
provider: 目标提供商
|
||||
|
||||
Returns:
|
||||
应用映射后的请求
|
||||
"""
|
||||
# 获取请求的模型名
|
||||
source_model = request.model
|
||||
|
||||
# 查找映射
|
||||
mapping = await self.get_mapping(source_model, provider.id)
|
||||
|
||||
if mapping:
|
||||
# 应用映射
|
||||
original_model = request.model
|
||||
request.model = mapping.model.provider_model_name
|
||||
|
||||
logger.debug(f"Applied model mapping for provider {provider.name}: "
|
||||
f"{original_model} -> {mapping.model.provider_model_name}")
|
||||
else:
|
||||
# 没有找到映射,使用原始模型名
|
||||
logger.debug(f"No model mapping found for {source_model} with provider {provider.name}, "
|
||||
f"forwarding with original model name")
|
||||
|
||||
return request
|
||||
|
||||
async def get_mapping(
|
||||
self, source_model: str, provider_id: str
|
||||
) -> Optional[ModelMapping]: # UUID
|
||||
"""
|
||||
获取模型映射
|
||||
|
||||
优化后逻辑:
|
||||
1. 使用统一的 ModelMappingResolver 解析别名(带缓存)
|
||||
2. 通过 GlobalModel 找到该 Provider 的 Model 实现
|
||||
3. 使用独立的映射缓存
|
||||
|
||||
Args:
|
||||
source_model: 用户请求的模型名或别名
|
||||
provider_id: 提供商ID (UUID)
|
||||
|
||||
Returns:
|
||||
模型映射对象(包含 model 字段),如果没有找到返回None
|
||||
"""
|
||||
# 检查缓存
|
||||
cache_key = f"{provider_id}:{source_model}"
|
||||
if cache_key in self._cache:
|
||||
return self._cache[cache_key]
|
||||
|
||||
mapping = None
|
||||
|
||||
# 步骤 1 & 2: 通过统一的模型映射解析服务
|
||||
mapping_resolver = get_model_mapping_resolver()
|
||||
global_model = await mapping_resolver.get_global_model_by_request(
|
||||
self.db, source_model, provider_id
|
||||
)
|
||||
|
||||
if not global_model:
|
||||
logger.debug(f"GlobalModel not found: {source_model} (provider={provider_id[:8]}...)")
|
||||
self._cache[cache_key] = None
|
||||
return None
|
||||
|
||||
# 步骤 3: 查找该 Provider 是否有实现这个 GlobalModel 的 Model(使用缓存)
|
||||
model = await ModelCacheService.get_model_by_provider_and_global_model(
|
||||
self.db, provider_id, global_model.id
|
||||
)
|
||||
|
||||
if model:
|
||||
# 只有当模型名发生变化时才返回映射
|
||||
if model.provider_model_name != source_model:
|
||||
mapping = type(
|
||||
"obj",
|
||||
(object,),
|
||||
{
|
||||
"source_model": source_model,
|
||||
"model": model,
|
||||
"is_active": True,
|
||||
"provider_id": provider_id,
|
||||
},
|
||||
)()
|
||||
|
||||
logger.debug(f"Found model mapping: {source_model} -> {model.provider_model_name} "
|
||||
f"(provider={provider_id[:8]}...)")
|
||||
else:
|
||||
logger.debug(f"Model found but no name change: {source_model} (provider={provider_id[:8]}...)")
|
||||
|
||||
# 缓存结果
|
||||
self._cache[cache_key] = mapping
|
||||
|
||||
return mapping
|
||||
|
||||
def get_all_mappings(self, provider_id: str) -> List[ModelMapping]: # UUID
|
||||
"""
|
||||
获取提供商的所有可用模型(通过 GlobalModel)
|
||||
|
||||
方案 A: 返回该 Provider 所有可用的 GlobalModel
|
||||
|
||||
Args:
|
||||
provider_id: 提供商ID (UUID)
|
||||
|
||||
Returns:
|
||||
模型映射列表(模拟的 ModelMapping 对象列表)
|
||||
"""
|
||||
# 查询该 Provider 的所有活跃 Model
|
||||
models = (
|
||||
self.db.query(Model)
|
||||
.join(GlobalModel)
|
||||
.filter(
|
||||
Model.provider_id == provider_id,
|
||||
Model.is_active == True,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 构造兼容的 ModelMapping 对象列表
|
||||
mappings = []
|
||||
for model in models:
|
||||
mapping = type(
|
||||
"obj",
|
||||
(object,),
|
||||
{
|
||||
"source_model": model.global_model.name,
|
||||
"model": model,
|
||||
"is_active": True,
|
||||
"provider_id": provider_id,
|
||||
},
|
||||
)()
|
||||
mappings.append(mapping)
|
||||
|
||||
return mappings
|
||||
|
||||
def get_supported_models(self, provider_id: str) -> List[str]: # UUID
|
||||
"""
|
||||
获取提供商支持的所有源模型名
|
||||
|
||||
Args:
|
||||
provider_id: 提供商ID (UUID)
|
||||
|
||||
Returns:
|
||||
支持的模型名列表
|
||||
"""
|
||||
mappings = self.get_all_mappings(provider_id)
|
||||
return [mapping.source_model for mapping in mappings]
|
||||
|
||||
async def validate_request(
|
||||
self, request: ClaudeMessagesRequest, provider: Provider
|
||||
) -> tuple[bool, Optional[str]]:
|
||||
"""
|
||||
验证请求是否符合映射的限制
|
||||
|
||||
Args:
|
||||
request: 请求对象
|
||||
provider: 提供商对象
|
||||
|
||||
Returns:
|
||||
(是否有效, 错误信息)
|
||||
"""
|
||||
mapping = await self.get_mapping(request.model, provider.id)
|
||||
|
||||
if not mapping:
|
||||
# 没有映射,可能是默认支持的模型
|
||||
return True, None
|
||||
|
||||
if not mapping.is_active:
|
||||
return False, f"Model mapping for {request.model} is disabled"
|
||||
|
||||
# 不限制max_tokens,作为中转服务不应该限制用户的请求
|
||||
# if request.max_tokens and request.max_tokens > mapping.max_output_tokens:
|
||||
# return False, (
|
||||
# f"Requested max_tokens {request.max_tokens} exceeds limit "
|
||||
# f"{mapping.max_output_tokens} for model {request.model}"
|
||||
# )
|
||||
|
||||
# 可以添加更多验证逻辑,比如检查输入长度等
|
||||
|
||||
return True, None
|
||||
|
||||
def clear_cache(self):
|
||||
"""清空缓存"""
|
||||
self._cache.clear()
|
||||
logger.debug("Model mapping cache cleared")
|
||||
|
||||
def refresh_cache(self, provider_id: Optional[str] = None): # UUID
|
||||
"""
|
||||
刷新缓存
|
||||
|
||||
Args:
|
||||
provider_id: 如果指定,只刷新该提供商的缓存 (UUID)
|
||||
"""
|
||||
if provider_id:
|
||||
# 清除特定提供商的缓存
|
||||
keys_to_remove = [
|
||||
key for key in self._cache.keys() if key.startswith(f"{provider_id}:")
|
||||
]
|
||||
for key in keys_to_remove:
|
||||
del self._cache[key]
|
||||
logger.debug(f"Refreshed cache for provider {provider_id}")
|
||||
else:
|
||||
# 清空所有缓存
|
||||
self.clear_cache()
|
||||
|
||||
|
||||
class ModelRoutingMiddleware:
|
||||
"""
|
||||
模型路由中间件
|
||||
根据模型名选择合适的提供商
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session):
|
||||
"""
|
||||
初始化模型路由中间件
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
"""
|
||||
self.db = db
|
||||
self.mapper = ModelMapperMiddleware(db)
|
||||
|
||||
def select_provider(
|
||||
self,
|
||||
model_name: str,
|
||||
preferred_provider: Optional[str] = None,
|
||||
allowed_api_formats: Optional[List[str]] = None,
|
||||
request_id: Optional[str] = None,
|
||||
) -> Optional[Provider]:
|
||||
"""
|
||||
根据模型名选择提供商
|
||||
|
||||
逻辑:
|
||||
1. 如果指定了提供商,使用指定的提供商
|
||||
2. 如果没指定,使用默认提供商
|
||||
3. 选定提供商后,会检查该提供商的模型映射(在apply_mapping中处理)
|
||||
4. 如果指定了allowed_api_formats,只选择符合格式的提供商
|
||||
|
||||
Args:
|
||||
model_name: 请求的模型名
|
||||
preferred_provider: 首选提供商名称
|
||||
allowed_api_formats: 允许的API格式列表(如 ['CLAUDE', 'CLAUDE_CLI'])
|
||||
request_id: 请求ID(用于日志关联)
|
||||
|
||||
Returns:
|
||||
选中的提供商,如果没有找到返回None
|
||||
"""
|
||||
request_prefix = f"ID:{request_id} | " if request_id else ""
|
||||
|
||||
# 1. 如果指定了提供商,直接使用
|
||||
if preferred_provider:
|
||||
provider = (
|
||||
self.db.query(Provider)
|
||||
.filter(Provider.name == preferred_provider, Provider.is_active == True)
|
||||
.first()
|
||||
)
|
||||
|
||||
if provider:
|
||||
# 检查API格式 - 从 endpoints 中检查
|
||||
if allowed_api_formats:
|
||||
# 检查是否有符合要求的活跃端点
|
||||
has_matching_endpoint = any(
|
||||
ep.is_active and ep.api_format and ep.api_format in allowed_api_formats
|
||||
for ep in provider.endpoints
|
||||
)
|
||||
if not has_matching_endpoint:
|
||||
logger.warning(f"Specified provider {provider.name} has no active endpoints with allowed API formats ({allowed_api_formats})")
|
||||
# 不返回该提供商,继续查找
|
||||
else:
|
||||
logger.debug(f" └─ {request_prefix}使用指定提供商: {provider.name} | 模型:{model_name}")
|
||||
return provider
|
||||
else:
|
||||
logger.debug(f" └─ {request_prefix}使用指定提供商: {provider.name} | 模型:{model_name}")
|
||||
return provider
|
||||
else:
|
||||
logger.warning(f"Specified provider {preferred_provider} not found or inactive")
|
||||
|
||||
# 2. 查找优先级最高的活动提供商(provider_priority 最小)
|
||||
query = self.db.query(Provider).filter(Provider.is_active == True)
|
||||
|
||||
# 如果指定了API格式过滤,添加过滤条件 - 检查是否有符合要求的 endpoint
|
||||
if allowed_api_formats:
|
||||
query = (
|
||||
query.join(ProviderEndpoint)
|
||||
.filter(
|
||||
ProviderEndpoint.is_active == True,
|
||||
ProviderEndpoint.api_format.in_(allowed_api_formats),
|
||||
)
|
||||
.distinct()
|
||||
)
|
||||
|
||||
# 按 provider_priority 排序,优先级最高(数字最小)的在前
|
||||
best_provider = query.order_by(Provider.provider_priority.asc(), Provider.id.asc()).first()
|
||||
|
||||
if best_provider:
|
||||
logger.debug(f" └─ {request_prefix}使用优先级最高提供商: {best_provider.name} (priority:{best_provider.provider_priority}) | 模型:{model_name}")
|
||||
return best_provider
|
||||
|
||||
# 3. 没有任何活动提供商
|
||||
if allowed_api_formats:
|
||||
logger.error(f"No active providers found with allowed API formats {allowed_api_formats}. Please configure at least one provider.")
|
||||
else:
|
||||
logger.error("No active providers found. Please configure at least one provider.")
|
||||
return None
|
||||
|
||||
def get_available_models(self) -> Dict[str, List[str]]:
|
||||
"""
|
||||
获取所有可用的模型及其提供商
|
||||
|
||||
方案 A: 基于 GlobalModel 查询
|
||||
|
||||
Returns:
|
||||
字典,键为 GlobalModel.name,值为支持该模型的提供商名列表
|
||||
"""
|
||||
result = {}
|
||||
|
||||
# 查询所有活跃的 GlobalModel 及其 Provider
|
||||
models = (
|
||||
self.db.query(GlobalModel.name, Provider.name)
|
||||
.join(Model, GlobalModel.id == Model.global_model_id)
|
||||
.join(Provider, Model.provider_id == Provider.id)
|
||||
.filter(
|
||||
GlobalModel.is_active == True, Model.is_active == True, Provider.is_active == True
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
for global_model_name, provider_name in models:
|
||||
if global_model_name not in result:
|
||||
result[global_model_name] = []
|
||||
if provider_name not in result[global_model_name]:
|
||||
result[global_model_name].append(provider_name)
|
||||
|
||||
return result
|
||||
|
||||
async def get_cheapest_provider(self, model_name: str) -> Optional[Provider]:
|
||||
"""
|
||||
获取某个模型最便宜的提供商
|
||||
|
||||
方案 A: 通过 GlobalModel 查找
|
||||
|
||||
Args:
|
||||
model_name: GlobalModel 名称或别名
|
||||
|
||||
Returns:
|
||||
最便宜的提供商
|
||||
"""
|
||||
# 步骤 1: 解析模型名
|
||||
global_model_name = await resolve_model_to_global_name(self.db, model_name)
|
||||
|
||||
# 步骤 2: 查找 GlobalModel
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(GlobalModel.name == global_model_name, GlobalModel.is_active == True)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not global_model:
|
||||
return None
|
||||
|
||||
# 步骤 3: 查询所有支持该模型的 Provider 及其价格
|
||||
models_with_providers = (
|
||||
self.db.query(Provider, Model)
|
||||
.join(Model, Provider.id == Model.provider_id)
|
||||
.filter(
|
||||
Model.global_model_id == global_model.id,
|
||||
Model.is_active == True,
|
||||
Provider.is_active == True,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not models_with_providers:
|
||||
return None
|
||||
|
||||
# 按总价格排序(输入+输出价格)
|
||||
cheapest = min(
|
||||
models_with_providers, key=lambda x: x[1].input_price_per_1m + x[1].output_price_per_1m
|
||||
)
|
||||
|
||||
provider = cheapest[0]
|
||||
model = cheapest[1]
|
||||
|
||||
logger.debug(f"Selected cheapest provider {provider.name} for model {model_name} "
|
||||
f"(input: ${model.input_price_per_1m}/M, output: ${model.output_price_per_1m}/M)")
|
||||
|
||||
return provider
|
||||
432
src/services/model/mapping_resolver.py
Normal file
432
src/services/model/mapping_resolver.py
Normal file
@@ -0,0 +1,432 @@
|
||||
"""
|
||||
模型映射解析服务
|
||||
|
||||
负责统一的模型别名/降级解析,按优先级顺序:
|
||||
1. 映射(mapping):Provider 特定 → 全局
|
||||
2. 别名(alias):Provider 特定 → 全局
|
||||
3. 直接匹配 GlobalModel.name
|
||||
|
||||
支持特性:
|
||||
- 带缓存(本地或 Redis),减少数据库访问
|
||||
- 提供模糊匹配能力,用于提示相似模型
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from src.core.logger import logger
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.constants import CacheSize, CacheTTL
|
||||
from src.core.logger import logger
|
||||
from src.models.database import GlobalModel, ModelMapping
|
||||
from src.services.cache.backend import BaseCacheBackend, get_cache_backend
|
||||
|
||||
|
||||
class ModelMappingResolver:
|
||||
"""统一的 ModelMapping 解析服务(可跨进程共享缓存)。"""
|
||||
|
||||
def __init__(self, cache_ttl: int = CacheTTL.MODEL_MAPPING, cache_backend_type: str = "auto"):
|
||||
self._cache_ttl = cache_ttl
|
||||
self._cache_backend_type = cache_backend_type
|
||||
self._mapping_cache: Optional[BaseCacheBackend] = None
|
||||
self._global_model_cache: Optional[BaseCacheBackend] = None
|
||||
self._initialized = False
|
||||
self._stats = {
|
||||
"mapping_hits": 0,
|
||||
"mapping_misses": 0,
|
||||
"global_hits": 0,
|
||||
"global_misses": 0,
|
||||
}
|
||||
|
||||
async def _ensure_initialized(self):
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
self._mapping_cache = await get_cache_backend(
|
||||
name="model_mapping_resolver:mapping",
|
||||
backend_type=self._cache_backend_type,
|
||||
max_size=CacheSize.MODEL_MAPPING,
|
||||
ttl=self._cache_ttl,
|
||||
)
|
||||
self._global_model_cache = await get_cache_backend(
|
||||
name="model_mapping_resolver:global",
|
||||
backend_type=self._cache_backend_type,
|
||||
max_size=CacheSize.MODEL_MAPPING,
|
||||
ttl=self._cache_ttl,
|
||||
)
|
||||
self._initialized = True
|
||||
logger.debug(f"[ModelMappingResolver] 缓存后端已初始化: {self._mapping_cache.get_stats()['backend']}")
|
||||
|
||||
def _cache_key(self, source_model: str, provider_id: Optional[str]) -> str:
|
||||
return f"{provider_id or 'global'}:{source_model}"
|
||||
|
||||
async def _lookup_target_global_model_id(
|
||||
self,
|
||||
db: Session,
|
||||
source_model: str,
|
||||
provider_id: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
按优先级查找目标 GlobalModel ID:
|
||||
1. 映射(mapping_type='mapping'):Provider 特定 → 全局
|
||||
2. 别名(mapping_type='alias'):Provider 特定 → 全局
|
||||
3. 直接匹配 GlobalModel.name
|
||||
"""
|
||||
await self._ensure_initialized()
|
||||
cache_key = self._cache_key(source_model, provider_id)
|
||||
cached = await self._mapping_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
self._stats["mapping_hits"] += 1
|
||||
return cached or None
|
||||
|
||||
self._stats["mapping_misses"] += 1
|
||||
|
||||
target_id: Optional[str] = None
|
||||
|
||||
# 优先级 1:查找映射(mapping_type='mapping')
|
||||
# 1.1 Provider 特定映射
|
||||
if provider_id:
|
||||
mapping = (
|
||||
db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == source_model,
|
||||
ModelMapping.provider_id == provider_id,
|
||||
ModelMapping.mapping_type == "mapping",
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if mapping:
|
||||
target_id = mapping.target_global_model_id
|
||||
logger.debug(f"[MappingResolver] 命中 Provider 映射: {source_model} -> {target_id[:8]}...")
|
||||
|
||||
# 1.2 全局映射
|
||||
if not target_id:
|
||||
mapping = (
|
||||
db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == source_model,
|
||||
ModelMapping.provider_id.is_(None),
|
||||
ModelMapping.mapping_type == "mapping",
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if mapping:
|
||||
target_id = mapping.target_global_model_id
|
||||
logger.debug(f"[MappingResolver] 命中全局映射: {source_model} -> {target_id[:8]}...")
|
||||
|
||||
# 优先级 2:查找别名(mapping_type='alias')
|
||||
# 2.1 Provider 特定别名
|
||||
if not target_id and provider_id:
|
||||
alias = (
|
||||
db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == source_model,
|
||||
ModelMapping.provider_id == provider_id,
|
||||
ModelMapping.mapping_type == "alias",
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if alias:
|
||||
target_id = alias.target_global_model_id
|
||||
logger.debug(f"[MappingResolver] 命中 Provider 别名: {source_model} -> {target_id[:8]}...")
|
||||
|
||||
# 2.2 全局别名
|
||||
if not target_id:
|
||||
alias = (
|
||||
db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == source_model,
|
||||
ModelMapping.provider_id.is_(None),
|
||||
ModelMapping.mapping_type == "alias",
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if alias:
|
||||
target_id = alias.target_global_model_id
|
||||
logger.debug(f"[MappingResolver] 命中全局别名: {source_model} -> {target_id[:8]}...")
|
||||
|
||||
# 优先级 3:直接匹配 GlobalModel.name
|
||||
if not target_id:
|
||||
global_model = (
|
||||
db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.name == source_model,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if global_model:
|
||||
target_id = global_model.id
|
||||
logger.debug(f"[MappingResolver] 直接匹配 GlobalModel: {source_model}")
|
||||
|
||||
cached_value = target_id if target_id is not None else ""
|
||||
await self._mapping_cache.set(cache_key, cached_value, self._cache_ttl)
|
||||
return target_id
|
||||
|
||||
async def resolve_to_global_model_name(
|
||||
self,
|
||||
db: Session,
|
||||
source_model: str,
|
||||
provider_id: Optional[str] = None,
|
||||
) -> str:
|
||||
"""解析模型名/别名为 GlobalModel.name。未找到时返回原始输入。"""
|
||||
target_id = await self._lookup_target_global_model_id(db, source_model, provider_id)
|
||||
if not target_id:
|
||||
return source_model
|
||||
|
||||
await self._ensure_initialized()
|
||||
cached_name = await self._global_model_cache.get(target_id)
|
||||
if cached_name:
|
||||
self._stats["global_hits"] += 1
|
||||
return cached_name
|
||||
|
||||
self._stats["global_misses"] += 1
|
||||
global_model = (
|
||||
db.query(GlobalModel)
|
||||
.filter(GlobalModel.id == target_id, GlobalModel.is_active == True)
|
||||
.first()
|
||||
)
|
||||
if global_model:
|
||||
await self._global_model_cache.set(target_id, global_model.name, self._cache_ttl)
|
||||
return global_model.name
|
||||
|
||||
return source_model
|
||||
|
||||
async def get_global_model_by_request(
|
||||
self,
|
||||
db: Session,
|
||||
source_model: str,
|
||||
provider_id: Optional[str] = None,
|
||||
) -> Optional[GlobalModel]:
|
||||
"""解析并返回 GlobalModel 对象(绑定当前 Session)。"""
|
||||
target_id = await self._lookup_target_global_model_id(db, source_model, provider_id)
|
||||
if not target_id:
|
||||
return None
|
||||
|
||||
global_model = (
|
||||
db.query(GlobalModel)
|
||||
.filter(GlobalModel.id == target_id, GlobalModel.is_active == True)
|
||||
.first()
|
||||
)
|
||||
return global_model
|
||||
|
||||
async def get_global_model_with_mapping_info(
|
||||
self,
|
||||
db: Session,
|
||||
source_model: str,
|
||||
provider_id: Optional[str] = None,
|
||||
) -> Tuple[Optional[GlobalModel], bool]:
|
||||
"""
|
||||
解析并返回 GlobalModel 对象,同时返回是否发生了映射。
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
source_model: 用户请求的模型名
|
||||
provider_id: Provider ID(可选)
|
||||
|
||||
Returns:
|
||||
(global_model, is_mapped) - GlobalModel 对象和是否发生了映射
|
||||
is_mapped=True 表示 source_model 通过 mapping 规则映射到了不同的模型
|
||||
is_mapped=False 表示 source_model 直接匹配或通过 alias 匹配
|
||||
"""
|
||||
await self._ensure_initialized()
|
||||
|
||||
# 先检查是否存在 mapping 类型的映射规则
|
||||
has_mapping = False
|
||||
|
||||
# 检查 Provider 特定映射
|
||||
if provider_id:
|
||||
mapping = (
|
||||
db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == source_model,
|
||||
ModelMapping.provider_id == provider_id,
|
||||
ModelMapping.mapping_type == "mapping",
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if mapping:
|
||||
has_mapping = True
|
||||
|
||||
# 检查全局映射
|
||||
if not has_mapping:
|
||||
mapping = (
|
||||
db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == source_model,
|
||||
ModelMapping.provider_id.is_(None),
|
||||
ModelMapping.mapping_type == "mapping",
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if mapping:
|
||||
has_mapping = True
|
||||
|
||||
# 获取 GlobalModel
|
||||
global_model = await self.get_global_model_by_request(db, source_model, provider_id)
|
||||
|
||||
return global_model, has_mapping
|
||||
|
||||
async def get_global_model_direct(
|
||||
self,
|
||||
db: Session,
|
||||
source_model: str,
|
||||
) -> Optional[GlobalModel]:
|
||||
"""
|
||||
直接通过模型名获取 GlobalModel,不应用任何映射规则。
|
||||
仅查找 alias 和直接匹配。
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
source_model: 模型名
|
||||
|
||||
Returns:
|
||||
GlobalModel 对象或 None
|
||||
"""
|
||||
# 优先级 1:查找别名(alias)
|
||||
# 全局别名
|
||||
alias = (
|
||||
db.query(ModelMapping)
|
||||
.filter(
|
||||
ModelMapping.source_model == source_model,
|
||||
ModelMapping.provider_id.is_(None),
|
||||
ModelMapping.mapping_type == "alias",
|
||||
ModelMapping.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if alias:
|
||||
global_model = (
|
||||
db.query(GlobalModel)
|
||||
.filter(GlobalModel.id == alias.target_global_model_id, GlobalModel.is_active == True)
|
||||
.first()
|
||||
)
|
||||
if global_model:
|
||||
return global_model
|
||||
|
||||
# 优先级 2:直接匹配 GlobalModel.name
|
||||
global_model = (
|
||||
db.query(GlobalModel)
|
||||
.filter(
|
||||
GlobalModel.name == source_model,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
return global_model
|
||||
|
||||
def find_similar_models(
|
||||
self,
|
||||
db: Session,
|
||||
invalid_model: str,
|
||||
limit: int = 3,
|
||||
threshold: float = 0.4,
|
||||
) -> List[Tuple[str, float]]:
|
||||
"""用于提示相似的 GlobalModel.name。"""
|
||||
from difflib import SequenceMatcher
|
||||
|
||||
all_models = db.query(GlobalModel.name).filter(GlobalModel.is_active == True).all()
|
||||
similarities: List[Tuple[str, float]] = []
|
||||
invalid_lower = invalid_model.lower()
|
||||
|
||||
for model in all_models:
|
||||
model_name = model.name
|
||||
ratio = SequenceMatcher(None, invalid_lower, model_name.lower()).ratio()
|
||||
if invalid_lower in model_name.lower() or model_name.lower() in invalid_lower:
|
||||
ratio += 0.2
|
||||
if ratio >= threshold:
|
||||
similarities.append((model_name, ratio))
|
||||
|
||||
similarities.sort(key=lambda item: item[1], reverse=True)
|
||||
return similarities[:limit]
|
||||
|
||||
async def invalidate_mapping_cache(self, source_model: str, provider_id: Optional[str] = None):
|
||||
await self._ensure_initialized()
|
||||
keys = [self._cache_key(source_model, provider_id)]
|
||||
if provider_id:
|
||||
keys.append(self._cache_key(source_model, None))
|
||||
for key in keys:
|
||||
await self._mapping_cache.delete(key)
|
||||
|
||||
async def invalidate_global_model_cache(self, global_model_id: Optional[str] = None):
|
||||
await self._ensure_initialized()
|
||||
if global_model_id:
|
||||
await self._global_model_cache.delete(global_model_id)
|
||||
else:
|
||||
await self._global_model_cache.clear()
|
||||
|
||||
async def clear_cache(self):
|
||||
await self._ensure_initialized()
|
||||
await self._mapping_cache.clear()
|
||||
await self._global_model_cache.clear()
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
total_mapping = self._stats["mapping_hits"] + self._stats["mapping_misses"]
|
||||
total_global = self._stats["global_hits"] + self._stats["global_misses"]
|
||||
stats = {
|
||||
"mapping_hit_rate": (
|
||||
self._stats["mapping_hits"] / total_mapping if total_mapping else 0.0
|
||||
),
|
||||
"global_hit_rate": self._stats["global_hits"] / total_global if total_global else 0.0,
|
||||
"stats": self._stats,
|
||||
}
|
||||
if self._initialized:
|
||||
stats["mapping_cache_backend"] = self._mapping_cache.get_stats()
|
||||
stats["global_cache_backend"] = self._global_model_cache.get_stats()
|
||||
return stats
|
||||
|
||||
|
||||
_model_mapping_resolver: Optional[ModelMappingResolver] = None
|
||||
|
||||
|
||||
def get_model_mapping_resolver(
|
||||
cache_ttl: int = 300, cache_backend_type: Optional[str] = None
|
||||
) -> ModelMappingResolver:
|
||||
global _model_mapping_resolver
|
||||
|
||||
if _model_mapping_resolver is None:
|
||||
if cache_backend_type is None:
|
||||
cache_backend_type = os.getenv("ALIAS_CACHE_BACKEND", "auto")
|
||||
_model_mapping_resolver = ModelMappingResolver(
|
||||
cache_ttl=cache_ttl,
|
||||
cache_backend_type=cache_backend_type,
|
||||
)
|
||||
logger.debug(f"[ModelMappingResolver] 初始化(cache_ttl={cache_ttl}s, backend={cache_backend_type})")
|
||||
|
||||
# 注册到缓存失效服务
|
||||
try:
|
||||
from src.services.cache.invalidation import get_cache_invalidation_service
|
||||
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.set_mapping_resolver(_model_mapping_resolver)
|
||||
except Exception as exc:
|
||||
logger.warning(f"[ModelMappingResolver] 注册缓存失效服务失败: {exc}")
|
||||
|
||||
return _model_mapping_resolver
|
||||
|
||||
|
||||
async def resolve_model_to_global_name(
|
||||
db: Session,
|
||||
source_model: str,
|
||||
provider_id: Optional[str] = None,
|
||||
) -> str:
|
||||
resolver = get_model_mapping_resolver()
|
||||
return await resolver.resolve_to_global_model_name(db, source_model, provider_id)
|
||||
|
||||
|
||||
async def get_global_model_by_request(
|
||||
db: Session,
|
||||
source_model: str,
|
||||
provider_id: Optional[str] = None,
|
||||
) -> Optional[GlobalModel]:
|
||||
resolver = get_model_mapping_resolver()
|
||||
return await resolver.get_global_model_by_request(db, source_model, provider_id)
|
||||
48
src/services/model/pricing_strategy.py
Normal file
48
src/services/model/pricing_strategy.py
Normal file
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
计费相关数据类
|
||||
|
||||
定义计费计算所需的数据结构。
|
||||
实际的计费逻辑已移至 ChatAdapterBase,每种 API 格式可以覆盖计费方法。
|
||||
|
||||
数据类:
|
||||
- UsageTokens: 请求的 token 使用量
|
||||
- PricingConfig: 价格配置
|
||||
- CostResult: 计费结果
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageTokens:
|
||||
"""请求的 token 使用量"""
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class PricingConfig:
|
||||
"""价格配置"""
|
||||
input_price_per_1m: float = 0.0
|
||||
output_price_per_1m: float = 0.0
|
||||
cache_creation_price_per_1m: Optional[float] = None
|
||||
cache_read_price_per_1m: Optional[float] = None
|
||||
price_per_request: Optional[float] = None
|
||||
tiered_pricing: Optional[dict] = None
|
||||
cache_ttl_minutes: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class CostResult:
|
||||
"""计费结果"""
|
||||
input_cost: float = 0.0
|
||||
output_cost: float = 0.0
|
||||
cache_creation_cost: float = 0.0
|
||||
cache_read_cost: float = 0.0
|
||||
cache_cost: float = 0.0
|
||||
request_cost: float = 0.0
|
||||
total_cost: float = 0.0
|
||||
tier_index: Optional[int] = None # 命中的阶梯索引
|
||||
356
src/services/model/service.py
Normal file
356
src/services/model/service.py
Normal file
@@ -0,0 +1,356 @@
|
||||
"""
|
||||
模型管理服务
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import and_
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.models.api import ModelCreate, ModelResponse, ModelUpdate
|
||||
from src.models.database import Model, Provider
|
||||
from src.services.cache.invalidation import get_cache_invalidation_service
|
||||
from src.services.cache.model_cache import ModelCacheService
|
||||
|
||||
|
||||
|
||||
class ModelService:
|
||||
"""模型管理服务"""
|
||||
|
||||
@staticmethod
|
||||
def create_model(db: Session, provider_id: str, model_data: ModelCreate) -> Model:
|
||||
"""创建模型"""
|
||||
# 检查提供商是否存在
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise NotFoundException(f"提供商 {provider_id} 不存在")
|
||||
|
||||
# 检查同一提供商下是否已存在同名模型
|
||||
existing = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
and_(
|
||||
Model.provider_id == provider_id,
|
||||
Model.provider_model_name == model_data.provider_model_name,
|
||||
)
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if existing:
|
||||
raise InvalidRequestException(
|
||||
f"提供商 {provider.name} 下已存在模型 {model_data.provider_model_name}"
|
||||
)
|
||||
|
||||
try:
|
||||
model = Model(
|
||||
provider_id=provider_id,
|
||||
global_model_id=model_data.global_model_id,
|
||||
provider_model_name=model_data.provider_model_name,
|
||||
price_per_request=model_data.price_per_request,
|
||||
tiered_pricing=model_data.tiered_pricing,
|
||||
supports_vision=model_data.supports_vision,
|
||||
supports_function_calling=model_data.supports_function_calling,
|
||||
supports_streaming=model_data.supports_streaming,
|
||||
supports_extended_thinking=model_data.supports_extended_thinking,
|
||||
is_active=model_data.is_active if model_data.is_active is not None else True,
|
||||
config=model_data.config,
|
||||
)
|
||||
db.add(model)
|
||||
db.commit()
|
||||
db.refresh(model)
|
||||
# 显式加载 global_model 关系
|
||||
if model.global_model_id:
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
model = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.global_model))
|
||||
.filter(Model.id == model.id)
|
||||
.first()
|
||||
)
|
||||
|
||||
logger.info(f"创建模型成功: provider={provider.name}, model={model.provider_model_name}, global_model_id={model.global_model_id}")
|
||||
return model
|
||||
|
||||
except IntegrityError as e:
|
||||
db.rollback()
|
||||
logger.error(f"创建模型失败: {str(e)}")
|
||||
raise InvalidRequestException("创建模型失败,请检查输入数据")
|
||||
|
||||
@staticmethod
|
||||
def get_model(db: Session, model_id: str) -> Model: # UUID
|
||||
"""获取模型详情"""
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
model = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.global_model))
|
||||
.filter(Model.id == model_id)
|
||||
.first()
|
||||
)
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def get_models_by_provider(
|
||||
db: Session,
|
||||
provider_id: str, # UUID
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: Optional[bool] = None,
|
||||
) -> List[Model]:
|
||||
"""获取提供商的模型列表"""
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
query = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.global_model))
|
||||
.filter(Model.provider_id == provider_id)
|
||||
)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(Model.is_active == is_active)
|
||||
|
||||
# 按创建时间排序
|
||||
query = query.order_by(Model.created_at.desc())
|
||||
|
||||
return query.offset(skip).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def get_all_models(
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: Optional[bool] = None,
|
||||
category: Optional[str] = None,
|
||||
) -> List[Model]:
|
||||
"""获取所有模型列表"""
|
||||
query = db.query(Model)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(Model.is_active == is_active)
|
||||
|
||||
# 按提供商和创建时间排序
|
||||
query = query.order_by(Model.provider_id, Model.created_at.desc())
|
||||
|
||||
return query.offset(skip).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def update_model(db: Session, model_id: str, model_data: ModelUpdate) -> Model: # UUID
|
||||
"""更新模型"""
|
||||
model = db.query(Model).filter(Model.id == model_id).first()
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
|
||||
# 更新字段
|
||||
update_data = model_data.model_dump(exclude_unset=True)
|
||||
|
||||
# 添加调试日志
|
||||
logger.debug(f"更新模型 {model_id} 收到的数据: {update_data}")
|
||||
logger.debug(f"更新前的 supports_vision: {model.supports_vision}, supports_function_calling: {model.supports_function_calling}, supports_extended_thinking: {model.supports_extended_thinking}")
|
||||
|
||||
for field, value in update_data.items():
|
||||
setattr(model, field, value)
|
||||
|
||||
logger.debug(f"更新后的 supports_vision: {model.supports_vision}, supports_function_calling: {model.supports_function_calling}, supports_extended_thinking: {model.supports_extended_thinking}")
|
||||
|
||||
try:
|
||||
db.commit()
|
||||
db.refresh(model)
|
||||
|
||||
# 清除 Redis 缓存(异步执行,不阻塞返回)
|
||||
asyncio.create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
)
|
||||
)
|
||||
|
||||
# 清除内存缓存(ModelMapperMiddleware 实例)
|
||||
if model.provider_id and model.global_model_id:
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.on_model_changed(model.provider_id, model.global_model_id)
|
||||
|
||||
logger.info(f"更新模型成功: id={model_id}, 最终 supports_vision: {model.supports_vision}, supports_function_calling: {model.supports_function_calling}, supports_extended_thinking: {model.supports_extended_thinking}")
|
||||
return model
|
||||
except IntegrityError as e:
|
||||
db.rollback()
|
||||
logger.error(f"更新模型失败: {str(e)}")
|
||||
raise InvalidRequestException("更新模型失败,请检查输入数据")
|
||||
|
||||
@staticmethod
|
||||
def delete_model(db: Session, model_id: str): # UUID
|
||||
"""删除模型
|
||||
|
||||
新架构删除逻辑:
|
||||
- Model 只是 Provider 对 GlobalModel 的实现,删除不影响 GlobalModel
|
||||
- 检查是否是该 GlobalModel 的最后一个实现(如果是,警告但允许删除)
|
||||
- 不检查 ModelMapping(映射是 GlobalModel 之间的关系,别名也统一存储在此表中)
|
||||
"""
|
||||
model = db.query(Model).filter(Model.id == model_id).first()
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
|
||||
# 检查这是否是该 GlobalModel 的最后一个关联提供商
|
||||
if model.global_model_id:
|
||||
other_implementations = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
Model.global_model_id == model.global_model_id,
|
||||
Model.id != model_id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.count()
|
||||
)
|
||||
|
||||
if other_implementations == 0:
|
||||
logger.warning(f"警告:删除模型 {model_id}(Provider: {model.provider_id[:8]}...)后,"
|
||||
f"GlobalModel '{model.global_model_id}' 将没有任何活跃的关联提供商")
|
||||
|
||||
try:
|
||||
db.delete(model)
|
||||
db.commit()
|
||||
logger.info(f"删除模型成功: id={model_id}, provider_model_name={model.provider_model_name}, "
|
||||
f"global_model_id={model.global_model_id[:8] if model.global_model_id else 'None'}...")
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"删除模型失败: {str(e)}")
|
||||
raise InvalidRequestException("删除模型失败")
|
||||
|
||||
@staticmethod
|
||||
def toggle_model_availability(db: Session, model_id: str, is_available: bool) -> Model: # UUID
|
||||
"""切换模型可用状态"""
|
||||
model = db.query(Model).filter(Model.id == model_id).first()
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
|
||||
model.is_available = is_available
|
||||
db.commit()
|
||||
db.refresh(model)
|
||||
|
||||
# 清除 Redis 缓存
|
||||
asyncio.create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
)
|
||||
)
|
||||
|
||||
# 清除内存缓存(ModelMapperMiddleware 实例)
|
||||
if model.provider_id and model.global_model_id:
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.on_model_changed(model.provider_id, model.global_model_id)
|
||||
|
||||
status = "可用" if is_available else "不可用"
|
||||
logger.info(f"更新模型可用状态: id={model_id}, status={status}")
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def get_model_by_name(db: Session, provider_id: str, model_name: str) -> Optional[Model]:
|
||||
"""根据 provider_model_name 获取模型"""
|
||||
return (
|
||||
db.query(Model)
|
||||
.filter(and_(Model.provider_id == provider_id, Model.provider_model_name == model_name))
|
||||
.first()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def batch_create_models(
|
||||
db: Session, provider_id: str, models_data: List[ModelCreate]
|
||||
) -> List[Model]: # UUID
|
||||
"""批量创建模型"""
|
||||
# 检查提供商是否存在
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise NotFoundException(f"提供商 {provider_id} 不存在")
|
||||
|
||||
created_models = []
|
||||
for model_data in models_data:
|
||||
# 检查是否已存在
|
||||
existing = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
and_(
|
||||
Model.provider_id == provider_id,
|
||||
Model.provider_model_name == model_data.provider_model_name,
|
||||
)
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if existing:
|
||||
logger.warning(f"模型 {model_data.provider_model_name} 已存在,跳过创建")
|
||||
continue
|
||||
|
||||
model = Model(
|
||||
provider_id=provider_id,
|
||||
global_model_id=model_data.global_model_id,
|
||||
provider_model_name=model_data.provider_model_name,
|
||||
price_per_request=model_data.price_per_request,
|
||||
tiered_pricing=model_data.tiered_pricing,
|
||||
supports_vision=model_data.supports_vision,
|
||||
supports_function_calling=model_data.supports_function_calling,
|
||||
supports_streaming=model_data.supports_streaming,
|
||||
supports_extended_thinking=model_data.supports_extended_thinking,
|
||||
is_active=model_data.is_active,
|
||||
config=model_data.config,
|
||||
)
|
||||
db.add(model)
|
||||
created_models.append(model)
|
||||
|
||||
if created_models:
|
||||
try:
|
||||
db.commit()
|
||||
for model in created_models:
|
||||
db.refresh(model)
|
||||
logger.info(f"批量创建 {len(created_models)} 个模型成功")
|
||||
except IntegrityError as e:
|
||||
db.rollback()
|
||||
logger.error(f"批量创建模型失败: {str(e)}")
|
||||
raise InvalidRequestException("批量创建模型失败")
|
||||
|
||||
return created_models
|
||||
|
||||
@staticmethod
|
||||
def convert_to_response(model: Model) -> ModelResponse:
|
||||
"""转换为响应模型(新架构:从 GlobalModel 获取显示信息和默认值)"""
|
||||
return ModelResponse(
|
||||
id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
provider_model_name=model.provider_model_name,
|
||||
# 原始配置值(可能为空)
|
||||
price_per_request=model.price_per_request,
|
||||
tiered_pricing=model.tiered_pricing,
|
||||
supports_vision=model.supports_vision,
|
||||
supports_function_calling=model.supports_function_calling,
|
||||
supports_streaming=model.supports_streaming,
|
||||
supports_extended_thinking=model.supports_extended_thinking,
|
||||
supports_image_generation=model.supports_image_generation,
|
||||
# 有效值(合并 Model 和 GlobalModel 默认值)
|
||||
effective_tiered_pricing=model.get_effective_tiered_pricing(),
|
||||
effective_input_price=model.get_effective_input_price(),
|
||||
effective_output_price=model.get_effective_output_price(),
|
||||
effective_price_per_request=model.get_effective_price_per_request(),
|
||||
effective_supports_vision=model.get_effective_supports_vision(),
|
||||
effective_supports_function_calling=model.get_effective_supports_function_calling(),
|
||||
effective_supports_streaming=model.get_effective_supports_streaming(),
|
||||
effective_supports_extended_thinking=model.get_effective_supports_extended_thinking(),
|
||||
effective_supports_image_generation=model.get_effective_supports_image_generation(),
|
||||
is_active=model.is_active,
|
||||
is_available=model.is_available if model.is_available is not None else True,
|
||||
created_at=model.created_at,
|
||||
updated_at=model.updated_at,
|
||||
# GlobalModel 信息(如果存在)
|
||||
global_model_name=model.global_model.name if model.global_model else None,
|
||||
global_model_display_name=(
|
||||
model.global_model.display_name if model.global_model else None
|
||||
),
|
||||
)
|
||||
22
src/services/orchestration/__init__.py
Normal file
22
src/services/orchestration/__init__.py
Normal file
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
Orchestration 模块
|
||||
|
||||
提供请求编排相关的组件:
|
||||
- FallbackOrchestrator: 故障转移编排器,协调请求的完整生命周期
|
||||
- CandidateResolver: 候选解析器,负责获取和排序可用的 Provider 组合
|
||||
- RequestDispatcher: 请求分发器,负责执行单个候选请求
|
||||
- ErrorClassifier: 错误分类器,负责错误分类和处理策略
|
||||
"""
|
||||
|
||||
from .candidate_resolver import CandidateResolver
|
||||
from .error_classifier import ErrorAction, ErrorClassifier
|
||||
from .fallback_orchestrator import FallbackOrchestrator
|
||||
from .request_dispatcher import RequestDispatcher
|
||||
|
||||
__all__ = [
|
||||
"FallbackOrchestrator",
|
||||
"CandidateResolver",
|
||||
"RequestDispatcher",
|
||||
"ErrorClassifier",
|
||||
"ErrorAction",
|
||||
]
|
||||
242
src/services/orchestration/candidate_resolver.py
Normal file
242
src/services/orchestration/candidate_resolver.py
Normal file
@@ -0,0 +1,242 @@
|
||||
"""
|
||||
候选解析器
|
||||
|
||||
负责获取和排序可用的 Provider/Endpoint/Key 组合
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.enums import APIFormat
|
||||
from src.core.exceptions import ProviderNotAvailableException
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey
|
||||
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
|
||||
|
||||
|
||||
|
||||
class CandidateResolver:
|
||||
"""
|
||||
候选解析器 - 负责获取和排序可用的 Provider 组合
|
||||
|
||||
职责:
|
||||
1. 从 CacheAwareScheduler 获取所有可用候选
|
||||
2. 创建候选记录(用于追踪)
|
||||
3. 提供候选的迭代和过滤功能
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: Session,
|
||||
cache_scheduler: CacheAwareScheduler,
|
||||
) -> None:
|
||||
"""
|
||||
初始化候选解析器
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
cache_scheduler: 缓存感知调度器
|
||||
"""
|
||||
self.db = db
|
||||
self.cache_scheduler = cache_scheduler
|
||||
|
||||
async def fetch_candidates(
|
||||
self,
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
user_api_key: Optional[ApiKey] = None,
|
||||
request_id: Optional[str] = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[List[ProviderCandidate], str]:
|
||||
"""
|
||||
获取所有可用候选
|
||||
|
||||
Args:
|
||||
api_format: API 格式
|
||||
model_name: 模型名称
|
||||
affinity_key: 亲和性标识符(通常为API Key ID,用于缓存亲和性)
|
||||
user_api_key: 用户 API Key(用于 allowed_providers/allowed_api_formats 过滤)
|
||||
request_id: 请求 ID(用于日志)
|
||||
is_stream: 是否是流式请求,如果为 True 则过滤不支持流式的 Provider
|
||||
capability_requirements: 能力需求(用于过滤不满足能力要求的 Key)
|
||||
|
||||
Returns:
|
||||
(所有候选组合的列表, global_model_id)
|
||||
|
||||
Raises:
|
||||
ProviderNotAvailableException: 没有找到任何可用候选时
|
||||
"""
|
||||
all_candidates: List[ProviderCandidate] = []
|
||||
provider_offset = 0
|
||||
provider_batch_size = 20
|
||||
global_model_id: Optional[str] = None
|
||||
|
||||
while True:
|
||||
candidates, resolved_global_model_id = await self.cache_scheduler.list_all_candidates(
|
||||
db=self.db,
|
||||
api_format=api_format,
|
||||
model_name=model_name,
|
||||
affinity_key=affinity_key,
|
||||
user_api_key=user_api_key,
|
||||
provider_offset=provider_offset,
|
||||
provider_limit=provider_batch_size,
|
||||
is_stream=is_stream,
|
||||
capability_requirements=capability_requirements,
|
||||
)
|
||||
|
||||
if resolved_global_model_id and global_model_id is None:
|
||||
global_model_id = resolved_global_model_id
|
||||
|
||||
if not candidates:
|
||||
break
|
||||
|
||||
all_candidates.extend(candidates)
|
||||
provider_offset += provider_batch_size
|
||||
|
||||
if not all_candidates:
|
||||
logger.error(f" [{request_id}] 没有找到任何可用的 Provider/Endpoint/Key 组合")
|
||||
request_type = "流式" if is_stream else "非流式"
|
||||
raise ProviderNotAvailableException(
|
||||
f"没有可用的 Provider 支持模型 {model_name} 的{request_type}请求"
|
||||
)
|
||||
|
||||
logger.debug(f" [{request_id}] 获取到 {len(all_candidates)} 个候选组合")
|
||||
|
||||
# 如果没有解析到 global_model_id,使用原始 model_name 作为后备
|
||||
return all_candidates, global_model_id or model_name
|
||||
|
||||
def create_candidate_records(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
request_id: Optional[str],
|
||||
user_id: str,
|
||||
user_api_key: ApiKey,
|
||||
required_capabilities: Optional[Dict[str, bool]] = None,
|
||||
) -> Dict[Tuple[int, int], str]:
|
||||
"""
|
||||
为所有候选预先创建 available 状态记录(批量插入优化)
|
||||
|
||||
Args:
|
||||
all_candidates: 所有候选组合
|
||||
request_id: 请求 ID
|
||||
user_id: 用户 ID
|
||||
user_api_key: 用户 API Key 对象
|
||||
required_capabilities: 请求需要的能力标签
|
||||
|
||||
Returns:
|
||||
candidate_record_map: {(candidate_index, retry_index): candidate_record_id}
|
||||
"""
|
||||
from src.models.database import RequestCandidate
|
||||
|
||||
candidate_records_to_insert: List[Dict[str, Any]] = []
|
||||
candidate_record_map: Dict[Tuple[int, int], str] = {}
|
||||
|
||||
# 只保存启用的能力(值为 True 的)
|
||||
active_capabilities = None
|
||||
if required_capabilities:
|
||||
active_capabilities = {k: v for k, v in required_capabilities.items() if v}
|
||||
if not active_capabilities:
|
||||
active_capabilities = None
|
||||
|
||||
for candidate_index, candidate in enumerate(all_candidates):
|
||||
provider = candidate.provider
|
||||
endpoint = candidate.endpoint
|
||||
key = candidate.key
|
||||
|
||||
if candidate.is_skipped:
|
||||
record_id = str(uuid.uuid4())
|
||||
candidate_records_to_insert.append(
|
||||
{
|
||||
"id": record_id,
|
||||
"request_id": request_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"user_id": user_id,
|
||||
"api_key_id": user_api_key.id if user_api_key else None,
|
||||
"provider_id": provider.id,
|
||||
"endpoint_id": endpoint.id,
|
||||
"key_id": key.id,
|
||||
"status": "skipped",
|
||||
"skip_reason": candidate.skip_reason,
|
||||
"is_cached": candidate.is_cached,
|
||||
"extra_data": {},
|
||||
"required_capabilities": active_capabilities,
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
}
|
||||
)
|
||||
candidate_record_map[(candidate_index, 0)] = record_id
|
||||
else:
|
||||
max_retries_for_candidate = endpoint.max_retries if candidate.is_cached else 1
|
||||
|
||||
for retry_index in range(max_retries_for_candidate):
|
||||
record_id = str(uuid.uuid4())
|
||||
candidate_records_to_insert.append(
|
||||
{
|
||||
"id": record_id,
|
||||
"request_id": request_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": retry_index,
|
||||
"user_id": user_id,
|
||||
"api_key_id": user_api_key.id if user_api_key else None,
|
||||
"provider_id": provider.id,
|
||||
"endpoint_id": endpoint.id,
|
||||
"key_id": key.id,
|
||||
"status": "available",
|
||||
"is_cached": candidate.is_cached,
|
||||
"extra_data": {},
|
||||
"required_capabilities": active_capabilities,
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
}
|
||||
)
|
||||
candidate_record_map[(candidate_index, retry_index)] = record_id
|
||||
|
||||
if candidate_records_to_insert:
|
||||
self.db.bulk_insert_mappings(
|
||||
RequestCandidate, candidate_records_to_insert # type: ignore
|
||||
)
|
||||
self.db.flush()
|
||||
|
||||
logger.debug(f" [{request_id}] 批量插入完成: {len(candidate_records_to_insert)} 条记录")
|
||||
|
||||
return candidate_record_map
|
||||
|
||||
def get_active_candidates(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
) -> List[Tuple[int, ProviderCandidate]]:
|
||||
"""
|
||||
获取所有非跳过的候选(带索引)
|
||||
|
||||
Args:
|
||||
all_candidates: 所有候选组合
|
||||
|
||||
Returns:
|
||||
List of (index, candidate) for non-skipped candidates
|
||||
"""
|
||||
return [(i, c) for i, c in enumerate(all_candidates) if not c.is_skipped]
|
||||
|
||||
def count_total_attempts(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
) -> int:
|
||||
"""
|
||||
计算总尝试次数
|
||||
|
||||
Args:
|
||||
all_candidates: 所有候选组合
|
||||
|
||||
Returns:
|
||||
总尝试次数
|
||||
"""
|
||||
total = 0
|
||||
for candidate in all_candidates:
|
||||
if not candidate.is_skipped:
|
||||
endpoint = candidate.endpoint
|
||||
max_retries = int(endpoint.max_retries) if candidate.is_cached else 1
|
||||
total += max_retries
|
||||
return total
|
||||
530
src/services/orchestration/error_classifier.py
Normal file
530
src/services/orchestration/error_classifier.py
Normal file
@@ -0,0 +1,530 @@
|
||||
"""
|
||||
错误分类器
|
||||
|
||||
负责错误分类和处理策略决定
|
||||
"""
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.enums import APIFormat
|
||||
from src.core.exceptions import (
|
||||
ConcurrencyLimitError,
|
||||
ProviderAuthException,
|
||||
ProviderException,
|
||||
ProviderNotAvailableException,
|
||||
ProviderRateLimitException,
|
||||
UpstreamClientException,
|
||||
)
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.cache.aware_scheduler import CacheAwareScheduler
|
||||
from src.services.health.monitor import health_monitor
|
||||
from src.services.provider.format import normalize_api_format
|
||||
from src.services.rate_limit.adaptive_concurrency import get_adaptive_manager
|
||||
from src.services.rate_limit.detector import RateLimitType, detect_rate_limit_type
|
||||
|
||||
|
||||
|
||||
class ErrorAction(Enum):
|
||||
"""错误处理动作"""
|
||||
|
||||
CONTINUE = "continue" # 继续重试当前候选
|
||||
BREAK = "break" # 跳到下一个候选
|
||||
RAISE = "raise" # 直接抛出异常
|
||||
|
||||
|
||||
class ErrorClassifier:
|
||||
"""
|
||||
错误分类器 - 负责错误分类和处理策略
|
||||
|
||||
职责:
|
||||
1. 将错误分类为可重试/不可重试
|
||||
2. 决定错误后的处理动作(重试/切换/放弃)
|
||||
3. 处理特定类型的错误(如 429 限流)
|
||||
4. 更新健康状态和缓存亲和性
|
||||
"""
|
||||
|
||||
# 需要触发故障转移的错误类型
|
||||
RETRIABLE_ERRORS: Tuple[type, ...] = (
|
||||
ProviderException, # 包含所有 Provider 异常子类
|
||||
ConnectionError, # Python 标准连接错误
|
||||
TimeoutError, # Python 标准超时错误
|
||||
httpx.TransportError, # HTTPX 传输错误
|
||||
)
|
||||
|
||||
# 不可重试的错误类型(直接抛出)
|
||||
NON_RETRIABLE_ERRORS: Tuple[type, ...] = (
|
||||
ValueError, # 参数错误
|
||||
TypeError, # 类型错误
|
||||
KeyError, # 键错误
|
||||
UpstreamClientException, # 上游客户端错误
|
||||
)
|
||||
|
||||
# 表示客户端请求错误的关键词(不区分大小写)
|
||||
# 这些错误是由用户请求本身导致的,换 Provider 也无济于事
|
||||
CLIENT_ERROR_PATTERNS: Tuple[str, ...] = (
|
||||
"could not process image", # 图片处理失败
|
||||
"image too large", # 图片过大
|
||||
"invalid image", # 无效图片
|
||||
"unsupported image", # 不支持的图片格式
|
||||
"invalid_request_error", # OpenAI/Claude 通用客户端错误类型
|
||||
"content_policy_violation", # 内容违规
|
||||
"invalid_api_key", # 无效的 API Key(不同于认证失败)
|
||||
"context_length_exceeded", # 上下文长度超限
|
||||
"max_tokens", # token 数超限
|
||||
"invalid_prompt", # 无效的提示词
|
||||
"content too long", # 内容过长
|
||||
"message is too long", # 消息过长
|
||||
"prompt is too long", # Prompt 超长(第三方代理常见格式)
|
||||
"image exceeds", # 图片超出限制
|
||||
"pdf too large", # PDF 过大
|
||||
"file too large", # 文件过大
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: Session,
|
||||
adaptive_manager: Any = None,
|
||||
cache_scheduler: Optional[CacheAwareScheduler] = None,
|
||||
) -> None:
|
||||
"""
|
||||
初始化错误分类器
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
adaptive_manager: 自适应并发管理器
|
||||
cache_scheduler: 缓存调度器(可选)
|
||||
"""
|
||||
self.db = db
|
||||
self.adaptive_manager = adaptive_manager or get_adaptive_manager()
|
||||
self.cache_scheduler = cache_scheduler
|
||||
|
||||
def _is_client_error(self, error_text: Optional[str]) -> bool:
|
||||
"""
|
||||
检测错误响应是否为客户端错误(不应重试)
|
||||
|
||||
Args:
|
||||
error_text: 错误响应文本
|
||||
|
||||
Returns:
|
||||
是否为客户端错误
|
||||
"""
|
||||
if not error_text:
|
||||
return False
|
||||
|
||||
error_lower = error_text.lower()
|
||||
return any(pattern.lower() in error_lower for pattern in self.CLIENT_ERROR_PATTERNS)
|
||||
|
||||
def _extract_error_message(self, error_text: Optional[str]) -> Optional[str]:
|
||||
"""
|
||||
从错误响应中提取错误消息
|
||||
|
||||
支持格式:
|
||||
- {"error": {"message": "..."}} (OpenAI/Claude)
|
||||
- {"error": {"type": "...", "message": "..."}}
|
||||
- {"error": "..."}
|
||||
- {"message": "..."}
|
||||
|
||||
Args:
|
||||
error_text: 错误响应文本
|
||||
|
||||
Returns:
|
||||
提取的错误消息,如果无法解析则返回原始文本
|
||||
"""
|
||||
if not error_text:
|
||||
return None
|
||||
|
||||
try:
|
||||
data = json.loads(error_text)
|
||||
|
||||
# {"error": {"message": "..."}} 或 {"error": {"type": "...", "message": "..."}}
|
||||
if isinstance(data.get("error"), dict):
|
||||
error_obj = data["error"]
|
||||
message = error_obj.get("message", "")
|
||||
error_type = error_obj.get("type", "")
|
||||
if message:
|
||||
if error_type:
|
||||
return f"{error_type}: {message}"
|
||||
return str(message)
|
||||
|
||||
# {"error": "..."}
|
||||
if isinstance(data.get("error"), str):
|
||||
return str(data["error"])
|
||||
|
||||
# {"message": "..."}
|
||||
if isinstance(data.get("message"), str):
|
||||
return str(data["message"])
|
||||
|
||||
except (json.JSONDecodeError, TypeError, KeyError):
|
||||
pass
|
||||
|
||||
# 无法解析,返回原始文本(截断)
|
||||
return error_text[:500] if len(error_text) > 500 else error_text
|
||||
|
||||
def classify(
|
||||
self,
|
||||
error: Exception,
|
||||
has_retry_left: bool = False,
|
||||
) -> ErrorAction:
|
||||
"""
|
||||
分类错误,返回处理动作
|
||||
|
||||
Args:
|
||||
error: 异常对象
|
||||
has_retry_left: 当前候选是否还有重试次数
|
||||
|
||||
Returns:
|
||||
ErrorAction: 处理动作
|
||||
"""
|
||||
if isinstance(error, ConcurrencyLimitError):
|
||||
return ErrorAction.BREAK
|
||||
|
||||
if isinstance(error, httpx.HTTPStatusError):
|
||||
# HTTP 错误根据状态码决定
|
||||
return ErrorAction.CONTINUE if has_retry_left else ErrorAction.BREAK
|
||||
|
||||
if isinstance(error, self.RETRIABLE_ERRORS):
|
||||
return ErrorAction.CONTINUE if has_retry_left else ErrorAction.BREAK
|
||||
|
||||
if isinstance(error, self.NON_RETRIABLE_ERRORS):
|
||||
return ErrorAction.RAISE
|
||||
|
||||
# 未知错误,直接抛出
|
||||
return ErrorAction.RAISE
|
||||
|
||||
async def handle_rate_limit(
|
||||
self,
|
||||
key: ProviderAPIKey,
|
||||
provider_name: str,
|
||||
current_concurrent: Optional[int],
|
||||
exception: ProviderRateLimitException,
|
||||
request_id: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
处理 429 速率限制错误的自适应调整
|
||||
|
||||
Args:
|
||||
key: API Key 对象
|
||||
provider_name: 提供商名称
|
||||
current_concurrent: 当前并发数
|
||||
exception: 速率限制异常
|
||||
request_id: 请求 ID(用于日志)
|
||||
|
||||
Returns:
|
||||
限制类型: "concurrent" 或 "rpm" 或 "unknown"
|
||||
"""
|
||||
try:
|
||||
# 提取响应头(如果有)
|
||||
response_headers = {}
|
||||
if hasattr(exception, "response_headers"):
|
||||
response_headers = exception.response_headers or {}
|
||||
|
||||
# 检测速率限制类型
|
||||
rate_limit_info = detect_rate_limit_type(
|
||||
headers=response_headers,
|
||||
provider_name=provider_name,
|
||||
current_concurrent=current_concurrent,
|
||||
)
|
||||
|
||||
logger.info(f" [{request_id}] 429错误分析: "
|
||||
f"类型={rate_limit_info.limit_type}, "
|
||||
f"retry_after={rate_limit_info.retry_after}s, "
|
||||
f"当前并发={current_concurrent}")
|
||||
|
||||
# 调用自适应管理器处理
|
||||
new_limit = self.adaptive_manager.handle_429_error(
|
||||
db=self.db,
|
||||
key=key,
|
||||
rate_limit_info=rate_limit_info,
|
||||
current_concurrent=current_concurrent,
|
||||
)
|
||||
|
||||
if rate_limit_info.limit_type == RateLimitType.CONCURRENT:
|
||||
logger.warning(f" [{request_id}] 自适应调整: " f"Key {key.id[:8]}... 并发限制 -> {new_limit}")
|
||||
return "concurrent"
|
||||
elif rate_limit_info.limit_type == RateLimitType.RPM:
|
||||
logger.info(f" [{request_id}] [RPM] RPM限制,需要切换Provider")
|
||||
return "rpm"
|
||||
else:
|
||||
return "unknown"
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f" [{request_id}] 处理429错误时异常: {e}")
|
||||
return "unknown"
|
||||
|
||||
def convert_http_error(
|
||||
self,
|
||||
error: httpx.HTTPStatusError,
|
||||
provider_name: str,
|
||||
error_response_text: Optional[str] = None,
|
||||
) -> Union[ProviderException, UpstreamClientException]:
|
||||
"""
|
||||
转换 HTTP 错误为 Provider 异常
|
||||
|
||||
Args:
|
||||
error: HTTP 状态错误
|
||||
provider_name: Provider 名称
|
||||
error_response_text: 错误响应文本(可选)
|
||||
|
||||
Returns:
|
||||
ProviderException 或 UpstreamClientException: 转换后的异常
|
||||
"""
|
||||
status = error.response.status_code if error.response else None
|
||||
|
||||
# 提取可读的错误消息
|
||||
extracted_message = self._extract_error_message(error_response_text)
|
||||
|
||||
# 构建详细错误信息
|
||||
if extracted_message:
|
||||
detailed_message = f"提供商 '{provider_name}' 返回错误 {status}: {extracted_message}"
|
||||
else:
|
||||
detailed_message = f"提供商 '{provider_name}' 返回错误: {status}"
|
||||
|
||||
if status == 401:
|
||||
return ProviderAuthException(provider_name=provider_name)
|
||||
|
||||
if status == 429:
|
||||
return ProviderRateLimitException(
|
||||
message=error_response_text or f"提供商 '{provider_name}' 速率限制",
|
||||
provider_name=provider_name,
|
||||
response_headers=dict(error.response.headers) if error.response else None,
|
||||
retry_after=(
|
||||
int(error.response.headers.get("retry-after", 0))
|
||||
if error.response and error.response.headers.get("retry-after")
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
# 400 错误:检查是否为客户端请求错误(不应重试)
|
||||
if status == 400 and self._is_client_error(error_response_text):
|
||||
logger.info(f"检测到客户端请求错误,不进行重试: {extracted_message}")
|
||||
return UpstreamClientException(
|
||||
message=extracted_message or "请求无效",
|
||||
provider_name=provider_name,
|
||||
status_code=400,
|
||||
upstream_error=error_response_text,
|
||||
)
|
||||
|
||||
if status and status >= 500:
|
||||
return ProviderNotAvailableException(
|
||||
message=detailed_message,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
|
||||
return ProviderNotAvailableException(
|
||||
message=detailed_message,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
|
||||
async def handle_http_error(
|
||||
self,
|
||||
http_error: httpx.HTTPStatusError,
|
||||
*,
|
||||
provider: Provider,
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
affinity_key: str,
|
||||
api_format: Union[str, APIFormat],
|
||||
global_model_id: str,
|
||||
request_id: Optional[str],
|
||||
captured_key_concurrent: Optional[int],
|
||||
elapsed_ms: Optional[int],
|
||||
attempt: int,
|
||||
max_attempts: int,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
处理 HTTP 错误,返回 extra_data
|
||||
|
||||
Args:
|
||||
http_error: HTTP 状态错误
|
||||
provider: Provider 对象
|
||||
endpoint: Endpoint 对象
|
||||
key: API Key 对象
|
||||
affinity_key: 亲和性标识符(通常为 API Key ID)
|
||||
api_format: API 格式
|
||||
global_model_id: GlobalModel ID(规范化的模型标识)
|
||||
request_id: 请求 ID
|
||||
captured_key_concurrent: 捕获的并发数
|
||||
elapsed_ms: 耗时(毫秒)
|
||||
attempt: 当前尝试次数
|
||||
max_attempts: 最大尝试次数
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: 额外数据,包含:
|
||||
- error_response: 错误响应文本(如有)
|
||||
- converted_error: 转换后的异常对象(用于判断是否应该重试)
|
||||
"""
|
||||
provider_name = str(provider.name)
|
||||
|
||||
# 尝试读取错误响应内容
|
||||
error_response_text = None
|
||||
try:
|
||||
if http_error.response and hasattr(http_error.response, "text"):
|
||||
error_response_text = http_error.response.text[:1000] # 限制长度
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.warning(f" [{request_id}] HTTP错误 (attempt={attempt}/{max_attempts}): "
|
||||
f"{http_error.response.status_code if http_error.response else 'unknown'}")
|
||||
|
||||
converted_error = self.convert_http_error(http_error, provider_name, error_response_text)
|
||||
|
||||
# 构建 extra_data,包含转换后的异常
|
||||
extra_data: Dict[str, Any] = {
|
||||
"converted_error": converted_error,
|
||||
}
|
||||
if error_response_text:
|
||||
extra_data["error_response"] = error_response_text
|
||||
|
||||
# 转换 api_format 为字符串
|
||||
api_format_str = (
|
||||
normalize_api_format(api_format).value
|
||||
if isinstance(api_format, (str, APIFormat))
|
||||
else str(api_format)
|
||||
)
|
||||
|
||||
# 处理客户端请求错误(不应重试,不失效缓存,不记录健康失败)
|
||||
if isinstance(converted_error, UpstreamClientException):
|
||||
logger.warning(f" [{request_id}] 客户端请求错误,不进行重试: {converted_error.message}")
|
||||
return extra_data
|
||||
|
||||
# 处理认证错误
|
||||
if isinstance(converted_error, ProviderAuthException):
|
||||
if endpoint and key and self.cache_scheduler is not None:
|
||||
await self.cache_scheduler.invalidate_cache(
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format_str,
|
||||
global_model_id=global_model_id,
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
)
|
||||
if key:
|
||||
health_monitor.record_failure(
|
||||
db=self.db,
|
||||
key_id=str(key.id),
|
||||
error_type="ProviderAuthException",
|
||||
)
|
||||
return extra_data
|
||||
|
||||
# 处理限流错误
|
||||
if isinstance(converted_error, ProviderRateLimitException) and key:
|
||||
await self.handle_rate_limit(
|
||||
key=key,
|
||||
provider_name=provider_name,
|
||||
current_concurrent=captured_key_concurrent,
|
||||
exception=converted_error,
|
||||
request_id=request_id,
|
||||
)
|
||||
if endpoint and self.cache_scheduler is not None:
|
||||
await self.cache_scheduler.invalidate_cache(
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format_str,
|
||||
global_model_id=global_model_id,
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
)
|
||||
else:
|
||||
# 其他错误也失效缓存
|
||||
if endpoint and key and self.cache_scheduler is not None:
|
||||
await self.cache_scheduler.invalidate_cache(
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format_str,
|
||||
global_model_id=global_model_id,
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
)
|
||||
|
||||
# 记录健康失败
|
||||
if key:
|
||||
health_monitor.record_failure(
|
||||
db=self.db,
|
||||
key_id=str(key.id),
|
||||
error_type=type(converted_error).__name__,
|
||||
)
|
||||
|
||||
return extra_data
|
||||
|
||||
async def handle_retriable_error(
|
||||
self,
|
||||
error: Exception,
|
||||
*,
|
||||
provider: Provider,
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
affinity_key: str,
|
||||
api_format: Union[str, APIFormat],
|
||||
global_model_id: str,
|
||||
captured_key_concurrent: Optional[int],
|
||||
elapsed_ms: Optional[int],
|
||||
request_id: Optional[str],
|
||||
attempt: int,
|
||||
max_attempts: int,
|
||||
) -> None:
|
||||
"""
|
||||
处理可重试错误
|
||||
|
||||
Args:
|
||||
error: 异常对象
|
||||
provider: Provider 对象
|
||||
endpoint: Endpoint 对象
|
||||
key: API Key 对象
|
||||
affinity_key: 亲和性标识符(通常为 API Key ID)
|
||||
api_format: API 格式
|
||||
global_model_id: GlobalModel ID(规范化的模型标识,用于缓存亲和性)
|
||||
captured_key_concurrent: 捕获的并发数
|
||||
elapsed_ms: 耗时(毫秒)
|
||||
request_id: 请求 ID
|
||||
attempt: 当前尝试次数
|
||||
max_attempts: 最大尝试次数
|
||||
"""
|
||||
provider_name = str(provider.name)
|
||||
|
||||
logger.warning(f" [{request_id}] 请求失败 (attempt={attempt}/{max_attempts}): "
|
||||
f"{type(error).__name__}: {str(error)}")
|
||||
|
||||
# 转换 api_format 为字符串
|
||||
api_format_str = (
|
||||
normalize_api_format(api_format).value
|
||||
if isinstance(api_format, (str, APIFormat))
|
||||
else str(api_format)
|
||||
)
|
||||
|
||||
# 处理限流错误
|
||||
if isinstance(error, ProviderRateLimitException) and key:
|
||||
await self.handle_rate_limit(
|
||||
key=key,
|
||||
provider_name=provider_name,
|
||||
current_concurrent=captured_key_concurrent,
|
||||
exception=error,
|
||||
request_id=request_id,
|
||||
)
|
||||
if endpoint and self.cache_scheduler is not None:
|
||||
await self.cache_scheduler.invalidate_cache(
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format_str,
|
||||
global_model_id=global_model_id,
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
)
|
||||
elif endpoint and key and self.cache_scheduler is not None:
|
||||
# 其他错误也失效缓存
|
||||
await self.cache_scheduler.invalidate_cache(
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format_str,
|
||||
global_model_id=global_model_id,
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
)
|
||||
|
||||
# 记录健康失败
|
||||
if key:
|
||||
health_monitor.record_failure(
|
||||
db=self.db,
|
||||
key_id=str(key.id),
|
||||
error_type=type(error).__name__,
|
||||
)
|
||||
757
src/services/orchestration/fallback_orchestrator.py
Normal file
757
src/services/orchestration/fallback_orchestrator.py
Normal file
@@ -0,0 +1,757 @@
|
||||
"""
|
||||
故障转移编排器(预取+顺序遍历策略)
|
||||
|
||||
功能:
|
||||
1. 预先获取所有可用的 Provider/Endpoint/Key 组合
|
||||
2. 按优先级顺序遍历组合(每个只尝试一次)
|
||||
3. 集成 HealthMonitor 记录成功/失败
|
||||
4. 集成 ConcurrencyManager 管理并发(支持缓存用户优先级)
|
||||
5. 缓存亲和性管理(自动失效失败的Key)
|
||||
|
||||
优化亮点:
|
||||
- 避免运行时重复查询数据库
|
||||
- 精确控制重试次数(=实际组合数)
|
||||
- 清晰的故障转移逻辑,易于维护和调试
|
||||
|
||||
重构说明:
|
||||
- 职责已拆分到独立组件(src/services/orchestration/):
|
||||
- CandidateResolver: 候选解析器,负责获取和排序可用的 Provider 组合
|
||||
- RequestDispatcher: 请求分发器,负责执行单个候选请求
|
||||
- ErrorClassifier: 错误分类器,负责错误分类和处理策略
|
||||
- 本类作为协调者,组合使用上述组件
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable, Dict, List, NoReturn, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
from redis import Redis
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.enums import APIFormat
|
||||
from src.core.exceptions import (
|
||||
ConcurrencyLimitError,
|
||||
ProviderNotAvailableException,
|
||||
UpstreamClientException,
|
||||
)
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.cache.aware_scheduler import (
|
||||
CacheAwareScheduler,
|
||||
ProviderCandidate,
|
||||
get_cache_aware_scheduler,
|
||||
)
|
||||
from src.services.provider.format import normalize_api_format
|
||||
from src.services.rate_limit.adaptive_concurrency import get_adaptive_manager
|
||||
from src.services.rate_limit.concurrency_manager import get_concurrency_manager
|
||||
from src.services.request.candidate import RequestCandidateService
|
||||
from src.services.request.executor import ExecutionError, RequestExecutor
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
from .candidate_resolver import CandidateResolver
|
||||
from .error_classifier import ErrorClassifier
|
||||
from .request_dispatcher import RequestDispatcher
|
||||
|
||||
|
||||
class FallbackOrchestrator:
|
||||
"""
|
||||
故障转移编排器(预取+顺序遍历策略)
|
||||
|
||||
负责协调请求的完整生命周期:
|
||||
1. 预先获取所有可用的 Provider+Endpoint+Key 组合(按优先级排序)
|
||||
2. 按顺序遍历每个组合,获取并发槽位(缓存用户优先)
|
||||
3. 发送请求
|
||||
4. 记录结果(成功/失败,更新健康度)
|
||||
5. 失败时自动切换到下一个组合,直到成功或全部失败
|
||||
|
||||
故障转移策略(V2 - 预取优化):
|
||||
- 启动时预先获取所有符合条件的 Provider/Endpoint/Key 组合
|
||||
- 按优先级排序:Provider.provider_priority → Key.internal_priority(Endpoint在Provider内唯一,无需排序)
|
||||
- 过滤条件:活跃状态、健康度、熔断器状态、模型支持、API格式匹配
|
||||
- 顺序遍历组合列表,每个组合只尝试一次
|
||||
- 重试次数 = 实际可用组合数(无固定上限,避免过度重试)
|
||||
- 优势:可预测、高效、公平、资源友好
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session, redis_client: Optional[Redis] = None) -> None:
|
||||
"""
|
||||
初始化编排器
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
redis_client: Redis客户端(可选,用于缓存和并发控制)
|
||||
"""
|
||||
self.db = db
|
||||
self.redis = redis_client
|
||||
self.cache_scheduler: Optional[CacheAwareScheduler] = None
|
||||
self.concurrency_manager: Any = None
|
||||
self.adaptive_manager = get_adaptive_manager() # 自适应并发管理器
|
||||
self.request_executor: Optional[RequestExecutor] = None
|
||||
|
||||
# 拆分后的组件(延迟初始化)
|
||||
self._candidate_resolver: Optional[CandidateResolver] = None
|
||||
self._request_dispatcher: Optional[RequestDispatcher] = None
|
||||
self._error_classifier: Optional[ErrorClassifier] = None
|
||||
|
||||
async def _ensure_initialized(self) -> None:
|
||||
"""确保异步组件已初始化"""
|
||||
if self.cache_scheduler is None:
|
||||
priority_mode = SystemConfigService.get_config(
|
||||
self.db,
|
||||
"provider_priority_mode",
|
||||
CacheAwareScheduler.PRIORITY_MODE_PROVIDER,
|
||||
)
|
||||
self.cache_scheduler = await get_cache_aware_scheduler(
|
||||
self.redis,
|
||||
priority_mode=priority_mode,
|
||||
)
|
||||
else:
|
||||
# 确保运行时配置变更能生效
|
||||
priority_mode = SystemConfigService.get_config(
|
||||
self.db,
|
||||
"provider_priority_mode",
|
||||
CacheAwareScheduler.PRIORITY_MODE_PROVIDER,
|
||||
)
|
||||
self.cache_scheduler.set_priority_mode(priority_mode)
|
||||
|
||||
# 确保 cache_scheduler 内部组件也已初始化
|
||||
await self.cache_scheduler._ensure_initialized()
|
||||
|
||||
if self.concurrency_manager is None:
|
||||
self.concurrency_manager = await get_concurrency_manager()
|
||||
|
||||
if self.request_executor is None and self.concurrency_manager is not None:
|
||||
self.request_executor = RequestExecutor(
|
||||
db=self.db,
|
||||
concurrency_manager=self.concurrency_manager,
|
||||
adaptive_manager=self.adaptive_manager,
|
||||
)
|
||||
|
||||
# 初始化拆分后的组件
|
||||
if self._candidate_resolver is None:
|
||||
self._candidate_resolver = CandidateResolver(
|
||||
db=self.db,
|
||||
cache_scheduler=self.cache_scheduler,
|
||||
)
|
||||
|
||||
if self._error_classifier is None:
|
||||
self._error_classifier = ErrorClassifier(
|
||||
db=self.db,
|
||||
cache_scheduler=self.cache_scheduler,
|
||||
adaptive_manager=self.adaptive_manager,
|
||||
)
|
||||
|
||||
if self._request_dispatcher is None and self.request_executor is not None:
|
||||
self._request_dispatcher = RequestDispatcher(
|
||||
db=self.db,
|
||||
request_executor=self.request_executor,
|
||||
cache_scheduler=self.cache_scheduler,
|
||||
)
|
||||
|
||||
async def _fetch_all_candidates(
|
||||
self,
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
user_api_key: Optional[ApiKey] = None,
|
||||
request_id: Optional[str] = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[List[ProviderCandidate], str]:
|
||||
"""
|
||||
收集所有可用的 Provider/Endpoint/Key 候选组合
|
||||
|
||||
委托给 CandidateResolver 处理。
|
||||
|
||||
Args:
|
||||
api_format: API 格式
|
||||
model_name: 模型名称
|
||||
affinity_key: 亲和性标识符(通常为API Key ID,用于缓存亲和性)
|
||||
user_api_key: 用户 API Key(用于 allowed_providers/allowed_api_formats 过滤)
|
||||
request_id: 请求 ID(用于日志)
|
||||
is_stream: 是否是流式请求,如果为 True 则过滤不支持流式的 Provider
|
||||
capability_requirements: 能力需求(用于过滤不满足能力要求的 Key)
|
||||
|
||||
Returns:
|
||||
(所有候选组合的列表, global_model_id)
|
||||
|
||||
Raises:
|
||||
ProviderNotAvailableException: 没有找到任何可用候选时
|
||||
"""
|
||||
assert self._candidate_resolver is not None
|
||||
return await self._candidate_resolver.fetch_candidates(
|
||||
api_format=api_format,
|
||||
model_name=model_name,
|
||||
affinity_key=affinity_key,
|
||||
user_api_key=user_api_key,
|
||||
request_id=request_id,
|
||||
is_stream=is_stream,
|
||||
capability_requirements=capability_requirements,
|
||||
)
|
||||
|
||||
def _create_candidate_records(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
request_id: Optional[str],
|
||||
user_id: str,
|
||||
user_api_key: ApiKey,
|
||||
required_capabilities: Optional[Dict[str, bool]] = None,
|
||||
) -> Dict[Tuple[int, int], str]:
|
||||
"""
|
||||
为所有候选预先创建 available 状态记录(批量插入优化)
|
||||
|
||||
委托给 CandidateResolver 处理。
|
||||
|
||||
Args:
|
||||
all_candidates: 所有候选组合
|
||||
request_id: 请求 ID
|
||||
user_id: 用户 ID
|
||||
user_api_key: 用户 API Key 对象
|
||||
required_capabilities: 请求需要的能力标签
|
||||
|
||||
Returns:
|
||||
candidate_record_map: {(candidate_index, retry_index): candidate_record_id}
|
||||
"""
|
||||
assert self._candidate_resolver is not None
|
||||
return self._candidate_resolver.create_candidate_records(
|
||||
all_candidates=all_candidates,
|
||||
request_id=request_id,
|
||||
user_id=user_id,
|
||||
user_api_key=user_api_key,
|
||||
required_capabilities=required_capabilities,
|
||||
)
|
||||
|
||||
async def _try_single_candidate(
|
||||
self,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
retry_index: int,
|
||||
candidate_record_id: str,
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
global_model_id: str,
|
||||
attempt_counter: int,
|
||||
max_attempts: int,
|
||||
is_stream: bool = False,
|
||||
) -> Tuple[Any, str, str, str, str, str]:
|
||||
"""
|
||||
尝试单个候选执行请求
|
||||
|
||||
委托给 RequestDispatcher 处理。
|
||||
|
||||
Args:
|
||||
candidate: 候选对象
|
||||
candidate_index: 候选索引
|
||||
retry_index: 重试索引
|
||||
candidate_record_id: 候选记录 ID
|
||||
user_api_key: 用户 API Key
|
||||
request_func: 请求函数
|
||||
request_id: 请求 ID
|
||||
api_format: API 格式
|
||||
model_name: 模型名称
|
||||
affinity_key: 亲和性标识符(通常为API Key ID)
|
||||
global_model_id: GlobalModel ID(规范化的模型标识,用于缓存亲和性)
|
||||
attempt_counter: 尝试计数
|
||||
max_attempts: 最大尝试次数
|
||||
is_stream: 是否为流式请求
|
||||
|
||||
Returns:
|
||||
(response, provider_name, candidate_record_id, provider_id, endpoint_id, key_id)
|
||||
|
||||
Raises:
|
||||
ExecutionError: 执行失败时
|
||||
"""
|
||||
assert self._request_dispatcher is not None
|
||||
return await self._request_dispatcher.dispatch(
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=retry_index,
|
||||
candidate_record_id=candidate_record_id,
|
||||
user_api_key=user_api_key,
|
||||
request_func=request_func,
|
||||
request_id=request_id,
|
||||
api_format=api_format,
|
||||
model_name=model_name,
|
||||
affinity_key=affinity_key,
|
||||
global_model_id=global_model_id,
|
||||
attempt_counter=attempt_counter,
|
||||
max_attempts=max_attempts,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
|
||||
async def _handle_candidate_error(
|
||||
self,
|
||||
exec_err: ExecutionError,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_record_id: str,
|
||||
retry_index: int,
|
||||
max_retries_for_candidate: int,
|
||||
affinity_key: str,
|
||||
api_format: APIFormat,
|
||||
global_model_id: str,
|
||||
request_id: Optional[str],
|
||||
attempt: int,
|
||||
max_attempts: int,
|
||||
) -> str:
|
||||
"""
|
||||
处理候选执行错误
|
||||
|
||||
Args:
|
||||
exec_err: 执行错误
|
||||
candidate: 候选对象
|
||||
candidate_record_id: 候选记录 ID
|
||||
retry_index: 当前重试索引
|
||||
max_retries_for_candidate: 该候选的最大重试次数
|
||||
affinity_key: 亲和性标识符(通常为API Key ID)
|
||||
api_format: API 格式
|
||||
global_model_id: GlobalModel ID(规范化的模型标识)
|
||||
request_id: 请求 ID
|
||||
attempt: 当前尝试次数
|
||||
max_attempts: 最大尝试次数
|
||||
|
||||
Returns:
|
||||
action: "continue" (继续重试), "break" (跳到下一个候选), "raise" (抛出异常)
|
||||
"""
|
||||
provider = candidate.provider
|
||||
endpoint = candidate.endpoint
|
||||
key = candidate.key
|
||||
|
||||
context = exec_err.context
|
||||
captured_key_concurrent = context.concurrent_requests
|
||||
elapsed_ms = context.elapsed_ms
|
||||
cause = exec_err.cause
|
||||
|
||||
has_retry_left = retry_index < (max_retries_for_candidate - 1)
|
||||
|
||||
# 确保 error_classifier 已初始化
|
||||
assert self._error_classifier is not None, "ErrorClassifier not initialized"
|
||||
|
||||
if isinstance(cause, ConcurrencyLimitError):
|
||||
logger.warning(f" [{request_id}] 并发限制 (attempt={attempt}/{max_attempts}): {cause}")
|
||||
RequestCandidateService.mark_candidate_skipped(
|
||||
db=self.db,
|
||||
candidate_id=candidate_record_id,
|
||||
skip_reason=f"并发限制: {str(cause)}",
|
||||
)
|
||||
return "break"
|
||||
|
||||
if isinstance(cause, httpx.HTTPStatusError):
|
||||
status_code = cause.response.status_code
|
||||
# 使用 ErrorClassifier 处理 HTTP 错误
|
||||
extra_data = await self._error_classifier.handle_http_error(
|
||||
http_error=cause,
|
||||
provider=provider,
|
||||
endpoint=endpoint,
|
||||
key=key,
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format,
|
||||
global_model_id=global_model_id,
|
||||
request_id=request_id,
|
||||
captured_key_concurrent=captured_key_concurrent,
|
||||
elapsed_ms=elapsed_ms,
|
||||
max_attempts=max_attempts,
|
||||
attempt=attempt,
|
||||
)
|
||||
|
||||
# 检查是否为客户端请求错误(不应重试)
|
||||
converted_error = extra_data.get("converted_error")
|
||||
# 从 extra_data 中移除 converted_error,避免序列化问题
|
||||
serializable_extra_data = {k: v for k, v in extra_data.items() if k != "converted_error"}
|
||||
|
||||
if isinstance(converted_error, UpstreamClientException):
|
||||
logger.warning(f" [{request_id}] 客户端请求错误,停止重试: {converted_error.message}")
|
||||
RequestCandidateService.mark_candidate_failed(
|
||||
db=self.db,
|
||||
candidate_id=candidate_record_id,
|
||||
error_type="UpstreamClientException",
|
||||
error_message=converted_error.message,
|
||||
status_code=status_code,
|
||||
latency_ms=elapsed_ms,
|
||||
concurrent_requests=captured_key_concurrent,
|
||||
extra_data=serializable_extra_data,
|
||||
)
|
||||
# 重新包装异常,附加 request_metadata 以便记录 usage
|
||||
converted_error.request_metadata = {
|
||||
"provider": provider.name,
|
||||
"provider_id": str(provider.id),
|
||||
"provider_endpoint_id": str(endpoint.id),
|
||||
"provider_api_key_id": str(key.id),
|
||||
"api_format": api_format.value if hasattr(api_format, "value") else str(api_format),
|
||||
}
|
||||
raise converted_error
|
||||
|
||||
RequestCandidateService.mark_candidate_failed(
|
||||
db=self.db,
|
||||
candidate_id=candidate_record_id,
|
||||
error_type="HTTPStatusError",
|
||||
error_message=f"HTTP {status_code}: {str(cause)}",
|
||||
status_code=status_code,
|
||||
latency_ms=elapsed_ms,
|
||||
concurrent_requests=captured_key_concurrent,
|
||||
extra_data=serializable_extra_data,
|
||||
)
|
||||
return "continue" if has_retry_left else "break"
|
||||
|
||||
if isinstance(cause, self._error_classifier.RETRIABLE_ERRORS):
|
||||
# 使用 ErrorClassifier 处理可重试错误
|
||||
await self._error_classifier.handle_retriable_error(
|
||||
error=cause,
|
||||
provider=provider,
|
||||
endpoint=endpoint,
|
||||
key=key,
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format,
|
||||
global_model_id=global_model_id,
|
||||
captured_key_concurrent=captured_key_concurrent,
|
||||
elapsed_ms=elapsed_ms,
|
||||
request_id=request_id,
|
||||
attempt=attempt,
|
||||
max_attempts=max_attempts,
|
||||
)
|
||||
# str(cause) 可能为空(如 httpx 超时异常),使用 repr() 作为备用
|
||||
error_msg = str(cause) or repr(cause)
|
||||
RequestCandidateService.mark_candidate_failed(
|
||||
db=self.db,
|
||||
candidate_id=candidate_record_id,
|
||||
error_type=type(cause).__name__,
|
||||
error_message=error_msg,
|
||||
latency_ms=elapsed_ms,
|
||||
concurrent_requests=captured_key_concurrent,
|
||||
)
|
||||
return "continue" if has_retry_left else "break"
|
||||
|
||||
# 未知错误:记录失败并抛出
|
||||
error_msg = str(cause) or repr(cause)
|
||||
RequestCandidateService.mark_candidate_failed(
|
||||
db=self.db,
|
||||
candidate_id=candidate_record_id,
|
||||
error_type=type(cause).__name__,
|
||||
error_message=error_msg,
|
||||
latency_ms=elapsed_ms,
|
||||
concurrent_requests=captured_key_concurrent,
|
||||
)
|
||||
return "raise"
|
||||
|
||||
def _create_pending_usage_record(
|
||||
self,
|
||||
request_id: Optional[str],
|
||||
user_api_key: ApiKey,
|
||||
model_name: str,
|
||||
is_stream: bool,
|
||||
api_format_enum: APIFormat,
|
||||
) -> None:
|
||||
"""创建 pending 状态的使用记录(用于实时状态追踪)"""
|
||||
if not request_id:
|
||||
return
|
||||
|
||||
from src.services.usage.service import UsageService
|
||||
|
||||
try:
|
||||
from src.models.database import User
|
||||
|
||||
user = self.db.query(User).filter(User.id == user_api_key.user_id).first()
|
||||
UsageService.create_pending_usage(
|
||||
db=self.db,
|
||||
request_id=request_id,
|
||||
user=user,
|
||||
api_key=user_api_key,
|
||||
model=model_name,
|
||||
is_stream=is_stream,
|
||||
api_format=api_format_enum.value,
|
||||
)
|
||||
except Exception as e:
|
||||
# 创建 pending 记录失败不应阻塞请求
|
||||
logger.warning(f"创建 pending 使用记录失败: {e}")
|
||||
|
||||
async def _execute_candidates_loop(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
candidate_record_map: Dict[Tuple[int, int], str],
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
api_format_enum: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
global_model_id: str,
|
||||
is_stream: bool = False,
|
||||
) -> Tuple[Any, str, Optional[str], Optional[str], Optional[str], Optional[str]]:
|
||||
"""遍历所有候选执行请求,返回第一个成功的结果或抛出异常"""
|
||||
attempt_counter = 0
|
||||
max_attempts = 0
|
||||
last_error: Optional[Exception] = None
|
||||
last_candidate: Optional[ProviderCandidate] = None
|
||||
|
||||
for candidate_index, candidate in enumerate(all_candidates):
|
||||
last_candidate = candidate
|
||||
|
||||
if candidate.is_skipped:
|
||||
logger.debug(f" [{request_id}] 跳过候选: Provider={candidate.provider.name}, "
|
||||
f"Reason={candidate.skip_reason}")
|
||||
continue
|
||||
|
||||
result = await self._try_candidate_with_retries(
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
candidate_record_map=candidate_record_map,
|
||||
user_api_key=user_api_key,
|
||||
request_func=request_func,
|
||||
request_id=request_id,
|
||||
api_format_enum=api_format_enum,
|
||||
model_name=model_name,
|
||||
affinity_key=affinity_key,
|
||||
global_model_id=global_model_id,
|
||||
attempt_counter=attempt_counter,
|
||||
max_attempts=max_attempts,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
|
||||
if result["success"]:
|
||||
response: Tuple[Any, str, Optional[str], Optional[str], Optional[str], Optional[str]] = result["response"]
|
||||
return response
|
||||
|
||||
# 更新计数器和错误信息
|
||||
attempt_counter = result["attempt_counter"]
|
||||
max_attempts = result["max_attempts"]
|
||||
if result.get("error"):
|
||||
last_error = result["error"]
|
||||
if result.get("should_raise") and last_error is not None:
|
||||
self._attach_metadata_to_error(last_error, last_candidate, model_name, api_format_enum)
|
||||
raise last_error
|
||||
|
||||
# 所有组合都已尝试完毕,全部失败
|
||||
self._raise_all_failed_exception(request_id, max_attempts, last_candidate, model_name, api_format_enum)
|
||||
|
||||
async def _try_candidate_with_retries(
|
||||
self,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
candidate_record_map: Dict[Tuple[int, int], str],
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
api_format_enum: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
global_model_id: str,
|
||||
attempt_counter: int,
|
||||
max_attempts: int,
|
||||
is_stream: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
"""尝试单个候选(含重试逻辑),返回执行结果"""
|
||||
provider = candidate.provider
|
||||
endpoint = candidate.endpoint
|
||||
max_retries_for_candidate = int(endpoint.max_retries) if candidate.is_cached else 1
|
||||
|
||||
for retry_index in range(max_retries_for_candidate):
|
||||
attempt_counter += 1
|
||||
max_attempts = max(max_attempts, attempt_counter)
|
||||
|
||||
if retry_index == 0:
|
||||
# 首次尝试该候选
|
||||
cache_hint = " (cached)" if candidate.is_cached else ""
|
||||
logger.info(f" [{request_id[:8] if request_id else 'N/A'}] -> {provider.name}{cache_hint}")
|
||||
else:
|
||||
logger.info(f" [{request_id[:8] if request_id else 'N/A'}] -> {provider.name} (retry {retry_index})")
|
||||
|
||||
candidate_record_id = candidate_record_map[(candidate_index, retry_index)]
|
||||
|
||||
try:
|
||||
response = await self._try_single_candidate(
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=retry_index,
|
||||
candidate_record_id=candidate_record_id,
|
||||
user_api_key=user_api_key,
|
||||
request_func=request_func,
|
||||
request_id=request_id,
|
||||
api_format=api_format_enum,
|
||||
model_name=model_name,
|
||||
affinity_key=affinity_key,
|
||||
global_model_id=global_model_id,
|
||||
attempt_counter=attempt_counter,
|
||||
max_attempts=max_attempts,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
return {"success": True, "response": response}
|
||||
|
||||
except ExecutionError as exec_err:
|
||||
action = await self._handle_candidate_error(
|
||||
exec_err=exec_err,
|
||||
candidate=candidate,
|
||||
candidate_record_id=candidate_record_id,
|
||||
retry_index=retry_index,
|
||||
max_retries_for_candidate=max_retries_for_candidate,
|
||||
affinity_key=affinity_key,
|
||||
api_format=api_format_enum,
|
||||
global_model_id=global_model_id,
|
||||
request_id=request_id,
|
||||
attempt=attempt_counter,
|
||||
max_attempts=max_attempts,
|
||||
)
|
||||
|
||||
if action == "continue":
|
||||
continue
|
||||
elif action == "break":
|
||||
break
|
||||
elif action == "raise":
|
||||
return {
|
||||
"success": False,
|
||||
"should_raise": True,
|
||||
"error": exec_err.cause,
|
||||
"attempt_counter": attempt_counter,
|
||||
"max_attempts": max_attempts,
|
||||
}
|
||||
|
||||
return {
|
||||
"success": False,
|
||||
"attempt_counter": attempt_counter,
|
||||
"max_attempts": max_attempts,
|
||||
}
|
||||
|
||||
def _attach_metadata_to_error(
|
||||
self,
|
||||
error: Optional[Exception],
|
||||
candidate: Optional[ProviderCandidate],
|
||||
model_name: str,
|
||||
api_format_enum: APIFormat,
|
||||
) -> None:
|
||||
"""附加 candidate 信息到异常,以便记录 usage"""
|
||||
if not error or not candidate:
|
||||
return
|
||||
|
||||
from src.services.request.result import RequestMetadata
|
||||
|
||||
existing_metadata = getattr(error, "request_metadata", None)
|
||||
if existing_metadata and getattr(existing_metadata, "api_format", None):
|
||||
return # 已有完整的 metadata
|
||||
|
||||
metadata = RequestMetadata(
|
||||
provider_request_headers=(
|
||||
getattr(existing_metadata, "provider_request_headers", {})
|
||||
if existing_metadata
|
||||
else {}
|
||||
),
|
||||
provider=getattr(existing_metadata, "provider", None) or str(candidate.provider.name),
|
||||
model=getattr(existing_metadata, "model", None) or model_name,
|
||||
provider_id=getattr(existing_metadata, "provider_id", None) or str(candidate.provider.id),
|
||||
provider_endpoint_id=(
|
||||
getattr(existing_metadata, "provider_endpoint_id", None)
|
||||
or str(candidate.endpoint.id)
|
||||
),
|
||||
provider_api_key_id=(
|
||||
getattr(existing_metadata, "provider_api_key_id", None)
|
||||
or str(candidate.key.id)
|
||||
),
|
||||
api_format=api_format_enum.value,
|
||||
)
|
||||
# 使用 setattr 避免类型检查错误
|
||||
setattr(error, "request_metadata", metadata)
|
||||
|
||||
def _raise_all_failed_exception(
|
||||
self,
|
||||
request_id: Optional[str],
|
||||
max_attempts: int,
|
||||
last_candidate: Optional[ProviderCandidate],
|
||||
model_name: str,
|
||||
api_format_enum: APIFormat,
|
||||
) -> NoReturn:
|
||||
"""所有组合都失败时抛出异常"""
|
||||
logger.error(f" [{request_id}] 所有 {max_attempts} 个组合均失败")
|
||||
|
||||
request_metadata = None
|
||||
if last_candidate:
|
||||
request_metadata = {
|
||||
"provider": last_candidate.provider.name,
|
||||
"model": model_name,
|
||||
"provider_id": str(last_candidate.provider.id),
|
||||
"provider_endpoint_id": str(last_candidate.endpoint.id),
|
||||
"provider_api_key_id": str(last_candidate.key.id),
|
||||
"api_format": api_format_enum.value,
|
||||
}
|
||||
|
||||
raise ProviderNotAvailableException(
|
||||
f"所有Provider均不可用,已尝试{max_attempts}个组合",
|
||||
request_metadata=request_metadata,
|
||||
)
|
||||
|
||||
async def execute_with_fallback(
|
||||
self,
|
||||
api_format: Union[str, APIFormat],
|
||||
model_name: str,
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[[Provider, ProviderEndpoint, ProviderAPIKey], Any],
|
||||
request_id: Optional[str] = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[Any, str, Optional[str], Optional[str], Optional[str], Optional[str]]:
|
||||
"""
|
||||
执行请求,并在失败时自动故障转移(缓存感知)
|
||||
|
||||
Args:
|
||||
api_format: API 格式(如 'CLAUDE', 'OPENAI')
|
||||
model_name: 模型名称
|
||||
user_api_key: 用户的 API Key对象
|
||||
request_func: 请求函数,接收 (provider, endpoint, key) 参数,返回响应
|
||||
request_id: 请求 ID(用于日志)
|
||||
is_stream: 是否是流式请求,如果为 True 则过滤不支持流式的 Provider
|
||||
capability_requirements: 能力需求(用于过滤不满足能力要求的 Key)
|
||||
|
||||
Returns:
|
||||
(请求响应, 实际Provider名称, RequestTraceAttempt ID, provider_id, endpoint_id, key_id)
|
||||
|
||||
Raises:
|
||||
ProviderNotAvailableException: 所有 Providers 都失败后抛出
|
||||
"""
|
||||
await self._ensure_initialized()
|
||||
|
||||
# 准备执行上下文
|
||||
affinity_key = str(user_api_key.id)
|
||||
user_id = str(user_api_key.user_id)
|
||||
api_format_enum = normalize_api_format(api_format)
|
||||
|
||||
logger.debug(f"[FallbackOrchestrator] execute_with_fallback 被调用: "
|
||||
f"api_format={api_format_enum.value}, model_name={model_name}, "
|
||||
f"request_id={request_id}, is_stream={is_stream}")
|
||||
|
||||
# 创建 pending 状态的使用记录
|
||||
self._create_pending_usage_record(request_id, user_api_key, model_name, is_stream, api_format_enum)
|
||||
|
||||
# 1. 收集所有候选(同时获取规范化的 global_model_id 用于缓存亲和性)
|
||||
all_candidates, global_model_id = await self._fetch_all_candidates(
|
||||
api_format=api_format_enum,
|
||||
model_name=model_name,
|
||||
affinity_key=affinity_key,
|
||||
user_api_key=user_api_key,
|
||||
request_id=request_id,
|
||||
is_stream=is_stream,
|
||||
capability_requirements=capability_requirements,
|
||||
)
|
||||
|
||||
# 2. 批量创建候选记录
|
||||
candidate_record_map = self._create_candidate_records(
|
||||
all_candidates=all_candidates,
|
||||
request_id=request_id,
|
||||
user_id=user_id,
|
||||
user_api_key=user_api_key,
|
||||
required_capabilities=capability_requirements,
|
||||
)
|
||||
|
||||
# 3. 遍历候选执行请求(使用 global_model_id 用于缓存亲和性)
|
||||
return await self._execute_candidates_loop(
|
||||
all_candidates=all_candidates,
|
||||
candidate_record_map=candidate_record_map,
|
||||
user_api_key=user_api_key,
|
||||
request_func=request_func,
|
||||
request_id=request_id,
|
||||
api_format_enum=api_format_enum,
|
||||
model_name=model_name,
|
||||
affinity_key=affinity_key,
|
||||
global_model_id=global_model_id,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
158
src/services/orchestration/request_dispatcher.py
Normal file
158
src/services/orchestration/request_dispatcher.py
Normal file
@@ -0,0 +1,158 @@
|
||||
"""
|
||||
请求分发器
|
||||
|
||||
负责执行单个候选请求
|
||||
"""
|
||||
|
||||
from typing import Any, Callable, Optional, Tuple
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.enums import APIFormat
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey
|
||||
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
|
||||
from src.services.request.candidate import RequestCandidateService
|
||||
from src.services.request.executor import RequestExecutor
|
||||
|
||||
|
||||
|
||||
class RequestDispatcher:
|
||||
"""
|
||||
请求分发器 - 负责执行单个候选请求
|
||||
|
||||
职责:
|
||||
1. 执行请求并返回结果
|
||||
2. 更新候选状态(pending -> success/failed)
|
||||
3. 设置缓存亲和性(成功时)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: Session,
|
||||
request_executor: RequestExecutor,
|
||||
cache_scheduler: Optional[CacheAwareScheduler] = None,
|
||||
) -> None:
|
||||
"""
|
||||
初始化请求分发器
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
request_executor: 请求执行器
|
||||
cache_scheduler: 缓存调度器(可选)
|
||||
"""
|
||||
self.db = db
|
||||
self.request_executor = request_executor
|
||||
self.cache_scheduler = cache_scheduler
|
||||
|
||||
async def dispatch(
|
||||
self,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
retry_index: int,
|
||||
candidate_record_id: str,
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
global_model_id: str,
|
||||
attempt_counter: int,
|
||||
max_attempts: int,
|
||||
is_stream: bool = False,
|
||||
) -> Tuple[Any, str, str, str, str, str]:
|
||||
"""
|
||||
执行请求并返回结果
|
||||
|
||||
Args:
|
||||
candidate: 候选对象
|
||||
candidate_index: 候选索引
|
||||
retry_index: 重试索引
|
||||
candidate_record_id: 候选记录 ID
|
||||
user_api_key: 用户 API Key
|
||||
request_func: 请求函数
|
||||
request_id: 请求 ID
|
||||
api_format: API 格式
|
||||
model_name: 模型名称
|
||||
affinity_key: 亲和性标识符(通常为API Key ID)
|
||||
global_model_id: GlobalModel ID(规范化的模型标识,用于缓存亲和性)
|
||||
attempt_counter: 尝试计数
|
||||
max_attempts: 最大尝试次数
|
||||
is_stream: 是否为流式请求
|
||||
|
||||
Returns:
|
||||
(response, provider_name, candidate_record_id, provider_id, endpoint_id, key_id)
|
||||
|
||||
Raises:
|
||||
ExecutionError: 执行失败时
|
||||
"""
|
||||
provider = candidate.provider
|
||||
endpoint = candidate.endpoint
|
||||
key = candidate.key
|
||||
|
||||
# 显式转换为 str
|
||||
provider_id = str(provider.id)
|
||||
provider_name = str(provider.name)
|
||||
endpoint_id = str(endpoint.id)
|
||||
key_id = str(key.id)
|
||||
cache_ttl_minutes = int(key.cache_ttl_minutes or 0)
|
||||
provider_supports_caching = cache_ttl_minutes > 0
|
||||
provider_cache_ttl_seconds: Optional[int] = (
|
||||
cache_ttl_minutes * 60 if cache_ttl_minutes > 0 else None
|
||||
)
|
||||
|
||||
# 更新状态为 pending
|
||||
RequestCandidateService.update_candidate_status(
|
||||
db=self.db, candidate_id=candidate_record_id, status="pending"
|
||||
)
|
||||
|
||||
# 执行请求
|
||||
execution_result = await self.request_executor.execute(
|
||||
candidate=candidate,
|
||||
candidate_id=candidate_record_id,
|
||||
candidate_index=candidate_index,
|
||||
user_api_key=user_api_key,
|
||||
request_func=request_func,
|
||||
request_id=request_id,
|
||||
api_format=api_format,
|
||||
model_name=model_name,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
|
||||
context = execution_result.context
|
||||
elapsed_ms = context.elapsed_ms or 0
|
||||
|
||||
# 流式请求:标记为 streaming 状态(请求尚未完成)
|
||||
# 非流式请求:标记为 success 状态
|
||||
# 注意:executor.execute() 内部已经处理了状态标记,这里不再重复
|
||||
# 流式请求的 success 状态会在流完成后由 _record_stream_stats 方法标记
|
||||
|
||||
# 设置缓存亲和性
|
||||
if provider_supports_caching and self.cache_scheduler is not None:
|
||||
try:
|
||||
api_format_str = (
|
||||
api_format.value if isinstance(api_format, APIFormat) else api_format
|
||||
)
|
||||
await self.cache_scheduler.set_cache_affinity(
|
||||
affinity_key=affinity_key,
|
||||
provider_id=provider_id,
|
||||
endpoint_id=endpoint_id,
|
||||
key_id=key_id,
|
||||
api_format=api_format_str,
|
||||
global_model_id=global_model_id,
|
||||
ttl=provider_cache_ttl_seconds,
|
||||
)
|
||||
except Exception as cache_exc:
|
||||
logger.warning(f" [{request_id}] 设置缓存亲和性失败: {cache_exc}")
|
||||
|
||||
logger.debug(f" [{request_id}] 请求成功: Provider={provider_name}, 耗时={elapsed_ms}ms")
|
||||
|
||||
return (
|
||||
execution_result.response,
|
||||
provider_name,
|
||||
candidate_record_id,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
)
|
||||
16
src/services/provider/__init__.py
Normal file
16
src/services/provider/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
Provider 服务模块
|
||||
|
||||
包含 Provider 管理、格式处理、传输层等功能。
|
||||
"""
|
||||
|
||||
from src.services.provider.format import normalize_api_format
|
||||
from src.services.provider.service import ProviderService
|
||||
from src.services.provider.transport import build_provider_headers, build_provider_url
|
||||
|
||||
__all__ = [
|
||||
"ProviderService",
|
||||
"normalize_api_format",
|
||||
"build_provider_headers",
|
||||
"build_provider_url",
|
||||
]
|
||||
21
src/services/provider/format.py
Normal file
21
src/services/provider/format.py
Normal file
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
API 格式辅助函数,确保在调度/编排链路中使用统一的枚举值。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
from src.core.api_format_metadata import resolve_api_format
|
||||
from src.core.enums import APIFormat
|
||||
|
||||
|
||||
def normalize_api_format(
|
||||
value: Union[str, APIFormat, None], default: APIFormat = APIFormat.CLAUDE
|
||||
) -> APIFormat:
|
||||
"""
|
||||
将任意字符串/枚举值归一化为 APIFormat。
|
||||
未识别的值回退到默认枚举(默认 CLAUDE)。
|
||||
"""
|
||||
resolved = resolve_api_format(value)
|
||||
return resolved or default
|
||||
61
src/services/provider/response_normalizer.py
Normal file
61
src/services/provider/response_normalizer.py
Normal file
@@ -0,0 +1,61 @@
|
||||
"""响应标准化服务,用于 STANDARD 模式下的响应格式验证和补全"""
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.claude import ClaudeResponse
|
||||
|
||||
|
||||
|
||||
class ResponseNormalizer:
|
||||
"""响应标准化器 - 用于标准模式下验证和补全响应字段"""
|
||||
|
||||
@staticmethod
|
||||
def normalize_claude_response(
|
||||
response_data: Dict[str, Any], request_id: Optional[str] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
标准化 Claude API 响应
|
||||
|
||||
Args:
|
||||
response_data: 原始响应数据
|
||||
request_id: 请求ID(用于日志)
|
||||
|
||||
Returns:
|
||||
标准化后的响应数据(失败时返回原始数据)
|
||||
"""
|
||||
if "error" in response_data:
|
||||
logger.debug(f"[ResponseNormalizer] 检测到错误响应,跳过标准化 | ID:{request_id}")
|
||||
return response_data
|
||||
|
||||
try:
|
||||
validated = ClaudeResponse.model_validate(response_data)
|
||||
normalized = validated.model_dump(mode="json", exclude_none=False)
|
||||
|
||||
logger.debug(f"[ResponseNormalizer] 响应标准化成功 | ID:{request_id}")
|
||||
return normalized
|
||||
|
||||
except Exception as e:
|
||||
logger.debug(f"[ResponseNormalizer] 响应验证失败,透传原始数据 | ID:{request_id}")
|
||||
return response_data
|
||||
|
||||
@staticmethod
|
||||
def should_normalize(response_data: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
判断是否需要进行标准化
|
||||
|
||||
Args:
|
||||
response_data: 响应数据
|
||||
|
||||
Returns:
|
||||
是否需要标准化
|
||||
"""
|
||||
# 错误响应不需要标准化
|
||||
if "error" in response_data:
|
||||
return False
|
||||
|
||||
# 已经包含新字段的响应不需要再次标准化
|
||||
if "context_management" in response_data and "container" in response_data:
|
||||
return False
|
||||
|
||||
return True
|
||||
159
src/services/provider/service.py
Normal file
159
src/services/provider/service.py
Normal file
@@ -0,0 +1,159 @@
|
||||
"""
|
||||
提供商服务
|
||||
负责提供商选择、模型映射和请求处理
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import GlobalModel, Model, Provider
|
||||
from src.services.model.cost import ModelCostService
|
||||
from src.services.model.mapper import ModelMapperMiddleware, ModelRoutingMiddleware
|
||||
|
||||
|
||||
|
||||
class ProviderService:
|
||||
"""提供商服务类"""
|
||||
|
||||
def __init__(self, db: Session):
|
||||
"""
|
||||
初始化提供商服务
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
"""
|
||||
self.db = db
|
||||
self.mapper = ModelMapperMiddleware(db)
|
||||
self.router = ModelRoutingMiddleware(db)
|
||||
self.cost_service = ModelCostService(db)
|
||||
|
||||
async def _check_model_availability(self, model_name: str):
|
||||
"""
|
||||
检查模型是否可用(严格白名单模式)
|
||||
|
||||
Args:
|
||||
model_name: 模型名称
|
||||
|
||||
Returns:
|
||||
Model对象如果存在且激活,否则None
|
||||
"""
|
||||
# 首先检查是否有直接的模型记录
|
||||
model = (
|
||||
self.db.query(Model)
|
||||
.filter(Model.provider_model_name == model_name, Model.is_active == True)
|
||||
.first()
|
||||
)
|
||||
|
||||
if model:
|
||||
return model
|
||||
|
||||
# 方案 A:检查是否是别名(全局别名系统)
|
||||
from src.services.model.mapping_resolver import resolve_model_to_global_name
|
||||
|
||||
global_model_name = await resolve_model_to_global_name(self.db, model_name)
|
||||
|
||||
# 查找 GlobalModel
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(GlobalModel.name == global_model_name, GlobalModel.is_active == True)
|
||||
.first()
|
||||
)
|
||||
|
||||
if global_model:
|
||||
# 查找任意 Provider 的 Model 实现
|
||||
model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(Model.global_model_id == global_model.id, Model.is_active == True)
|
||||
.first()
|
||||
)
|
||||
if model_obj:
|
||||
return model_obj
|
||||
|
||||
return None
|
||||
|
||||
async def _check_provider_model_availability(self, provider_id: str, model_name: str):
|
||||
"""
|
||||
检查特定提供商是否支持特定模型
|
||||
|
||||
Args:
|
||||
provider_id: 提供商ID
|
||||
model_name: 模型名称
|
||||
|
||||
Returns:
|
||||
Model对象如果该提供商支持该模型且激活,否则None
|
||||
"""
|
||||
# 首先检查该提供商下是否有直接的模型记录
|
||||
model = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_id,
|
||||
Model.provider_model_name == model_name,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if model:
|
||||
return model
|
||||
|
||||
# 方案 A:检查是否是别名
|
||||
from src.services.model.mapping_resolver import resolve_model_to_global_name
|
||||
|
||||
global_model_name = await resolve_model_to_global_name(self.db, model_name, provider_id)
|
||||
|
||||
# 查找 GlobalModel
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(GlobalModel.name == global_model_name, GlobalModel.is_active == True)
|
||||
.first()
|
||||
)
|
||||
|
||||
if global_model:
|
||||
# 查找该 Provider 是否有实现该 GlobalModel
|
||||
model_obj = (
|
||||
self.db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_id,
|
||||
Model.global_model_id == global_model.id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if model_obj:
|
||||
return model_obj
|
||||
|
||||
return None
|
||||
|
||||
def calculate_cost(
|
||||
self, provider: Provider, model: str, input_tokens: int, output_tokens: int
|
||||
) -> Dict[str, float]:
|
||||
"""
|
||||
计算使用成本
|
||||
|
||||
Args:
|
||||
provider: 提供商对象
|
||||
model: 模型名
|
||||
input_tokens: 输入tokens
|
||||
output_tokens: 输出tokens
|
||||
|
||||
Returns:
|
||||
成本信息
|
||||
"""
|
||||
return self.mapper.calculate_cost(model, provider.id, input_tokens, output_tokens)
|
||||
|
||||
def get_available_models(self) -> Dict[str, list]:
|
||||
"""
|
||||
获取所有可用的模型
|
||||
|
||||
Returns:
|
||||
模型和支持的提供商映射
|
||||
"""
|
||||
return self.router.get_available_models()
|
||||
|
||||
def clear_cache(self):
|
||||
"""清空缓存"""
|
||||
self.mapper.clear_cache()
|
||||
self.cost_service.clear_cache()
|
||||
logger.info("Provider service cache cleared")
|
||||
146
src/services/provider/transport.py
Normal file
146
src/services/provider/transport.py
Normal file
@@ -0,0 +1,146 @@
|
||||
"""
|
||||
统一的 Provider 请求构建工具。
|
||||
|
||||
负责:
|
||||
- 根据 endpoint/key 构建标准请求头
|
||||
- 根据 API 格式或端点配置生成请求 URL
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.core.api_format_metadata import get_auth_config, get_default_path, resolve_api_format
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.enums import APIFormat
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
|
||||
def build_provider_headers(
|
||||
endpoint,
|
||||
key,
|
||||
original_headers: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
根据 endpoint/key 构建请求头,并透传客户端自定义头。
|
||||
"""
|
||||
headers: Dict[str, str] = {}
|
||||
|
||||
decrypted_key = crypto_service.decrypt(key.api_key)
|
||||
|
||||
# 根据 API 格式自动选择认证头
|
||||
api_format = getattr(endpoint, "api_format", None)
|
||||
resolved_format = resolve_api_format(api_format)
|
||||
auth_header, auth_type = (
|
||||
get_auth_config(resolved_format) if resolved_format else ("Authorization", "bearer")
|
||||
)
|
||||
|
||||
if auth_type == "bearer":
|
||||
headers[auth_header] = f"Bearer {decrypted_key}"
|
||||
else:
|
||||
headers[auth_header] = decrypted_key
|
||||
|
||||
if endpoint.headers:
|
||||
headers.update(endpoint.headers)
|
||||
|
||||
excluded_headers = {
|
||||
"host",
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"x-goog-api-key",
|
||||
"content-length",
|
||||
"transfer-encoding",
|
||||
}
|
||||
|
||||
if original_headers:
|
||||
for name, value in original_headers.items():
|
||||
if name.lower() not in excluded_headers:
|
||||
headers[name] = value
|
||||
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
|
||||
if "Content-Type" not in headers and "content-type" not in headers:
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
return headers
|
||||
|
||||
|
||||
def build_provider_url(
|
||||
endpoint,
|
||||
*,
|
||||
query_params: Optional[Dict[str, Any]] = None,
|
||||
path_params: Optional[Dict[str, Any]] = None,
|
||||
is_stream: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
根据 endpoint 配置生成请求 URL
|
||||
|
||||
优先级:
|
||||
1. endpoint.custom_path - 自定义路径(支持模板变量如 {model})
|
||||
2. API 格式默认路径 - 根据 api_format 自动选择
|
||||
|
||||
Args:
|
||||
endpoint: 端点配置
|
||||
query_params: 查询参数
|
||||
path_params: 路径模板参数 (如 {model})
|
||||
is_stream: 是否为流式请求,用于 Gemini API 选择正确的操作方法
|
||||
"""
|
||||
base = endpoint.base_url.rstrip("/")
|
||||
|
||||
# 准备路径参数,添加 Gemini API 所需的 action 参数
|
||||
effective_path_params = dict(path_params) if path_params else {}
|
||||
|
||||
# 为 Gemini API 格式自动添加 action 参数
|
||||
resolved_format = resolve_api_format(endpoint.api_format)
|
||||
if resolved_format in (APIFormat.GEMINI, APIFormat.GEMINI_CLI):
|
||||
if "action" not in effective_path_params:
|
||||
effective_path_params["action"] = (
|
||||
"streamGenerateContent" if is_stream else "generateContent"
|
||||
)
|
||||
|
||||
# 优先使用 custom_path 字段
|
||||
if endpoint.custom_path:
|
||||
path = endpoint.custom_path
|
||||
if effective_path_params:
|
||||
try:
|
||||
path = path.format(**effective_path_params)
|
||||
except KeyError:
|
||||
# 如果模板变量不匹配,保持原路径
|
||||
pass
|
||||
else:
|
||||
# 使用 API 格式的默认路径
|
||||
path = _resolve_default_path(endpoint.api_format)
|
||||
if effective_path_params:
|
||||
try:
|
||||
path = path.format(**effective_path_params)
|
||||
except KeyError:
|
||||
# 如果模板变量不匹配,保持原路径
|
||||
pass
|
||||
|
||||
if not path.startswith("/"):
|
||||
path = f"/{path}"
|
||||
|
||||
url = f"{base}{path}"
|
||||
|
||||
# 添加查询参数
|
||||
if query_params:
|
||||
query_string = urlencode(query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
|
||||
return url
|
||||
|
||||
|
||||
def _resolve_default_path(api_format) -> str:
|
||||
"""
|
||||
根据 API 格式返回默认路径
|
||||
"""
|
||||
resolved = resolve_api_format(api_format)
|
||||
if resolved:
|
||||
return get_default_path(resolved)
|
||||
|
||||
logger.warning(f"Unknown api_format '{api_format}' for endpoint, fallback to '/'")
|
||||
return "/"
|
||||
19
src/services/rate_limit/__init__.py
Normal file
19
src/services/rate_limit/__init__.py
Normal file
@@ -0,0 +1,19 @@
|
||||
"""
|
||||
限流服务模块
|
||||
|
||||
包含自适应并发控制、RPM限流、IP限流等功能。
|
||||
"""
|
||||
|
||||
from src.services.rate_limit.adaptive_concurrency import AdaptiveConcurrencyManager
|
||||
from src.services.rate_limit.concurrency_manager import ConcurrencyManager
|
||||
from src.services.rate_limit.detector import RateLimitDetector
|
||||
from src.services.rate_limit.ip_limiter import IPRateLimiter
|
||||
from src.services.rate_limit.rpm_limiter import RPMLimiter
|
||||
|
||||
__all__ = [
|
||||
"AdaptiveConcurrencyManager",
|
||||
"ConcurrencyManager",
|
||||
"IPRateLimiter",
|
||||
"RPMLimiter",
|
||||
"RateLimitDetector",
|
||||
]
|
||||
558
src/services/rate_limit/adaptive_concurrency.py
Normal file
558
src/services/rate_limit/adaptive_concurrency.py
Normal file
@@ -0,0 +1,558 @@
|
||||
"""
|
||||
自适应并发调整器 - 基于滑动窗口利用率的并发限制调整
|
||||
|
||||
核心改进(相对于旧版基于"持续高利用率"的方案):
|
||||
- 使用滑动窗口采样,容忍并发波动
|
||||
- 基于窗口内高利用率采样比例决策,而非要求连续高利用率
|
||||
- 增加探测性扩容机制,长时间稳定时主动尝试扩容
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional, cast
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.constants import ConcurrencyDefaults
|
||||
from src.core.batch_committer import get_batch_committer
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.rate_limit.detector import RateLimitInfo, RateLimitType
|
||||
|
||||
|
||||
class AdaptiveStrategy:
|
||||
"""自适应策略类型"""
|
||||
|
||||
AIMD = "aimd" # 加性增-乘性减 (Additive Increase Multiplicative Decrease)
|
||||
CONSERVATIVE = "conservative" # 保守策略(只减不增)
|
||||
AGGRESSIVE = "aggressive" # 激进策略(快速探测)
|
||||
|
||||
|
||||
class AdaptiveConcurrencyManager:
|
||||
"""
|
||||
自适应并发管理器
|
||||
|
||||
核心算法:基于滑动窗口利用率的 AIMD
|
||||
- 滑动窗口记录最近 N 次请求的利用率
|
||||
- 当窗口内高利用率采样比例 >= 60% 时触发扩容
|
||||
- 遇到 429 错误时乘性减少 (*0.7)
|
||||
- 长时间无 429 且有流量时触发探测性扩容
|
||||
|
||||
扩容条件(满足任一即可):
|
||||
1. 滑动窗口扩容:窗口内 >= 60% 的采样利用率 >= 70%,且不在冷却期
|
||||
2. 探测性扩容:距上次 429 超过 30 分钟,且期间有足够请求量
|
||||
|
||||
关键特性:
|
||||
1. 滑动窗口容忍并发波动,不会因单次低利用率重置
|
||||
2. 区分并发限制和 RPM 限制
|
||||
3. 探测性扩容避免长期卡在低限制
|
||||
4. 记录调整历史
|
||||
"""
|
||||
|
||||
# 默认配置 - 使用统一常量
|
||||
DEFAULT_INITIAL_LIMIT = ConcurrencyDefaults.INITIAL_LIMIT
|
||||
MIN_CONCURRENT_LIMIT = ConcurrencyDefaults.MIN_CONCURRENT_LIMIT
|
||||
MAX_CONCURRENT_LIMIT = ConcurrencyDefaults.MAX_CONCURRENT_LIMIT
|
||||
|
||||
# AIMD 参数
|
||||
INCREASE_STEP = ConcurrencyDefaults.INCREASE_STEP
|
||||
DECREASE_MULTIPLIER = ConcurrencyDefaults.DECREASE_MULTIPLIER
|
||||
|
||||
# 滑动窗口参数
|
||||
UTILIZATION_WINDOW_SIZE = ConcurrencyDefaults.UTILIZATION_WINDOW_SIZE
|
||||
UTILIZATION_WINDOW_SECONDS = ConcurrencyDefaults.UTILIZATION_WINDOW_SECONDS
|
||||
UTILIZATION_THRESHOLD = ConcurrencyDefaults.UTILIZATION_THRESHOLD
|
||||
HIGH_UTILIZATION_RATIO = ConcurrencyDefaults.HIGH_UTILIZATION_RATIO
|
||||
MIN_SAMPLES_FOR_DECISION = ConcurrencyDefaults.MIN_SAMPLES_FOR_DECISION
|
||||
|
||||
# 探测性扩容参数
|
||||
PROBE_INCREASE_INTERVAL_MINUTES = ConcurrencyDefaults.PROBE_INCREASE_INTERVAL_MINUTES
|
||||
PROBE_INCREASE_MIN_REQUESTS = ConcurrencyDefaults.PROBE_INCREASE_MIN_REQUESTS
|
||||
|
||||
# 记录历史数量
|
||||
MAX_HISTORY_RECORDS = 20
|
||||
|
||||
def __init__(self, strategy: str = AdaptiveStrategy.AIMD):
|
||||
"""
|
||||
初始化自适应并发管理器
|
||||
|
||||
Args:
|
||||
strategy: 调整策略
|
||||
"""
|
||||
self.strategy = strategy
|
||||
|
||||
def handle_429_error(
|
||||
self,
|
||||
db: Session,
|
||||
key: ProviderAPIKey,
|
||||
rate_limit_info: RateLimitInfo,
|
||||
current_concurrent: Optional[int] = None,
|
||||
) -> int:
|
||||
"""
|
||||
处理429错误,调整并发限制
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
key: API Key对象
|
||||
rate_limit_info: 速率限制信息
|
||||
current_concurrent: 当前并发数
|
||||
|
||||
Returns:
|
||||
调整后的并发限制
|
||||
"""
|
||||
# max_concurrent=NULL 表示启用自适应,max_concurrent=数字 表示固定限制
|
||||
is_adaptive = key.max_concurrent is None
|
||||
|
||||
if not is_adaptive:
|
||||
logger.debug(
|
||||
f"Key {key.id} 设置了固定并发限制 ({key.max_concurrent}),跳过自适应调整"
|
||||
)
|
||||
return int(key.max_concurrent) # type: ignore[arg-type]
|
||||
|
||||
# 更新429统计
|
||||
key.last_429_at = datetime.now(timezone.utc) # type: ignore[assignment]
|
||||
key.last_429_type = rate_limit_info.limit_type # type: ignore[assignment]
|
||||
key.last_concurrent_peak = current_concurrent # type: ignore[assignment]
|
||||
|
||||
# 遇到 429 错误,清空利用率采样窗口(重新开始收集)
|
||||
key.utilization_samples = [] # type: ignore[assignment]
|
||||
|
||||
if rate_limit_info.limit_type == RateLimitType.CONCURRENT:
|
||||
# 并发限制:减少并发数
|
||||
key.concurrent_429_count = int(key.concurrent_429_count or 0) + 1 # type: ignore[assignment]
|
||||
|
||||
# 获取当前有效限制(自适应模式使用 learned_max_concurrent)
|
||||
old_limit = int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
|
||||
new_limit = self._decrease_limit(old_limit, current_concurrent)
|
||||
|
||||
logger.warning(
|
||||
f"[CONCURRENT] 并发限制触发: Key {key.id[:8]}... | "
|
||||
f"当前并发: {current_concurrent} | "
|
||||
f"调整: {old_limit} -> {new_limit}"
|
||||
)
|
||||
|
||||
# 记录调整历史
|
||||
self._record_adjustment(
|
||||
key,
|
||||
old_limit=old_limit,
|
||||
new_limit=new_limit,
|
||||
reason="concurrent_429",
|
||||
current_concurrent=current_concurrent,
|
||||
)
|
||||
|
||||
# 更新学习到的并发限制
|
||||
key.learned_max_concurrent = new_limit # type: ignore[assignment]
|
||||
|
||||
elif rate_limit_info.limit_type == RateLimitType.RPM:
|
||||
# RPM限制:不调整并发,只记录
|
||||
key.rpm_429_count = int(key.rpm_429_count or 0) + 1 # type: ignore[assignment]
|
||||
|
||||
logger.info(
|
||||
f"[RPM] RPM限制触发: Key {key.id[:8]}... | "
|
||||
f"retry_after: {rate_limit_info.retry_after}s | "
|
||||
f"不调整并发限制"
|
||||
)
|
||||
|
||||
else:
|
||||
# 未知类型:保守处理,轻微减少
|
||||
logger.warning(
|
||||
f"[UNKNOWN] 未知429类型: Key {key.id[:8]}... | "
|
||||
f"当前并发: {current_concurrent} | "
|
||||
f"保守减少并发"
|
||||
)
|
||||
|
||||
old_limit = int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
|
||||
new_limit = max(int(old_limit * 0.9), self.MIN_CONCURRENT_LIMIT) # 减少10%
|
||||
|
||||
self._record_adjustment(
|
||||
key,
|
||||
old_limit=old_limit,
|
||||
new_limit=new_limit,
|
||||
reason="unknown_429",
|
||||
current_concurrent=current_concurrent,
|
||||
)
|
||||
|
||||
key.learned_max_concurrent = new_limit # type: ignore[assignment]
|
||||
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
return int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
|
||||
|
||||
def handle_success(
|
||||
self,
|
||||
db: Session,
|
||||
key: ProviderAPIKey,
|
||||
current_concurrent: int,
|
||||
) -> Optional[int]:
|
||||
"""
|
||||
处理成功请求,基于滑动窗口利用率考虑增加并发限制
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
key: API Key对象
|
||||
current_concurrent: 当前并发数(必需,用于计算利用率)
|
||||
|
||||
Returns:
|
||||
调整后的并发限制(如果有调整),否则返回 None
|
||||
"""
|
||||
# max_concurrent=NULL 表示启用自适应
|
||||
is_adaptive = key.max_concurrent is None
|
||||
|
||||
if not is_adaptive:
|
||||
return None
|
||||
|
||||
current_limit = int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
|
||||
|
||||
# 计算当前利用率
|
||||
utilization = float(current_concurrent / current_limit) if current_limit > 0 else 0.0
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
now_ts = now.timestamp()
|
||||
|
||||
# 更新滑动窗口
|
||||
samples = self._update_utilization_window(key, now_ts, utilization)
|
||||
|
||||
# 检查是否满足扩容条件
|
||||
increase_reason = self._check_increase_conditions(key, samples, now)
|
||||
|
||||
if increase_reason and current_limit < self.MAX_CONCURRENT_LIMIT:
|
||||
old_limit = current_limit
|
||||
new_limit = self._increase_limit(current_limit)
|
||||
|
||||
# 计算窗口统计用于日志
|
||||
avg_util = sum(s["util"] for s in samples) / len(samples) if samples else 0
|
||||
high_util_count = sum(1 for s in samples if s["util"] >= self.UTILIZATION_THRESHOLD)
|
||||
high_util_ratio = high_util_count / len(samples) if samples else 0
|
||||
|
||||
logger.info(
|
||||
f"[INCREASE] {increase_reason}: Key {key.id[:8]}... | "
|
||||
f"窗口采样: {len(samples)} | "
|
||||
f"平均利用率: {avg_util:.1%} | "
|
||||
f"高利用率比例: {high_util_ratio:.1%} | "
|
||||
f"调整: {old_limit} -> {new_limit}"
|
||||
)
|
||||
|
||||
# 记录调整历史
|
||||
self._record_adjustment(
|
||||
key,
|
||||
old_limit=old_limit,
|
||||
new_limit=new_limit,
|
||||
reason=increase_reason,
|
||||
avg_utilization=round(avg_util, 2),
|
||||
high_util_ratio=round(high_util_ratio, 2),
|
||||
sample_count=len(samples),
|
||||
current_concurrent=current_concurrent,
|
||||
)
|
||||
|
||||
# 更新限制
|
||||
key.learned_max_concurrent = new_limit # type: ignore[assignment]
|
||||
|
||||
# 如果是探测性扩容,更新探测时间
|
||||
if increase_reason == "probe_increase":
|
||||
key.last_probe_increase_at = now # type: ignore[assignment]
|
||||
|
||||
# 扩容后清空采样窗口,重新开始收集
|
||||
key.utilization_samples = [] # type: ignore[assignment]
|
||||
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
return new_limit
|
||||
|
||||
# 定期持久化采样数据(每5个采样保存一次)
|
||||
if len(samples) % 5 == 0:
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
return None
|
||||
|
||||
def _update_utilization_window(
|
||||
self, key: ProviderAPIKey, now_ts: float, utilization: float
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
更新利用率滑动窗口
|
||||
|
||||
Args:
|
||||
key: API Key对象
|
||||
now_ts: 当前时间戳
|
||||
utilization: 当前利用率
|
||||
|
||||
Returns:
|
||||
更新后的采样列表
|
||||
"""
|
||||
samples: List[Dict[str, Any]] = list(key.utilization_samples or [])
|
||||
|
||||
# 添加新采样
|
||||
samples.append({"ts": now_ts, "util": round(utilization, 3)})
|
||||
|
||||
# 移除过期采样(超过时间窗口)
|
||||
cutoff_ts = now_ts - self.UTILIZATION_WINDOW_SECONDS
|
||||
samples = [s for s in samples if s["ts"] > cutoff_ts]
|
||||
|
||||
# 限制采样数量
|
||||
if len(samples) > self.UTILIZATION_WINDOW_SIZE:
|
||||
samples = samples[-self.UTILIZATION_WINDOW_SIZE:]
|
||||
|
||||
# 更新到 key 对象
|
||||
key.utilization_samples = samples # type: ignore[assignment]
|
||||
|
||||
return samples
|
||||
|
||||
def _check_increase_conditions(
|
||||
self, key: ProviderAPIKey, samples: List[Dict[str, Any]], now: datetime
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
检查是否满足扩容条件
|
||||
|
||||
Args:
|
||||
key: API Key对象
|
||||
samples: 利用率采样列表
|
||||
now: 当前时间
|
||||
|
||||
Returns:
|
||||
扩容原因(如果满足条件),否则返回 None
|
||||
"""
|
||||
# 检查是否在冷却期
|
||||
if self._is_in_cooldown(key):
|
||||
return None
|
||||
|
||||
# 条件1:滑动窗口扩容
|
||||
if len(samples) >= self.MIN_SAMPLES_FOR_DECISION:
|
||||
high_util_count = sum(1 for s in samples if s["util"] >= self.UTILIZATION_THRESHOLD)
|
||||
high_util_ratio = high_util_count / len(samples)
|
||||
|
||||
if high_util_ratio >= self.HIGH_UTILIZATION_RATIO:
|
||||
return "high_utilization"
|
||||
|
||||
# 条件2:探测性扩容(长时间无 429 且有流量)
|
||||
if self._should_probe_increase(key, samples, now):
|
||||
return "probe_increase"
|
||||
|
||||
return None
|
||||
|
||||
def _should_probe_increase(
|
||||
self, key: ProviderAPIKey, samples: List[Dict[str, Any]], now: datetime
|
||||
) -> bool:
|
||||
"""
|
||||
检查是否应该进行探测性扩容
|
||||
|
||||
条件:
|
||||
1. 距上次 429 超过 PROBE_INCREASE_INTERVAL_MINUTES 分钟
|
||||
2. 距上次探测性扩容超过 PROBE_INCREASE_INTERVAL_MINUTES 分钟
|
||||
3. 期间有足够的请求量(采样数 >= PROBE_INCREASE_MIN_REQUESTS)
|
||||
4. 平均利用率 > 30%(说明确实有使用需求)
|
||||
|
||||
Args:
|
||||
key: API Key对象
|
||||
samples: 利用率采样列表
|
||||
now: 当前时间
|
||||
|
||||
Returns:
|
||||
是否应该探测性扩容
|
||||
"""
|
||||
probe_interval_seconds = self.PROBE_INCREASE_INTERVAL_MINUTES * 60
|
||||
|
||||
# 检查距上次 429 的时间
|
||||
if key.last_429_at:
|
||||
last_429_at = cast(datetime, key.last_429_at)
|
||||
time_since_429 = (now - last_429_at).total_seconds()
|
||||
if time_since_429 < probe_interval_seconds:
|
||||
return False
|
||||
|
||||
# 检查距上次探测性扩容的时间
|
||||
if key.last_probe_increase_at:
|
||||
last_probe = cast(datetime, key.last_probe_increase_at)
|
||||
time_since_probe = (now - last_probe).total_seconds()
|
||||
if time_since_probe < probe_interval_seconds:
|
||||
return False
|
||||
|
||||
# 检查请求量
|
||||
if len(samples) < self.PROBE_INCREASE_MIN_REQUESTS:
|
||||
return False
|
||||
|
||||
# 检查平均利用率(确保确实有使用需求)
|
||||
avg_util = sum(s["util"] for s in samples) / len(samples)
|
||||
if avg_util < 0.3: # 至少 30% 利用率
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _is_in_cooldown(self, key: ProviderAPIKey) -> bool:
|
||||
"""
|
||||
检查是否在 429 错误后的冷却期内
|
||||
|
||||
Args:
|
||||
key: API Key对象
|
||||
|
||||
Returns:
|
||||
True 如果在冷却期内,否则 False
|
||||
"""
|
||||
if key.last_429_at is None:
|
||||
return False
|
||||
|
||||
last_429_at = cast(datetime, key.last_429_at)
|
||||
time_since_429 = (datetime.now(timezone.utc) - last_429_at).total_seconds()
|
||||
cooldown_seconds = ConcurrencyDefaults.COOLDOWN_AFTER_429_MINUTES * 60
|
||||
|
||||
return bool(time_since_429 < cooldown_seconds)
|
||||
|
||||
def _decrease_limit(
|
||||
self,
|
||||
current_limit: int,
|
||||
current_concurrent: Optional[int] = None,
|
||||
) -> int:
|
||||
"""
|
||||
减少并发限制
|
||||
|
||||
策略:
|
||||
- 如果知道当前并发数,设置为当前并发的70%
|
||||
- 否则,使用乘性减少
|
||||
"""
|
||||
if current_concurrent:
|
||||
# 基于当前并发数减少
|
||||
new_limit = max(
|
||||
int(current_concurrent * self.DECREASE_MULTIPLIER), self.MIN_CONCURRENT_LIMIT
|
||||
)
|
||||
else:
|
||||
# 乘性减少
|
||||
new_limit = max(
|
||||
int(current_limit * self.DECREASE_MULTIPLIER), self.MIN_CONCURRENT_LIMIT
|
||||
)
|
||||
|
||||
return new_limit
|
||||
|
||||
def _increase_limit(self, current_limit: int) -> int:
|
||||
"""
|
||||
增加并发限制
|
||||
|
||||
策略:加性增加 (+1)
|
||||
"""
|
||||
new_limit = min(current_limit + self.INCREASE_STEP, self.MAX_CONCURRENT_LIMIT)
|
||||
return new_limit
|
||||
|
||||
def _record_adjustment(
|
||||
self,
|
||||
key: ProviderAPIKey,
|
||||
old_limit: int,
|
||||
new_limit: int,
|
||||
reason: str,
|
||||
**extra_data: Any,
|
||||
) -> None:
|
||||
"""
|
||||
记录并发调整历史
|
||||
|
||||
Args:
|
||||
key: API Key对象
|
||||
old_limit: 原限制
|
||||
new_limit: 新限制
|
||||
reason: 调整原因
|
||||
**extra_data: 额外数据
|
||||
"""
|
||||
history: List[Dict[str, Any]] = list(key.adjustment_history or [])
|
||||
|
||||
record = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"old_limit": old_limit,
|
||||
"new_limit": new_limit,
|
||||
"reason": reason,
|
||||
**extra_data,
|
||||
}
|
||||
history.append(record)
|
||||
|
||||
# 保留最近N条记录
|
||||
if len(history) > self.MAX_HISTORY_RECORDS:
|
||||
history = history[-self.MAX_HISTORY_RECORDS:]
|
||||
|
||||
key.adjustment_history = history # type: ignore[assignment]
|
||||
|
||||
def get_adjustment_stats(self, key: ProviderAPIKey) -> Dict[str, Any]:
|
||||
"""
|
||||
获取调整统计信息
|
||||
|
||||
Args:
|
||||
key: API Key对象
|
||||
|
||||
Returns:
|
||||
统计信息
|
||||
"""
|
||||
history: List[Dict[str, Any]] = list(key.adjustment_history or [])
|
||||
samples: List[Dict[str, Any]] = list(key.utilization_samples or [])
|
||||
|
||||
# max_concurrent=NULL 表示自适应,否则为固定限制
|
||||
is_adaptive = key.max_concurrent is None
|
||||
current_limit = int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
|
||||
effective_limit = current_limit if is_adaptive else int(key.max_concurrent) # type: ignore
|
||||
|
||||
# 计算窗口统计
|
||||
avg_utilization: Optional[float] = None
|
||||
high_util_ratio: Optional[float] = None
|
||||
if samples:
|
||||
avg_utilization = sum(s["util"] for s in samples) / len(samples)
|
||||
high_util_count = sum(1 for s in samples if s["util"] >= self.UTILIZATION_THRESHOLD)
|
||||
high_util_ratio = high_util_count / len(samples)
|
||||
|
||||
last_429_at_str: Optional[str] = None
|
||||
if key.last_429_at:
|
||||
last_429_at_str = cast(datetime, key.last_429_at).isoformat()
|
||||
|
||||
last_probe_at_str: Optional[str] = None
|
||||
if key.last_probe_increase_at:
|
||||
last_probe_at_str = cast(datetime, key.last_probe_increase_at).isoformat()
|
||||
|
||||
return {
|
||||
"adaptive_mode": is_adaptive,
|
||||
"max_concurrent": key.max_concurrent, # NULL=自适应,数字=固定限制
|
||||
"effective_limit": effective_limit, # 当前有效限制
|
||||
"learned_limit": key.learned_max_concurrent, # 学习到的限制
|
||||
"concurrent_429_count": int(key.concurrent_429_count or 0),
|
||||
"rpm_429_count": int(key.rpm_429_count or 0),
|
||||
"last_429_at": last_429_at_str,
|
||||
"last_429_type": key.last_429_type,
|
||||
"adjustment_count": len(history),
|
||||
"recent_adjustments": history[-5:] if history else [],
|
||||
# 滑动窗口相关
|
||||
"window_sample_count": len(samples),
|
||||
"window_avg_utilization": round(avg_utilization, 3) if avg_utilization else None,
|
||||
"window_high_util_ratio": round(high_util_ratio, 3) if high_util_ratio else None,
|
||||
"utilization_threshold": self.UTILIZATION_THRESHOLD,
|
||||
"high_util_ratio_threshold": self.HIGH_UTILIZATION_RATIO,
|
||||
"min_samples_for_decision": self.MIN_SAMPLES_FOR_DECISION,
|
||||
# 探测性扩容相关
|
||||
"last_probe_increase_at": last_probe_at_str,
|
||||
"probe_increase_interval_minutes": self.PROBE_INCREASE_INTERVAL_MINUTES,
|
||||
}
|
||||
|
||||
def reset_learning(self, db: Session, key: ProviderAPIKey) -> None:
|
||||
"""
|
||||
重置学习状态(管理员功能)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
key: API Key对象
|
||||
"""
|
||||
logger.info(f"[RESET] 重置学习状态: Key {key.id[:8]}...")
|
||||
|
||||
key.learned_max_concurrent = None # type: ignore[assignment]
|
||||
key.concurrent_429_count = 0 # type: ignore[assignment]
|
||||
key.rpm_429_count = 0 # type: ignore[assignment]
|
||||
key.last_429_at = None # type: ignore[assignment]
|
||||
key.last_429_type = None # type: ignore[assignment]
|
||||
key.last_concurrent_peak = None # type: ignore[assignment]
|
||||
key.adjustment_history = [] # type: ignore[assignment]
|
||||
key.utilization_samples = [] # type: ignore[assignment]
|
||||
key.last_probe_increase_at = None # type: ignore[assignment]
|
||||
|
||||
db.flush()
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
|
||||
# 全局单例
|
||||
_adaptive_manager: Optional[AdaptiveConcurrencyManager] = None
|
||||
|
||||
|
||||
def get_adaptive_manager() -> AdaptiveConcurrencyManager:
|
||||
"""获取全局自适应管理器单例"""
|
||||
global _adaptive_manager
|
||||
if _adaptive_manager is None:
|
||||
_adaptive_manager = AdaptiveConcurrencyManager()
|
||||
return _adaptive_manager
|
||||
340
src/services/rate_limit/adaptive_reservation.py
Normal file
340
src/services/rate_limit/adaptive_reservation.py
Normal file
@@ -0,0 +1,340 @@
|
||||
"""
|
||||
自适应预留比例管理器
|
||||
|
||||
根据学习置信度和当前负载动态计算缓存用户预留比例,
|
||||
解决固定 30% 预留在学习初期和负载变化时的不适应问题。
|
||||
|
||||
核心思路:
|
||||
1. 探测阶段:使用低预留,让系统快速学习真实并发限制
|
||||
2. 稳定阶段:根据置信度和负载动态调整预留比例
|
||||
3. 置信度计算:综合考虑连续成功次数、429冷却时间、调整历史稳定性
|
||||
"""
|
||||
|
||||
import statistics
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
from src.config.constants import AdaptiveReservationDefaults
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReservationConfig:
|
||||
"""预留比例配置(使用统一常量作为默认值)"""
|
||||
|
||||
# 探测阶段配置
|
||||
probe_phase_requests: int = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.PROBE_PHASE_REQUESTS
|
||||
)
|
||||
probe_reservation: float = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.PROBE_RESERVATION
|
||||
)
|
||||
|
||||
# 稳定阶段配置
|
||||
stable_min_reservation: float = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.STABLE_MIN_RESERVATION
|
||||
)
|
||||
stable_max_reservation: float = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.STABLE_MAX_RESERVATION
|
||||
)
|
||||
|
||||
# 置信度计算参数
|
||||
success_count_for_full_confidence: int = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.SUCCESS_COUNT_FOR_FULL_CONFIDENCE
|
||||
)
|
||||
cooldown_hours_for_full_confidence: int = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.COOLDOWN_HOURS_FOR_FULL_CONFIDENCE
|
||||
)
|
||||
|
||||
# 负载阈值
|
||||
low_load_threshold: float = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.LOW_LOAD_THRESHOLD
|
||||
)
|
||||
high_load_threshold: float = field(
|
||||
default_factory=lambda: AdaptiveReservationDefaults.HIGH_LOAD_THRESHOLD
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReservationResult:
|
||||
"""预留比例计算结果"""
|
||||
|
||||
ratio: float # 最终预留比例
|
||||
phase: str # 当前阶段: "probe" | "stable"
|
||||
confidence: float # 置信度 (0-1)
|
||||
load_factor: float # 负载因子 (0-1)
|
||||
details: Dict[str, Any] # 详细信息
|
||||
|
||||
|
||||
class AdaptiveReservationManager:
|
||||
"""
|
||||
自适应预留比例管理器
|
||||
|
||||
工作原理:
|
||||
1. 探测阶段(请求数 < 阈值):
|
||||
- 使用低预留比例(10%),不浪费资源
|
||||
- 让系统快速探测真实并发限制
|
||||
|
||||
2. 稳定阶段(请求数 >= 阈值):
|
||||
- 根据置信度和负载动态计算预留比例
|
||||
- 置信度高 + 负载高 = 高预留(保护缓存用户)
|
||||
- 置信度低或负载低 = 低预留(避免浪费)
|
||||
|
||||
置信度因素:
|
||||
- 连续成功次数:越多说明当前限制越准确
|
||||
- 429冷却时间:距离上次429越久越稳定
|
||||
- 调整历史稳定性:最近调整的方差越小越稳定
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[ReservationConfig] = None):
|
||||
self.config = config or ReservationConfig()
|
||||
self._cache: Dict[str, ReservationResult] = {} # 简单的内存缓存
|
||||
|
||||
def calculate_reservation(
|
||||
self,
|
||||
key: "ProviderAPIKey",
|
||||
current_concurrent: int = 0,
|
||||
effective_limit: Optional[int] = None,
|
||||
) -> ReservationResult:
|
||||
"""
|
||||
计算当前应使用的预留比例
|
||||
|
||||
Args:
|
||||
key: ProviderAPIKey 对象
|
||||
current_concurrent: 当前并发数
|
||||
effective_limit: 有效并发限制(学习值或配置值)
|
||||
|
||||
Returns:
|
||||
ReservationResult 包含预留比例和详细信息
|
||||
"""
|
||||
# 计算总请求数(用于判断阶段)
|
||||
total_requests = self._get_total_requests(key)
|
||||
|
||||
# 计算负载率
|
||||
load_ratio = self._calculate_load_ratio(current_concurrent, effective_limit)
|
||||
|
||||
# 阶段1: 探测阶段
|
||||
if total_requests < self.config.probe_phase_requests:
|
||||
return ReservationResult(
|
||||
ratio=self.config.probe_reservation,
|
||||
phase="probe",
|
||||
confidence=0.0,
|
||||
load_factor=load_ratio,
|
||||
details={
|
||||
"total_requests": total_requests,
|
||||
"probe_threshold": self.config.probe_phase_requests,
|
||||
"reason": "探测阶段,使用低预留让系统学习真实限制",
|
||||
},
|
||||
)
|
||||
|
||||
# 阶段2: 稳定阶段
|
||||
confidence = self._calculate_confidence(key)
|
||||
ratio = self._calculate_stable_ratio(confidence, load_ratio)
|
||||
|
||||
return ReservationResult(
|
||||
ratio=ratio,
|
||||
phase="stable",
|
||||
confidence=confidence,
|
||||
load_factor=load_ratio,
|
||||
details={
|
||||
"total_requests": total_requests,
|
||||
"confidence_factors": self._get_confidence_breakdown(key),
|
||||
"reason": self._get_ratio_reason(confidence, load_ratio),
|
||||
},
|
||||
)
|
||||
|
||||
def _get_total_requests(self, key: "ProviderAPIKey") -> int:
|
||||
"""获取总请求数(用于判断是否过了探测阶段)"""
|
||||
# 使用总请求计数作为基准
|
||||
request_count = key.request_count or 0
|
||||
|
||||
# 如果 request_count 为 0,使用 429 计数 + 成功计数作为近似值
|
||||
if request_count == 0:
|
||||
concurrent_429 = key.concurrent_429_count or 0
|
||||
rpm_429 = key.rpm_429_count or 0
|
||||
success_count = key.success_count or 0
|
||||
# 调整历史中的记录数也可以参考
|
||||
history_count = len(key.adjustment_history or []) * 10
|
||||
return concurrent_429 + rpm_429 + success_count + history_count
|
||||
|
||||
return request_count
|
||||
|
||||
def _calculate_load_ratio(
|
||||
self, current_concurrent: int, effective_limit: Optional[int]
|
||||
) -> float:
|
||||
"""计算当前负载率"""
|
||||
if not effective_limit or effective_limit <= 0:
|
||||
return 0.0
|
||||
return min(current_concurrent / effective_limit, 1.0)
|
||||
|
||||
def _calculate_confidence(self, key: "ProviderAPIKey") -> float:
|
||||
"""
|
||||
计算学习值的置信度 (0-1)
|
||||
|
||||
三个因素各占一定权重:
|
||||
- 成功率:40%(基于总成功数/总请求数)
|
||||
- 429冷却时间:30%
|
||||
- 调整历史稳定性:30%
|
||||
"""
|
||||
scores = self._get_confidence_breakdown(key)
|
||||
return min(
|
||||
scores["success_score"] + scores["cooldown_score"] + scores["stability_score"], 1.0
|
||||
)
|
||||
|
||||
def _get_confidence_breakdown(self, key: "ProviderAPIKey") -> Dict[str, float]:
|
||||
"""获取置信度各因素的详细分数"""
|
||||
# 因素1: 成功率(权重 40%)
|
||||
# 使用成功率而非连续成功次数,更准确反映 Key 的稳定性
|
||||
request_count = key.request_count or 0
|
||||
success_count = key.success_count or 0
|
||||
|
||||
if request_count >= self.config.success_count_for_full_confidence:
|
||||
# 请求数足够时,根据成功率计算
|
||||
success_rate = success_count / request_count if request_count > 0 else 0
|
||||
success_score = success_rate * 0.4
|
||||
elif request_count > 0:
|
||||
# 请求数不足时,按比例折算
|
||||
progress_ratio = request_count / self.config.success_count_for_full_confidence
|
||||
success_rate = success_count / request_count
|
||||
success_score = success_rate * progress_ratio * 0.4
|
||||
else:
|
||||
success_score = 0.0
|
||||
|
||||
# 因素2: 429冷却时间(权重 30%)
|
||||
if key.last_429_at:
|
||||
now = datetime.now(timezone.utc)
|
||||
# 确保 last_429_at 有时区信息
|
||||
last_429 = key.last_429_at
|
||||
if last_429.tzinfo is None:
|
||||
last_429 = last_429.replace(tzinfo=timezone.utc)
|
||||
hours_since_429 = (now - last_429).total_seconds() / 3600
|
||||
cooldown_ratio = min(
|
||||
hours_since_429 / self.config.cooldown_hours_for_full_confidence, 1.0
|
||||
)
|
||||
cooldown_score = cooldown_ratio * 0.3
|
||||
else:
|
||||
# 从未触发 429,给满分
|
||||
cooldown_score = 0.3
|
||||
|
||||
# 因素3: 调整历史稳定性(权重 30%)
|
||||
history = key.adjustment_history or []
|
||||
if len(history) >= 3:
|
||||
# 取最近的调整记录
|
||||
recent = history[-5:] if len(history) >= 5 else history
|
||||
limits = [h.get("new_limit", 0) for h in recent if h.get("new_limit")]
|
||||
|
||||
if len(limits) >= 2:
|
||||
try:
|
||||
variance = statistics.variance(limits)
|
||||
# 方差越小越稳定,方差为10时分数接近0
|
||||
stability_ratio = max(0, 1 - variance / 10)
|
||||
stability_score = stability_ratio * 0.3
|
||||
except statistics.StatisticsError:
|
||||
stability_score = 0.15
|
||||
else:
|
||||
stability_score = 0.15
|
||||
else:
|
||||
# 历史数据不足,给一半分
|
||||
stability_score = 0.15
|
||||
|
||||
# 计算成功率用于返回
|
||||
success_rate_pct = (success_count / request_count * 100) if request_count > 0 else None
|
||||
|
||||
return {
|
||||
"success_score": round(success_score, 3),
|
||||
"cooldown_score": round(cooldown_score, 3),
|
||||
"stability_score": round(stability_score, 3),
|
||||
"request_count": request_count,
|
||||
"success_count": success_count,
|
||||
"success_rate": round(success_rate_pct, 1) if success_rate_pct is not None else None,
|
||||
"hours_since_429": (
|
||||
round(
|
||||
(
|
||||
datetime.now(timezone.utc) - key.last_429_at.replace(tzinfo=timezone.utc)
|
||||
).total_seconds()
|
||||
/ 3600,
|
||||
1,
|
||||
)
|
||||
if key.last_429_at
|
||||
else None
|
||||
),
|
||||
"history_count": len(history),
|
||||
}
|
||||
|
||||
def _calculate_stable_ratio(self, confidence: float, load_ratio: float) -> float:
|
||||
"""
|
||||
计算稳定阶段的预留比例
|
||||
|
||||
策略:
|
||||
- 低负载(<50%):使用最小预留,槽位充足无需过多预留
|
||||
- 中等负载(50-80%):根据置信度线性增加预留
|
||||
- 高负载(>80%):根据置信度使用较高预留保护缓存用户
|
||||
"""
|
||||
min_r = self.config.stable_min_reservation
|
||||
max_r = self.config.stable_max_reservation
|
||||
|
||||
if load_ratio < self.config.low_load_threshold:
|
||||
# 低负载:使用最小预留
|
||||
return min_r
|
||||
|
||||
if load_ratio < self.config.high_load_threshold:
|
||||
# 中等负载:根据置信度和负载线性插值
|
||||
# 负载越高、置信度越高,预留越多
|
||||
load_factor = (load_ratio - self.config.low_load_threshold) / (
|
||||
self.config.high_load_threshold - self.config.low_load_threshold
|
||||
)
|
||||
return min_r + confidence * load_factor * (max_r - min_r)
|
||||
|
||||
# 高负载:根据置信度决定预留比例
|
||||
# 置信度高 → 接近最大预留
|
||||
# 置信度低 → 保守预留(避免基于不准确的学习值过度预留)
|
||||
return min_r + confidence * (max_r - min_r)
|
||||
|
||||
def _get_ratio_reason(self, confidence: float, load_ratio: float) -> str:
|
||||
"""生成预留比例的解释"""
|
||||
if load_ratio < self.config.low_load_threshold:
|
||||
return f"低负载({load_ratio:.0%}),使用最小预留"
|
||||
|
||||
if confidence < 0.3:
|
||||
return f"置信度低({confidence:.0%}),保守预留避免浪费"
|
||||
|
||||
if confidence > 0.7 and load_ratio > self.config.high_load_threshold:
|
||||
return f"高置信度({confidence:.0%})+高负载({load_ratio:.0%}),使用较高预留保护缓存用户"
|
||||
|
||||
return f"置信度{confidence:.0%},负载{load_ratio:.0%},动态计算预留"
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""获取管理器统计信息"""
|
||||
return {
|
||||
"config": {
|
||||
"probe_phase_requests": self.config.probe_phase_requests,
|
||||
"probe_reservation": self.config.probe_reservation,
|
||||
"stable_min_reservation": self.config.stable_min_reservation,
|
||||
"stable_max_reservation": self.config.stable_max_reservation,
|
||||
"low_load_threshold": self.config.low_load_threshold,
|
||||
"high_load_threshold": self.config.high_load_threshold,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# 全局单例
|
||||
_reservation_manager: Optional[AdaptiveReservationManager] = None
|
||||
|
||||
|
||||
def get_adaptive_reservation_manager() -> AdaptiveReservationManager:
|
||||
"""获取全局自适应预留管理器单例"""
|
||||
global _reservation_manager
|
||||
if _reservation_manager is None:
|
||||
_reservation_manager = AdaptiveReservationManager()
|
||||
return _reservation_manager
|
||||
|
||||
|
||||
def reset_adaptive_reservation_manager():
|
||||
"""重置全局单例(用于测试)"""
|
||||
global _reservation_manager
|
||||
_reservation_manager = None
|
||||
582
src/services/rate_limit/concurrency_manager.py
Normal file
582
src/services/rate_limit/concurrency_manager.py
Normal file
@@ -0,0 +1,582 @@
|
||||
"""
|
||||
并发管理器 - 支持 Redis 或内存的并发控制
|
||||
|
||||
功能:
|
||||
1. Endpoint 级别的并发限制
|
||||
2. ProviderAPIKey 级别的并发限制
|
||||
3. 分布式环境下优先使用 Redis,多实例共享
|
||||
4. 在开发/单实例场景下自动降级为内存计数
|
||||
5. 自动释放和异常处理(Redis 提供 TTL,内存模式请确保手动释放)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import timedelta
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
class ConcurrencyManager:
|
||||
"""分布式并发管理器"""
|
||||
|
||||
_instance: Optional["ConcurrencyManager"] = None
|
||||
_redis: Optional[aioredis.Redis] = None
|
||||
|
||||
def __new__(cls):
|
||||
"""单例模式"""
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
"""初始化内存后端结构(只执行一次)"""
|
||||
if hasattr(self, "_memory_initialized"):
|
||||
return
|
||||
|
||||
self._memory_lock: asyncio.Lock = asyncio.Lock()
|
||||
self._memory_endpoint_counts: dict[str, int] = {}
|
||||
self._memory_key_counts: dict[str, int] = {}
|
||||
self._memory_initialized = True
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""初始化 Redis 连接"""
|
||||
if self._redis is not None:
|
||||
return
|
||||
|
||||
# 优先使用 REDIS_URL,如果没有则根据密码构建 URL
|
||||
redis_url = os.getenv("REDIS_URL")
|
||||
|
||||
if not redis_url:
|
||||
# 本地开发模式:从 REDIS_PASSWORD 构建 URL
|
||||
redis_password = os.getenv("REDIS_PASSWORD")
|
||||
if redis_password:
|
||||
redis_url = f"redis://:{redis_password}@localhost:6379/0"
|
||||
else:
|
||||
redis_url = "redis://localhost:6379/0"
|
||||
|
||||
try:
|
||||
self._redis = await aioredis.from_url(
|
||||
redis_url,
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
socket_timeout=5.0,
|
||||
socket_connect_timeout=5.0,
|
||||
)
|
||||
# 测试连接
|
||||
await self._redis.ping()
|
||||
# 脱敏显示(隐藏密码)
|
||||
safe_url = redis_url.split("@")[-1] if "@" in redis_url else redis_url
|
||||
logger.info(f"[OK] Redis 连接成功: {safe_url}")
|
||||
except Exception as e:
|
||||
logger.error(f"[ERROR] Redis 连接失败: {e}")
|
||||
logger.warning("[WARN] 并发控制将被禁用(仅在单实例环境下安全)")
|
||||
self._redis = None
|
||||
|
||||
async def close(self) -> None:
|
||||
"""关闭 Redis 连接"""
|
||||
if self._redis:
|
||||
await self._redis.close()
|
||||
self._redis = None
|
||||
logger.info("Redis 连接已关闭")
|
||||
|
||||
def _get_endpoint_key(self, endpoint_id: str) -> str:
|
||||
"""获取 Endpoint 并发计数的 Redis Key"""
|
||||
return f"concurrency:endpoint:{endpoint_id}"
|
||||
|
||||
def _get_key_key(self, key_id: str) -> str:
|
||||
"""获取 ProviderAPIKey 并发计数的 Redis Key"""
|
||||
return f"concurrency:key:{key_id}"
|
||||
|
||||
async def get_current_concurrency(
|
||||
self, endpoint_id: Optional[str] = None, key_id: Optional[str] = None
|
||||
) -> Tuple[int, int]:
|
||||
"""
|
||||
获取当前并发数
|
||||
|
||||
Args:
|
||||
endpoint_id: Endpoint ID(可选)
|
||||
key_id: ProviderAPIKey ID(可选)
|
||||
|
||||
Returns:
|
||||
(endpoint_concurrency, key_concurrency)
|
||||
"""
|
||||
if self._redis is None:
|
||||
async with self._memory_lock:
|
||||
endpoint_count = (
|
||||
self._memory_endpoint_counts.get(endpoint_id, 0) if endpoint_id else 0
|
||||
)
|
||||
key_count = self._memory_key_counts.get(key_id, 0) if key_id else 0
|
||||
return endpoint_count, key_count
|
||||
|
||||
endpoint_count = 0
|
||||
key_count = 0
|
||||
|
||||
try:
|
||||
if endpoint_id:
|
||||
endpoint_key = self._get_endpoint_key(endpoint_id)
|
||||
result = await self._redis.get(endpoint_key)
|
||||
endpoint_count = int(result) if result else 0
|
||||
|
||||
if key_id:
|
||||
key_key = self._get_key_key(key_id)
|
||||
result = await self._redis.get(key_key)
|
||||
key_count = int(result) if result else 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取并发数失败: {e}")
|
||||
|
||||
return endpoint_count, key_count
|
||||
|
||||
async def check_available(
|
||||
self,
|
||||
endpoint_id: str,
|
||||
endpoint_max_concurrent: Optional[int],
|
||||
key_id: str,
|
||||
key_max_concurrent: Optional[int],
|
||||
) -> bool:
|
||||
"""
|
||||
检查是否可以获取并发槽位(不实际获取)
|
||||
|
||||
Args:
|
||||
endpoint_id: Endpoint ID
|
||||
endpoint_max_concurrent: Endpoint 最大并发数(None 表示不限制)
|
||||
key_id: ProviderAPIKey ID
|
||||
key_max_concurrent: Key 最大并发数(None 表示不限制)
|
||||
|
||||
Returns:
|
||||
是否可用(True/False)
|
||||
"""
|
||||
if self._redis is None:
|
||||
async with self._memory_lock:
|
||||
endpoint_count = self._memory_endpoint_counts.get(endpoint_id, 0)
|
||||
key_count = self._memory_key_counts.get(key_id, 0)
|
||||
|
||||
if (
|
||||
endpoint_max_concurrent is not None
|
||||
and endpoint_count >= endpoint_max_concurrent
|
||||
):
|
||||
return False
|
||||
|
||||
if key_max_concurrent is not None and key_count >= key_max_concurrent:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
endpoint_count, key_count = await self.get_current_concurrency(endpoint_id, key_id)
|
||||
|
||||
# 检查 Endpoint 级别限制
|
||||
if endpoint_max_concurrent is not None and endpoint_count >= endpoint_max_concurrent:
|
||||
return False
|
||||
|
||||
# 检查 Key 级别限制
|
||||
if key_max_concurrent is not None and key_count >= key_max_concurrent:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def acquire_slot(
|
||||
self,
|
||||
endpoint_id: str,
|
||||
endpoint_max_concurrent: Optional[int],
|
||||
key_id: str,
|
||||
key_max_concurrent: Optional[int],
|
||||
is_cached_user: bool = False, # 新增:是否是缓存用户
|
||||
cache_reservation_ratio: float = 0.3, # 新增:缓存预留比例
|
||||
ttl_seconds: int = 600, # 10分钟 TTL,防止死锁
|
||||
) -> bool:
|
||||
"""
|
||||
尝试获取并发槽位(支持缓存用户优先级)
|
||||
|
||||
Args:
|
||||
endpoint_id: Endpoint ID
|
||||
endpoint_max_concurrent: Endpoint 最大并发数(None 表示不限制)
|
||||
key_id: ProviderAPIKey ID
|
||||
key_max_concurrent: Key 最大并发数(None 表示不限制)
|
||||
is_cached_user: 是否是缓存用户(缓存用户可使用全部槽位)
|
||||
cache_reservation_ratio: 缓存预留比例(默认30%,只对新用户生效)
|
||||
ttl_seconds: TTL 秒数,防止异常情况下的死锁
|
||||
|
||||
Returns:
|
||||
是否成功获取(True/False)
|
||||
|
||||
缓存预留机制说明:
|
||||
- 假设 key_max_concurrent = 10, cache_reservation_ratio = 0.3
|
||||
- 新用户最多使用: 7个槽位 (10 * (1 - 0.3))
|
||||
- 缓存用户最多使用: 10个槽位(全部)
|
||||
- 预留的3个槽位专门给缓存用户,保证他们的请求优先
|
||||
"""
|
||||
if self._redis is None:
|
||||
async with self._memory_lock:
|
||||
endpoint_count = self._memory_endpoint_counts.get(endpoint_id, 0)
|
||||
key_count = self._memory_key_counts.get(key_id, 0)
|
||||
|
||||
# Endpoint 限制
|
||||
if (
|
||||
endpoint_max_concurrent is not None
|
||||
and endpoint_count >= endpoint_max_concurrent
|
||||
):
|
||||
return False
|
||||
|
||||
# Key 限制,包含缓存预留
|
||||
if key_max_concurrent is not None:
|
||||
if is_cached_user:
|
||||
if key_count >= key_max_concurrent:
|
||||
return False
|
||||
else:
|
||||
available_for_new = max(
|
||||
1, math.ceil(key_max_concurrent * (1 - cache_reservation_ratio))
|
||||
)
|
||||
if key_count >= available_for_new:
|
||||
return False
|
||||
|
||||
# 通过限制,更新计数
|
||||
self._memory_endpoint_counts[endpoint_id] = endpoint_count + 1
|
||||
self._memory_key_counts[key_id] = key_count + 1
|
||||
return True
|
||||
|
||||
endpoint_key = self._get_endpoint_key(endpoint_id)
|
||||
key_key = self._get_key_key(key_id)
|
||||
|
||||
try:
|
||||
# 使用 Lua 脚本保证原子性(新增缓存预留逻辑)
|
||||
lua_script = """
|
||||
local endpoint_key = KEYS[1]
|
||||
local key_key = KEYS[2]
|
||||
local endpoint_max = tonumber(ARGV[1])
|
||||
local key_max = tonumber(ARGV[2])
|
||||
local ttl = tonumber(ARGV[3])
|
||||
local is_cached = tonumber(ARGV[4]) -- 0=新用户, 1=缓存用户
|
||||
local cache_ratio = tonumber(ARGV[5]) -- 缓存预留比例
|
||||
|
||||
-- 获取当前值
|
||||
local endpoint_count = tonumber(redis.call('GET', endpoint_key) or '0')
|
||||
local key_count = tonumber(redis.call('GET', key_key) or '0')
|
||||
|
||||
-- 检查 endpoint 限制(-1 表示不限制)
|
||||
if endpoint_max >= 0 and endpoint_count >= endpoint_max then
|
||||
return 0 -- 失败:endpoint 已满
|
||||
end
|
||||
|
||||
-- 检查 key 限制(支持缓存预留)
|
||||
if key_max >= 0 then
|
||||
if is_cached == 0 then
|
||||
-- 新用户:只能使用 (1 - cache_ratio) 的槽位
|
||||
local available_for_new = math.floor(key_max * (1 - cache_ratio))
|
||||
if key_count >= available_for_new then
|
||||
return 0 -- 失败:新用户配额已满
|
||||
end
|
||||
else
|
||||
-- 缓存用户:可以使用全部槽位
|
||||
if key_count >= key_max then
|
||||
return 0 -- 失败:总配额已满
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
-- 增加计数
|
||||
redis.call('INCR', endpoint_key)
|
||||
redis.call('EXPIRE', endpoint_key, ttl)
|
||||
redis.call('INCR', key_key)
|
||||
redis.call('EXPIRE', key_key, ttl)
|
||||
|
||||
return 1 -- 成功
|
||||
"""
|
||||
|
||||
# 执行脚本
|
||||
result = await self._redis.eval(
|
||||
lua_script,
|
||||
2, # 2 个 KEYS
|
||||
endpoint_key,
|
||||
key_key,
|
||||
endpoint_max_concurrent if endpoint_max_concurrent is not None else -1,
|
||||
key_max_concurrent if key_max_concurrent is not None else -1,
|
||||
ttl_seconds,
|
||||
1 if is_cached_user else 0, # 缓存用户标志
|
||||
cache_reservation_ratio, # 预留比例
|
||||
)
|
||||
|
||||
success = result == 1
|
||||
|
||||
if success:
|
||||
user_type = "缓存用户" if is_cached_user else "新用户"
|
||||
logger.debug(
|
||||
f"[OK] 获取并发槽位成功: endpoint={endpoint_id}, key={key_id}, "
|
||||
f"类型={user_type}"
|
||||
)
|
||||
else:
|
||||
endpoint_count, key_count = await self.get_current_concurrency(endpoint_id, key_id)
|
||||
|
||||
# 计算新用户可用槽位
|
||||
if key_max_concurrent and not is_cached_user:
|
||||
available_for_new = int(key_max_concurrent * (1 - cache_reservation_ratio))
|
||||
user_info = f"新用户配额={available_for_new}, 当前={key_count}"
|
||||
else:
|
||||
user_info = f"缓存用户, 当前={key_count}/{key_max_concurrent}"
|
||||
|
||||
logger.warning(
|
||||
f"[WARN] 并发槽位已满: endpoint={endpoint_id}({endpoint_count}/{endpoint_max_concurrent}), "
|
||||
f"key={key_id}({user_info})"
|
||||
)
|
||||
|
||||
return success
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取并发槽位失败,降级到内存模式: {e}")
|
||||
# Redis 异常时降级到内存模式进行保守限流
|
||||
# 使用较低的限制值(原限制的 50%)避免上游 API 被打爆
|
||||
async with self._memory_lock:
|
||||
endpoint_count = self._memory_endpoint_counts.get(endpoint_id, 0)
|
||||
key_count = self._memory_key_counts.get(key_id, 0)
|
||||
|
||||
# 降级模式下使用更保守的限制(50%)
|
||||
fallback_endpoint_limit = (
|
||||
max(1, endpoint_max_concurrent // 2)
|
||||
if endpoint_max_concurrent is not None
|
||||
else None
|
||||
)
|
||||
fallback_key_limit = (
|
||||
max(1, key_max_concurrent // 2) if key_max_concurrent is not None else None
|
||||
)
|
||||
|
||||
if (
|
||||
fallback_endpoint_limit is not None
|
||||
and endpoint_count >= fallback_endpoint_limit
|
||||
):
|
||||
logger.warning(
|
||||
f"[FALLBACK] Endpoint 并发达到降级限制: {endpoint_count}/{fallback_endpoint_limit}"
|
||||
)
|
||||
return False
|
||||
|
||||
if fallback_key_limit is not None and key_count >= fallback_key_limit:
|
||||
logger.warning(
|
||||
f"[FALLBACK] Key 并发达到降级限制: {key_count}/{fallback_key_limit}"
|
||||
)
|
||||
return False
|
||||
|
||||
# 更新内存计数
|
||||
self._memory_endpoint_counts[endpoint_id] = endpoint_count + 1
|
||||
self._memory_key_counts[key_id] = key_count + 1
|
||||
logger.debug(
|
||||
f"[FALLBACK] 使用内存模式获取槽位: endpoint={endpoint_id}, key={key_id}"
|
||||
)
|
||||
return True
|
||||
|
||||
async def release_slot(self, endpoint_id: str, key_id: str) -> None:
|
||||
"""
|
||||
释放并发槽位
|
||||
|
||||
Args:
|
||||
endpoint_id: Endpoint ID
|
||||
key_id: ProviderAPIKey ID
|
||||
"""
|
||||
if self._redis is None:
|
||||
async with self._memory_lock:
|
||||
if endpoint_id in self._memory_endpoint_counts:
|
||||
self._memory_endpoint_counts[endpoint_id] = max(
|
||||
0, self._memory_endpoint_counts[endpoint_id] - 1
|
||||
)
|
||||
if self._memory_endpoint_counts[endpoint_id] == 0:
|
||||
self._memory_endpoint_counts.pop(endpoint_id, None)
|
||||
|
||||
if key_id in self._memory_key_counts:
|
||||
self._memory_key_counts[key_id] = max(0, self._memory_key_counts[key_id] - 1)
|
||||
if self._memory_key_counts[key_id] == 0:
|
||||
self._memory_key_counts.pop(key_id, None)
|
||||
return
|
||||
|
||||
endpoint_key = self._get_endpoint_key(endpoint_id)
|
||||
key_key = self._get_key_key(key_id)
|
||||
|
||||
try:
|
||||
# 使用 Lua 脚本保证原子性(不会减到负数)
|
||||
lua_script = """
|
||||
local endpoint_key = KEYS[1]
|
||||
local key_key = KEYS[2]
|
||||
|
||||
local endpoint_count = tonumber(redis.call('GET', endpoint_key) or '0')
|
||||
local key_count = tonumber(redis.call('GET', key_key) or '0')
|
||||
|
||||
if endpoint_count > 0 then
|
||||
redis.call('DECR', endpoint_key)
|
||||
end
|
||||
|
||||
if key_count > 0 then
|
||||
redis.call('DECR', key_key)
|
||||
end
|
||||
|
||||
return 1
|
||||
"""
|
||||
|
||||
await self._redis.eval(lua_script, 2, endpoint_key, key_key)
|
||||
|
||||
logger.debug(f"[OK] 释放并发槽位: endpoint={endpoint_id}, key={key_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"释放并发槽位失败: {e}")
|
||||
|
||||
@asynccontextmanager
|
||||
async def concurrency_guard(
|
||||
self,
|
||||
endpoint_id: str,
|
||||
endpoint_max_concurrent: Optional[int],
|
||||
key_id: str,
|
||||
key_max_concurrent: Optional[int],
|
||||
is_cached_user: bool = False, # 新增:是否是缓存用户
|
||||
cache_reservation_ratio: float = 0.3, # 新增:缓存预留比例
|
||||
):
|
||||
"""
|
||||
并发控制上下文管理器(支持缓存用户优先级)
|
||||
|
||||
用法:
|
||||
async with manager.concurrency_guard(
|
||||
endpoint_id, endpoint_max, key_id, key_max,
|
||||
is_cached_user=True # 缓存用户
|
||||
):
|
||||
# 执行请求
|
||||
response = await send_request(...)
|
||||
|
||||
如果获取失败,会抛出 ConcurrencyLimitError 异常
|
||||
"""
|
||||
# 尝试获取槽位(传递缓存用户参数)
|
||||
acquired = await self.acquire_slot(
|
||||
endpoint_id,
|
||||
endpoint_max_concurrent,
|
||||
key_id,
|
||||
key_max_concurrent,
|
||||
is_cached_user,
|
||||
cache_reservation_ratio,
|
||||
)
|
||||
|
||||
if not acquired:
|
||||
from src.core.exceptions import ConcurrencyLimitError
|
||||
|
||||
user_type = "缓存用户" if is_cached_user else "新用户"
|
||||
raise ConcurrencyLimitError(
|
||||
f"并发限制已达上限: endpoint={endpoint_id}, key={key_id}, 类型={user_type}"
|
||||
)
|
||||
|
||||
# 记录开始时间和状态
|
||||
import time
|
||||
|
||||
slot_acquired_at = time.time()
|
||||
exception_occurred = False
|
||||
|
||||
try:
|
||||
yield # 执行请求
|
||||
except Exception as e:
|
||||
# 记录异常
|
||||
exception_occurred = True
|
||||
raise
|
||||
finally:
|
||||
# 计算槽位占用时长
|
||||
slot_duration = time.time() - slot_acquired_at
|
||||
|
||||
# 记录 Prometheus 指标
|
||||
try:
|
||||
from src.core.metrics import (
|
||||
concurrency_slot_duration_seconds,
|
||||
concurrency_slot_release_total,
|
||||
)
|
||||
|
||||
# 记录槽位占用时长分布
|
||||
concurrency_slot_duration_seconds.labels(
|
||||
key_id=key_id[:8] if key_id else "unknown", # 只记录前8位
|
||||
exception=str(exception_occurred),
|
||||
).observe(slot_duration)
|
||||
|
||||
# 记录槽位释放计数
|
||||
concurrency_slot_release_total.labels(
|
||||
key_id=key_id[:8] if key_id else "unknown", exception=str(exception_occurred)
|
||||
).inc()
|
||||
|
||||
# 告警:槽位占用时间过长(超过 60 秒)
|
||||
if slot_duration > 60:
|
||||
logger.warning(
|
||||
f"[WARN] 并发槽位占用时间过长: "
|
||||
f"key_id={key_id[:8] if key_id else 'unknown'}..., "
|
||||
f"duration={slot_duration:.1f}s, "
|
||||
f"exception={exception_occurred}"
|
||||
)
|
||||
|
||||
except Exception as metric_error:
|
||||
# 指标记录失败不应影响业务逻辑
|
||||
logger.debug(f"记录并发指标失败: {metric_error}")
|
||||
|
||||
# 自动释放槽位(即使发生异常)
|
||||
await self.release_slot(endpoint_id, key_id)
|
||||
|
||||
async def reset_concurrency(
|
||||
self, endpoint_id: Optional[str] = None, key_id: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
重置并发计数(管理功能,慎用)
|
||||
|
||||
Args:
|
||||
endpoint_id: Endpoint ID(可选,None 表示重置所有 endpoint)
|
||||
key_id: ProviderAPIKey ID(可选,None 表示重置所有 key)
|
||||
"""
|
||||
if self._redis is None:
|
||||
async with self._memory_lock:
|
||||
if endpoint_id:
|
||||
self._memory_endpoint_counts.pop(endpoint_id, None)
|
||||
logger.info(f"[RESET] 重置 Endpoint 并发计数(内存): {endpoint_id}")
|
||||
else:
|
||||
count = len(self._memory_endpoint_counts)
|
||||
self._memory_endpoint_counts.clear()
|
||||
if count:
|
||||
logger.info(f"[RESET] 重置所有 Endpoint 并发计数(内存): {count} 个")
|
||||
|
||||
if key_id:
|
||||
self._memory_key_counts.pop(key_id, None)
|
||||
logger.info(f"[RESET] 重置 Key 并发计数(内存): {key_id}")
|
||||
else:
|
||||
count = len(self._memory_key_counts)
|
||||
self._memory_key_counts.clear()
|
||||
if count:
|
||||
logger.info(f"[RESET] 重置所有 Key 并发计数(内存): {count} 个")
|
||||
return
|
||||
|
||||
try:
|
||||
if endpoint_id:
|
||||
endpoint_key = self._get_endpoint_key(endpoint_id)
|
||||
await self._redis.delete(endpoint_key)
|
||||
logger.info(f"[RESET] 重置 Endpoint 并发计数: {endpoint_id}")
|
||||
else:
|
||||
# 重置所有 endpoint
|
||||
keys = await self._redis.keys("concurrency:endpoint:*")
|
||||
if keys:
|
||||
await self._redis.delete(*keys)
|
||||
logger.info(f"[RESET] 重置所有 Endpoint 并发计数: {len(keys)} 个")
|
||||
|
||||
if key_id:
|
||||
key_key = self._get_key_key(key_id)
|
||||
await self._redis.delete(key_key)
|
||||
logger.info(f"[RESET] 重置 Key 并发计数: {key_id}")
|
||||
else:
|
||||
# 重置所有 key
|
||||
keys = await self._redis.keys("concurrency:key:*")
|
||||
if keys:
|
||||
await self._redis.delete(*keys)
|
||||
logger.info(f"[RESET] 重置所有 Key 并发计数: {len(keys)} 个")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"重置并发计数失败: {e}")
|
||||
|
||||
|
||||
# 全局单例
|
||||
_concurrency_manager: Optional[ConcurrencyManager] = None
|
||||
|
||||
|
||||
async def get_concurrency_manager() -> ConcurrencyManager:
|
||||
"""获取全局 ConcurrencyManager 实例"""
|
||||
global _concurrency_manager
|
||||
|
||||
if _concurrency_manager is None:
|
||||
_concurrency_manager = ConcurrencyManager()
|
||||
await _concurrency_manager.initialize()
|
||||
|
||||
return _concurrency_manager
|
||||
333
src/services/rate_limit/detector.py
Normal file
333
src/services/rate_limit/detector.py
Normal file
@@ -0,0 +1,333 @@
|
||||
"""
|
||||
速率限制检测器 - 解析429响应头,区分并发限制和RPM限制
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
class RateLimitType:
|
||||
"""速率限制类型"""
|
||||
|
||||
CONCURRENT = "concurrent" # 并发限制
|
||||
RPM = "rpm" # 每分钟请求数限制
|
||||
DAILY = "daily" # 每日限制
|
||||
MONTHLY = "monthly" # 每月限制
|
||||
UNKNOWN = "unknown" # 未知类型
|
||||
|
||||
|
||||
class RateLimitInfo:
|
||||
"""速率限制信息"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
limit_type: str,
|
||||
retry_after: Optional[int] = None,
|
||||
limit_value: Optional[int] = None,
|
||||
remaining: Optional[int] = None,
|
||||
reset_at: Optional[datetime] = None,
|
||||
current_usage: Optional[int] = None,
|
||||
raw_headers: Optional[Dict[str, str]] = None,
|
||||
):
|
||||
self.limit_type = limit_type
|
||||
self.retry_after = retry_after # 需要等待的秒数
|
||||
self.limit_value = limit_value # 限制值
|
||||
self.remaining = remaining # 剩余配额
|
||||
self.reset_at = reset_at # 重置时间
|
||||
self.current_usage = current_usage # 当前使用量
|
||||
self.raw_headers = raw_headers or {}
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"RateLimitInfo(type={self.limit_type}, "
|
||||
f"retry_after={self.retry_after}, "
|
||||
f"limit={self.limit_value}, "
|
||||
f"remaining={self.remaining})"
|
||||
)
|
||||
|
||||
|
||||
class RateLimitDetector:
|
||||
"""
|
||||
速率限制检测器
|
||||
|
||||
支持的提供商:
|
||||
- Anthropic Claude API
|
||||
- OpenAI API
|
||||
- 通用 HTTP 标准头
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def detect_from_headers(
|
||||
headers: Dict[str, str],
|
||||
provider_name: str = "unknown",
|
||||
current_concurrent: Optional[int] = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
从响应头中检测速率限制类型
|
||||
|
||||
Args:
|
||||
headers: 429响应的HTTP头
|
||||
provider_name: 提供商名称(用于选择解析策略)
|
||||
current_concurrent: 当前并发数(用于判断是否为并发限制)
|
||||
|
||||
Returns:
|
||||
RateLimitInfo对象
|
||||
"""
|
||||
# 标准化header key (转小写)
|
||||
headers_lower = {k.lower(): v for k, v in headers.items()}
|
||||
|
||||
# 根据提供商选择解析策略
|
||||
if "anthropic" in provider_name.lower() or "claude" in provider_name.lower():
|
||||
return RateLimitDetector._parse_anthropic_headers(headers_lower, current_concurrent)
|
||||
elif "openai" in provider_name.lower():
|
||||
return RateLimitDetector._parse_openai_headers(headers_lower, current_concurrent)
|
||||
else:
|
||||
return RateLimitDetector._parse_generic_headers(headers_lower, current_concurrent)
|
||||
|
||||
@staticmethod
|
||||
def _parse_anthropic_headers(
|
||||
headers: Dict[str, str],
|
||||
current_concurrent: Optional[int] = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
解析 Anthropic Claude API 的速率限制头
|
||||
|
||||
常见头部:
|
||||
- anthropic-ratelimit-requests-limit: 50
|
||||
- anthropic-ratelimit-requests-remaining: 0
|
||||
- anthropic-ratelimit-requests-reset: 2024-01-01T00:00:00Z
|
||||
- anthropic-ratelimit-tokens-limit: 100000
|
||||
- anthropic-ratelimit-tokens-remaining: 50000
|
||||
- retry-after: 60
|
||||
"""
|
||||
retry_after = RateLimitDetector._parse_retry_after(headers)
|
||||
|
||||
# 获取请求限制信息
|
||||
requests_limit = RateLimitDetector._parse_int(
|
||||
headers.get("anthropic-ratelimit-requests-limit")
|
||||
)
|
||||
requests_remaining = RateLimitDetector._parse_int(
|
||||
headers.get("anthropic-ratelimit-requests-remaining")
|
||||
)
|
||||
requests_reset = RateLimitDetector._parse_datetime(
|
||||
headers.get("anthropic-ratelimit-requests-reset")
|
||||
)
|
||||
|
||||
# 判断限制类型
|
||||
# 1. 明确的 RPM 限制:请求数剩余为 0
|
||||
if requests_remaining is not None and requests_remaining == 0:
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.RPM,
|
||||
retry_after=retry_after,
|
||||
limit_value=requests_limit,
|
||||
remaining=requests_remaining,
|
||||
reset_at=requests_reset,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
# 2. 可能的并发限制判断(多条件综合)
|
||||
# 条件:当前并发数存在,且 remaining > 0(说明不是 RPM 耗尽)
|
||||
# 同时 retry_after 较短(并发限制通常 retry_after 较短,如 1-10 秒)
|
||||
is_likely_concurrent = (
|
||||
current_concurrent is not None
|
||||
and current_concurrent >= 2 # 至少有 2 个并发
|
||||
and (requests_remaining is None or requests_remaining > 0) # RPM 未耗尽
|
||||
and (retry_after is None or retry_after <= 30) # 短暂等待
|
||||
)
|
||||
|
||||
if is_likely_concurrent:
|
||||
logger.info(
|
||||
f"检测到可能的并发限制: current_concurrent={current_concurrent}, "
|
||||
f"remaining={requests_remaining}, retry_after={retry_after}"
|
||||
)
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.CONCURRENT,
|
||||
retry_after=retry_after,
|
||||
current_usage=current_concurrent,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
# 3. 未知类型
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.UNKNOWN,
|
||||
retry_after=retry_after,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_openai_headers(
|
||||
headers: Dict[str, str],
|
||||
current_concurrent: Optional[int] = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
解析 OpenAI API 的速率限制头
|
||||
|
||||
常见头部:
|
||||
- x-ratelimit-limit-requests: 3500
|
||||
- x-ratelimit-remaining-requests: 0
|
||||
- x-ratelimit-reset-requests: 2024-01-01T00:00:00Z
|
||||
- x-ratelimit-limit-tokens: 90000
|
||||
- x-ratelimit-remaining-tokens: 50000
|
||||
- retry-after: 60
|
||||
"""
|
||||
retry_after = RateLimitDetector._parse_retry_after(headers)
|
||||
|
||||
# 获取请求限制信息
|
||||
requests_limit = RateLimitDetector._parse_int(headers.get("x-ratelimit-limit-requests"))
|
||||
requests_remaining = RateLimitDetector._parse_int(
|
||||
headers.get("x-ratelimit-remaining-requests")
|
||||
)
|
||||
requests_reset = RateLimitDetector._parse_datetime(
|
||||
headers.get("x-ratelimit-reset-requests")
|
||||
)
|
||||
|
||||
# 判断限制类型
|
||||
# 1. 明确的 RPM 限制
|
||||
if requests_remaining is not None and requests_remaining == 0:
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.RPM,
|
||||
retry_after=retry_after,
|
||||
limit_value=requests_limit,
|
||||
remaining=requests_remaining,
|
||||
reset_at=requests_reset,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
# 2. 可能的并发限制(多条件综合判断)
|
||||
is_likely_concurrent = (
|
||||
current_concurrent is not None
|
||||
and current_concurrent >= 2
|
||||
and (requests_remaining is None or requests_remaining > 0)
|
||||
and (retry_after is None or retry_after <= 30)
|
||||
)
|
||||
|
||||
if is_likely_concurrent:
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.CONCURRENT,
|
||||
retry_after=retry_after,
|
||||
current_usage=current_concurrent,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
# 3. 未知类型
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.UNKNOWN,
|
||||
retry_after=retry_after,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_generic_headers(
|
||||
headers: Dict[str, str],
|
||||
current_concurrent: Optional[int] = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
解析通用的速率限制头
|
||||
|
||||
标准头部:
|
||||
- retry-after: 60
|
||||
- x-ratelimit-limit: 100
|
||||
- x-ratelimit-remaining: 0
|
||||
- x-ratelimit-reset: 1609459200
|
||||
"""
|
||||
retry_after = RateLimitDetector._parse_retry_after(headers)
|
||||
|
||||
limit_value = RateLimitDetector._parse_int(headers.get("x-ratelimit-limit"))
|
||||
remaining = RateLimitDetector._parse_int(headers.get("x-ratelimit-remaining"))
|
||||
|
||||
# 1. 明确的 RPM 限制
|
||||
if remaining is not None and remaining == 0:
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.RPM,
|
||||
retry_after=retry_after,
|
||||
limit_value=limit_value,
|
||||
remaining=remaining,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
# 2. 可能的并发限制
|
||||
is_likely_concurrent = (
|
||||
current_concurrent is not None
|
||||
and current_concurrent >= 2
|
||||
and (remaining is None or remaining > 0)
|
||||
and (retry_after is None or retry_after <= 30)
|
||||
)
|
||||
|
||||
if is_likely_concurrent:
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.CONCURRENT,
|
||||
retry_after=retry_after,
|
||||
current_usage=current_concurrent,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
# 3. 未知类型
|
||||
return RateLimitInfo(
|
||||
limit_type=RateLimitType.UNKNOWN,
|
||||
retry_after=retry_after,
|
||||
raw_headers=headers,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_retry_after(headers: Dict[str, str]) -> Optional[int]:
|
||||
"""解析 Retry-After 头"""
|
||||
retry_after_str = headers.get("retry-after")
|
||||
if not retry_after_str:
|
||||
return None
|
||||
|
||||
try:
|
||||
# 尝试解析为整数(秒数)
|
||||
return int(retry_after_str)
|
||||
except ValueError:
|
||||
# 尝试解析为HTTP日期格式
|
||||
try:
|
||||
retry_date = datetime.strptime(retry_after_str, "%a, %d %b %Y %H:%M:%S %Z")
|
||||
delta = retry_date - datetime.now(timezone.utc)
|
||||
return max(int(delta.total_seconds()), 0)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _parse_int(value: Optional[str]) -> Optional[int]:
|
||||
"""安全解析整数"""
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _parse_datetime(value: Optional[str]) -> Optional[datetime]:
|
||||
"""安全解析ISO 8601日期时间"""
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
# 尝试解析 ISO 8601 格式
|
||||
if value.endswith("Z"):
|
||||
value = value[:-1] + "+00:00"
|
||||
return datetime.fromisoformat(value)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
# 便捷函数
|
||||
def detect_rate_limit_type(
|
||||
headers: Dict[str, str],
|
||||
provider_name: str = "unknown",
|
||||
current_concurrent: Optional[int] = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
检测速率限制类型(便捷函数)
|
||||
|
||||
Args:
|
||||
headers: 429响应头
|
||||
provider_name: 提供商名称
|
||||
current_concurrent: 当前并发数
|
||||
|
||||
Returns:
|
||||
RateLimitInfo对象
|
||||
"""
|
||||
return RateLimitDetector.detect_from_headers(headers, provider_name, current_concurrent)
|
||||
351
src/services/rate_limit/ip_limiter.py
Normal file
351
src/services/rate_limit/ip_limiter.py
Normal file
@@ -0,0 +1,351 @@
|
||||
"""
|
||||
IP 级别的速率限制服务
|
||||
|
||||
提供基于 IP 地址的速率限制,防止暴力破解和 DDoS 攻击
|
||||
"""
|
||||
|
||||
import ipaddress
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, Optional, Set
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
|
||||
class IPRateLimiter:
|
||||
"""IP 速率限制服务"""
|
||||
|
||||
# Redis key 前缀
|
||||
RATE_LIMIT_PREFIX = "ip:rate_limit:"
|
||||
BLACKLIST_PREFIX = "ip:blacklist:"
|
||||
WHITELIST_KEY = "ip:whitelist"
|
||||
|
||||
# 默认限制配置(每分钟)
|
||||
DEFAULT_LIMITS = {
|
||||
"default": 100, # 默认限制
|
||||
"login": 5, # 登录接口
|
||||
"register": 3, # 注册接口
|
||||
"api": 60, # API 接口
|
||||
"public": 60, # 公共接口
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
async def check_limit(
|
||||
ip_address: str, endpoint_type: str = "default", limit: Optional[int] = None
|
||||
) -> tuple[bool, int, int]:
|
||||
"""
|
||||
检查 IP 是否超过速率限制
|
||||
|
||||
Args:
|
||||
ip_address: IP 地址
|
||||
endpoint_type: 端点类型(default, login, register, api, public)
|
||||
limit: 自定义限制值,None 则使用默认值
|
||||
|
||||
Returns:
|
||||
(是否允许, 剩余次数, 重置时间秒数)
|
||||
"""
|
||||
# 检查白名单
|
||||
if await IPRateLimiter.is_whitelisted(ip_address):
|
||||
return True, 999999, 60
|
||||
|
||||
# 检查黑名单
|
||||
if await IPRateLimiter.is_blacklisted(ip_address):
|
||||
logger.warning(f"黑名单 IP 尝试访问: {ip_address}, 类型: {endpoint_type}")
|
||||
return False, 0, 0
|
||||
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
# Redis 不可用时降级:允许访问但记录警告
|
||||
logger.warning("Redis 不可用,跳过 IP 速率限制(降级模式)")
|
||||
return True, 0, 60
|
||||
|
||||
# 确定限制值
|
||||
rate_limit = (
|
||||
limit if limit is not None else IPRateLimiter.DEFAULT_LIMITS.get(endpoint_type, 100)
|
||||
)
|
||||
|
||||
try:
|
||||
# Redis key: ip:rate_limit:{type}:{ip}
|
||||
redis_key = f"{IPRateLimiter.RATE_LIMIT_PREFIX}{endpoint_type}:{ip_address}"
|
||||
|
||||
# 使用 Redis 的滑动窗口计数器
|
||||
# INCR 并设置过期时间
|
||||
count = await redis_client.incr(redis_key)
|
||||
|
||||
# 第一次访问时设置过期时间
|
||||
if count == 1:
|
||||
await redis_client.expire(redis_key, 60) # 60秒窗口
|
||||
|
||||
# 获取 TTL(剩余过期时间)
|
||||
ttl = await redis_client.ttl(redis_key)
|
||||
if ttl < 0:
|
||||
# 如果没有过期时间,重新设置
|
||||
await redis_client.expire(redis_key, 60)
|
||||
ttl = 60
|
||||
|
||||
remaining = max(0, rate_limit - count)
|
||||
allowed = count <= rate_limit
|
||||
|
||||
if not allowed:
|
||||
logger.warning(f"IP 速率限制触发: {ip_address}, 类型: {endpoint_type}, 计数: {count}/{rate_limit}")
|
||||
|
||||
return allowed, remaining, ttl
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"检查 IP 速率限制失败: {e}")
|
||||
# 发生错误时允许访问,避免误杀
|
||||
return True, 0, 60
|
||||
|
||||
@staticmethod
|
||||
async def add_to_blacklist(
|
||||
ip_address: str, reason: str = "manual", ttl: Optional[int] = None
|
||||
) -> bool:
|
||||
"""
|
||||
将 IP 加入黑名单
|
||||
|
||||
Args:
|
||||
ip_address: IP 地址
|
||||
reason: 加入黑名单的原因
|
||||
ttl: 过期时间(秒),None 表示永久
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning("Redis 不可用,无法将 IP 加入黑名单")
|
||||
return False
|
||||
|
||||
try:
|
||||
redis_key = f"{IPRateLimiter.BLACKLIST_PREFIX}{ip_address}"
|
||||
|
||||
if ttl is not None:
|
||||
await redis_client.setex(redis_key, ttl, reason)
|
||||
else:
|
||||
await redis_client.set(redis_key, reason)
|
||||
|
||||
logger.warning(f"IP 已加入黑名单: {ip_address}, 原因: {reason}, TTL: {ttl or '永久'}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"添加 IP 到黑名单失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def remove_from_blacklist(ip_address: str) -> bool:
|
||||
"""
|
||||
从黑名单移除 IP
|
||||
|
||||
Args:
|
||||
ip_address: IP 地址
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning("Redis 不可用,无法从黑名单移除 IP")
|
||||
return False
|
||||
|
||||
try:
|
||||
redis_key = f"{IPRateLimiter.BLACKLIST_PREFIX}{ip_address}"
|
||||
deleted = await redis_client.delete(redis_key)
|
||||
|
||||
if deleted:
|
||||
logger.info(f"IP 已从黑名单移除: {ip_address}")
|
||||
|
||||
return bool(deleted)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"从黑名单移除 IP 失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def is_blacklisted(ip_address: str) -> bool:
|
||||
"""
|
||||
检查 IP 是否在黑名单中
|
||||
|
||||
Args:
|
||||
ip_address: IP 地址
|
||||
|
||||
Returns:
|
||||
是否在黑名单中
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
redis_key = f"{IPRateLimiter.BLACKLIST_PREFIX}{ip_address}"
|
||||
exists = await redis_client.exists(redis_key)
|
||||
return bool(exists)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"检查 IP 黑名单状态失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def add_to_whitelist(ip_address: str) -> bool:
|
||||
"""
|
||||
将 IP 加入白名单
|
||||
|
||||
Args:
|
||||
ip_address: IP 地址或 CIDR 格式(如 192.168.1.0/24)
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning("Redis 不可用,无法将 IP 加入白名单")
|
||||
return False
|
||||
|
||||
try:
|
||||
# 验证 IP 格式
|
||||
try:
|
||||
ipaddress.ip_network(ip_address, strict=False)
|
||||
except ValueError as e:
|
||||
logger.error(f"无效的 IP 地址格式: {ip_address}, 错误: {e}")
|
||||
return False
|
||||
|
||||
# 使用 Redis Set 存储白名单
|
||||
await redis_client.sadd(IPRateLimiter.WHITELIST_KEY, ip_address)
|
||||
|
||||
logger.info(f"IP 已加入白名单: {ip_address}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"添加 IP 到白名单失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def remove_from_whitelist(ip_address: str) -> bool:
|
||||
"""
|
||||
从白名单移除 IP
|
||||
|
||||
Args:
|
||||
ip_address: IP 地址
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
logger.warning("Redis 不可用,无法从白名单移除 IP")
|
||||
return False
|
||||
|
||||
try:
|
||||
removed = await redis_client.srem(IPRateLimiter.WHITELIST_KEY, ip_address)
|
||||
|
||||
if removed:
|
||||
logger.info(f"IP 已从白名单移除: {ip_address}")
|
||||
|
||||
return bool(removed)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"从白名单移除 IP 失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def is_whitelisted(ip_address: str) -> bool:
|
||||
"""
|
||||
检查 IP 是否在白名单中(支持 CIDR 匹配)
|
||||
|
||||
Args:
|
||||
ip_address: IP 地址
|
||||
|
||||
Returns:
|
||||
是否在白名单中
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
# 获取所有白名单条目
|
||||
whitelist = await redis_client.smembers(IPRateLimiter.WHITELIST_KEY)
|
||||
|
||||
if not whitelist:
|
||||
return False
|
||||
|
||||
# 将 IP 地址转换为 ip_address 对象
|
||||
try:
|
||||
ip_obj = ipaddress.ip_address(ip_address)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
# 检查是否匹配白名单中的任何条目
|
||||
for entry in whitelist:
|
||||
try:
|
||||
network = ipaddress.ip_network(entry, strict=False)
|
||||
if ip_obj in network:
|
||||
return True
|
||||
except ValueError:
|
||||
# 如果条目格式无效,跳过
|
||||
continue
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"检查 IP 白名单状态失败: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def get_blacklist_stats() -> Dict:
|
||||
"""
|
||||
获取黑名单统计信息
|
||||
|
||||
Returns:
|
||||
统计信息字典
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
return {"available": False, "total": 0, "error": "Redis 不可用"}
|
||||
|
||||
try:
|
||||
pattern = f"{IPRateLimiter.BLACKLIST_PREFIX}*"
|
||||
cursor = 0
|
||||
total = 0
|
||||
|
||||
while True:
|
||||
cursor, keys = await redis_client.scan(cursor=cursor, match=pattern, count=100)
|
||||
total += len(keys)
|
||||
|
||||
if cursor == 0:
|
||||
break
|
||||
|
||||
return {"available": True, "total": total}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取黑名单统计失败: {e}")
|
||||
return {"available": False, "total": 0, "error": str(e)}
|
||||
|
||||
@staticmethod
|
||||
async def get_whitelist() -> Set[str]:
|
||||
"""
|
||||
获取白名单列表
|
||||
|
||||
Returns:
|
||||
白名单 IP 集合
|
||||
"""
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
if redis_client is None:
|
||||
return set()
|
||||
|
||||
try:
|
||||
whitelist = await redis_client.smembers(IPRateLimiter.WHITELIST_KEY)
|
||||
return whitelist if whitelist else set()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取白名单失败: {e}")
|
||||
return set()
|
||||
139
src/services/rate_limit/rpm_limiter.py
Normal file
139
src/services/rate_limit/rpm_limiter.py
Normal file
@@ -0,0 +1,139 @@
|
||||
"""
|
||||
RPM (Requests Per Minute) 限流服务
|
||||
"""
|
||||
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, Tuple
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.batch_committer import get_batch_committer
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider
|
||||
from src.models.database_extensions import ProviderUsageTracking
|
||||
|
||||
|
||||
|
||||
class RPMLimiter:
|
||||
"""RPM限流器"""
|
||||
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
# 内存中的RPM计数器 {provider_id: (count, window_start)}
|
||||
self._rpm_counters: Dict[str, Tuple[int, float]] = {}
|
||||
|
||||
def check_and_increment(self, provider_id: str) -> bool:
|
||||
"""
|
||||
检查并递增RPM计数
|
||||
|
||||
Returns:
|
||||
True if allowed, False if rate limited
|
||||
"""
|
||||
provider = self.db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
return True
|
||||
|
||||
rpm_limit = provider.rpm_limit
|
||||
if rpm_limit is None:
|
||||
# 未设置限制
|
||||
return True
|
||||
|
||||
if rpm_limit == 0:
|
||||
logger.warning(f"Provider {provider.name} is fully restricted by RPM limit=0")
|
||||
return False
|
||||
|
||||
current_time = time.time()
|
||||
|
||||
# 检查是否需要重置
|
||||
if provider.rpm_reset_at and provider.rpm_reset_at < datetime.now(timezone.utc):
|
||||
provider.rpm_used = 0
|
||||
provider.rpm_reset_at = datetime.fromtimestamp(current_time + 60, tz=timezone.utc)
|
||||
self.db.commit() # 立即提交事务,释放数据库锁
|
||||
|
||||
# 检查是否超限
|
||||
if provider.rpm_used >= rpm_limit:
|
||||
logger.warning(f"Provider {provider.name} RPM limit exceeded")
|
||||
return False
|
||||
|
||||
# 递增计数
|
||||
provider.rpm_used += 1
|
||||
if not provider.rpm_reset_at:
|
||||
provider.rpm_reset_at = datetime.fromtimestamp(current_time + 60, tz=timezone.utc)
|
||||
|
||||
self.db.commit() # 立即提交事务,释放数据库锁
|
||||
return True
|
||||
|
||||
def record_usage(
|
||||
self, provider_id: str, success: bool, response_time_ms: float, cost_usd: float
|
||||
):
|
||||
"""记录使用情况到追踪表"""
|
||||
|
||||
# 获取当前分钟窗口
|
||||
now = datetime.now(timezone.utc)
|
||||
window_start = now.replace(second=0, microsecond=0)
|
||||
window_end = (
|
||||
window_start.replace(minute=window_start.minute + 1)
|
||||
if window_start.minute < 59
|
||||
else window_start.replace(hour=window_start.hour + 1, minute=0)
|
||||
)
|
||||
|
||||
# 查找或创建追踪记录
|
||||
tracking = (
|
||||
self.db.query(ProviderUsageTracking)
|
||||
.filter(
|
||||
ProviderUsageTracking.provider_id == provider_id,
|
||||
ProviderUsageTracking.window_start == window_start,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not tracking:
|
||||
tracking = ProviderUsageTracking(
|
||||
provider_id=provider_id, window_start=window_start, window_end=window_end
|
||||
)
|
||||
self.db.add(tracking)
|
||||
|
||||
# 更新统计
|
||||
tracking.total_requests += 1
|
||||
if success:
|
||||
tracking.successful_requests += 1
|
||||
else:
|
||||
tracking.failed_requests += 1
|
||||
|
||||
tracking.total_response_time_ms += response_time_ms
|
||||
tracking.avg_response_time_ms = tracking.total_response_time_ms / tracking.total_requests
|
||||
tracking.total_cost_usd += cost_usd
|
||||
|
||||
self.db.flush() # 只 flush,不立即 commit
|
||||
# RPM 使用统计是非关键数据,使用批量提交
|
||||
get_batch_committer().mark_dirty(self.db)
|
||||
|
||||
logger.debug(f"Recorded usage for provider {provider_id}")
|
||||
|
||||
def get_rpm_status(self, provider_id: str) -> Dict:
|
||||
"""获取提供商的RPM状态"""
|
||||
provider = self.db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
return {"error": "Provider not found"}
|
||||
|
||||
return {
|
||||
"provider_id": provider_id,
|
||||
"provider_name": provider.name,
|
||||
"rpm_limit": provider.rpm_limit,
|
||||
"rpm_used": provider.rpm_used,
|
||||
"rpm_reset_at": provider.rpm_reset_at.isoformat() if provider.rpm_reset_at else None,
|
||||
"available": (
|
||||
provider.rpm_limit - provider.rpm_used if provider.rpm_limit is not None else None
|
||||
),
|
||||
}
|
||||
|
||||
def reset_rpm_counter(self, provider_id: str):
|
||||
"""手动重置RPM计数器"""
|
||||
provider = self.db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if provider:
|
||||
provider.rpm_used = 0
|
||||
provider.rpm_reset_at = None
|
||||
self.db.commit() # 立即提交事务,释放数据库锁
|
||||
|
||||
logger.info(f"Reset RPM counter for provider {provider.name}")
|
||||
17
src/services/request/__init__.py
Normal file
17
src/services/request/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
请求处理服务模块
|
||||
|
||||
包含候选选择、执行等功能。
|
||||
|
||||
注意:
|
||||
- RequestBuilder 已移至 src.api.handlers.base.request_builder,请直接从该模块导入
|
||||
- record_failed_request 已移至 src.services.usage.recorder,请直接从该模块导入
|
||||
"""
|
||||
|
||||
from src.services.request.candidate import RequestCandidateService
|
||||
from src.services.request.executor import RequestExecutor
|
||||
|
||||
__all__ = [
|
||||
"RequestCandidateService",
|
||||
"RequestExecutor",
|
||||
]
|
||||
291
src/services/request/candidate.py
Normal file
291
src/services/request/candidate.py
Normal file
@@ -0,0 +1,291 @@
|
||||
"""
|
||||
请求候选记录服务 - 管理候选队列
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.batch_committer import get_batch_committer
|
||||
from src.models.database import RequestCandidate
|
||||
|
||||
|
||||
class RequestCandidateService:
|
||||
"""请求候选记录服务"""
|
||||
|
||||
@staticmethod
|
||||
def create_candidate(
|
||||
db: Session,
|
||||
request_id: str,
|
||||
candidate_index: int,
|
||||
retry_index: int = 0, # 新增:重试序号
|
||||
user_id: Optional[str] = None,
|
||||
api_key_id: Optional[str] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
endpoint_id: Optional[str] = None,
|
||||
key_id: Optional[str] = None,
|
||||
status: str = "available",
|
||||
skip_reason: Optional[str] = None,
|
||||
is_cached: bool = False,
|
||||
extra_data: Optional[dict] = None,
|
||||
required_capabilities: Optional[dict] = None,
|
||||
) -> RequestCandidate:
|
||||
"""
|
||||
创建候选记录
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
request_id: 请求ID
|
||||
candidate_index: 候选序号
|
||||
retry_index: 重试序号(从0开始)
|
||||
user_id: 用户ID
|
||||
api_key_id: API Key ID
|
||||
provider_id: Provider ID
|
||||
endpoint_id: Endpoint ID
|
||||
key_id: API Key ID
|
||||
status: 候选状态 ('available', 'used', 'skipped', 'success', 'failed')
|
||||
skip_reason: 跳过原因
|
||||
is_cached: 是否为缓存亲和性候选
|
||||
extra_data: 额外数据
|
||||
required_capabilities: 请求需要的能力标签
|
||||
"""
|
||||
candidate = RequestCandidate(
|
||||
id=str(uuid.uuid4()),
|
||||
request_id=request_id,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=retry_index, # 新增
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
provider_id=provider_id,
|
||||
endpoint_id=endpoint_id,
|
||||
key_id=key_id,
|
||||
status=status,
|
||||
skip_reason=skip_reason,
|
||||
is_cached=is_cached,
|
||||
extra_data=extra_data or {},
|
||||
required_capabilities=required_capabilities,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
db.add(candidate)
|
||||
db.flush() # 只flush,不立即 commit
|
||||
# 标记为批量提交(非关键数据,可延迟)
|
||||
get_batch_committer().mark_dirty(db)
|
||||
return candidate
|
||||
|
||||
@staticmethod
|
||||
def mark_candidate_started(db: Session, candidate_id: str) -> None:
|
||||
"""
|
||||
标记候选开始执行
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
candidate_id: 候选ID
|
||||
"""
|
||||
candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
|
||||
if candidate:
|
||||
candidate.status = "pending"
|
||||
candidate.started_at = datetime.now(timezone.utc)
|
||||
# 关键状态更新:立即提交,不使用批量提交
|
||||
# 原因:前端需要实时看到请求开始执行
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def update_candidate_status(db: Session, candidate_id: str, status: str) -> None:
|
||||
"""
|
||||
更新候选状态(通用方法)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
candidate_id: 候选ID
|
||||
status: 新状态(pending, available, success, failed, skipped)
|
||||
"""
|
||||
candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
|
||||
if candidate:
|
||||
candidate.status = status
|
||||
# 如果状态变更为 pending,记录开始时间
|
||||
if status == "pending" and not candidate.started_at:
|
||||
candidate.started_at = datetime.now(timezone.utc)
|
||||
# 立即提交,确保前端能实时看到状态变化
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def mark_candidate_streaming(
|
||||
db: Session,
|
||||
candidate_id: str,
|
||||
status_code: int = 200,
|
||||
concurrent_requests: Optional[int] = None,
|
||||
) -> None:
|
||||
"""
|
||||
标记候选为流式传输中
|
||||
|
||||
用于流式请求:连接建立成功后,流开始传输时调用。
|
||||
此时请求尚未完成,需要等流传输完毕后再调用 mark_candidate_success。
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
candidate_id: 候选ID
|
||||
status_code: HTTP 状态码(通常是 200)
|
||||
concurrent_requests: 并发请求数
|
||||
"""
|
||||
candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
|
||||
if candidate:
|
||||
candidate.status = "streaming"
|
||||
candidate.status_code = status_code
|
||||
candidate.concurrent_requests = concurrent_requests
|
||||
# streaming 状态不设置 finished_at,因为请求还在进行中
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def mark_candidate_success(
|
||||
db: Session,
|
||||
candidate_id: str,
|
||||
status_code: int,
|
||||
latency_ms: int,
|
||||
concurrent_requests: Optional[int] = None,
|
||||
extra_data: Optional[dict] = None,
|
||||
) -> None:
|
||||
"""
|
||||
标记候选执行成功
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
candidate_id: 候选ID
|
||||
status_code: HTTP 状态码
|
||||
latency_ms: 延迟(毫秒)
|
||||
concurrent_requests: 并发请求数
|
||||
extra_data: 额外数据
|
||||
"""
|
||||
candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
|
||||
if candidate:
|
||||
candidate.status = "success"
|
||||
candidate.status_code = status_code
|
||||
candidate.latency_ms = latency_ms
|
||||
candidate.concurrent_requests = concurrent_requests
|
||||
candidate.finished_at = datetime.now(timezone.utc)
|
||||
if extra_data:
|
||||
candidate.extra_data = {**(candidate.extra_data or {}), **extra_data}
|
||||
# 关键状态更新:立即提交,不使用批量提交
|
||||
# 原因:前端需要实时看到请求成功/失败状态
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def mark_candidate_failed(
|
||||
db: Session,
|
||||
candidate_id: str,
|
||||
error_type: str,
|
||||
error_message: str,
|
||||
status_code: Optional[int] = None,
|
||||
latency_ms: Optional[int] = None,
|
||||
concurrent_requests: Optional[int] = None,
|
||||
extra_data: Optional[dict] = None,
|
||||
) -> None:
|
||||
"""
|
||||
标记候选执行失败
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
candidate_id: 候选ID
|
||||
error_type: 错误类型
|
||||
error_message: 错误消息
|
||||
status_code: HTTP 状态码(如果有)
|
||||
latency_ms: 延迟(毫秒)
|
||||
concurrent_requests: 并发请求数
|
||||
extra_data: 额外数据
|
||||
"""
|
||||
candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
|
||||
if candidate:
|
||||
candidate.status = "failed"
|
||||
candidate.error_type = error_type
|
||||
candidate.error_message = error_message
|
||||
candidate.status_code = status_code
|
||||
candidate.latency_ms = latency_ms
|
||||
candidate.concurrent_requests = concurrent_requests
|
||||
candidate.finished_at = datetime.now(timezone.utc)
|
||||
if extra_data:
|
||||
candidate.extra_data = {**(candidate.extra_data or {}), **extra_data}
|
||||
# 关键状态更新:立即提交,不使用批量提交
|
||||
# 原因:前端需要实时看到请求成功/失败状态
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def mark_candidate_skipped(
|
||||
db: Session, candidate_id: str, skip_reason: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
标记候选为已跳过
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
candidate_id: 候选ID
|
||||
skip_reason: 跳过原因
|
||||
"""
|
||||
candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
|
||||
if candidate:
|
||||
candidate.status = "skipped"
|
||||
candidate.skip_reason = skip_reason
|
||||
candidate.finished_at = datetime.now(timezone.utc)
|
||||
db.flush() # 只 flush,不立即 commit
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
@staticmethod
|
||||
def get_candidates_by_request_id(db: Session, request_id: str) -> List[RequestCandidate]:
|
||||
"""
|
||||
获取请求的所有候选记录
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
request_id: 请求ID
|
||||
|
||||
Returns:
|
||||
候选记录列表,按 candidate_index 排序
|
||||
"""
|
||||
return (
|
||||
db.query(RequestCandidate)
|
||||
.filter(RequestCandidate.request_id == request_id)
|
||||
.order_by(RequestCandidate.candidate_index)
|
||||
.all()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_candidate_stats_by_provider(db: Session, provider_id: str, limit: int = 100) -> dict:
|
||||
"""
|
||||
获取 Provider 的候选统计
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
provider_id: Provider ID
|
||||
limit: 最近记录数量限制
|
||||
|
||||
Returns:
|
||||
统计信息字典
|
||||
"""
|
||||
candidates = (
|
||||
db.query(RequestCandidate)
|
||||
.filter(RequestCandidate.provider_id == provider_id)
|
||||
.order_by(RequestCandidate.created_at.desc())
|
||||
.limit(limit)
|
||||
.all()
|
||||
)
|
||||
|
||||
total_candidates = len(candidates)
|
||||
success_count = sum(1 for c in candidates if c.status == "success")
|
||||
failed_count = sum(1 for c in candidates if c.status == "failed")
|
||||
skipped_count = sum(1 for c in candidates if c.status == "skipped")
|
||||
pending_count = sum(1 for c in candidates if c.status == "pending")
|
||||
available_count = sum(1 for c in candidates if c.status == "available")
|
||||
|
||||
# 计算失败率(只统计已完成的候选,即成功或失败的)
|
||||
completed_count = success_count + failed_count
|
||||
failure_rate = (failed_count / completed_count * 100) if completed_count > 0 else 0
|
||||
|
||||
return {
|
||||
"total_attempts": total_candidates, # 前端使用 total_attempts 字段
|
||||
"success_count": success_count,
|
||||
"failed_count": failed_count,
|
||||
"skipped_count": skipped_count,
|
||||
"pending_count": pending_count,
|
||||
"available_count": available_count, # 新增:尚未被调度的候选数
|
||||
"failure_rate": round(failure_rate, 2),
|
||||
}
|
||||
193
src/services/request/executor.py
Normal file
193
src/services/request/executor.py
Normal file
@@ -0,0 +1,193 @@
|
||||
"""
|
||||
封装请求执行逻辑,包含并发控制与链路追踪。
|
||||
"""
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Optional, Union
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.enums import APIFormat
|
||||
from src.core.exceptions import ConcurrencyLimitError
|
||||
from src.core.logger import logger
|
||||
from src.services.health.monitor import health_monitor
|
||||
from src.services.rate_limit.adaptive_reservation import get_adaptive_reservation_manager
|
||||
from src.services.request.candidate import RequestCandidateService
|
||||
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutionContext:
|
||||
candidate_id: str
|
||||
candidate_index: int
|
||||
provider_id: str
|
||||
endpoint_id: str
|
||||
key_id: str
|
||||
user_id: Optional[str]
|
||||
api_key_id: Optional[str]
|
||||
is_cached_user: bool
|
||||
start_time: Optional[float] = None
|
||||
elapsed_ms: Optional[int] = None
|
||||
concurrent_requests: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutionResult:
|
||||
response: Any
|
||||
context: ExecutionContext
|
||||
|
||||
|
||||
class ExecutionError(Exception):
|
||||
def __init__(self, cause: Exception, context: ExecutionContext):
|
||||
super().__init__(str(cause))
|
||||
self.cause = cause
|
||||
self.context = context
|
||||
|
||||
|
||||
class RequestExecutor:
|
||||
def __init__(self, db: Session, concurrency_manager, adaptive_manager):
|
||||
self.db = db
|
||||
self.concurrency_manager = concurrency_manager
|
||||
self.adaptive_manager = adaptive_manager
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
*,
|
||||
candidate,
|
||||
candidate_id: str,
|
||||
candidate_index: int,
|
||||
user_api_key,
|
||||
request_func: Callable,
|
||||
request_id: Optional[str],
|
||||
api_format: Union[str, APIFormat],
|
||||
model_name: str,
|
||||
is_stream: bool = False,
|
||||
) -> ExecutionResult:
|
||||
provider = candidate.provider
|
||||
endpoint = candidate.endpoint
|
||||
key = candidate.key
|
||||
is_cached_user = bool(candidate.is_cached)
|
||||
|
||||
# 标记候选开始执行
|
||||
RequestCandidateService.mark_candidate_started(
|
||||
db=self.db,
|
||||
candidate_id=candidate_id,
|
||||
)
|
||||
|
||||
context = ExecutionContext(
|
||||
candidate_id=candidate_id,
|
||||
candidate_index=candidate_index,
|
||||
provider_id=provider.id,
|
||||
endpoint_id=endpoint.id,
|
||||
key_id=key.id,
|
||||
user_id=user_api_key.user_id,
|
||||
api_key_id=user_api_key.id,
|
||||
is_cached_user=is_cached_user,
|
||||
)
|
||||
|
||||
try:
|
||||
# 计算动态预留比例
|
||||
reservation_manager = get_adaptive_reservation_manager()
|
||||
# 获取当前并发数用于计算负载
|
||||
try:
|
||||
_, current_key_concurrent = await self.concurrency_manager.get_current_concurrency(
|
||||
endpoint_id=endpoint.id,
|
||||
key_id=key.id,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"获取并发数失败(用于预留计算): {e}")
|
||||
current_key_concurrent = 0
|
||||
|
||||
# 获取有效的并发限制(自适应或固定)
|
||||
effective_key_limit = (
|
||||
key.learned_max_concurrent if key.max_concurrent is None else key.max_concurrent
|
||||
)
|
||||
|
||||
reservation_result = reservation_manager.calculate_reservation(
|
||||
key=key,
|
||||
current_concurrent=current_key_concurrent,
|
||||
effective_limit=effective_key_limit,
|
||||
)
|
||||
dynamic_reservation_ratio = reservation_result.ratio
|
||||
|
||||
logger.debug(f"[Executor] 动态预留: key={key.id[:8]}..., "
|
||||
f"ratio={dynamic_reservation_ratio:.0%}, phase={reservation_result.phase}, "
|
||||
f"confidence={reservation_result.confidence:.0%}")
|
||||
|
||||
async with self.concurrency_manager.concurrency_guard(
|
||||
endpoint_id=endpoint.id,
|
||||
endpoint_max_concurrent=endpoint.max_concurrent,
|
||||
key_id=key.id,
|
||||
key_max_concurrent=effective_key_limit,
|
||||
is_cached_user=is_cached_user,
|
||||
cache_reservation_ratio=dynamic_reservation_ratio,
|
||||
):
|
||||
try:
|
||||
_, key_concurrent = await self.concurrency_manager.get_current_concurrency(
|
||||
endpoint_id=endpoint.id,
|
||||
key_id=key.id,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"获取并发数失败(guard 内): {e}")
|
||||
key_concurrent = None
|
||||
|
||||
context.concurrent_requests = key_concurrent
|
||||
context.start_time = time.time()
|
||||
|
||||
response = await request_func(provider, endpoint, key)
|
||||
|
||||
context.elapsed_ms = int((time.time() - context.start_time) * 1000)
|
||||
|
||||
health_monitor.record_success(
|
||||
db=self.db,
|
||||
key_id=key.id,
|
||||
response_time_ms=context.elapsed_ms,
|
||||
)
|
||||
|
||||
# 自适应模式:max_concurrent = NULL
|
||||
if key.max_concurrent is None and key_concurrent is not None:
|
||||
self.adaptive_manager.handle_success(
|
||||
db=self.db,
|
||||
key=key,
|
||||
current_concurrent=key_concurrent,
|
||||
)
|
||||
|
||||
# 根据是否为流式请求,标记不同状态
|
||||
if is_stream:
|
||||
# 流式请求:标记为 streaming 状态
|
||||
# 此时连接已建立但流传输尚未完成
|
||||
# success 状态会在流完成后由 _record_stream_stats 方法标记
|
||||
RequestCandidateService.mark_candidate_streaming(
|
||||
db=self.db,
|
||||
candidate_id=candidate_id,
|
||||
status_code=200,
|
||||
concurrent_requests=key_concurrent,
|
||||
)
|
||||
else:
|
||||
# 非流式请求:标记为 success 状态
|
||||
RequestCandidateService.mark_candidate_success(
|
||||
db=self.db,
|
||||
candidate_id=candidate_id,
|
||||
status_code=200,
|
||||
latency_ms=context.elapsed_ms,
|
||||
concurrent_requests=key_concurrent,
|
||||
extra_data={
|
||||
"is_cached_user": is_cached_user,
|
||||
"model_name": model_name,
|
||||
"api_format": (
|
||||
api_format.value if isinstance(api_format, APIFormat) else api_format
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
return ExecutionResult(response=response, context=context)
|
||||
except ConcurrencyLimitError as exc:
|
||||
raise ExecutionError(exc, context) from exc
|
||||
except Exception as exc:
|
||||
context.elapsed_ms = (
|
||||
int((time.time() - context.start_time) * 1000)
|
||||
if context.start_time is not None
|
||||
else None
|
||||
)
|
||||
raise ExecutionError(exc, context) from exc
|
||||
330
src/services/request/result.py
Normal file
330
src/services/request/result.py
Normal file
@@ -0,0 +1,330 @@
|
||||
"""
|
||||
统一的请求结果和元数据结构
|
||||
|
||||
设计原则:
|
||||
1. RequestMetadata: 描述请求执行的上下文(Provider、Endpoint、Key、API格式等)
|
||||
2. RequestResult: 封装请求的完整结果(成功/失败、响应、元数据、费用等)
|
||||
3. 确保 api_format 在整个链路中始终可用
|
||||
|
||||
使用场景:
|
||||
- ProviderService 创建 RequestMetadata
|
||||
- FallbackOrchestrator 在异常时补充 RequestMetadata
|
||||
- ChatHandlerBase 使用 RequestResult 记录 Usage
|
||||
- ChatAdapterBase 使用 RequestResult 处理异常响应
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
|
||||
|
||||
class RequestStatus(Enum):
|
||||
"""请求状态"""
|
||||
|
||||
SUCCESS = "success"
|
||||
FAILED = "failed"
|
||||
PARTIAL = "partial" # 流式请求部分成功
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestMetadata:
|
||||
"""
|
||||
请求元数据 - 描述请求执行的上下文
|
||||
|
||||
必填字段:
|
||||
- api_format: API 格式,必须在请求开始时就确定
|
||||
- provider: Provider 名称
|
||||
- model: 模型名称
|
||||
|
||||
可选字段:
|
||||
- provider_id, provider_endpoint_id, provider_api_key_id: Provider 追踪信息
|
||||
- provider_request_headers, provider_response_headers: 请求/响应头
|
||||
- attempt_id: 请求尝试 ID
|
||||
- original_model: 用户请求的原始模型名(映射前)
|
||||
"""
|
||||
|
||||
# 必填字段 - 在请求开始时就应该确定
|
||||
api_format: str
|
||||
provider: str = "unknown"
|
||||
model: str = "unknown"
|
||||
|
||||
# Provider 追踪信息
|
||||
provider_id: Optional[str] = None
|
||||
provider_endpoint_id: Optional[str] = None
|
||||
provider_api_key_id: Optional[str] = None
|
||||
|
||||
# 请求/响应头
|
||||
provider_request_headers: Dict[str, str] = field(default_factory=dict)
|
||||
provider_response_headers: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
# 其他元数据
|
||||
attempt_id: Optional[str] = None
|
||||
original_model: Optional[str] = None # 用户请求的原始模型名(用于价格计算)
|
||||
|
||||
# Provider 响应元数据(存储 provider 返回的额外信息,如 Gemini 的 modelVersion)
|
||||
response_metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def with_provider_info(
|
||||
self,
|
||||
provider: str,
|
||||
provider_id: str,
|
||||
provider_endpoint_id: str,
|
||||
provider_api_key_id: str,
|
||||
) -> "RequestMetadata":
|
||||
"""返回包含 Provider 信息的新 RequestMetadata"""
|
||||
return RequestMetadata(
|
||||
api_format=self.api_format,
|
||||
provider=provider,
|
||||
model=self.model,
|
||||
provider_id=provider_id,
|
||||
provider_endpoint_id=provider_endpoint_id,
|
||||
provider_api_key_id=provider_api_key_id,
|
||||
provider_request_headers=self.provider_request_headers,
|
||||
provider_response_headers=self.provider_response_headers,
|
||||
attempt_id=self.attempt_id,
|
||||
original_model=self.original_model,
|
||||
response_metadata=self.response_metadata,
|
||||
)
|
||||
|
||||
def with_response_headers(self, headers: Dict[str, str]) -> "RequestMetadata":
|
||||
"""返回包含响应头的新 RequestMetadata"""
|
||||
return RequestMetadata(
|
||||
api_format=self.api_format,
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
provider_id=self.provider_id,
|
||||
provider_endpoint_id=self.provider_endpoint_id,
|
||||
provider_api_key_id=self.provider_api_key_id,
|
||||
provider_request_headers=self.provider_request_headers,
|
||||
provider_response_headers=headers,
|
||||
attempt_id=self.attempt_id,
|
||||
original_model=self.original_model,
|
||||
response_metadata=self.response_metadata,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageInfo:
|
||||
"""Token 使用量信息"""
|
||||
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class CostInfo:
|
||||
"""费用信息"""
|
||||
|
||||
input_cost_usd: float = 0.0
|
||||
output_cost_usd: float = 0.0
|
||||
cache_creation_cost_usd: float = 0.0
|
||||
cache_read_cost_usd: float = 0.0
|
||||
cache_cost_usd: float = 0.0
|
||||
total_cost_usd: float = 0.0
|
||||
|
||||
# 实际费用(乘以 rate_multiplier 后)
|
||||
actual_input_cost_usd: float = 0.0
|
||||
actual_output_cost_usd: float = 0.0
|
||||
actual_cache_creation_cost_usd: float = 0.0
|
||||
actual_cache_read_cost_usd: float = 0.0
|
||||
actual_total_cost_usd: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
"""
|
||||
请求结果 - 封装请求的完整结果
|
||||
|
||||
用于:
|
||||
- 成功请求:包含响应数据、使用量、费用
|
||||
- 失败请求:包含错误信息、状态码
|
||||
- 流式请求:包含流生成器和元数据
|
||||
"""
|
||||
|
||||
# 状态
|
||||
status: RequestStatus
|
||||
|
||||
# 元数据(必须存在)
|
||||
metadata: RequestMetadata
|
||||
|
||||
# 响应相关
|
||||
response_data: Optional[Any] = None # 成功时的响应数据
|
||||
stream: Optional[AsyncIterator[str]] = None # 流式响应
|
||||
|
||||
# 使用量和费用
|
||||
usage: UsageInfo = field(default_factory=UsageInfo)
|
||||
cost: CostInfo = field(default_factory=CostInfo)
|
||||
|
||||
# 错误信息
|
||||
status_code: int = 200
|
||||
error_message: Optional[str] = None
|
||||
error_type: Optional[str] = None
|
||||
|
||||
# 计时
|
||||
response_time_ms: int = 0
|
||||
|
||||
# 请求信息(用于记录)
|
||||
is_stream: bool = False
|
||||
request_headers: Dict[str, str] = field(default_factory=dict)
|
||||
request_body: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def is_success(self) -> bool:
|
||||
return self.status == RequestStatus.SUCCESS
|
||||
|
||||
@property
|
||||
def is_failed(self) -> bool:
|
||||
return self.status == RequestStatus.FAILED
|
||||
|
||||
@classmethod
|
||||
def success(
|
||||
cls,
|
||||
metadata: RequestMetadata,
|
||||
response_data: Any,
|
||||
usage: UsageInfo,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
) -> "RequestResult":
|
||||
"""创建成功的请求结果"""
|
||||
return cls(
|
||||
status=RequestStatus.SUCCESS,
|
||||
metadata=metadata,
|
||||
response_data=response_data,
|
||||
usage=usage,
|
||||
status_code=200,
|
||||
response_time_ms=response_time_ms,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def failed(
|
||||
cls,
|
||||
metadata: RequestMetadata,
|
||||
status_code: int,
|
||||
error_message: str,
|
||||
error_type: str,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
) -> "RequestResult":
|
||||
"""创建失败的请求结果"""
|
||||
return cls(
|
||||
status=RequestStatus.FAILED,
|
||||
metadata=metadata,
|
||||
status_code=status_code,
|
||||
error_message=error_message,
|
||||
error_type=error_type,
|
||||
response_time_ms=response_time_ms,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_exception(
|
||||
cls,
|
||||
exception: Exception,
|
||||
api_format: str,
|
||||
model: str,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
) -> "RequestResult":
|
||||
"""从异常创建失败的请求结果"""
|
||||
# 尝试从异常中提取 metadata
|
||||
existing_metadata = getattr(exception, "request_metadata", None)
|
||||
|
||||
def get_meta_value(meta, key, default=None):
|
||||
"""从 metadata 中提取值,支持字典和对象两种形式"""
|
||||
if meta is None:
|
||||
return default
|
||||
if isinstance(meta, dict):
|
||||
return meta.get(key, default)
|
||||
return getattr(meta, key, default)
|
||||
|
||||
if existing_metadata:
|
||||
# 如果异常已有 metadata,使用它但确保 api_format 存在
|
||||
metadata = RequestMetadata(
|
||||
api_format=get_meta_value(existing_metadata, "api_format") or api_format,
|
||||
provider=get_meta_value(existing_metadata, "provider", "unknown") or "unknown",
|
||||
model=get_meta_value(existing_metadata, "model", model) or model,
|
||||
provider_id=get_meta_value(existing_metadata, "provider_id"),
|
||||
provider_endpoint_id=get_meta_value(existing_metadata, "provider_endpoint_id"),
|
||||
provider_api_key_id=get_meta_value(existing_metadata, "provider_api_key_id"),
|
||||
provider_request_headers=get_meta_value(existing_metadata, "provider_request_headers", {}),
|
||||
provider_response_headers=get_meta_value(
|
||||
existing_metadata, "provider_response_headers", {}
|
||||
),
|
||||
attempt_id=get_meta_value(existing_metadata, "attempt_id"),
|
||||
original_model=get_meta_value(existing_metadata, "original_model"),
|
||||
response_metadata=get_meta_value(existing_metadata, "response_metadata", {}),
|
||||
)
|
||||
else:
|
||||
# 创建最小的 metadata
|
||||
metadata = RequestMetadata(
|
||||
api_format=api_format,
|
||||
provider="unknown",
|
||||
model=model,
|
||||
)
|
||||
|
||||
# 确定状态码和错误类型
|
||||
from src.core.exceptions import (
|
||||
ProviderAuthException,
|
||||
ProviderNotAvailableException,
|
||||
ProviderRateLimitException,
|
||||
ProviderTimeoutException,
|
||||
)
|
||||
|
||||
if isinstance(exception, ProviderAuthException):
|
||||
status_code = 503
|
||||
error_type = "provider_auth_error"
|
||||
elif isinstance(exception, ProviderRateLimitException):
|
||||
status_code = 429
|
||||
error_type = "rate_limit_exceeded"
|
||||
elif isinstance(exception, ProviderTimeoutException):
|
||||
status_code = 504
|
||||
error_type = "timeout_error"
|
||||
elif isinstance(exception, ProviderNotAvailableException):
|
||||
status_code = 503
|
||||
error_type = "provider_unavailable"
|
||||
else:
|
||||
status_code = 500
|
||||
error_type = "internal_error"
|
||||
|
||||
return cls(
|
||||
status=RequestStatus.FAILED,
|
||||
metadata=metadata,
|
||||
status_code=status_code,
|
||||
error_message=str(exception),
|
||||
error_type=error_type,
|
||||
response_time_ms=response_time_ms,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
|
||||
|
||||
class StreamWithMetadata:
|
||||
"""带元数据的流式响应包装器"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream: AsyncIterator[str],
|
||||
metadata: RequestMetadata,
|
||||
response_headers_container: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self.stream = stream
|
||||
self.metadata = metadata
|
||||
self.response_headers_container = response_headers_container
|
||||
self._metadata_updated = False
|
||||
|
||||
def update_metadata_with_response_headers(self):
|
||||
"""使用实际的响应头更新元数据"""
|
||||
if self.response_headers_container and "headers" in self.response_headers_container:
|
||||
if not self._metadata_updated:
|
||||
self.metadata = self.metadata.with_response_headers(
|
||||
self.response_headers_container["headers"]
|
||||
)
|
||||
self._metadata_updated = True
|
||||
|
||||
def __aiter__(self):
|
||||
return self.stream
|
||||
|
||||
async def __anext__(self):
|
||||
return await self.stream.__anext__()
|
||||
23
src/services/system/__init__.py
Normal file
23
src/services/system/__init__.py
Normal file
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
系统服务模块
|
||||
|
||||
包含系统配置、审计日志、公告等功能。
|
||||
"""
|
||||
|
||||
from src.services.system.announcement import AnnouncementService
|
||||
from src.services.system.audit import AuditService
|
||||
from src.services.system.cleanup_scheduler import CleanupScheduler
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.system.scheduler import APP_TIMEZONE, TaskScheduler, get_scheduler
|
||||
from src.services.system.sync_stats import SyncStatsService
|
||||
|
||||
__all__ = [
|
||||
"SystemConfigService",
|
||||
"AuditService",
|
||||
"AnnouncementService",
|
||||
"CleanupScheduler",
|
||||
"SyncStatsService",
|
||||
"TaskScheduler",
|
||||
"get_scheduler",
|
||||
"APP_TIMEZONE",
|
||||
]
|
||||
241
src/services/system/announcement.py
Normal file
241
src/services/system/announcement.py
Normal file
@@ -0,0 +1,241 @@
|
||||
"""
|
||||
公告系统服务
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import and_, or_
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.exceptions import ForbiddenException, NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Announcement, AnnouncementRead, User, UserRole
|
||||
|
||||
|
||||
|
||||
class AnnouncementService:
|
||||
"""公告系统服务"""
|
||||
|
||||
@staticmethod
|
||||
def create_announcement(
|
||||
db: Session,
|
||||
author_id: str, # UUID
|
||||
title: str,
|
||||
content: str,
|
||||
type: str = "info",
|
||||
priority: int = 0,
|
||||
is_pinned: bool = False,
|
||||
start_time: Optional[datetime] = None,
|
||||
end_time: Optional[datetime] = None,
|
||||
) -> Announcement:
|
||||
"""创建公告"""
|
||||
# 验证作者是否为管理员
|
||||
author = db.query(User).filter(User.id == author_id).first()
|
||||
if not author or author.role != UserRole.ADMIN:
|
||||
raise ForbiddenException("Only administrators can create announcements")
|
||||
|
||||
# 验证类型
|
||||
if type not in ["info", "warning", "maintenance", "important"]:
|
||||
raise ValueError("Invalid announcement type")
|
||||
|
||||
announcement = Announcement(
|
||||
title=title,
|
||||
content=content,
|
||||
type=type,
|
||||
priority=priority,
|
||||
author_id=author_id,
|
||||
is_pinned=is_pinned,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
is_active=True,
|
||||
)
|
||||
|
||||
db.add(announcement)
|
||||
db.commit()
|
||||
db.refresh(announcement)
|
||||
|
||||
logger.info(f"Created announcement: {announcement.id} - {title}")
|
||||
return announcement
|
||||
|
||||
@staticmethod
|
||||
def get_announcements(
|
||||
db: Session,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
active_only: bool = True,
|
||||
include_read_status: bool = False,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
) -> dict:
|
||||
"""获取公告列表"""
|
||||
query = db.query(Announcement)
|
||||
|
||||
# 筛选条件
|
||||
if active_only:
|
||||
now = datetime.now(timezone.utc)
|
||||
query = query.filter(
|
||||
Announcement.is_active == True,
|
||||
or_(Announcement.start_time == None, Announcement.start_time <= now),
|
||||
or_(Announcement.end_time == None, Announcement.end_time >= now),
|
||||
)
|
||||
|
||||
# 排序:置顶优先,然后按优先级和创建时间
|
||||
query = query.order_by(
|
||||
Announcement.is_pinned.desc(),
|
||||
Announcement.priority.desc(),
|
||||
Announcement.created_at.desc(),
|
||||
)
|
||||
|
||||
# 分页
|
||||
total = query.count()
|
||||
announcements = query.offset(offset).limit(limit).all()
|
||||
|
||||
# 获取已读状态
|
||||
read_announcement_ids = set()
|
||||
unread_count = 0
|
||||
|
||||
if user_id and include_read_status:
|
||||
read_records = (
|
||||
db.query(AnnouncementRead.announcement_id)
|
||||
.filter(AnnouncementRead.user_id == user_id)
|
||||
.all()
|
||||
)
|
||||
read_announcement_ids = {r[0] for r in read_records}
|
||||
unread_count = total - len(read_announcement_ids)
|
||||
|
||||
# 构建返回数据
|
||||
items = []
|
||||
for announcement in announcements:
|
||||
item = {
|
||||
"id": announcement.id,
|
||||
"title": announcement.title,
|
||||
"content": announcement.content,
|
||||
"type": announcement.type,
|
||||
"priority": announcement.priority,
|
||||
"is_pinned": announcement.is_pinned,
|
||||
"is_active": announcement.is_active,
|
||||
"author": {"id": announcement.author.id, "username": announcement.author.username},
|
||||
"start_time": announcement.start_time,
|
||||
"end_time": announcement.end_time,
|
||||
"created_at": announcement.created_at,
|
||||
"updated_at": announcement.updated_at,
|
||||
}
|
||||
|
||||
if include_read_status and user_id:
|
||||
item["is_read"] = announcement.id in read_announcement_ids
|
||||
|
||||
items.append(item)
|
||||
|
||||
result = {"items": items, "total": total}
|
||||
|
||||
if include_read_status and user_id:
|
||||
result["unread_count"] = unread_count
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def get_announcement(db: Session, announcement_id: str) -> Announcement: # UUID
|
||||
"""获取单个公告"""
|
||||
announcement = db.query(Announcement).filter(Announcement.id == announcement_id).first()
|
||||
|
||||
if not announcement:
|
||||
raise NotFoundException("Announcement not found")
|
||||
|
||||
return announcement
|
||||
|
||||
@staticmethod
|
||||
def update_announcement(
|
||||
db: Session,
|
||||
announcement_id: str, # UUID
|
||||
user_id: str, # UUID
|
||||
title: Optional[str] = None,
|
||||
content: Optional[str] = None,
|
||||
type: Optional[str] = None,
|
||||
priority: Optional[int] = None,
|
||||
is_active: Optional[bool] = None,
|
||||
is_pinned: Optional[bool] = None,
|
||||
start_time: Optional[datetime] = None,
|
||||
end_time: Optional[datetime] = None,
|
||||
) -> Announcement:
|
||||
"""更新公告"""
|
||||
# 验证用户是否为管理员
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user or user.role != UserRole.ADMIN:
|
||||
raise ForbiddenException("Only administrators can update announcements")
|
||||
|
||||
announcement = AnnouncementService.get_announcement(db, announcement_id)
|
||||
|
||||
# 更新提供的字段
|
||||
if title is not None:
|
||||
announcement.title = title
|
||||
if content is not None:
|
||||
announcement.content = content
|
||||
if type is not None:
|
||||
if type not in ["info", "warning", "maintenance", "important"]:
|
||||
raise ValueError("Invalid announcement type")
|
||||
announcement.type = type
|
||||
if priority is not None:
|
||||
announcement.priority = priority
|
||||
if is_active is not None:
|
||||
announcement.is_active = is_active
|
||||
if is_pinned is not None:
|
||||
announcement.is_pinned = is_pinned
|
||||
if start_time is not None:
|
||||
announcement.start_time = start_time
|
||||
if end_time is not None:
|
||||
announcement.end_time = end_time
|
||||
|
||||
db.commit()
|
||||
db.refresh(announcement)
|
||||
|
||||
logger.info(f"Updated announcement: {announcement_id}")
|
||||
return announcement
|
||||
|
||||
@staticmethod
|
||||
def delete_announcement(db: Session, announcement_id: str, user_id: str) -> None: # UUID
|
||||
"""删除公告"""
|
||||
# 验证用户是否为管理员
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user or user.role != UserRole.ADMIN:
|
||||
raise ForbiddenException("Only administrators can delete announcements")
|
||||
|
||||
announcement = AnnouncementService.get_announcement(db, announcement_id)
|
||||
|
||||
db.delete(announcement)
|
||||
db.commit()
|
||||
|
||||
logger.info(f"Deleted announcement: {announcement_id}")
|
||||
|
||||
@staticmethod
|
||||
def mark_as_read(db: Session, announcement_id: str, user_id: str) -> None: # UUID
|
||||
"""标记公告为已读"""
|
||||
# 检查公告是否存在
|
||||
announcement = AnnouncementService.get_announcement(db, announcement_id)
|
||||
|
||||
# 检查是否已经标记为已读
|
||||
existing = (
|
||||
db.query(AnnouncementRead)
|
||||
.filter(
|
||||
AnnouncementRead.user_id == user_id,
|
||||
AnnouncementRead.announcement_id == announcement_id,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not existing:
|
||||
read_record = AnnouncementRead(user_id=user_id, announcement_id=announcement_id)
|
||||
db.add(read_record)
|
||||
db.commit()
|
||||
|
||||
logger.info(f"User {user_id} marked announcement {announcement_id} as read")
|
||||
|
||||
@staticmethod
|
||||
def get_active_announcements(db: Session, user_id: Optional[str] = None) -> dict: # UUID
|
||||
"""获取当前有效的公告(首页展示用)"""
|
||||
return AnnouncementService.get_announcements(
|
||||
db=db,
|
||||
user_id=user_id,
|
||||
active_only=True,
|
||||
include_read_status=True if user_id else False,
|
||||
limit=10,
|
||||
)
|
||||
459
src/services/system/audit.py
Normal file
459
src/services/system/audit.py
Normal file
@@ -0,0 +1,459 @@
|
||||
"""
|
||||
审计日志服务
|
||||
记录所有重要操作和安全事件
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.database import get_db
|
||||
from src.models.database import AuditEventType, AuditLog
|
||||
from src.utils.transaction_manager import transactional
|
||||
|
||||
|
||||
|
||||
# 审计模型已移至 src/models/database.py
|
||||
|
||||
|
||||
class AuditService:
|
||||
"""审计服务"""
|
||||
|
||||
@staticmethod
|
||||
@transactional(commit=False) # 不自动提交,让调用方决定
|
||||
def log_event(
|
||||
db: Session,
|
||||
event_type: AuditEventType,
|
||||
description: str,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
api_key_id: Optional[str] = None, # UUID
|
||||
ip_address: Optional[str] = None,
|
||||
user_agent: Optional[str] = None,
|
||||
request_id: Optional[str] = None,
|
||||
status_code: Optional[int] = None,
|
||||
error_message: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> AuditLog:
|
||||
"""
|
||||
记录审计事件
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
event_type: 事件类型
|
||||
description: 事件描述
|
||||
user_id: 用户ID
|
||||
api_key_id: API密钥ID
|
||||
ip_address: IP地址
|
||||
user_agent: 用户代理
|
||||
request_id: 请求ID
|
||||
status_code: 状态码
|
||||
error_message: 错误消息
|
||||
metadata: 额外元数据
|
||||
|
||||
Returns:
|
||||
审计日志记录
|
||||
"""
|
||||
try:
|
||||
audit_log = AuditLog(
|
||||
event_type=event_type.value,
|
||||
description=description,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
ip_address=ip_address,
|
||||
user_agent=user_agent,
|
||||
request_id=request_id,
|
||||
status_code=status_code,
|
||||
error_message=error_message,
|
||||
event_metadata=metadata,
|
||||
)
|
||||
|
||||
db.add(audit_log)
|
||||
db.commit() # 立即提交事务,释放数据库锁
|
||||
db.refresh(audit_log)
|
||||
|
||||
# 同时记录到系统日志
|
||||
log_message = (
|
||||
f"AUDIT [{event_type.value}] - {description} | "
|
||||
f"user_id={user_id}, ip={ip_address}"
|
||||
)
|
||||
|
||||
if event_type in [
|
||||
AuditEventType.UNAUTHORIZED_ACCESS,
|
||||
AuditEventType.SUSPICIOUS_ACTIVITY,
|
||||
]:
|
||||
logger.warning(log_message)
|
||||
elif event_type in [AuditEventType.LOGIN_FAILED, AuditEventType.REQUEST_FAILED]:
|
||||
logger.info(log_message)
|
||||
else:
|
||||
logger.debug(log_message)
|
||||
|
||||
return audit_log
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to log audit event: {e}")
|
||||
db.rollback()
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def log_login_attempt(
|
||||
db: Session,
|
||||
email: str,
|
||||
success: bool,
|
||||
ip_address: str,
|
||||
user_agent: str,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
error_reason: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
记录登录尝试
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
email: 登录邮箱
|
||||
success: 是否成功
|
||||
ip_address: IP地址
|
||||
user_agent: 用户代理
|
||||
user_id: 用户ID(成功时)
|
||||
error_reason: 失败原因
|
||||
"""
|
||||
event_type = AuditEventType.LOGIN_SUCCESS if success else AuditEventType.LOGIN_FAILED
|
||||
description = f"Login attempt for {email}"
|
||||
if not success and error_reason:
|
||||
description += f": {error_reason}"
|
||||
|
||||
AuditService.log_event(
|
||||
db=db,
|
||||
event_type=event_type,
|
||||
description=description,
|
||||
user_id=user_id,
|
||||
ip_address=ip_address,
|
||||
user_agent=user_agent,
|
||||
metadata={"email": email},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def log_api_request(
|
||||
db: Session,
|
||||
user_id: str, # UUID
|
||||
api_key_id: str, # UUID
|
||||
request_id: str,
|
||||
model: str,
|
||||
provider: str,
|
||||
success: bool,
|
||||
ip_address: str,
|
||||
status_code: int,
|
||||
error_message: Optional[str] = None,
|
||||
input_tokens: Optional[int] = None,
|
||||
output_tokens: Optional[int] = None,
|
||||
cost_usd: Optional[float] = None,
|
||||
):
|
||||
"""
|
||||
记录API请求
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID
|
||||
api_key_id: API密钥ID
|
||||
request_id: 请求ID
|
||||
model: 模型名称
|
||||
provider: 提供商名称
|
||||
success: 是否成功
|
||||
ip_address: IP地址
|
||||
status_code: 状态码
|
||||
error_message: 错误消息
|
||||
input_tokens: 输入tokens
|
||||
output_tokens: 输出tokens
|
||||
cost_usd: 成本(美元)
|
||||
"""
|
||||
event_type = AuditEventType.REQUEST_SUCCESS if success else AuditEventType.REQUEST_FAILED
|
||||
description = f"API request to {provider}/{model}"
|
||||
|
||||
metadata = {"model": model, "provider": provider}
|
||||
|
||||
if input_tokens:
|
||||
metadata["input_tokens"] = input_tokens
|
||||
if output_tokens:
|
||||
metadata["output_tokens"] = output_tokens
|
||||
if cost_usd:
|
||||
metadata["cost_usd"] = cost_usd
|
||||
|
||||
AuditService.log_event(
|
||||
db=db,
|
||||
event_type=event_type,
|
||||
description=description,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
request_id=request_id,
|
||||
ip_address=ip_address,
|
||||
status_code=status_code,
|
||||
error_message=error_message,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def log_security_event(
|
||||
db: Session,
|
||||
event_type: AuditEventType,
|
||||
description: str,
|
||||
ip_address: str,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
severity: str = "medium",
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
"""
|
||||
记录安全事件
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
event_type: 事件类型
|
||||
description: 事件描述
|
||||
ip_address: IP地址
|
||||
user_id: 用户ID
|
||||
severity: 严重程度 (low, medium, high, critical)
|
||||
details: 详细信息
|
||||
"""
|
||||
event_metadata = {"severity": severity}
|
||||
if details:
|
||||
event_metadata.update(details)
|
||||
|
||||
AuditService.log_event(
|
||||
db=db,
|
||||
event_type=event_type,
|
||||
description=description,
|
||||
user_id=user_id,
|
||||
ip_address=ip_address,
|
||||
metadata=event_metadata,
|
||||
)
|
||||
|
||||
# 对于高严重性事件,简化日志输出
|
||||
if severity in ["high", "critical"]:
|
||||
logger.error(f"安全告警 [{severity.upper()}]: {description}")
|
||||
|
||||
@staticmethod
|
||||
def get_user_audit_logs(
|
||||
db: Session,
|
||||
user_id: str, # UUID
|
||||
event_types: Optional[List[AuditEventType]] = None,
|
||||
limit: int = 100,
|
||||
) -> List[AuditLog]:
|
||||
"""
|
||||
获取用户的审计日志
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID
|
||||
event_types: 事件类型过滤
|
||||
limit: 返回数量限制
|
||||
|
||||
Returns:
|
||||
审计日志列表
|
||||
"""
|
||||
query = db.query(AuditLog).filter(AuditLog.user_id == user_id)
|
||||
|
||||
if event_types:
|
||||
event_type_values = [et.value for et in event_types]
|
||||
query = query.filter(AuditLog.event_type.in_(event_type_values))
|
||||
|
||||
return query.order_by(AuditLog.created_at.desc()).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def get_suspicious_activities(db: Session, hours: int = 24, limit: int = 100) -> List[AuditLog]:
|
||||
"""
|
||||
获取可疑活动
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
hours: 时间范围(小时)
|
||||
limit: 返回数量限制
|
||||
|
||||
Returns:
|
||||
可疑活动列表
|
||||
"""
|
||||
cutoff_time = datetime.now(timezone.utc) - timedelta(hours=hours)
|
||||
|
||||
suspicious_types = [
|
||||
AuditEventType.SUSPICIOUS_ACTIVITY.value,
|
||||
AuditEventType.UNAUTHORIZED_ACCESS.value,
|
||||
AuditEventType.LOGIN_FAILED.value,
|
||||
AuditEventType.REQUEST_RATE_LIMITED.value,
|
||||
]
|
||||
|
||||
return (
|
||||
db.query(AuditLog)
|
||||
.filter(AuditLog.event_type.in_(suspicious_types), AuditLog.created_at >= cutoff_time)
|
||||
.order_by(AuditLog.created_at.desc())
|
||||
.limit(limit)
|
||||
.all()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def analyze_user_behavior(db: Session, user_id: str, days: int = 30) -> Dict[str, Any]: # UUID
|
||||
"""
|
||||
分析用户行为
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID
|
||||
days: 分析天数
|
||||
|
||||
Returns:
|
||||
行为分析结果
|
||||
"""
|
||||
from sqlalchemy import func
|
||||
|
||||
cutoff_time = datetime.now(timezone.utc) - timedelta(days=days)
|
||||
|
||||
# 统计各种事件类型
|
||||
event_counts = (
|
||||
db.query(AuditLog.event_type, func.count(AuditLog.id).label("count"))
|
||||
.filter(AuditLog.user_id == user_id, AuditLog.created_at >= cutoff_time)
|
||||
.group_by(AuditLog.event_type)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 统计失败请求
|
||||
failed_requests = (
|
||||
db.query(func.count(AuditLog.id))
|
||||
.filter(
|
||||
AuditLog.user_id == user_id,
|
||||
AuditLog.event_type == AuditEventType.REQUEST_FAILED.value,
|
||||
AuditLog.created_at >= cutoff_time,
|
||||
)
|
||||
.scalar()
|
||||
)
|
||||
|
||||
# 统计成功请求
|
||||
success_requests = (
|
||||
db.query(func.count(AuditLog.id))
|
||||
.filter(
|
||||
AuditLog.user_id == user_id,
|
||||
AuditLog.event_type == AuditEventType.REQUEST_SUCCESS.value,
|
||||
AuditLog.created_at >= cutoff_time,
|
||||
)
|
||||
.scalar()
|
||||
)
|
||||
|
||||
# 获取最近的可疑活动
|
||||
recent_suspicious = (
|
||||
db.query(AuditLog)
|
||||
.filter(
|
||||
AuditLog.user_id == user_id,
|
||||
AuditLog.event_type.in_(
|
||||
[
|
||||
AuditEventType.SUSPICIOUS_ACTIVITY.value,
|
||||
AuditEventType.UNAUTHORIZED_ACCESS.value,
|
||||
]
|
||||
),
|
||||
AuditLog.created_at >= cutoff_time,
|
||||
)
|
||||
.count()
|
||||
)
|
||||
|
||||
return {
|
||||
"user_id": user_id,
|
||||
"period_days": days,
|
||||
"event_counts": {event: count for event, count in event_counts},
|
||||
"failed_requests": failed_requests or 0,
|
||||
"success_requests": success_requests or 0,
|
||||
"success_rate": (
|
||||
success_requests / (success_requests + failed_requests)
|
||||
if (success_requests + failed_requests) > 0
|
||||
else 0
|
||||
),
|
||||
"suspicious_activities": recent_suspicious,
|
||||
"analysis_time": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def log_event_auto(
|
||||
event_type: AuditEventType,
|
||||
description: str,
|
||||
user_id: Optional[str] = None,
|
||||
api_key_id: Optional[str] = None,
|
||||
ip_address: Optional[str] = None,
|
||||
user_agent: Optional[str] = None,
|
||||
request_id: Optional[str] = None,
|
||||
status_code: Optional[int] = None,
|
||||
error_message: Optional[str] = None,
|
||||
event_metadata: Optional[Dict[str, Any]] = None,
|
||||
db: Optional[Session] = None,
|
||||
) -> Optional[AuditLog]:
|
||||
"""
|
||||
自动管理数据库会话的审计日志记录方法
|
||||
适用于中间件等无法直接获取数据库会话的场景
|
||||
|
||||
Args:
|
||||
event_type: 事件类型
|
||||
description: 事件描述
|
||||
user_id: 用户ID
|
||||
api_key_id: API密钥ID
|
||||
ip_address: IP地址
|
||||
user_agent: 用户代理
|
||||
request_id: 请求ID
|
||||
status_code: 状态码
|
||||
error_message: 错误消息
|
||||
event_metadata: 额外元数据
|
||||
db: 数据库会话(可选,如不提供则自动创建)
|
||||
|
||||
Returns:
|
||||
审计日志记录
|
||||
"""
|
||||
# 如果提供了数据库会话,使用它(不自动提交)
|
||||
if db is not None:
|
||||
try:
|
||||
audit_log = AuditService.log_event(
|
||||
db=db,
|
||||
event_type=event_type,
|
||||
description=description,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
ip_address=ip_address,
|
||||
user_agent=user_agent,
|
||||
request_id=request_id,
|
||||
status_code=status_code,
|
||||
error_message=error_message,
|
||||
metadata=event_metadata,
|
||||
)
|
||||
# 注意:不在这里提交,让调用方决定何时提交
|
||||
return audit_log
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to log audit event: {e}")
|
||||
return None
|
||||
|
||||
# 如果没有提供会话,自动创建并管理
|
||||
db_session = None
|
||||
try:
|
||||
db_session = next(get_db())
|
||||
|
||||
audit_log = AuditService.log_event(
|
||||
db=db_session,
|
||||
event_type=event_type,
|
||||
description=description,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
ip_address=ip_address,
|
||||
user_agent=user_agent,
|
||||
request_id=request_id,
|
||||
status_code=status_code,
|
||||
error_message=error_message,
|
||||
metadata=event_metadata,
|
||||
)
|
||||
|
||||
db_session.commit()
|
||||
return audit_log
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to log audit event with auto session: {e}")
|
||||
if db_session is not None:
|
||||
db_session.rollback()
|
||||
return None
|
||||
finally:
|
||||
if db_session is not None:
|
||||
db_session.close()
|
||||
|
||||
|
||||
# 全局审计服务实例
|
||||
audit_service = AuditService()
|
||||
597
src/services/system/cleanup_scheduler.py
Normal file
597
src/services/system/cleanup_scheduler.py
Normal file
@@ -0,0 +1,597 @@
|
||||
"""
|
||||
使用记录清理定时任务
|
||||
|
||||
分级清理策略:
|
||||
- detail_log_retention_days: 压缩 request_body 和 response_body 到压缩字段
|
||||
- header_retention_days: 清空 request_headers 和 response_headers
|
||||
- log_retention_days: 删除整条记录
|
||||
|
||||
统计聚合任务:
|
||||
- 每天凌晨聚合前一天的统计数据
|
||||
- 更新全局统计汇总
|
||||
|
||||
使用 APScheduler 进行任务调度,支持时区配置。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from sqlalchemy import delete
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import Usage
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.system.scheduler import get_scheduler
|
||||
from src.services.system.stats_aggregator import StatsAggregatorService
|
||||
from src.services.user.apikey import ApiKeyService
|
||||
from src.utils.compression import compress_json
|
||||
|
||||
|
||||
class CleanupScheduler:
|
||||
"""使用记录清理调度器"""
|
||||
|
||||
def __init__(self):
|
||||
self.running = False
|
||||
self._interval_tasks = []
|
||||
|
||||
async def start(self):
|
||||
"""启动调度器"""
|
||||
if self.running:
|
||||
logger.warning("Cleanup scheduler already running")
|
||||
return
|
||||
|
||||
self.running = True
|
||||
logger.info("使用记录清理调度器已启动")
|
||||
|
||||
scheduler = get_scheduler()
|
||||
|
||||
# 注册定时任务(使用业务时区)
|
||||
# 统计聚合任务 - 凌晨 1 点执行
|
||||
scheduler.add_cron_job(
|
||||
self._scheduled_stats_aggregation,
|
||||
hour=1,
|
||||
minute=0,
|
||||
job_id="stats_aggregation",
|
||||
name="统计数据聚合",
|
||||
)
|
||||
|
||||
# 清理任务 - 凌晨 3 点执行
|
||||
scheduler.add_cron_job(
|
||||
self._scheduled_cleanup,
|
||||
hour=3,
|
||||
minute=0,
|
||||
job_id="usage_cleanup",
|
||||
name="使用记录清理",
|
||||
)
|
||||
|
||||
# 连接池监控 - 每 5 分钟
|
||||
scheduler.add_interval_job(
|
||||
self._scheduled_monitor,
|
||||
minutes=5,
|
||||
job_id="pool_monitor",
|
||||
name="连接池监控",
|
||||
)
|
||||
|
||||
# Pending 状态清理 - 每 5 分钟
|
||||
scheduler.add_interval_job(
|
||||
self._scheduled_pending_cleanup,
|
||||
minutes=5,
|
||||
job_id="pending_cleanup",
|
||||
name="Pending状态清理",
|
||||
)
|
||||
|
||||
# 启动时执行一次初始化任务
|
||||
asyncio.create_task(self._run_startup_tasks())
|
||||
|
||||
async def _run_startup_tasks(self):
|
||||
"""启动时执行的初始化任务"""
|
||||
# 延迟一点执行,确保系统完全启动
|
||||
await asyncio.sleep(2)
|
||||
|
||||
try:
|
||||
logger.info("启动时执行首次清理任务...")
|
||||
await self._perform_cleanup()
|
||||
except Exception as e:
|
||||
logger.exception(f"启动时清理任务执行出错: {e}")
|
||||
|
||||
try:
|
||||
logger.info("启动时检查统计数据...")
|
||||
await self._perform_stats_aggregation(backfill=True)
|
||||
except Exception as e:
|
||||
logger.exception(f"启动时统计聚合任务出错: {e}")
|
||||
|
||||
async def stop(self):
|
||||
"""停止调度器"""
|
||||
if not self.running:
|
||||
return
|
||||
|
||||
self.running = False
|
||||
scheduler = get_scheduler()
|
||||
scheduler.stop()
|
||||
|
||||
logger.info("使用记录清理调度器已停止")
|
||||
|
||||
# ========== 任务函数(APScheduler 直接调用异步函数) ==========
|
||||
|
||||
async def _scheduled_stats_aggregation(self):
|
||||
"""统计聚合任务(定时调用)"""
|
||||
await self._perform_stats_aggregation()
|
||||
|
||||
async def _scheduled_cleanup(self):
|
||||
"""清理任务(定时调用)"""
|
||||
await self._perform_cleanup()
|
||||
|
||||
async def _scheduled_monitor(self):
|
||||
"""监控任务(定时调用)"""
|
||||
try:
|
||||
from src.database import log_pool_status
|
||||
|
||||
log_pool_status()
|
||||
except Exception as e:
|
||||
logger.exception(f"连接池监控任务出错: {e}")
|
||||
|
||||
async def _scheduled_pending_cleanup(self):
|
||||
"""Pending 清理任务(定时调用)"""
|
||||
await self._perform_pending_cleanup()
|
||||
|
||||
# ========== 实际任务实现 ==========
|
||||
|
||||
async def _perform_stats_aggregation(self, backfill: bool = False):
|
||||
"""执行统计聚合任务
|
||||
|
||||
Args:
|
||||
backfill: 是否回填历史数据(首次启动时使用)
|
||||
"""
|
||||
db = create_session()
|
||||
try:
|
||||
# 检查是否启用统计聚合
|
||||
if not SystemConfigService.get_config(db, "enable_stats_aggregation", True):
|
||||
logger.info("统计聚合已禁用,跳过聚合任务")
|
||||
return
|
||||
|
||||
logger.info("开始执行统计数据聚合...")
|
||||
|
||||
if backfill:
|
||||
# 首次启动时回填历史数据
|
||||
from src.models.database import StatsSummary
|
||||
|
||||
summary = db.query(StatsSummary).first()
|
||||
if not summary:
|
||||
logger.info("检测到首次运行,开始回填历史统计数据...")
|
||||
days_to_backfill = SystemConfigService.get_config(
|
||||
db, "stats_backfill_days", 365
|
||||
)
|
||||
count = StatsAggregatorService.backfill_historical_data(
|
||||
db, days=days_to_backfill
|
||||
)
|
||||
logger.info(f"历史数据回填完成,共 {count} 天")
|
||||
return
|
||||
|
||||
# 聚合昨天的数据
|
||||
now = datetime.now(timezone.utc)
|
||||
yesterday = (now - timedelta(days=1)).replace(
|
||||
hour=0, minute=0, second=0, microsecond=0
|
||||
)
|
||||
|
||||
StatsAggregatorService.aggregate_daily_stats(db, yesterday)
|
||||
|
||||
# 聚合所有用户的昨日数据
|
||||
from src.models.database import User as DBUser
|
||||
|
||||
users = db.query(DBUser.id).filter(DBUser.is_active.is_(True)).all()
|
||||
for (user_id,) in users:
|
||||
try:
|
||||
StatsAggregatorService.aggregate_user_daily_stats(db, user_id, yesterday)
|
||||
except Exception as e:
|
||||
logger.warning(f"聚合用户 {user_id} 统计数据失败: {e}")
|
||||
# 回滚当前用户的失败操作,继续处理其他用户
|
||||
try:
|
||||
db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 更新全局汇总
|
||||
StatsAggregatorService.update_summary(db)
|
||||
|
||||
logger.info("统计数据聚合完成")
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"统计聚合任务执行失败: {e}")
|
||||
db.rollback()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def _perform_pending_cleanup(self):
|
||||
"""执行 pending 状态清理"""
|
||||
db = create_session()
|
||||
try:
|
||||
from src.services.usage.service import UsageService
|
||||
|
||||
# 获取配置的超时时间(默认 10 分钟)
|
||||
timeout_minutes = SystemConfigService.get_config(
|
||||
db, "pending_request_timeout_minutes", 10
|
||||
)
|
||||
|
||||
# 执行清理
|
||||
cleaned_count = UsageService.cleanup_stale_pending_requests(
|
||||
db, timeout_minutes=timeout_minutes
|
||||
)
|
||||
|
||||
if cleaned_count > 0:
|
||||
logger.info(f"清理了 {cleaned_count} 条超时的 pending/streaming 请求")
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"清理 pending 请求失败: {e}")
|
||||
db.rollback()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def _perform_cleanup(self):
|
||||
"""执行清理任务"""
|
||||
db = create_session()
|
||||
try:
|
||||
# 检查是否启用自动清理
|
||||
if not SystemConfigService.get_config(db, "enable_auto_cleanup", True):
|
||||
logger.info("自动清理已禁用,跳过清理任务")
|
||||
return
|
||||
|
||||
logger.info("开始执行使用记录分级清理...")
|
||||
|
||||
# 获取配置参数
|
||||
detail_retention = SystemConfigService.get_config(db, "detail_log_retention_days", 7)
|
||||
compressed_retention = SystemConfigService.get_config(
|
||||
db, "compressed_log_retention_days", 90
|
||||
)
|
||||
header_retention = SystemConfigService.get_config(db, "header_retention_days", 90)
|
||||
log_retention = SystemConfigService.get_config(db, "log_retention_days", 365)
|
||||
batch_size = SystemConfigService.get_config(db, "cleanup_batch_size", 1000)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
# 1. 压缩详细日志 (body 字段 -> 压缩字段)
|
||||
detail_cutoff = now - timedelta(days=detail_retention)
|
||||
body_compressed = await self._cleanup_body_fields(db, detail_cutoff, batch_size)
|
||||
|
||||
# 2. 清理压缩字段(90天后)
|
||||
compressed_cutoff = now - timedelta(days=compressed_retention)
|
||||
compressed_cleaned = await self._cleanup_compressed_fields(
|
||||
db, compressed_cutoff, batch_size
|
||||
)
|
||||
|
||||
# 3. 清理请求头
|
||||
header_cutoff = now - timedelta(days=header_retention)
|
||||
header_cleaned = await self._cleanup_header_fields(db, header_cutoff, batch_size)
|
||||
|
||||
# 4. 删除过期记录
|
||||
log_cutoff = now - timedelta(days=log_retention)
|
||||
records_deleted = await self._delete_old_records(db, log_cutoff, batch_size)
|
||||
|
||||
# 5. 清理过期的API Keys
|
||||
auto_delete = SystemConfigService.get_config(db, "auto_delete_expired_keys", False)
|
||||
keys_cleaned = ApiKeyService.cleanup_expired_keys(db, auto_delete=auto_delete)
|
||||
|
||||
logger.info(
|
||||
f"清理完成: 压缩 {body_compressed} 条, "
|
||||
f"清理压缩字段 {compressed_cleaned} 条, "
|
||||
f"清理header {header_cleaned} 条, "
|
||||
f"删除记录 {records_deleted} 条, "
|
||||
f"清理过期Keys {keys_cleaned} 条"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"清理任务执行失败: {e}")
|
||||
db.rollback()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def _cleanup_body_fields(
|
||||
self, db: Session, cutoff_time: datetime, batch_size: int
|
||||
) -> int:
|
||||
"""压缩 request_body 和 response_body 字段到压缩字段
|
||||
|
||||
逐条处理,确保每条记录都正确更新
|
||||
"""
|
||||
from sqlalchemy import null, update
|
||||
|
||||
total_compressed = 0
|
||||
no_progress_count = 0 # 连续无进展计数
|
||||
processed_ids: set = set() # 记录已处理的 ID,防止重复处理
|
||||
|
||||
while True:
|
||||
batch_db = create_session()
|
||||
try:
|
||||
# 1. 查询需要压缩的记录
|
||||
# 注意:排除已经是 NULL 或 JSON null 的记录
|
||||
records = (
|
||||
batch_db.query(Usage.id, Usage.request_body, Usage.response_body)
|
||||
.filter(Usage.created_at < cutoff_time)
|
||||
.filter((Usage.request_body.isnot(None)) | (Usage.response_body.isnot(None)))
|
||||
.limit(batch_size)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not records:
|
||||
break
|
||||
|
||||
# 过滤掉实际值为 None 的记录(JSON null 被解析为 Python None)
|
||||
valid_records = [
|
||||
(rid, req, resp)
|
||||
for rid, req, resp in records
|
||||
if req is not None or resp is not None
|
||||
]
|
||||
|
||||
if not valid_records:
|
||||
# 所有记录都是 JSON null,需要清理它们
|
||||
logger.warning(
|
||||
f"检测到 {len(records)} 条记录的 body 字段为 JSON null,进行清理"
|
||||
)
|
||||
for record_id, _, _ in records:
|
||||
batch_db.execute(
|
||||
update(Usage)
|
||||
.where(Usage.id == record_id)
|
||||
.values(request_body=null(), response_body=null())
|
||||
)
|
||||
batch_db.commit()
|
||||
continue
|
||||
|
||||
# 检测是否有重复的 ID(说明更新未生效)
|
||||
current_ids = {r[0] for r in valid_records}
|
||||
repeated_ids = current_ids & processed_ids
|
||||
if repeated_ids:
|
||||
logger.error(
|
||||
f"检测到重复处理的记录 ID: {list(repeated_ids)[:5]}...,"
|
||||
"说明数据库更新未生效,终止循环"
|
||||
)
|
||||
break
|
||||
|
||||
batch_success = 0
|
||||
|
||||
# 2. 逐条更新(确保每条都正确处理)
|
||||
for record_id, req_body, resp_body in valid_records:
|
||||
try:
|
||||
# 使用 null() 确保设置的是 SQL NULL 而不是 JSON null
|
||||
result = batch_db.execute(
|
||||
update(Usage)
|
||||
.where(Usage.id == record_id)
|
||||
.values(
|
||||
request_body=null(),
|
||||
response_body=null(),
|
||||
request_body_compressed=compress_json(req_body)
|
||||
if req_body
|
||||
else None,
|
||||
response_body_compressed=compress_json(resp_body)
|
||||
if resp_body
|
||||
else None,
|
||||
)
|
||||
)
|
||||
if result.rowcount > 0:
|
||||
batch_success += 1
|
||||
processed_ids.add(record_id)
|
||||
except Exception as e:
|
||||
logger.warning(f"压缩记录 {record_id} 失败: {e}")
|
||||
continue
|
||||
|
||||
batch_db.commit()
|
||||
|
||||
# 3. 检查是否有实际进展
|
||||
if batch_success == 0:
|
||||
no_progress_count += 1
|
||||
if no_progress_count >= 3:
|
||||
logger.error(
|
||||
f"压缩 body 字段连续 {no_progress_count} 批无进展,"
|
||||
"终止循环以避免死循环"
|
||||
)
|
||||
break
|
||||
else:
|
||||
no_progress_count = 0 # 重置计数
|
||||
|
||||
total_compressed += batch_success
|
||||
logger.debug(
|
||||
f"已压缩 {batch_success} 条记录的 body 字段,累计 {total_compressed} 条"
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"压缩 body 字段失败: {e}")
|
||||
try:
|
||||
batch_db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
finally:
|
||||
batch_db.close()
|
||||
|
||||
return total_compressed
|
||||
|
||||
async def _cleanup_compressed_fields(
|
||||
self, db: Session, cutoff_time: datetime, batch_size: int
|
||||
) -> int:
|
||||
"""清理压缩字段(90天后删除压缩的body)
|
||||
|
||||
每批使用短生命周期 session,避免 ORM 缓存问题
|
||||
"""
|
||||
from sqlalchemy import null, update
|
||||
|
||||
total_cleaned = 0
|
||||
|
||||
while True:
|
||||
batch_db = create_session()
|
||||
try:
|
||||
# 查询需要清理压缩字段的记录
|
||||
records_to_clean = (
|
||||
batch_db.query(Usage.id)
|
||||
.filter(Usage.created_at < cutoff_time)
|
||||
.filter(
|
||||
(Usage.request_body_compressed.isnot(None))
|
||||
| (Usage.response_body_compressed.isnot(None))
|
||||
)
|
||||
.limit(batch_size)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not records_to_clean:
|
||||
break
|
||||
|
||||
record_ids = [r.id for r in records_to_clean]
|
||||
|
||||
# 批量更新,使用 null() 确保设置 SQL NULL
|
||||
result = batch_db.execute(
|
||||
update(Usage)
|
||||
.where(Usage.id.in_(record_ids))
|
||||
.values(
|
||||
request_body_compressed=null(),
|
||||
response_body_compressed=null(),
|
||||
)
|
||||
)
|
||||
|
||||
rows_updated = result.rowcount
|
||||
batch_db.commit()
|
||||
|
||||
if rows_updated == 0:
|
||||
logger.warning("清理压缩字段: rowcount=0,可能存在问题")
|
||||
break
|
||||
|
||||
total_cleaned += rows_updated
|
||||
logger.debug(f"已清理 {rows_updated} 条记录的压缩字段,累计 {total_cleaned} 条")
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"清理压缩字段失败: {e}")
|
||||
try:
|
||||
batch_db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
finally:
|
||||
batch_db.close()
|
||||
|
||||
return total_cleaned
|
||||
|
||||
async def _cleanup_header_fields(
|
||||
self, db: Session, cutoff_time: datetime, batch_size: int
|
||||
) -> int:
|
||||
"""清理 request_headers, response_headers 和 provider_request_headers 字段
|
||||
|
||||
每批使用短生命周期 session,避免 ORM 缓存问题
|
||||
"""
|
||||
from sqlalchemy import null, update
|
||||
|
||||
total_cleaned = 0
|
||||
|
||||
while True:
|
||||
batch_db = create_session()
|
||||
try:
|
||||
# 先查询需要清理的记录ID(分批)
|
||||
records_to_clean = (
|
||||
batch_db.query(Usage.id)
|
||||
.filter(Usage.created_at < cutoff_time)
|
||||
.filter(
|
||||
(Usage.request_headers.isnot(None))
|
||||
| (Usage.response_headers.isnot(None))
|
||||
| (Usage.provider_request_headers.isnot(None))
|
||||
)
|
||||
.limit(batch_size)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not records_to_clean:
|
||||
break
|
||||
|
||||
record_ids = [r.id for r in records_to_clean]
|
||||
|
||||
# 批量更新,使用 null() 确保设置 SQL NULL
|
||||
result = batch_db.execute(
|
||||
update(Usage)
|
||||
.where(Usage.id.in_(record_ids))
|
||||
.values(
|
||||
request_headers=null(),
|
||||
response_headers=null(),
|
||||
provider_request_headers=null(),
|
||||
)
|
||||
)
|
||||
|
||||
rows_updated = result.rowcount
|
||||
batch_db.commit()
|
||||
|
||||
if rows_updated == 0:
|
||||
logger.warning("清理 header 字段: rowcount=0,可能存在问题")
|
||||
break
|
||||
|
||||
total_cleaned += rows_updated
|
||||
logger.debug(
|
||||
f"已清理 {rows_updated} 条记录的 header 字段,累计 {total_cleaned} 条"
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"清理 header 字段失败: {e}")
|
||||
try:
|
||||
batch_db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
finally:
|
||||
batch_db.close()
|
||||
|
||||
return total_cleaned
|
||||
|
||||
async def _delete_old_records(self, db: Session, cutoff_time: datetime, batch_size: int) -> int:
|
||||
"""删除过期的完整记录"""
|
||||
total_deleted = 0
|
||||
|
||||
while True:
|
||||
try:
|
||||
# 查询要删除的记录ID(分批)
|
||||
records_to_delete = (
|
||||
db.query(Usage.id)
|
||||
.filter(Usage.created_at < cutoff_time)
|
||||
.limit(batch_size)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not records_to_delete:
|
||||
break
|
||||
|
||||
record_ids = [r.id for r in records_to_delete]
|
||||
|
||||
# 执行删除
|
||||
result = db.execute(
|
||||
delete(Usage)
|
||||
.where(Usage.id.in_(record_ids))
|
||||
.execution_options(synchronize_session=False)
|
||||
)
|
||||
|
||||
rows_deleted = result.rowcount
|
||||
db.commit()
|
||||
|
||||
total_deleted += rows_deleted
|
||||
logger.debug(f"已删除 {rows_deleted} 条过期记录,累计 {total_deleted} 条")
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"删除过期记录失败: {e}")
|
||||
try:
|
||||
db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
|
||||
return total_deleted
|
||||
|
||||
|
||||
# 全局单例
|
||||
_cleanup_scheduler = None
|
||||
|
||||
|
||||
def get_cleanup_scheduler() -> CleanupScheduler:
|
||||
"""获取清理调度器单例"""
|
||||
global _cleanup_scheduler
|
||||
if _cleanup_scheduler is None:
|
||||
_cleanup_scheduler = CleanupScheduler()
|
||||
return _cleanup_scheduler
|
||||
257
src/services/system/config.py
Normal file
257
src/services/system/config.py
Normal file
@@ -0,0 +1,257 @@
|
||||
"""
|
||||
系统配置服务
|
||||
"""
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, SystemConfig
|
||||
|
||||
|
||||
|
||||
class LogLevel(str, Enum):
|
||||
"""日志记录级别"""
|
||||
|
||||
BASIC = "basic" # 仅记录基本信息(tokens、成本等)
|
||||
HEADERS = "headers" # 记录基本信息+请求/响应头(敏感信息会脱敏)
|
||||
FULL = "full" # 记录完整请求和响应(包含body,敏感信息会脱敏)
|
||||
|
||||
|
||||
class SystemConfigService:
|
||||
"""系统配置服务类"""
|
||||
|
||||
# 默认配置
|
||||
DEFAULT_CONFIGS = {
|
||||
"request_log_level": {
|
||||
"value": LogLevel.BASIC.value,
|
||||
"description": "请求记录级别:basic(基本信息), headers(含请求头), full(完整请求响应)",
|
||||
},
|
||||
"max_request_body_size": {
|
||||
"value": 1048576, # 1MB
|
||||
"description": "最大请求体记录大小(字节),超过此大小的请求体将被截断(仅影响数据库记录,不影响真实API请求)",
|
||||
},
|
||||
"max_response_body_size": {
|
||||
"value": 1048576, # 1MB
|
||||
"description": "最大响应体记录大小(字节),超过此大小的响应体将被截断(仅影响数据库记录,不影响真实API响应)",
|
||||
},
|
||||
"sensitive_headers": {
|
||||
"value": ["authorization", "x-api-key", "api-key", "cookie", "set-cookie"],
|
||||
"description": "敏感请求头列表,这些请求头会被脱敏处理",
|
||||
},
|
||||
# 分级清理策略
|
||||
"detail_log_retention_days": {
|
||||
"value": 7,
|
||||
"description": "详细日志保留天数,超过此天数后压缩 request_body 和 response_body 到压缩字段",
|
||||
},
|
||||
"compressed_log_retention_days": {
|
||||
"value": 90,
|
||||
"description": "压缩日志保留天数,超过此天数后删除压缩的 body 字段(保留headers和统计)",
|
||||
},
|
||||
"header_retention_days": {
|
||||
"value": 90,
|
||||
"description": "请求头保留天数,超过此天数后清空 request_headers 和 response_headers 字段",
|
||||
},
|
||||
"log_retention_days": {
|
||||
"value": 365,
|
||||
"description": "完整日志保留天数,超过此天数后删除整条记录(保留核心统计)",
|
||||
},
|
||||
"enable_auto_cleanup": {
|
||||
"value": True,
|
||||
"description": "是否启用自动清理任务,每天凌晨执行分级清理",
|
||||
},
|
||||
"cleanup_batch_size": {
|
||||
"value": 1000,
|
||||
"description": "每批次清理的记录数,避免单次操作过大影响数据库性能",
|
||||
},
|
||||
"provider_priority_mode": {
|
||||
"value": "provider",
|
||||
"description": "优先级策略:provider(提供商优先模式) 或 global_key(全局Key优先模式)",
|
||||
},
|
||||
"auto_delete_expired_keys": {
|
||||
"value": False,
|
||||
"description": "是否自动删除过期的API Key(True=物理删除,False=仅禁用),仅管理员可配置",
|
||||
},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_config(cls, db: Session, key: str, default: Any = None) -> Optional[Any]:
|
||||
"""获取系统配置值"""
|
||||
config = db.query(SystemConfig).filter(SystemConfig.key == key).first()
|
||||
if config:
|
||||
return config.value
|
||||
|
||||
# 如果配置不存在,检查默认值
|
||||
if key in cls.DEFAULT_CONFIGS:
|
||||
return cls.DEFAULT_CONFIGS[key]["value"]
|
||||
|
||||
return default
|
||||
|
||||
@staticmethod
|
||||
def set_config(db: Session, key: str, value: Any, description: str = None) -> SystemConfig:
|
||||
"""设置系统配置值"""
|
||||
config = db.query(SystemConfig).filter(SystemConfig.key == key).first()
|
||||
|
||||
if config:
|
||||
# 更新现有配置
|
||||
config.value = value
|
||||
if description:
|
||||
config.description = description
|
||||
else:
|
||||
# 创建新配置
|
||||
config = SystemConfig(key=key, value=value, description=description)
|
||||
db.add(config)
|
||||
|
||||
db.commit()
|
||||
db.refresh(config)
|
||||
return config
|
||||
|
||||
@staticmethod
|
||||
def get_default_provider(db: Session) -> Optional[str]:
|
||||
"""
|
||||
获取系统默认提供商
|
||||
优先级:1. 管理员设置的默认提供商 2. 数据库中第一个可用提供商
|
||||
"""
|
||||
# 首先尝试获取管理员设置的默认提供商
|
||||
default_provider = SystemConfigService.get_config(db, "default_provider")
|
||||
if default_provider:
|
||||
return default_provider
|
||||
|
||||
# 如果没有设置,fallback到数据库中第一个可用提供商
|
||||
first_provider = db.query(Provider).filter(Provider.is_active == True).first()
|
||||
|
||||
if first_provider:
|
||||
return first_provider.name
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def set_default_provider(db: Session, provider_name: str) -> SystemConfig:
|
||||
"""设置系统默认提供商"""
|
||||
return SystemConfigService.set_config(
|
||||
db, "default_provider", provider_name, "系统默认提供商,当用户未设置个人提供商时使用"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_all_configs(db: Session) -> list:
|
||||
"""获取所有系统配置"""
|
||||
configs = db.query(SystemConfig).all()
|
||||
return [
|
||||
{
|
||||
"key": config.key,
|
||||
"value": config.value,
|
||||
"description": config.description,
|
||||
"updated_at": config.updated_at.isoformat(),
|
||||
}
|
||||
for config in configs
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def delete_config(db: Session, key: str) -> bool:
|
||||
"""删除系统配置"""
|
||||
config = db.query(SystemConfig).filter(SystemConfig.key == key).first()
|
||||
if config:
|
||||
db.delete(config)
|
||||
db.commit()
|
||||
return True
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def init_default_configs(cls, db: Session):
|
||||
"""初始化默认配置"""
|
||||
for key, default_config in cls.DEFAULT_CONFIGS.items():
|
||||
if not db.query(SystemConfig).filter(SystemConfig.key == key).first():
|
||||
config = SystemConfig(
|
||||
key=key,
|
||||
value=default_config["value"],
|
||||
description=default_config["description"],
|
||||
)
|
||||
db.add(config)
|
||||
|
||||
db.commit()
|
||||
logger.info("初始化默认系统配置完成")
|
||||
|
||||
@classmethod
|
||||
def get_log_level(cls, db: Session) -> LogLevel:
|
||||
"""获取日志记录级别"""
|
||||
level = cls.get_config(db, "request_log_level", LogLevel.BASIC.value)
|
||||
if isinstance(level, str):
|
||||
return LogLevel(level)
|
||||
return level
|
||||
|
||||
@classmethod
|
||||
def should_log_headers(cls, db: Session) -> bool:
|
||||
"""是否应该记录请求头"""
|
||||
log_level = cls.get_log_level(db)
|
||||
return log_level in [LogLevel.HEADERS, LogLevel.FULL]
|
||||
|
||||
@classmethod
|
||||
def should_log_body(cls, db: Session) -> bool:
|
||||
"""是否应该记录请求体和响应体"""
|
||||
log_level = cls.get_log_level(db)
|
||||
return log_level == LogLevel.FULL
|
||||
|
||||
@classmethod
|
||||
def should_mask_sensitive_data(cls, db: Session) -> bool:
|
||||
"""是否应该脱敏敏感数据(始终脱敏)"""
|
||||
_ = db # 保持接口一致性
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def get_sensitive_headers(cls, db: Session) -> list:
|
||||
"""获取敏感请求头列表"""
|
||||
return cls.get_config(db, "sensitive_headers", [])
|
||||
|
||||
@classmethod
|
||||
def mask_sensitive_headers(cls, db: Session, headers: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""脱敏敏感请求头"""
|
||||
if not cls.should_mask_sensitive_data(db):
|
||||
return headers
|
||||
|
||||
sensitive_headers = cls.get_sensitive_headers(db)
|
||||
masked_headers = {}
|
||||
|
||||
for key, value in headers.items():
|
||||
if key.lower() in [h.lower() for h in sensitive_headers]:
|
||||
# 保留前后各4个字符,中间用星号替换
|
||||
if len(str(value)) > 8:
|
||||
masked_value = str(value)[:4] + "****" + str(value)[-4:]
|
||||
else:
|
||||
masked_value = "****"
|
||||
masked_headers[key] = masked_value
|
||||
else:
|
||||
masked_headers[key] = value
|
||||
|
||||
return masked_headers
|
||||
|
||||
@classmethod
|
||||
def truncate_body(cls, db: Session, body: Any, is_request: bool = True) -> Any:
|
||||
"""截断过大的请求体或响应体"""
|
||||
max_size_key = "max_request_body_size" if is_request else "max_response_body_size"
|
||||
max_size = cls.get_config(db, max_size_key, 102400)
|
||||
|
||||
if not body:
|
||||
return body
|
||||
|
||||
# 转换为字符串以计算大小
|
||||
body_str = json.dumps(body) if isinstance(body, (dict, list)) else str(body)
|
||||
|
||||
if len(body_str) > max_size:
|
||||
# 截断并添加提示
|
||||
truncated_str = body_str[:max_size]
|
||||
if isinstance(body, (dict, list)):
|
||||
try:
|
||||
# 尝试保持JSON格式
|
||||
return {
|
||||
"_truncated": True,
|
||||
"_original_size": len(body_str),
|
||||
"_content": truncated_str,
|
||||
}
|
||||
except:
|
||||
pass
|
||||
return truncated_str + f"\n... (truncated, original size: {len(body_str)} bytes)"
|
||||
|
||||
return body
|
||||
187
src/services/system/scheduler.py
Normal file
187
src/services/system/scheduler.py
Normal file
@@ -0,0 +1,187 @@
|
||||
"""
|
||||
统一定时任务调度器
|
||||
|
||||
使用 APScheduler 管理所有定时任务,支持时区配置。
|
||||
所有定时任务使用应用时区(APP_TIMEZONE)配置执行时间,
|
||||
数据存储仍然使用 UTC。
|
||||
"""
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from apscheduler.triggers.cron import CronTrigger
|
||||
from apscheduler.triggers.interval import IntervalTrigger
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
# 应用时区配置,默认为 Asia/Shanghai
|
||||
APP_TIMEZONE = os.getenv("APP_TIMEZONE", "Asia/Shanghai")
|
||||
|
||||
|
||||
class TaskScheduler:
|
||||
"""统一定时任务调度器"""
|
||||
|
||||
_instance: Optional["TaskScheduler"] = None
|
||||
|
||||
def __init__(self):
|
||||
self.scheduler = AsyncIOScheduler(timezone=APP_TIMEZONE)
|
||||
self._started = False
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> "TaskScheduler":
|
||||
"""获取调度器单例"""
|
||||
if cls._instance is None:
|
||||
cls._instance = TaskScheduler()
|
||||
return cls._instance
|
||||
|
||||
def add_cron_job(
|
||||
self,
|
||||
func,
|
||||
hour: int,
|
||||
minute: int = 0,
|
||||
job_id: str = None,
|
||||
name: str = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
添加 cron 定时任务
|
||||
|
||||
Args:
|
||||
func: 要执行的函数
|
||||
hour: 执行时间(小时),使用业务时区
|
||||
minute: 执行时间(分钟)
|
||||
job_id: 任务ID
|
||||
name: 任务名称(用于日志)
|
||||
**kwargs: 传递给任务函数的参数
|
||||
"""
|
||||
trigger = CronTrigger(hour=hour, minute=minute, timezone=APP_TIMEZONE)
|
||||
|
||||
job_id = job_id or func.__name__
|
||||
display_name = name or job_id
|
||||
|
||||
self.scheduler.add_job(
|
||||
func,
|
||||
trigger,
|
||||
id=job_id,
|
||||
name=display_name,
|
||||
replace_existing=True,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"已注册定时任务: {display_name}, "
|
||||
f"执行时间: {hour:02d}:{minute:02d} ({APP_TIMEZONE})"
|
||||
)
|
||||
|
||||
def add_interval_job(
|
||||
self,
|
||||
func,
|
||||
seconds: int = None,
|
||||
minutes: int = None,
|
||||
hours: int = None,
|
||||
job_id: str = None,
|
||||
name: str = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
添加间隔执行任务
|
||||
|
||||
Args:
|
||||
func: 要执行的函数
|
||||
seconds: 间隔秒数
|
||||
minutes: 间隔分钟数
|
||||
hours: 间隔小时数
|
||||
job_id: 任务ID
|
||||
name: 任务名称
|
||||
**kwargs: 传递给任务函数的参数
|
||||
"""
|
||||
# 构建 trigger 参数,过滤掉 None 值
|
||||
trigger_kwargs = {}
|
||||
if seconds is not None:
|
||||
trigger_kwargs["seconds"] = seconds
|
||||
if minutes is not None:
|
||||
trigger_kwargs["minutes"] = minutes
|
||||
if hours is not None:
|
||||
trigger_kwargs["hours"] = hours
|
||||
|
||||
trigger = IntervalTrigger(**trigger_kwargs)
|
||||
|
||||
job_id = job_id or func.__name__
|
||||
display_name = name or job_id
|
||||
|
||||
# 计算间隔描述
|
||||
interval_parts = []
|
||||
if hours:
|
||||
interval_parts.append(f"{hours}小时")
|
||||
if minutes:
|
||||
interval_parts.append(f"{minutes}分钟")
|
||||
if seconds:
|
||||
interval_parts.append(f"{seconds}秒")
|
||||
interval_desc = "".join(interval_parts) or "未知间隔"
|
||||
|
||||
self.scheduler.add_job(
|
||||
func,
|
||||
trigger,
|
||||
id=job_id,
|
||||
name=display_name,
|
||||
replace_existing=True,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
logger.info(f"已注册间隔任务: {display_name}, 执行间隔: {interval_desc}")
|
||||
|
||||
def start(self):
|
||||
"""启动调度器"""
|
||||
if self._started:
|
||||
logger.warning("调度器已在运行中")
|
||||
return
|
||||
|
||||
self.scheduler.start()
|
||||
self._started = True
|
||||
logger.info(f"定时任务调度器已启动,应用时区: {APP_TIMEZONE}")
|
||||
|
||||
# 打印下次执行时间
|
||||
self._log_next_run_times()
|
||||
|
||||
def stop(self):
|
||||
"""停止调度器"""
|
||||
if not self._started:
|
||||
return
|
||||
|
||||
self.scheduler.shutdown(wait=False)
|
||||
self._started = False
|
||||
logger.info("定时任务调度器已停止")
|
||||
|
||||
def _log_next_run_times(self):
|
||||
"""记录所有任务的下次执行时间"""
|
||||
jobs = self.scheduler.get_jobs()
|
||||
if not jobs:
|
||||
return
|
||||
|
||||
logger.info("已注册的定时任务:")
|
||||
for job in jobs:
|
||||
next_run = job.next_run_time
|
||||
if next_run:
|
||||
# 计算距离下次执行的时间
|
||||
now = datetime.now(next_run.tzinfo)
|
||||
delta = next_run - now
|
||||
hours, remainder = divmod(int(delta.total_seconds()), 3600)
|
||||
minutes = remainder // 60
|
||||
|
||||
logger.info(
|
||||
f" - {job.name}: 下次执行 {next_run.strftime('%Y-%m-%d %H:%M')} "
|
||||
f"({hours}小时{minutes}分钟后)"
|
||||
)
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""调度器是否在运行"""
|
||||
return self._started
|
||||
|
||||
|
||||
# 便捷函数
|
||||
def get_scheduler() -> TaskScheduler:
|
||||
"""获取调度器单例"""
|
||||
return TaskScheduler.get_instance()
|
||||
436
src/services/system/stats_aggregator.py
Normal file
436
src/services/system/stats_aggregator.py
Normal file
@@ -0,0 +1,436 @@
|
||||
"""统计数据聚合服务
|
||||
|
||||
实现预聚合统计,避免每次请求都全表扫描。
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import (
|
||||
ApiKey,
|
||||
RequestCandidate,
|
||||
StatsDaily,
|
||||
StatsSummary,
|
||||
StatsUserDaily,
|
||||
Usage,
|
||||
)
|
||||
from src.models.database import User as DBUser
|
||||
|
||||
|
||||
class StatsAggregatorService:
|
||||
"""统计数据聚合服务"""
|
||||
|
||||
@staticmethod
|
||||
def aggregate_daily_stats(db: Session, date: datetime) -> StatsDaily:
|
||||
"""聚合指定日期的统计数据
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
date: 要聚合的日期 (会自动转为 UTC 当天开始)
|
||||
|
||||
Returns:
|
||||
StatsDaily 记录
|
||||
"""
|
||||
# 确保日期是 UTC 当天开始
|
||||
day_start = date.replace(hour=0, minute=0, second=0, microsecond=0, tzinfo=timezone.utc)
|
||||
day_end = day_start + timedelta(days=1)
|
||||
|
||||
# 检查是否已存在该日期的记录
|
||||
existing = db.query(StatsDaily).filter(StatsDaily.date == day_start).first()
|
||||
if existing:
|
||||
stats = existing
|
||||
else:
|
||||
stats = StatsDaily(id=str(uuid.uuid4()), date=day_start)
|
||||
|
||||
# 基础请求统计
|
||||
base_query = db.query(Usage).filter(
|
||||
and_(Usage.created_at >= day_start, Usage.created_at < day_end)
|
||||
)
|
||||
|
||||
total_requests = base_query.count()
|
||||
|
||||
# 如果没有请求,直接返回空记录
|
||||
if total_requests == 0:
|
||||
stats.total_requests = 0
|
||||
stats.success_requests = 0
|
||||
stats.error_requests = 0
|
||||
stats.input_tokens = 0
|
||||
stats.output_tokens = 0
|
||||
stats.cache_creation_tokens = 0
|
||||
stats.cache_read_tokens = 0
|
||||
stats.total_cost = 0.0
|
||||
stats.actual_total_cost = 0.0
|
||||
stats.input_cost = 0.0
|
||||
stats.output_cost = 0.0
|
||||
stats.cache_creation_cost = 0.0
|
||||
stats.cache_read_cost = 0.0
|
||||
stats.avg_response_time_ms = 0.0
|
||||
stats.fallback_count = 0
|
||||
|
||||
if not existing:
|
||||
db.add(stats)
|
||||
db.commit()
|
||||
return stats
|
||||
|
||||
# 错误请求数
|
||||
error_requests = (
|
||||
base_query.filter(
|
||||
(Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
).count()
|
||||
)
|
||||
|
||||
# Token 和成本聚合
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost"),
|
||||
func.sum(Usage.input_cost_usd).label("input_cost"),
|
||||
func.sum(Usage.output_cost_usd).label("output_cost"),
|
||||
func.sum(Usage.cache_creation_cost_usd).label("cache_creation_cost"),
|
||||
func.sum(Usage.cache_read_cost_usd).label("cache_read_cost"),
|
||||
func.avg(Usage.response_time_ms).label("avg_response_time"),
|
||||
)
|
||||
.filter(and_(Usage.created_at >= day_start, Usage.created_at < day_end))
|
||||
.first()
|
||||
)
|
||||
|
||||
# Fallback 统计 (执行候选数 > 1 的请求数)
|
||||
fallback_subquery = (
|
||||
db.query(
|
||||
RequestCandidate.request_id,
|
||||
func.count(RequestCandidate.id).label("executed_count"),
|
||||
)
|
||||
.filter(
|
||||
and_(
|
||||
RequestCandidate.created_at >= day_start,
|
||||
RequestCandidate.created_at < day_end,
|
||||
RequestCandidate.status.in_(["success", "failed"]),
|
||||
)
|
||||
)
|
||||
.group_by(RequestCandidate.request_id)
|
||||
.subquery()
|
||||
)
|
||||
fallback_count = (
|
||||
db.query(func.count())
|
||||
.select_from(fallback_subquery)
|
||||
.filter(fallback_subquery.c.executed_count > 1)
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
|
||||
# 使用维度统计
|
||||
unique_models = (
|
||||
db.query(func.count(func.distinct(Usage.model)))
|
||||
.filter(and_(Usage.created_at >= day_start, Usage.created_at < day_end))
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
unique_providers = (
|
||||
db.query(func.count(func.distinct(Usage.provider)))
|
||||
.filter(and_(Usage.created_at >= day_start, Usage.created_at < day_end))
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
|
||||
# 更新统计记录
|
||||
stats.total_requests = total_requests
|
||||
stats.success_requests = total_requests - error_requests
|
||||
stats.error_requests = error_requests
|
||||
stats.input_tokens = int(aggregated.input_tokens or 0)
|
||||
stats.output_tokens = int(aggregated.output_tokens or 0)
|
||||
stats.cache_creation_tokens = int(aggregated.cache_creation_tokens or 0)
|
||||
stats.cache_read_tokens = int(aggregated.cache_read_tokens or 0)
|
||||
stats.total_cost = float(aggregated.total_cost or 0)
|
||||
stats.actual_total_cost = float(aggregated.actual_total_cost or 0)
|
||||
stats.input_cost = float(aggregated.input_cost or 0)
|
||||
stats.output_cost = float(aggregated.output_cost or 0)
|
||||
stats.cache_creation_cost = float(aggregated.cache_creation_cost or 0)
|
||||
stats.cache_read_cost = float(aggregated.cache_read_cost or 0)
|
||||
stats.avg_response_time_ms = float(aggregated.avg_response_time or 0)
|
||||
stats.fallback_count = fallback_count
|
||||
stats.unique_models = unique_models
|
||||
stats.unique_providers = unique_providers
|
||||
|
||||
if not existing:
|
||||
db.add(stats)
|
||||
db.commit()
|
||||
|
||||
logger.info(f"[StatsAggregator] 聚合日期 {day_start.date()} 完成: {total_requests} 请求")
|
||||
return stats
|
||||
|
||||
@staticmethod
|
||||
def aggregate_user_daily_stats(
|
||||
db: Session, user_id: str, date: datetime
|
||||
) -> StatsUserDaily:
|
||||
"""聚合指定用户指定日期的统计数据"""
|
||||
day_start = date.replace(hour=0, minute=0, second=0, microsecond=0, tzinfo=timezone.utc)
|
||||
day_end = day_start + timedelta(days=1)
|
||||
|
||||
existing = (
|
||||
db.query(StatsUserDaily)
|
||||
.filter(and_(StatsUserDaily.user_id == user_id, StatsUserDaily.date == day_start))
|
||||
.first()
|
||||
)
|
||||
|
||||
if existing:
|
||||
stats = existing
|
||||
else:
|
||||
stats = StatsUserDaily(id=str(uuid.uuid4()), user_id=user_id, date=day_start)
|
||||
|
||||
# 用户请求统计
|
||||
base_query = db.query(Usage).filter(
|
||||
and_(
|
||||
Usage.user_id == user_id,
|
||||
Usage.created_at >= day_start,
|
||||
Usage.created_at < day_end,
|
||||
)
|
||||
)
|
||||
|
||||
total_requests = base_query.count()
|
||||
|
||||
if total_requests == 0:
|
||||
stats.total_requests = 0
|
||||
stats.success_requests = 0
|
||||
stats.error_requests = 0
|
||||
stats.input_tokens = 0
|
||||
stats.output_tokens = 0
|
||||
stats.cache_creation_tokens = 0
|
||||
stats.cache_read_tokens = 0
|
||||
stats.total_cost = 0.0
|
||||
|
||||
if not existing:
|
||||
db.add(stats)
|
||||
db.commit()
|
||||
return stats
|
||||
|
||||
error_requests = (
|
||||
base_query.filter(
|
||||
(Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
).count()
|
||||
)
|
||||
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost"),
|
||||
)
|
||||
.filter(
|
||||
and_(
|
||||
Usage.user_id == user_id,
|
||||
Usage.created_at >= day_start,
|
||||
Usage.created_at < day_end,
|
||||
)
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
stats.total_requests = total_requests
|
||||
stats.success_requests = total_requests - error_requests
|
||||
stats.error_requests = error_requests
|
||||
stats.input_tokens = int(aggregated.input_tokens or 0)
|
||||
stats.output_tokens = int(aggregated.output_tokens or 0)
|
||||
stats.cache_creation_tokens = int(aggregated.cache_creation_tokens or 0)
|
||||
stats.cache_read_tokens = int(aggregated.cache_read_tokens or 0)
|
||||
stats.total_cost = float(aggregated.total_cost or 0)
|
||||
|
||||
if not existing:
|
||||
db.add(stats)
|
||||
db.commit()
|
||||
return stats
|
||||
|
||||
@staticmethod
|
||||
def update_summary(db: Session) -> StatsSummary:
|
||||
"""更新全局统计汇总
|
||||
|
||||
汇总截止到昨天的所有数据。
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
today = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
cutoff_date = today # 不含今天
|
||||
|
||||
# 获取或创建 summary 记录
|
||||
summary = db.query(StatsSummary).first()
|
||||
if not summary:
|
||||
summary = StatsSummary(id=str(uuid.uuid4()), cutoff_date=cutoff_date)
|
||||
|
||||
# 从 stats_daily 聚合历史数据
|
||||
daily_aggregated = (
|
||||
db.query(
|
||||
func.sum(StatsDaily.total_requests).label("total_requests"),
|
||||
func.sum(StatsDaily.success_requests).label("success_requests"),
|
||||
func.sum(StatsDaily.error_requests).label("error_requests"),
|
||||
func.sum(StatsDaily.input_tokens).label("input_tokens"),
|
||||
func.sum(StatsDaily.output_tokens).label("output_tokens"),
|
||||
func.sum(StatsDaily.cache_creation_tokens).label("cache_creation_tokens"),
|
||||
func.sum(StatsDaily.cache_read_tokens).label("cache_read_tokens"),
|
||||
func.sum(StatsDaily.total_cost).label("total_cost"),
|
||||
func.sum(StatsDaily.actual_total_cost).label("actual_total_cost"),
|
||||
)
|
||||
.filter(StatsDaily.date < cutoff_date)
|
||||
.first()
|
||||
)
|
||||
|
||||
# 用户/API Key 统计
|
||||
total_users = db.query(func.count(DBUser.id)).scalar() or 0
|
||||
active_users = (
|
||||
db.query(func.count(DBUser.id)).filter(DBUser.is_active.is_(True)).scalar() or 0
|
||||
)
|
||||
total_api_keys = db.query(func.count(ApiKey.id)).scalar() or 0
|
||||
active_api_keys = (
|
||||
db.query(func.count(ApiKey.id)).filter(ApiKey.is_active.is_(True)).scalar() or 0
|
||||
)
|
||||
|
||||
# 更新 summary
|
||||
summary.cutoff_date = cutoff_date
|
||||
summary.all_time_requests = int(daily_aggregated.total_requests or 0)
|
||||
summary.all_time_success_requests = int(daily_aggregated.success_requests or 0)
|
||||
summary.all_time_error_requests = int(daily_aggregated.error_requests or 0)
|
||||
summary.all_time_input_tokens = int(daily_aggregated.input_tokens or 0)
|
||||
summary.all_time_output_tokens = int(daily_aggregated.output_tokens or 0)
|
||||
summary.all_time_cache_creation_tokens = int(daily_aggregated.cache_creation_tokens or 0)
|
||||
summary.all_time_cache_read_tokens = int(daily_aggregated.cache_read_tokens or 0)
|
||||
summary.all_time_cost = float(daily_aggregated.total_cost or 0)
|
||||
summary.all_time_actual_cost = float(daily_aggregated.actual_total_cost or 0)
|
||||
summary.total_users = total_users
|
||||
summary.active_users = active_users
|
||||
summary.total_api_keys = total_api_keys
|
||||
summary.active_api_keys = active_api_keys
|
||||
|
||||
db.add(summary)
|
||||
db.commit()
|
||||
|
||||
logger.info(f"[StatsAggregator] 更新全局汇总完成,截止日期: {cutoff_date.date()}")
|
||||
return summary
|
||||
|
||||
@staticmethod
|
||||
def get_today_realtime_stats(db: Session) -> dict:
|
||||
"""获取今日实时统计(用于与预聚合数据合并)"""
|
||||
now = datetime.now(timezone.utc)
|
||||
today = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
|
||||
base_query = db.query(Usage).filter(Usage.created_at >= today)
|
||||
|
||||
total_requests = base_query.count()
|
||||
|
||||
if total_requests == 0:
|
||||
return {
|
||||
"total_requests": 0,
|
||||
"success_requests": 0,
|
||||
"error_requests": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"cache_creation_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"total_cost": 0.0,
|
||||
"actual_total_cost": 0.0,
|
||||
}
|
||||
|
||||
error_requests = (
|
||||
base_query.filter(
|
||||
(Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
).count()
|
||||
)
|
||||
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost"),
|
||||
)
|
||||
.filter(Usage.created_at >= today)
|
||||
.first()
|
||||
)
|
||||
|
||||
return {
|
||||
"total_requests": total_requests,
|
||||
"success_requests": total_requests - error_requests,
|
||||
"error_requests": error_requests,
|
||||
"input_tokens": int(aggregated.input_tokens or 0),
|
||||
"output_tokens": int(aggregated.output_tokens or 0),
|
||||
"cache_creation_tokens": int(aggregated.cache_creation_tokens or 0),
|
||||
"cache_read_tokens": int(aggregated.cache_read_tokens or 0),
|
||||
"total_cost": float(aggregated.total_cost or 0),
|
||||
"actual_total_cost": float(aggregated.actual_total_cost or 0),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_combined_stats(db: Session) -> dict:
|
||||
"""获取合并后的统计数据(预聚合 + 今日实时)"""
|
||||
summary = db.query(StatsSummary).first()
|
||||
today_stats = StatsAggregatorService.get_today_realtime_stats(db)
|
||||
|
||||
if not summary:
|
||||
# 如果没有预聚合数据,返回今日数据
|
||||
return today_stats
|
||||
|
||||
return {
|
||||
"total_requests": summary.all_time_requests + today_stats["total_requests"],
|
||||
"success_requests": summary.all_time_success_requests
|
||||
+ today_stats["success_requests"],
|
||||
"error_requests": summary.all_time_error_requests + today_stats["error_requests"],
|
||||
"input_tokens": summary.all_time_input_tokens + today_stats["input_tokens"],
|
||||
"output_tokens": summary.all_time_output_tokens + today_stats["output_tokens"],
|
||||
"cache_creation_tokens": summary.all_time_cache_creation_tokens
|
||||
+ today_stats["cache_creation_tokens"],
|
||||
"cache_read_tokens": summary.all_time_cache_read_tokens
|
||||
+ today_stats["cache_read_tokens"],
|
||||
"total_cost": summary.all_time_cost + today_stats["total_cost"],
|
||||
"actual_total_cost": summary.all_time_actual_cost + today_stats["actual_total_cost"],
|
||||
"total_users": summary.total_users,
|
||||
"active_users": summary.active_users,
|
||||
"total_api_keys": summary.total_api_keys,
|
||||
"active_api_keys": summary.active_api_keys,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def backfill_historical_data(db: Session, days: int = 365) -> int:
|
||||
"""回填历史数据(首次部署时使用)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
days: 要回填的天数
|
||||
|
||||
Returns:
|
||||
回填的天数
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
today = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
|
||||
# 找到最早的 Usage 记录
|
||||
earliest = db.query(func.min(Usage.created_at)).scalar()
|
||||
if not earliest:
|
||||
logger.info("[StatsAggregator] 没有历史数据需要回填")
|
||||
return 0
|
||||
|
||||
# 计算需要回填的日期范围
|
||||
earliest_date = earliest.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
start_date = max(earliest_date, today - timedelta(days=days))
|
||||
|
||||
count = 0
|
||||
current_date = start_date
|
||||
while current_date < today:
|
||||
StatsAggregatorService.aggregate_daily_stats(db, current_date)
|
||||
count += 1
|
||||
current_date += timedelta(days=1)
|
||||
|
||||
# 更新汇总
|
||||
if count > 0:
|
||||
StatsAggregatorService.update_summary(db)
|
||||
|
||||
logger.info(f"[StatsAggregator] 回填历史数据完成,共 {count} 天")
|
||||
return count
|
||||
142
src/services/system/sync_stats.py
Normal file
142
src/services/system/sync_stats.py
Normal file
@@ -0,0 +1,142 @@
|
||||
"""
|
||||
API密钥统计同步服务
|
||||
定期同步API密钥的统计数据,确保与实际使用记录一致
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, Usage
|
||||
|
||||
|
||||
|
||||
class SyncStatsService:
|
||||
"""API密钥统计同步服务"""
|
||||
|
||||
# 分页批量大小
|
||||
BATCH_SIZE = 100
|
||||
|
||||
@staticmethod
|
||||
def sync_api_key_stats(db: Session, api_key_id: Optional[str] = None) -> dict: # UUID
|
||||
"""
|
||||
同步API密钥的统计数据
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
api_key_id: 指定要同步的API密钥ID,如果不指定则同步所有
|
||||
|
||||
Returns:
|
||||
同步结果统计
|
||||
"""
|
||||
result = {"synced": 0, "updated": 0, "errors": 0}
|
||||
|
||||
try:
|
||||
# 获取要同步的API密钥(使用分页避免大数据量问题)
|
||||
if api_key_id:
|
||||
api_keys = db.query(ApiKey).filter(ApiKey.id == api_key_id).all()
|
||||
else:
|
||||
# 分页处理,避免一次加载所有数据
|
||||
offset = 0
|
||||
api_keys = []
|
||||
while True:
|
||||
batch = db.query(ApiKey).offset(offset).limit(SyncStatsService.BATCH_SIZE).all()
|
||||
if not batch:
|
||||
break
|
||||
api_keys.extend(batch)
|
||||
offset += SyncStatsService.BATCH_SIZE
|
||||
|
||||
for api_key in api_keys:
|
||||
try:
|
||||
# 计算实际的使用统计
|
||||
stats = (
|
||||
db.query(
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_cost_usd).label("cost"),
|
||||
)
|
||||
.filter(Usage.api_key_id == api_key.id)
|
||||
.first()
|
||||
)
|
||||
|
||||
actual_requests = stats.requests or 0
|
||||
actual_cost = float(stats.cost or 0)
|
||||
|
||||
# 获取最后使用时间
|
||||
last_usage = (
|
||||
db.query(Usage.created_at)
|
||||
.filter(Usage.api_key_id == api_key.id)
|
||||
.order_by(Usage.created_at.desc())
|
||||
.first()
|
||||
)
|
||||
|
||||
# 检查是否需要更新
|
||||
needs_update = False
|
||||
if api_key.total_requests != actual_requests:
|
||||
logger.info(f"API密钥 {api_key.id} 请求数不一致: {api_key.total_requests} -> {actual_requests}")
|
||||
api_key.total_requests = actual_requests
|
||||
needs_update = True
|
||||
|
||||
if abs(api_key.total_cost_usd - actual_cost) > 0.0001:
|
||||
logger.info(f"API密钥 {api_key.id} 费用不一致: {api_key.total_cost_usd} -> {actual_cost}")
|
||||
api_key.total_cost_usd = actual_cost
|
||||
needs_update = True
|
||||
|
||||
if last_usage and api_key.last_used_at != last_usage[0]:
|
||||
api_key.last_used_at = last_usage[0]
|
||||
needs_update = True
|
||||
|
||||
result["synced"] += 1
|
||||
if needs_update:
|
||||
result["updated"] += 1
|
||||
logger.info(f"已更新API密钥 {api_key.id} 的统计数据")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"同步API密钥 {api_key.id} 统计时出错: {e}")
|
||||
result["errors"] += 1
|
||||
# 回滚当前失败的操作,继续处理其他密钥
|
||||
try:
|
||||
db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 提交所有更改
|
||||
db.commit()
|
||||
logger.info(f"同步完成: 处理 {result['synced']} 个密钥, 更新 {result['updated']} 个, 错误 {result['errors']} 个")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"同步统计数据时出错: {e}")
|
||||
db.rollback()
|
||||
raise
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def get_api_key_real_stats(db: Session, api_key_id: str) -> dict: # UUID
|
||||
"""
|
||||
获取API密钥的实际统计数据(直接从使用记录计算)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
api_key_id: API密钥ID
|
||||
|
||||
Returns:
|
||||
实际的统计数据
|
||||
"""
|
||||
# 计算实际的使用统计
|
||||
stats = (
|
||||
db.query(
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_cost_usd).label("cost"),
|
||||
func.max(Usage.created_at).label("last_used"),
|
||||
)
|
||||
.filter(Usage.api_key_id == api_key_id)
|
||||
.first()
|
||||
)
|
||||
|
||||
return {
|
||||
"total_requests": stats.requests or 0,
|
||||
"total_cost_usd": float(stats.cost or 0),
|
||||
"last_used_at": stats.last_used,
|
||||
}
|
||||
15
src/services/usage/__init__.py
Normal file
15
src/services/usage/__init__.py
Normal file
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
使用量服务模块
|
||||
|
||||
包含使用量追踪、流式使用量、配额调度等功能。
|
||||
"""
|
||||
|
||||
from src.services.usage.quota_scheduler import QuotaScheduler
|
||||
from src.services.usage.service import UsageService
|
||||
from src.services.usage.stream import StreamUsageTracker
|
||||
|
||||
__all__ = [
|
||||
"UsageService",
|
||||
"StreamUsageTracker",
|
||||
"QuotaScheduler",
|
||||
]
|
||||
161
src/services/usage/quota_scheduler.py
Normal file
161
src/services/usage/quota_scheduler.py
Normal file
@@ -0,0 +1,161 @@
|
||||
"""
|
||||
额度周期重置定时任务
|
||||
|
||||
支持按天数周期重置额度:
|
||||
- quota_reset_day: 重置周期(天数),例如7=每周,30=每月
|
||||
- quota_last_reset_at: 上次重置时间,用于计算下次重置
|
||||
|
||||
使用统一的 TaskScheduler 进行调度。
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from src.core.enums import ProviderBillingType
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import Provider
|
||||
from src.services.system.scheduler import get_scheduler
|
||||
|
||||
|
||||
class QuotaScheduler:
|
||||
"""额度周期重置调度器"""
|
||||
|
||||
def __init__(self):
|
||||
self.running = False
|
||||
|
||||
async def start(self):
|
||||
"""启动调度器"""
|
||||
if self.running:
|
||||
logger.warning("Quota scheduler already running")
|
||||
return
|
||||
|
||||
self.running = True
|
||||
logger.info("Quota scheduler started")
|
||||
|
||||
scheduler = get_scheduler()
|
||||
|
||||
# 每小时检查一次额度重置
|
||||
scheduler.add_interval_job(
|
||||
self._scheduled_quota_check,
|
||||
hours=1,
|
||||
job_id="quota_reset_check",
|
||||
name="额度周期重置检查",
|
||||
)
|
||||
|
||||
# 启动时立即执行一次检查
|
||||
await self._check_and_reset_quotas()
|
||||
|
||||
async def stop(self):
|
||||
"""停止调度器"""
|
||||
if not self.running:
|
||||
return
|
||||
|
||||
self.running = False
|
||||
logger.info("Quota scheduler stopped")
|
||||
|
||||
async def _scheduled_quota_check(self):
|
||||
"""额度检查任务(定时调用)"""
|
||||
await self._check_and_reset_quotas()
|
||||
|
||||
async def _check_and_reset_quotas(self):
|
||||
"""检查并重置周期额度"""
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
# 获取所有定额类型的提供商
|
||||
providers = (
|
||||
db.query(Provider)
|
||||
.filter(
|
||||
Provider.billing_type == ProviderBillingType.MONTHLY_QUOTA,
|
||||
Provider.is_active == True,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not providers:
|
||||
logger.debug("No quota providers to check")
|
||||
return
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
reset_count = 0
|
||||
|
||||
for provider in providers:
|
||||
try:
|
||||
# 如果没有上次重置时间,初始化为当前时间
|
||||
if provider.quota_last_reset_at is None:
|
||||
provider.quota_last_reset_at = now
|
||||
db.commit()
|
||||
logger.info(f"Initialized quota_last_reset_at for provider {provider.name}")
|
||||
continue
|
||||
|
||||
# 计算距离上次重置的天数
|
||||
days_since_reset = (now - provider.quota_last_reset_at).days
|
||||
|
||||
# 如果达到或超过重置周期,执行重置
|
||||
if days_since_reset >= provider.quota_reset_day:
|
||||
logger.info(f"Resetting quota for provider {provider.name}")
|
||||
|
||||
provider.monthly_used_usd = 0.0
|
||||
provider.rpm_used = 0 # 同时重置RPM计数
|
||||
provider.rpm_reset_at = None
|
||||
provider.quota_last_reset_at = now
|
||||
reset_count += 1
|
||||
|
||||
# 检查是否过期
|
||||
if provider.quota_expires_at and provider.quota_expires_at < now:
|
||||
logger.warning(f"Provider {provider.name} quota expired")
|
||||
# 可以选择禁用过期的提供商
|
||||
# provider.is_active = False
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error processing provider {provider.name}: {e}")
|
||||
|
||||
if reset_count > 0:
|
||||
db.commit()
|
||||
logger.info(f"Reset quotas for {reset_count} providers")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def force_reset(self, provider_id: str = None):
|
||||
"""手动强制重置额度"""
|
||||
db = create_session()
|
||||
try:
|
||||
now = datetime.now(timezone.utc)
|
||||
if provider_id:
|
||||
# 重置指定提供商
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if provider and provider.billing_type == ProviderBillingType.MONTHLY_QUOTA:
|
||||
provider.monthly_used_usd = 0.0
|
||||
provider.rpm_used = 0
|
||||
provider.rpm_reset_at = None
|
||||
provider.quota_last_reset_at = now
|
||||
db.commit()
|
||||
logger.info(f"Force reset quota for provider {provider.name}")
|
||||
else:
|
||||
# 重置所有定额提供商
|
||||
providers = (
|
||||
db.query(Provider)
|
||||
.filter(Provider.billing_type == ProviderBillingType.MONTHLY_QUOTA)
|
||||
.all()
|
||||
)
|
||||
for provider in providers:
|
||||
provider.monthly_used_usd = 0.0
|
||||
provider.rpm_used = 0
|
||||
provider.rpm_reset_at = None
|
||||
provider.quota_last_reset_at = now
|
||||
db.commit()
|
||||
logger.info(f"Force reset quotas for {len(providers)} providers")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
# 全局单例
|
||||
_quota_scheduler = None
|
||||
|
||||
|
||||
def get_quota_scheduler() -> QuotaScheduler:
|
||||
"""获取全局调度器实例"""
|
||||
global _quota_scheduler
|
||||
if _quota_scheduler is None:
|
||||
_quota_scheduler = QuotaScheduler()
|
||||
return _quota_scheduler
|
||||
244
src/services/usage/recorder.py
Normal file
244
src/services/usage/recorder.py
Normal file
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
统一的 Usage 记录器
|
||||
|
||||
设计原则:
|
||||
1. 单一入口:所有 Usage 记录都通过 UsageRecorder
|
||||
2. 自动处理:根据 RequestResult 自动判断成功/失败
|
||||
3. 完整记录:确保所有必要字段都被记录
|
||||
4. 异步友好:支持后台异步记录,不阻塞主流程
|
||||
|
||||
使用方式:
|
||||
```python
|
||||
recorder = UsageRecorder(db, user, api_key)
|
||||
|
||||
# 记录成功请求
|
||||
await recorder.record_success(result)
|
||||
|
||||
# 记录失败请求
|
||||
await recorder.record_failure(result)
|
||||
|
||||
# 或者自动判断
|
||||
await recorder.record(result)
|
||||
```
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, User
|
||||
from src.services.request.result import RequestResult
|
||||
from src.services.system.audit import audit_service
|
||||
from src.services.usage.service import UsageService
|
||||
|
||||
|
||||
|
||||
class UsageRecorder:
|
||||
"""
|
||||
统一的 Usage 记录器
|
||||
|
||||
职责:
|
||||
1. 记录成功请求的 Usage(包含 token 使用量和费用)
|
||||
2. 记录失败请求的 Usage(token=0,记录错误信息)
|
||||
3. 记录审计日志
|
||||
4. 更新用户配额
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: Session,
|
||||
user: User,
|
||||
api_key: ApiKey,
|
||||
client_ip: str = "unknown",
|
||||
request_id: Optional[str] = None,
|
||||
):
|
||||
self.db = db
|
||||
self.user = user
|
||||
self.api_key = api_key
|
||||
self.client_ip = client_ip
|
||||
self.request_id = request_id
|
||||
|
||||
async def record(self, result: RequestResult) -> None:
|
||||
"""
|
||||
根据 RequestResult 自动判断并记录 Usage
|
||||
|
||||
Args:
|
||||
result: 请求结果
|
||||
"""
|
||||
if result.is_success:
|
||||
await self.record_success(result)
|
||||
else:
|
||||
await self.record_failure(result)
|
||||
|
||||
async def record_success(
|
||||
self,
|
||||
result: RequestResult,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
记录成功请求的 Usage
|
||||
|
||||
Args:
|
||||
result: 成功的请求结果
|
||||
request_headers: 原始请求头(可选,用于调试)
|
||||
request_body: 原始请求体(可选,用于调试)
|
||||
"""
|
||||
metadata = result.metadata
|
||||
usage = result.usage
|
||||
|
||||
# 确定 target_model:当存在 original_model 且与 model 不同时,说明发生了映射
|
||||
target_model = None
|
||||
if metadata.original_model and metadata.original_model != metadata.model:
|
||||
target_model = metadata.model
|
||||
|
||||
await UsageService.record_usage(
|
||||
db=self.db,
|
||||
user=self.user,
|
||||
api_key=self.api_key,
|
||||
provider=metadata.provider,
|
||||
model=metadata.original_model or metadata.model,
|
||||
target_model=target_model,
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
cache_creation_input_tokens=usage.cache_creation_input_tokens,
|
||||
cache_read_input_tokens=usage.cache_read_input_tokens,
|
||||
request_type="chat",
|
||||
api_format=metadata.api_format,
|
||||
is_stream=result.is_stream,
|
||||
response_time_ms=result.response_time_ms,
|
||||
status_code=200,
|
||||
error_message=None,
|
||||
metadata=metadata.response_metadata if metadata.response_metadata else None,
|
||||
request_headers=request_headers or result.request_headers,
|
||||
request_body=request_body or result.request_body,
|
||||
provider_request_headers=metadata.provider_request_headers,
|
||||
response_headers=metadata.provider_response_headers,
|
||||
response_body=result.response_data if isinstance(result.response_data, dict) else {},
|
||||
request_id=self.request_id,
|
||||
provider_id=metadata.provider_id,
|
||||
provider_endpoint_id=metadata.provider_endpoint_id,
|
||||
provider_api_key_id=metadata.provider_api_key_id,
|
||||
status="completed", # 成功请求
|
||||
)
|
||||
|
||||
# 记录审计日志
|
||||
audit_service.log_api_request(
|
||||
db=self.db,
|
||||
user_id=self.user.id,
|
||||
api_key_id=self.api_key.id,
|
||||
request_id=self.request_id or "",
|
||||
model=metadata.original_model or metadata.model,
|
||||
provider=metadata.provider,
|
||||
success=True,
|
||||
ip_address=self.client_ip,
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
logger.debug(f"[UsageRecorder] 成功记录: provider={metadata.provider}, "
|
||||
f"model={metadata.model}, api_format={metadata.api_format}, "
|
||||
f"tokens={usage.input_tokens}+{usage.output_tokens}")
|
||||
|
||||
async def record_failure(
|
||||
self,
|
||||
result: RequestResult,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
记录失败请求的 Usage
|
||||
|
||||
Args:
|
||||
result: 失败的请求结果
|
||||
request_headers: 原始请求头
|
||||
request_body: 原始请求体
|
||||
"""
|
||||
metadata = result.metadata
|
||||
|
||||
# 确定 target_model:当存在 original_model 且与 model 不同时,说明发生了映射
|
||||
target_model = None
|
||||
if metadata.original_model and metadata.original_model != metadata.model:
|
||||
target_model = metadata.model
|
||||
|
||||
await UsageService.record_usage(
|
||||
db=self.db,
|
||||
user=self.user,
|
||||
api_key=self.api_key,
|
||||
provider=metadata.provider,
|
||||
model=metadata.original_model or metadata.model,
|
||||
target_model=target_model,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
request_type="chat",
|
||||
api_format=metadata.api_format,
|
||||
is_stream=result.is_stream,
|
||||
response_time_ms=result.response_time_ms,
|
||||
status_code=result.status_code,
|
||||
error_message=result.error_message,
|
||||
metadata=metadata.response_metadata if metadata.response_metadata else None,
|
||||
request_headers=request_headers or result.request_headers,
|
||||
request_body=request_body or result.request_body,
|
||||
provider_request_headers=metadata.provider_request_headers,
|
||||
response_headers={},
|
||||
response_body={"error": result.error_message} if result.error_message else {},
|
||||
request_id=self.request_id,
|
||||
provider_id=metadata.provider_id,
|
||||
provider_endpoint_id=metadata.provider_endpoint_id,
|
||||
provider_api_key_id=metadata.provider_api_key_id,
|
||||
status="failed", # 失败请求
|
||||
)
|
||||
|
||||
# 记录审计日志
|
||||
audit_service.log_api_request(
|
||||
db=self.db,
|
||||
user_id=self.user.id,
|
||||
api_key_id=self.api_key.id,
|
||||
request_id=self.request_id or "",
|
||||
model=metadata.original_model or metadata.model,
|
||||
provider=metadata.provider,
|
||||
success=False,
|
||||
ip_address=self.client_ip,
|
||||
status_code=result.status_code,
|
||||
error_message=result.error_message,
|
||||
)
|
||||
|
||||
logger.debug(f"[UsageRecorder] 失败记录: provider={metadata.provider}, "
|
||||
f"model={metadata.model}, api_format={metadata.api_format}, "
|
||||
f"status={result.status_code}, error={result.error_message[:100] if result.error_message else 'N/A'}")
|
||||
|
||||
async def record_from_exception(
|
||||
self,
|
||||
exception: Exception,
|
||||
api_format: str,
|
||||
model: str,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
从异常创建 RequestResult 并记录失败
|
||||
|
||||
这是一个便捷方法,用于在异常处理中快速记录失败请求。
|
||||
|
||||
Args:
|
||||
exception: 捕获的异常
|
||||
api_format: API 格式(必须提供,确保始终有值)
|
||||
model: 模型名称
|
||||
response_time_ms: 响应时间
|
||||
is_stream: 是否流式请求
|
||||
request_headers: 原始请求头
|
||||
request_body: 原始请求体
|
||||
"""
|
||||
result = RequestResult.from_exception(
|
||||
exception=exception,
|
||||
api_format=api_format,
|
||||
model=model,
|
||||
response_time_ms=response_time_ms,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
result.request_headers = request_headers or {}
|
||||
result.request_body = request_body or {}
|
||||
|
||||
await self.record_failure(result, request_headers, request_body)
|
||||
1306
src/services/usage/service.py
Normal file
1306
src/services/usage/service.py
Normal file
File diff suppressed because it is too large
Load Diff
1077
src/services/usage/stream.py
Normal file
1077
src/services/usage/stream.py
Normal file
File diff suppressed because it is too large
Load Diff
15
src/services/user/__init__.py
Normal file
15
src/services/user/__init__.py
Normal file
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
用户服务模块
|
||||
|
||||
包含用户管理、API Key 管理等功能。
|
||||
"""
|
||||
|
||||
from src.services.user.apikey import ApiKeyService
|
||||
from src.services.user.preference import PreferenceService
|
||||
from src.services.user.service import UserService
|
||||
|
||||
__all__ = [
|
||||
"UserService",
|
||||
"ApiKeyService",
|
||||
"PreferenceService",
|
||||
]
|
||||
393
src/services/user/apikey.py
Normal file
393
src/services/user/apikey.py
Normal file
@@ -0,0 +1,393 @@
|
||||
"""
|
||||
API密钥管理服务
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, Usage, User
|
||||
|
||||
|
||||
|
||||
class ApiKeyService:
|
||||
"""API密钥管理服务"""
|
||||
|
||||
@staticmethod
|
||||
def create_api_key(
|
||||
db: Session,
|
||||
user_id: str, # UUID
|
||||
name: Optional[str] = None,
|
||||
allowed_providers: Optional[List[str]] = None,
|
||||
allowed_api_formats: Optional[List[str]] = None,
|
||||
allowed_models: Optional[List[str]] = None,
|
||||
rate_limit: int = 100,
|
||||
concurrent_limit: int = 5,
|
||||
expire_days: Optional[int] = None,
|
||||
initial_balance_usd: Optional[float] = None,
|
||||
is_standalone: bool = False,
|
||||
auto_delete_on_expiry: bool = False,
|
||||
) -> tuple[ApiKey, str]:
|
||||
"""创建新的API密钥,返回密钥对象和明文密钥
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID
|
||||
name: 密钥名称
|
||||
allowed_providers: 允许的提供商列表
|
||||
allowed_api_formats: 允许的 API 格式列表
|
||||
allowed_models: 允许的模型列表
|
||||
rate_limit: 速率限制
|
||||
concurrent_limit: 并发限制
|
||||
expire_days: 过期天数,None = 永不过期
|
||||
initial_balance_usd: 初始余额(USD),仅用于独立Key,None = 无限制
|
||||
is_standalone: 是否为独立余额Key(仅管理员可创建)
|
||||
auto_delete_on_expiry: 过期后是否自动删除(True=物理删除,False=仅禁用)
|
||||
"""
|
||||
|
||||
# 生成密钥
|
||||
key = ApiKey.generate_key()
|
||||
key_hash = ApiKey.hash_key(key)
|
||||
key_encrypted = crypto_service.encrypt(key) # 加密存储密钥
|
||||
|
||||
# 计算过期时间
|
||||
expires_at = None
|
||||
if expire_days:
|
||||
expires_at = datetime.now(timezone.utc) + timedelta(days=expire_days)
|
||||
|
||||
api_key = ApiKey(
|
||||
user_id=user_id,
|
||||
key_hash=key_hash,
|
||||
key_encrypted=key_encrypted,
|
||||
name=name or f"API Key {datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}",
|
||||
allowed_providers=allowed_providers,
|
||||
allowed_api_formats=allowed_api_formats,
|
||||
allowed_models=allowed_models,
|
||||
rate_limit=rate_limit,
|
||||
concurrent_limit=concurrent_limit,
|
||||
expires_at=expires_at,
|
||||
balance_used_usd=0.0,
|
||||
current_balance_usd=initial_balance_usd, # 直接使用初始余额,None = 无限制
|
||||
is_standalone=is_standalone,
|
||||
auto_delete_on_expiry=auto_delete_on_expiry,
|
||||
is_active=True,
|
||||
)
|
||||
|
||||
db.add(api_key)
|
||||
db.commit()
|
||||
db.refresh(api_key)
|
||||
|
||||
logger.info(f"创建API密钥: 用户ID {user_id}, 密钥名 {api_key.name}, "
|
||||
f"独立Key={is_standalone}, 初始余额={initial_balance_usd}")
|
||||
return api_key, key # 返回密钥对象和明文密钥
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(db: Session, key_id: str) -> Optional[ApiKey]: # UUID
|
||||
"""获取API密钥"""
|
||||
return db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
|
||||
@staticmethod
|
||||
def get_api_key_by_key(db: Session, key: str) -> Optional[ApiKey]:
|
||||
"""通过密钥字符串获取API密钥"""
|
||||
key_hash = ApiKey.hash_key(key)
|
||||
return db.query(ApiKey).filter(ApiKey.key_hash == key_hash).first()
|
||||
|
||||
@staticmethod
|
||||
def list_user_api_keys(
|
||||
db: Session, user_id: str, is_active: Optional[bool] = None # UUID
|
||||
) -> List[ApiKey]:
|
||||
"""列出用户的所有API密钥(不包括独立Key)"""
|
||||
query = db.query(ApiKey).filter(
|
||||
ApiKey.user_id == user_id, ApiKey.is_standalone == False # 排除独立Key
|
||||
)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(ApiKey.is_active == is_active)
|
||||
|
||||
return query.order_by(ApiKey.created_at.desc()).all()
|
||||
|
||||
@staticmethod
|
||||
def list_standalone_api_keys(db: Session, is_active: Optional[bool] = None) -> List[ApiKey]:
|
||||
"""列出所有独立余额Key(仅管理员可用)"""
|
||||
query = db.query(ApiKey).filter(ApiKey.is_standalone == True)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(ApiKey.is_active == is_active)
|
||||
|
||||
return query.order_by(ApiKey.created_at.desc()).all()
|
||||
|
||||
@staticmethod
|
||||
def update_api_key(db: Session, key_id: str, **kwargs) -> Optional[ApiKey]: # UUID
|
||||
"""更新API密钥"""
|
||||
api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
if not api_key:
|
||||
return None
|
||||
|
||||
# 可更新的字段
|
||||
updatable_fields = [
|
||||
"name",
|
||||
"allowed_providers",
|
||||
"allowed_api_formats",
|
||||
"allowed_models",
|
||||
"rate_limit",
|
||||
"concurrent_limit",
|
||||
"is_active",
|
||||
"expires_at",
|
||||
"balance_limit_usd",
|
||||
"auto_delete_on_expiry",
|
||||
]
|
||||
|
||||
for field, value in kwargs.items():
|
||||
if field in updatable_fields and value is not None:
|
||||
setattr(api_key, field, value)
|
||||
|
||||
api_key.updated_at = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
db.refresh(api_key)
|
||||
|
||||
logger.debug(f"更新API密钥: ID {key_id}")
|
||||
return api_key
|
||||
|
||||
@staticmethod
|
||||
def delete_api_key(db: Session, key_id: str) -> bool: # UUID
|
||||
"""删除API密钥(禁用)"""
|
||||
api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
if not api_key:
|
||||
return False
|
||||
|
||||
api_key.is_active = False
|
||||
api_key.updated_at = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
|
||||
logger.info(f"删除API密钥: ID {key_id}")
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def get_remaining_balance(api_key: ApiKey) -> Optional[float]:
|
||||
"""计算剩余余额(仅用于独立Key)
|
||||
|
||||
Returns:
|
||||
剩余余额,None 表示无限制或非独立Key
|
||||
"""
|
||||
if not api_key.is_standalone:
|
||||
return None
|
||||
|
||||
if api_key.current_balance_usd is None:
|
||||
return None
|
||||
|
||||
# 剩余余额 = 当前余额 - 已使用余额
|
||||
remaining = api_key.current_balance_usd - (api_key.balance_used_usd or 0)
|
||||
return max(0, remaining) # 不能为负数
|
||||
|
||||
@staticmethod
|
||||
def check_balance(api_key: ApiKey) -> tuple[bool, Optional[float]]:
|
||||
"""检查余额限制(仅用于独立Key)
|
||||
|
||||
Returns:
|
||||
(is_allowed, remaining_balance): 是否允许请求,剩余余额(None表示无限制)
|
||||
"""
|
||||
if not api_key.is_standalone:
|
||||
# 非独立Key不检查余额
|
||||
return True, None
|
||||
|
||||
# 使用新的预付费模式: current_balance_usd
|
||||
if api_key.current_balance_usd is None:
|
||||
# 无余额限制
|
||||
return True, None
|
||||
|
||||
# 使用统一的余额计算方法
|
||||
remaining = ApiKeyService.get_remaining_balance(api_key)
|
||||
is_allowed = remaining > 0 if remaining is not None else True
|
||||
|
||||
if not is_allowed:
|
||||
logger.warning(f"API密钥余额不足: Key ID {api_key.id}, " f"剩余余额 ${remaining:.4f}")
|
||||
|
||||
return is_allowed, remaining
|
||||
|
||||
@staticmethod
|
||||
def check_rate_limit(db: Session, api_key: ApiKey, window_minutes: int = 1) -> tuple[bool, int]:
|
||||
"""检查速率限制"""
|
||||
|
||||
# 计算时间窗口
|
||||
window_start = datetime.now(timezone.utc) - timedelta(minutes=window_minutes)
|
||||
|
||||
# 统计窗口内的请求数
|
||||
request_count = (
|
||||
db.query(func.count(Usage.id))
|
||||
.filter(Usage.api_key_id == api_key.id, Usage.created_at >= window_start)
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
|
||||
# 检查是否超限
|
||||
is_allowed = request_count < api_key.rate_limit
|
||||
|
||||
if not is_allowed:
|
||||
logger.warning(f"API密钥速率限制: Key ID {api_key.id}, 请求数 {request_count}/{api_key.rate_limit}")
|
||||
|
||||
return is_allowed, api_key.rate_limit - request_count
|
||||
|
||||
@staticmethod
|
||||
def add_balance(db: Session, key_id: str, amount_usd: float) -> Optional[ApiKey]:
|
||||
"""为独立余额Key调整余额
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
key_id: API Key ID
|
||||
amount_usd: 要调整的余额金额(USD),正数为增加,负数为扣除
|
||||
|
||||
Returns:
|
||||
更新后的API Key对象,如果Key不存在或不是独立Key则返回None
|
||||
"""
|
||||
api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
if not api_key:
|
||||
logger.warning(f"余额调整失败: Key ID {key_id} 不存在")
|
||||
return None
|
||||
|
||||
if not api_key.is_standalone:
|
||||
logger.warning(f"余额调整失败: Key ID {key_id} 不是独立余额Key")
|
||||
return None
|
||||
|
||||
if amount_usd == 0:
|
||||
logger.warning(f"余额调整失败: 调整金额不能为0,当前值 ${amount_usd}")
|
||||
return None
|
||||
|
||||
# 如果是扣除(负数),检查是否超过当前余额
|
||||
if amount_usd < 0:
|
||||
current = api_key.current_balance_usd or 0
|
||||
if abs(amount_usd) > current:
|
||||
logger.warning(f"余额扣除失败: 扣除金额 ${abs(amount_usd):.4f} 超过当前余额 ${current:.4f}")
|
||||
return None
|
||||
|
||||
# 调整当前余额
|
||||
if api_key.current_balance_usd is None:
|
||||
api_key.current_balance_usd = amount_usd if amount_usd > 0 else 0
|
||||
else:
|
||||
api_key.current_balance_usd = max(0, api_key.current_balance_usd + amount_usd)
|
||||
|
||||
api_key.updated_at = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
db.refresh(api_key)
|
||||
|
||||
action = "增加" if amount_usd > 0 else "扣除"
|
||||
logger.info(f"余额调整成功: Key ID {key_id}, {action} ${abs(amount_usd):.4f}, "
|
||||
f"新余额 ${api_key.current_balance_usd:.4f}")
|
||||
return api_key
|
||||
|
||||
@staticmethod
|
||||
def cleanup_expired_keys(db: Session, auto_delete: bool = False) -> int:
|
||||
"""清理过期的API密钥
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
auto_delete: 全局默认行为(True=物理删除,False=仅禁用)
|
||||
单个Key的 auto_delete_on_expiry 字段会覆盖此设置
|
||||
|
||||
Returns:
|
||||
int: 清理的密钥数量
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
expired_keys = (
|
||||
db.query(ApiKey)
|
||||
.filter(ApiKey.expires_at <= now, ApiKey.is_active == True) # 只处理仍然活跃的
|
||||
.all()
|
||||
)
|
||||
|
||||
count = 0
|
||||
for api_key in expired_keys:
|
||||
# 优先使用Key自身的auto_delete_on_expiry设置,否则使用全局设置
|
||||
should_delete = (
|
||||
api_key.auto_delete_on_expiry
|
||||
if api_key.auto_delete_on_expiry is not None
|
||||
else auto_delete
|
||||
)
|
||||
|
||||
if should_delete:
|
||||
# 物理删除(Usage记录会保留,因为是 SET NULL)
|
||||
db.delete(api_key)
|
||||
logger.info(f"删除过期API密钥: ID {api_key.id}, 名称 {api_key.name}, "
|
||||
f"过期时间 {api_key.expires_at}")
|
||||
else:
|
||||
# 仅禁用
|
||||
api_key.is_active = False
|
||||
api_key.updated_at = now
|
||||
logger.info(f"禁用过期API密钥: ID {api_key.id}, 名称 {api_key.name}, "
|
||||
f"过期时间 {api_key.expires_at}")
|
||||
count += 1
|
||||
|
||||
if count > 0:
|
||||
db.commit()
|
||||
|
||||
return count
|
||||
|
||||
@staticmethod
|
||||
def get_api_key_stats(
|
||||
db: Session,
|
||||
key_id: str, # UUID
|
||||
start_date: Optional[datetime] = None,
|
||||
end_date: Optional[datetime] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""获取API密钥使用统计"""
|
||||
|
||||
api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
if not api_key:
|
||||
return {}
|
||||
|
||||
query = db.query(Usage).filter(Usage.api_key_id == key_id)
|
||||
|
||||
if start_date:
|
||||
query = query.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
query = query.filter(Usage.created_at <= end_date)
|
||||
|
||||
# 统计数据
|
||||
stats = db.query(
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_tokens).label("tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("cost_usd"),
|
||||
func.avg(Usage.response_time_ms).label("avg_response_time"),
|
||||
).filter(Usage.api_key_id == key_id)
|
||||
|
||||
if start_date:
|
||||
stats = stats.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
stats = stats.filter(Usage.created_at <= end_date)
|
||||
|
||||
result = stats.first()
|
||||
|
||||
# 按天统计
|
||||
daily_stats = db.query(
|
||||
func.date(Usage.created_at).label("date"),
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_tokens).label("tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("cost_usd"),
|
||||
).filter(Usage.api_key_id == key_id)
|
||||
|
||||
if start_date:
|
||||
daily_stats = daily_stats.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
daily_stats = daily_stats.filter(Usage.created_at <= end_date)
|
||||
|
||||
daily_stats = daily_stats.group_by(func.date(Usage.created_at)).all()
|
||||
|
||||
return {
|
||||
"key_id": key_id,
|
||||
"key_name": api_key.name,
|
||||
"total_requests": result.requests or 0,
|
||||
"total_tokens": result.tokens or 0,
|
||||
"total_cost_usd": float(result.cost_usd or 0),
|
||||
"avg_response_time_ms": float(result.avg_response_time or 0),
|
||||
"daily_stats": [
|
||||
{
|
||||
"date": stat.date.isoformat() if stat.date else None,
|
||||
"requests": stat.requests,
|
||||
"tokens": stat.tokens,
|
||||
"cost_usd": float(stat.cost_usd),
|
||||
}
|
||||
for stat in daily_stats
|
||||
],
|
||||
}
|
||||
137
src/services/user/preference.py
Normal file
137
src/services/user/preference.py
Normal file
@@ -0,0 +1,137 @@
|
||||
"""
|
||||
用户偏好设置服务
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.exceptions import NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, User, UserPreference
|
||||
|
||||
|
||||
|
||||
class PreferenceService:
|
||||
"""用户偏好设置服务"""
|
||||
|
||||
@staticmethod
|
||||
def get_or_create_preferences(db: Session, user_id: str) -> UserPreference: # UUID
|
||||
"""获取或创建用户偏好设置"""
|
||||
preferences = db.query(UserPreference).filter(UserPreference.user_id == user_id).first()
|
||||
|
||||
if not preferences:
|
||||
# 创建默认偏好设置
|
||||
preferences = UserPreference(
|
||||
user_id=user_id,
|
||||
theme="light",
|
||||
language="zh-CN",
|
||||
timezone="Asia/Shanghai",
|
||||
email_notifications=True,
|
||||
usage_alerts=True,
|
||||
announcement_notifications=True,
|
||||
)
|
||||
db.add(preferences)
|
||||
db.commit()
|
||||
db.refresh(preferences)
|
||||
logger.info(f"Created default preferences for user {user_id}")
|
||||
|
||||
return preferences
|
||||
|
||||
@staticmethod
|
||||
def update_preferences(
|
||||
db: Session,
|
||||
user_id: str, # UUID
|
||||
avatar_url: Optional[str] = None,
|
||||
bio: Optional[str] = None,
|
||||
default_provider_id: Optional[str] = None, # UUID
|
||||
theme: Optional[str] = None,
|
||||
language: Optional[str] = None,
|
||||
timezone: Optional[str] = None,
|
||||
email_notifications: Optional[bool] = None,
|
||||
usage_alerts: Optional[bool] = None,
|
||||
announcement_notifications: Optional[bool] = None,
|
||||
) -> UserPreference:
|
||||
"""更新用户偏好设置"""
|
||||
preferences = PreferenceService.get_or_create_preferences(db, user_id)
|
||||
|
||||
# 更新提供的字段
|
||||
if avatar_url is not None:
|
||||
preferences.avatar_url = avatar_url
|
||||
if bio is not None:
|
||||
preferences.bio = bio
|
||||
if default_provider_id is not None:
|
||||
# 验证提供商是否存在且活跃
|
||||
provider = (
|
||||
db.query(Provider)
|
||||
.filter(Provider.id == default_provider_id, Provider.is_active == True)
|
||||
.first()
|
||||
)
|
||||
if not provider:
|
||||
raise NotFoundException("Provider not found or inactive")
|
||||
preferences.default_provider_id = default_provider_id
|
||||
if theme is not None:
|
||||
if theme not in ["light", "dark", "auto"]:
|
||||
raise ValueError("Invalid theme. Must be 'light', 'dark', or 'auto'")
|
||||
preferences.theme = theme
|
||||
if language is not None:
|
||||
preferences.language = language
|
||||
if timezone is not None:
|
||||
preferences.timezone = timezone
|
||||
if email_notifications is not None:
|
||||
preferences.email_notifications = email_notifications
|
||||
if usage_alerts is not None:
|
||||
preferences.usage_alerts = usage_alerts
|
||||
if announcement_notifications is not None:
|
||||
preferences.announcement_notifications = announcement_notifications
|
||||
|
||||
db.commit()
|
||||
db.refresh(preferences)
|
||||
logger.info(f"Updated preferences for user {user_id}")
|
||||
|
||||
return preferences
|
||||
|
||||
@staticmethod
|
||||
def get_user_with_preferences(db: Session, user_id: str) -> dict: # UUID
|
||||
"""获取用户信息及其偏好设置"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
raise NotFoundException("User not found")
|
||||
|
||||
preferences = PreferenceService.get_or_create_preferences(db, user_id)
|
||||
|
||||
# 构建返回数据
|
||||
user_data = {
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"username": user.username,
|
||||
"role": user.role.value,
|
||||
"is_active": user.is_active,
|
||||
"created_at": user.created_at,
|
||||
"last_login_at": user.last_login_at,
|
||||
"preferences": {
|
||||
"avatar_url": preferences.avatar_url,
|
||||
"bio": preferences.bio,
|
||||
"default_provider": (
|
||||
preferences.default_provider.name if preferences.default_provider else None
|
||||
),
|
||||
"theme": preferences.theme,
|
||||
"language": preferences.language,
|
||||
"timezone": preferences.timezone,
|
||||
"notifications": {
|
||||
"email": preferences.email_notifications,
|
||||
"usage_alerts": preferences.usage_alerts,
|
||||
"announcements": preferences.announcement_notifications,
|
||||
},
|
||||
},
|
||||
# 配额信息
|
||||
"quota_usd": user.quota_usd,
|
||||
"used_usd": user.used_usd,
|
||||
"stats": {
|
||||
"total_cost": user.used_usd,
|
||||
"total_cost_all_time": user.total_usd,
|
||||
"api_keys_count": len(user.api_keys),
|
||||
},
|
||||
}
|
||||
|
||||
return user_data
|
||||
433
src/services/user/service.py
Normal file
433
src/services/user/service.py
Normal file
@@ -0,0 +1,433 @@
|
||||
"""
|
||||
用户管理服务
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.core.validators import EmailValidator, PasswordValidator, UsernameValidator
|
||||
from src.models.database import ApiKey, GlobalModel, Model, Provider, Usage, User, UserRole
|
||||
from src.services.cache.user_cache import UserCacheService
|
||||
from src.utils.transaction_manager import retry_on_database_error, transactional
|
||||
|
||||
|
||||
|
||||
class UserService:
|
||||
"""用户管理服务"""
|
||||
|
||||
@staticmethod
|
||||
@transactional()
|
||||
@retry_on_database_error(max_retries=3)
|
||||
def create_user(
|
||||
db: Session,
|
||||
email: str,
|
||||
username: str,
|
||||
password: str,
|
||||
role: UserRole = UserRole.USER,
|
||||
quota_usd: Optional[float] = 10.0,
|
||||
) -> User:
|
||||
"""创建新用户,quota_usd 为 None 表示无限制"""
|
||||
|
||||
# 验证邮箱格式
|
||||
valid, error_msg = EmailValidator.validate(email)
|
||||
if not valid:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 验证用户名格式
|
||||
valid, error_msg = UsernameValidator.validate(username)
|
||||
if not valid:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 验证密码复杂度
|
||||
valid, error_msg = PasswordValidator.validate(password)
|
||||
if not valid:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# 检查邮箱是否已存在
|
||||
if db.query(User).filter(User.email == email).first():
|
||||
raise ValueError(f"邮箱已存在: {email}")
|
||||
|
||||
# 检查用户名是否已存在
|
||||
if db.query(User).filter(User.username == username).first():
|
||||
raise ValueError(f"用户名已存在: {username}")
|
||||
|
||||
user = User(
|
||||
email=email,
|
||||
username=username,
|
||||
role=role,
|
||||
quota_usd=quota_usd,
|
||||
is_active=True,
|
||||
)
|
||||
user.set_password(password)
|
||||
|
||||
db.add(user)
|
||||
db.commit() # 立即提交事务,释放数据库锁
|
||||
db.refresh(user)
|
||||
|
||||
logger.info(f"创建新用户: {email} (ID: {user.id}, 角色: {role.value})")
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
@transactional()
|
||||
def create_user_with_api_key(
|
||||
db: Session,
|
||||
email: str,
|
||||
username: str,
|
||||
password: str,
|
||||
api_key_name: str = "默认密钥",
|
||||
role: UserRole = UserRole.USER,
|
||||
quota_usd: Optional[float] = 10.0,
|
||||
concurrent_limit: int = 5,
|
||||
) -> tuple[User, ApiKey]:
|
||||
"""
|
||||
创建用户并同时创建API密钥(原子操作)
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
email: 邮箱
|
||||
username: 用户名
|
||||
password: 密码
|
||||
api_key_name: API密钥名称
|
||||
role: 用户角色
|
||||
quota_usd: USD配额,None 表示无限制
|
||||
concurrent_limit: 并发限制
|
||||
|
||||
Returns:
|
||||
tuple[User, ApiKey]: 用户对象和API密钥对象
|
||||
|
||||
Raises:
|
||||
ValueError: 当验证失败时
|
||||
"""
|
||||
# 创建用户
|
||||
user = UserService.create_user(
|
||||
db=db, email=email, username=username, password=password, role=role, quota_usd=quota_usd
|
||||
)
|
||||
|
||||
# 导入API密钥服务(避免循环导入)
|
||||
from .apikey import ApiKeyService
|
||||
|
||||
# 创建API密钥(返回值是 (api_key, plain_key))
|
||||
api_key, plain_key = ApiKeyService.create_api_key(
|
||||
db=db, user_id=user.id, name=api_key_name, concurrent_limit=concurrent_limit
|
||||
)
|
||||
|
||||
logger.info(f"创建用户和API密钥完成: {email} (用户ID: {user.id}, 密钥ID: {api_key.id})")
|
||||
|
||||
# 返回用户对象、API Key对象和明文密钥
|
||||
return user, api_key, plain_key
|
||||
|
||||
@staticmethod
|
||||
def get_user(db: Session, user_id: str) -> Optional[User]:
|
||||
"""获取用户"""
|
||||
import random
|
||||
import time
|
||||
|
||||
# 添加重试机制处理数据库并发问题
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
return user
|
||||
except Exception as e:
|
||||
if attempt < max_retries - 1:
|
||||
# 添加随机延迟避免并发冲突
|
||||
time.sleep(random.uniform(0.01, 0.05))
|
||||
db.rollback() # 回滚事务
|
||||
continue
|
||||
else:
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def get_user_by_email(db: Session, email: str) -> Optional[User]:
|
||||
"""通过邮箱获取用户"""
|
||||
return db.query(User).filter(User.email == email).first()
|
||||
|
||||
@staticmethod
|
||||
def list_users(
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
role: Optional[UserRole] = None,
|
||||
is_active: Optional[bool] = None,
|
||||
) -> List[User]:
|
||||
"""列出用户"""
|
||||
query = db.query(User)
|
||||
|
||||
if role:
|
||||
query = query.filter(User.role == role)
|
||||
if is_active is not None:
|
||||
query = query.filter(User.is_active == is_active)
|
||||
|
||||
return query.offset(skip).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
@transactional()
|
||||
def update_user(db: Session, user_id: str, **kwargs) -> Optional[User]:
|
||||
"""更新用户信息"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
return None
|
||||
|
||||
# 可更新的字段
|
||||
updatable_fields = [
|
||||
"email",
|
||||
"username",
|
||||
"quota_usd",
|
||||
"is_active",
|
||||
"role",
|
||||
# 访问限制字段
|
||||
"allowed_providers",
|
||||
"allowed_endpoints",
|
||||
"allowed_models",
|
||||
]
|
||||
|
||||
# 允许设置为 None 的字段(表示无限制)
|
||||
nullable_fields = ["quota_usd", "allowed_providers", "allowed_endpoints", "allowed_models"]
|
||||
|
||||
for field, value in kwargs.items():
|
||||
if field not in updatable_fields:
|
||||
continue
|
||||
# nullable_fields 中的字段允许设置为 None
|
||||
if field in nullable_fields:
|
||||
setattr(user, field, value)
|
||||
elif value is not None:
|
||||
setattr(user, field, value)
|
||||
|
||||
# 如果提供了新密码
|
||||
if "password" in kwargs and kwargs["password"]:
|
||||
# 验证新密码复杂度
|
||||
valid, error_msg = PasswordValidator.validate(kwargs["password"])
|
||||
if not valid:
|
||||
raise ValueError(error_msg)
|
||||
user.set_password(kwargs["password"])
|
||||
|
||||
user.updated_at = datetime.now(timezone.utc)
|
||||
db.commit() # 立即提交事务,释放数据库锁
|
||||
db.refresh(user)
|
||||
|
||||
# 清除用户缓存
|
||||
asyncio.create_task(UserCacheService.invalidate_user_cache(user.id, user.email))
|
||||
|
||||
logger.debug(f"更新用户信息: {user.email} (ID: {user_id})")
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
@transactional()
|
||||
def delete_user(db: Session, user_id: str) -> bool:
|
||||
"""删除用户(硬删除)
|
||||
|
||||
删除流程:
|
||||
1. 手动删除关联的子记录(避免 SQLAlchemy ORM 与数据库 CASCADE 冲突)
|
||||
2. 删除用户记录
|
||||
3. 历史 Usage 记录保留,user_id 会被数据库设为 NULL
|
||||
4. 新用户注册时会有新的 UUID,看不到旧用户的记录
|
||||
"""
|
||||
from src.models.database import (
|
||||
AnnouncementRead,
|
||||
ApiKey,
|
||||
UserPreference,
|
||||
UserQuota,
|
||||
)
|
||||
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
return False
|
||||
|
||||
# 记录删除信息用于日志
|
||||
email = user.email
|
||||
|
||||
# 手动删除子记录,避免 SQLAlchemy 的 ORM cascade 与数据库 CASCADE 冲突
|
||||
# 这些表的数据库外键已经设置了 ON DELETE CASCADE,但 SQLAlchemy 会先尝试 UPDATE 设置为 NULL
|
||||
# 所以我们手动删除来避免这个问题
|
||||
db.query(UserPreference).filter(UserPreference.user_id == user_id).delete(
|
||||
synchronize_session=False
|
||||
)
|
||||
db.query(UserQuota).filter(UserQuota.user_id == user_id).delete(synchronize_session=False)
|
||||
db.query(AnnouncementRead).filter(AnnouncementRead.user_id == user_id).delete(
|
||||
synchronize_session=False
|
||||
)
|
||||
api_key_count = db.query(ApiKey).filter(ApiKey.user_id == user_id).count()
|
||||
db.query(ApiKey).filter(ApiKey.user_id == user_id).delete(synchronize_session=False)
|
||||
|
||||
# 现在删除用户(Usage, AuditLog, RequestAttempt 会通过数据库 SET NULL 保留)
|
||||
db.delete(user)
|
||||
db.commit() # 立即提交事务,释放数据库锁
|
||||
|
||||
# 清除用户缓存
|
||||
asyncio.create_task(UserCacheService.invalidate_user_cache(user_id, email))
|
||||
|
||||
logger.info(f"删除用户: {email} (ID: {user_id}), 同时删除 {api_key_count} 个API密钥")
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
@transactional()
|
||||
def change_password(
|
||||
db: Session, user_id: str, old_password: str, new_password: str
|
||||
) -> tuple[bool, str]:
|
||||
"""
|
||||
更改用户密码
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID
|
||||
old_password: 旧密码
|
||||
new_password: 新密码
|
||||
|
||||
Returns:
|
||||
(是否成功, 消息)
|
||||
"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
return False, "用户不存在"
|
||||
|
||||
# 验证旧密码
|
||||
if not user.verify_password(old_password):
|
||||
logger.warning(f"密码更改失败 - 旧密码错误: 用户ID {user_id}")
|
||||
return False, "旧密码错误"
|
||||
|
||||
# 验证新密码复杂度
|
||||
valid, error_msg = PasswordValidator.validate(new_password)
|
||||
if not valid:
|
||||
return False, error_msg
|
||||
|
||||
# 检查新密码不能与旧密码相同
|
||||
if old_password == new_password:
|
||||
return False, "新密码不能与旧密码相同"
|
||||
|
||||
# 设置新密码
|
||||
user.set_password(new_password)
|
||||
user.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
# 清除用户缓存
|
||||
asyncio.create_task(UserCacheService.invalidate_user_cache(user.id, user.email))
|
||||
|
||||
logger.info(f"密码更改成功: 用户ID {user_id}")
|
||||
return True, "密码更改成功"
|
||||
|
||||
@staticmethod
|
||||
def update_user_quota(
|
||||
db: Session,
|
||||
user_id: str,
|
||||
quota_usd: Optional[float] = None,
|
||||
) -> Optional[User]:
|
||||
"""更新用户配额"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
return None
|
||||
|
||||
if quota_usd is not None:
|
||||
user.quota_usd = quota_usd
|
||||
|
||||
db.commit()
|
||||
db.refresh(user)
|
||||
|
||||
# 清除用户缓存
|
||||
asyncio.create_task(UserCacheService.invalidate_user_cache(user.id, user.email))
|
||||
|
||||
logger.debug(f"更新用户配额: {user.email} (USD: {quota_usd})")
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
def get_user_usage_stats(
|
||||
db: Session,
|
||||
user_id: str,
|
||||
start_date: Optional[datetime] = None,
|
||||
end_date: Optional[datetime] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""获取用户使用统计"""
|
||||
|
||||
query = db.query(Usage).filter(Usage.user_id == user_id)
|
||||
|
||||
if start_date:
|
||||
query = query.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
query = query.filter(Usage.created_at <= end_date)
|
||||
|
||||
# 统计数据
|
||||
stats = db.query(
|
||||
func.count(Usage.id).label("total_requests"),
|
||||
func.sum(Usage.total_tokens).label("total_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost_usd"),
|
||||
func.avg(Usage.response_time_ms).label("avg_response_time"),
|
||||
).filter(Usage.user_id == user_id)
|
||||
|
||||
if start_date:
|
||||
stats = stats.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
stats = stats.filter(Usage.created_at <= end_date)
|
||||
|
||||
result = stats.first()
|
||||
|
||||
# 按模型分组统计
|
||||
model_stats = db.query(
|
||||
Usage.model,
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_tokens).label("tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("cost_usd"),
|
||||
).filter(Usage.user_id == user_id)
|
||||
|
||||
if start_date:
|
||||
model_stats = model_stats.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
model_stats = model_stats.filter(Usage.created_at <= end_date)
|
||||
|
||||
model_stats = model_stats.group_by(Usage.model).all()
|
||||
|
||||
return {
|
||||
"total_requests": result.total_requests or 0,
|
||||
"total_tokens": result.total_tokens or 0,
|
||||
"total_cost_usd": float(result.total_cost_usd or 0),
|
||||
"avg_response_time_ms": float(result.avg_response_time or 0),
|
||||
"by_model": [
|
||||
{
|
||||
"model": stat.model,
|
||||
"requests": stat.requests,
|
||||
"tokens": stat.tokens,
|
||||
"cost_usd": float(stat.cost_usd),
|
||||
}
|
||||
for stat in model_stats
|
||||
],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_user_available_models(db: Session, user: User) -> List[Model]:
|
||||
"""获取用户可用的模型
|
||||
|
||||
新架构:通过 GlobalModel + Model 关联查询用户可用模型
|
||||
逻辑:用户可用提供商 → Provider 的 Model 实现 → 关联的 GlobalModel
|
||||
"""
|
||||
# 获取用户可用的提供商
|
||||
if user.role == UserRole.ADMIN:
|
||||
# 管理员可以使用所有活动提供商
|
||||
provider_ids = [
|
||||
p.id for p in db.query(Provider.id).filter(Provider.is_active == True).all()
|
||||
]
|
||||
else:
|
||||
# 普通用户使用关联的提供商
|
||||
provider_ids = [p.id for p in user.providers]
|
||||
|
||||
if not provider_ids:
|
||||
return []
|
||||
|
||||
# 查询这些提供商的所有活跃 Model(关联 GlobalModel)
|
||||
models = (
|
||||
db.query(Model)
|
||||
.join(GlobalModel, Model.global_model_id == GlobalModel.id)
|
||||
.filter(
|
||||
and_(
|
||||
Model.provider_id.in_(provider_ids),
|
||||
Model.is_active == True,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
logger.debug(f"用户 {user.email} 可用模型: {len(models)} 个 (提供商数: {len(provider_ids)})")
|
||||
|
||||
return models
|
||||
Reference in New Issue
Block a user