Initial commit

This commit is contained in:
fawney19
2025-12-10 20:52:44 +08:00
commit f784106826
485 changed files with 110993 additions and 0 deletions

3
src/services/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
"""
服务层模块
"""

View File

@@ -0,0 +1,13 @@
"""
认证服务模块
包含认证服务、JWT 黑名单等功能。
"""
from src.services.auth.jwt_blacklist import JWTBlacklistService
from src.services.auth.service import AuthService
__all__ = [
"AuthService",
"JWTBlacklistService",
]

View 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}"
# 计算 TTLToken 过期前的剩余时间)
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)}

View 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
View 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
View 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 IDapi_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

File diff suppressed because it is too large Load Diff

330
src/services/cache/backend.py vendored Normal file
View 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
View 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: 相关 ProviderNone 表示全局)
"""
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
View 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
View 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
View 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
View 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

View File

@@ -0,0 +1,7 @@
"""
能力服务模块
"""
from .resolver import CapabilityResolver
__all__ = ["CapabilityResolver"]

View 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

View File

@@ -0,0 +1,14 @@
"""
健康监控服务模块
包含健康监控相关功能:
- health_monitor: 健康度监控单例
- HealthMonitor: 健康监控类
"""
from .monitor import HealthMonitor, health_monitor
__all__ = [
"health_monitor",
"HealthMonitor",
]

View 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}")

View 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)

View 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
View 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,
)
)

View 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

View 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

View File

@@ -0,0 +1,432 @@
"""
模型映射解析服务
负责统一的模型别名/降级解析,按优先级顺序:
1. 映射mappingProvider 特定 → 全局
2. 别名aliasProvider 特定 → 全局
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)

View 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 # 命中的阶梯索引

View 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
),
)

View 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",
]

View 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

View 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__,
)

View 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_priorityEndpoint在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,
)

View 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,
)

View 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",
]

View 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

View 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

View 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")

View 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 "/"

View 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",
]

View 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

View 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

View 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

View 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)

View 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()

View 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}")

View 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",
]

View 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),
}

View 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

View 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__()

View 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",
]

View 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,
)

View 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()

View 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

View 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 KeyTrue=物理删除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

View 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()

View 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

View 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,
}

View 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",
]

View 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

View 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. 记录失败请求的 Usagetoken=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)

File diff suppressed because it is too large Load Diff

1077
src/services/usage/stream.py Normal file

File diff suppressed because it is too large Load Diff

View 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
View 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仅用于独立KeyNone = 无限制
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
],
}

View 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

View 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