i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 重构限流系统和健康监控，支持按 API 格式区分

Browse Source 
- 将 adaptive_concurrency 重命名为 adaptive_rpm，从并发控制改为 RPM 控制
- 健康监控器支持按 API 格式独立管理健康度和熔断器状态
- 新增 model_permissions 模块，支持按格式配置允许的模型
- 重构前端提供商相关表单组件，新增 Collapsible UI 组件
- 新增数据库迁移脚本支持新的数据结构
This commit is contained in:
fawney19
2026-01-10 18:43:53 +08:00
parent dd2fbf4424
commit 09e0f594ff
97 changed files with 6642 additions and 4169 deletions
Download Patch File Download Diff File Expand all files Collapse all files

224
src/services/cache/aware_scheduler.py vendored
View File

@@ -34,7 +34,7 @@ import hashlib
import random
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from sqlalchemy.orm import Session, selectinload
@@ -80,8 +80,6 @@ class ProviderCandidate:
@dataclass
class ConcurrencySnapshot:
endpoint_current: int
endpoint_limit: Optional[int]
key_current: int
key_limit: Optional[int]
is_cached_user: bool = False
@@ -91,11 +89,9 @@ class ConcurrencySnapshot:
reservation_confidence: float = 0.0
def describe(self) -> str:
endpoint_limit_text = str(self.endpoint_limit) if self.endpoint_limit is not None else "inf"
key_limit_text = str(self.key_limit) if self.key_limit is not None else "inf"
reservation_text = f"{self.reservation_ratio:.0%}" if self.reservation_ratio > 0 else "N/A"
return (
f"endpoint={self.endpoint_current}/{endpoint_limit_text}, "
f"key={self.key_current}/{key_limit_text}, "
f"cached={self.is_cached_user}, "
f"reserve={reservation_text}({self.reservation_phase})"
@@ -246,9 +242,8 @@ class CacheAwareScheduler:
if not candidates:
if provider_offset == 0:
# 没有找到任何候选，提供友好的错误提示
error_msg = f"模型 '{model_name}' 不可用"
raise ProviderNotAvailableException(error_msg)
# 没有找到任何候选，提供友好的错误提示（不暴露内部信息）
raise ProviderNotAvailableException("请求的模型当前不可用")
break
self._metrics["total_batches"] += 1
@@ -270,7 +265,6 @@ class CacheAwareScheduler:
is_cached_user = bool(candidate.is_cached)
can_use, snapshot = await self._check_concurrent_available(
endpoint,
key,
is_cached_user=is_cached_user,
)
@@ -312,47 +306,51 @@ class CacheAwareScheduler:
provider_offset += provider_batch_size
raise ProviderNotAvailableException(f"所有Provider的资源当前不可用 (model={model_name})")
raise ProviderNotAvailableException("服务暂时繁忙，请稍后重试")
def _get_effective_concurrent_limit(self, key: ProviderAPIKey) -> Optional[int]:
def _get_effective_rpm_limit(self, key: ProviderAPIKey) -> Optional[int]:
"""
获取有效的并发限制
获取有效的 RPM 限制
新逻辑：
- max_concurrent=NULL: 启用自适应，使用 learned_max_concurrent（如无学习记录则为 None）
- max_concurrent=数字: 固定限制，直接使用该值
- rpm_limit=NULL: 启用自适应，使用 learned_rpm_limit（如无学习记录则使用默认初始值）
- rpm_limit=数字: 固定限制，直接使用该值
Args:
key: API Key对象
Returns:
有效的并发限制（None 表示不限制）
有效的 RPM 限制（None 表示不限制）
"""
if key.max_concurrent is None:
if key.rpm_limit is None:
# 自适应模式：使用学习到的值
learned = key.learned_max_concurrent
return int(learned) if learned is not None else None
learned = key.learned_rpm_limit
if learned is not None:
return int(learned)
# 未学习到值时，使用默认初始限制，避免无限制打爆上游
from src.config.constants import RPMDefaults
return int(RPMDefaults.INITIAL_LIMIT)
else:
# 固定限制模式
return int(key.max_concurrent)
return int(key.rpm_limit)
async def _check_concurrent_available(
self,
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
is_cached_user: bool = False,
) -> Tuple[bool, ConcurrencySnapshot]:
"""
检查并发是否可用（使用动态预留机制）
检查 RPM 限制是否可用（使用动态预留机制）
核心逻辑 - 动态缓存预留机制:
- 总槽位: 有效并发限制（固定值或学习到的值）
- 总槽位: 有效 RPM 限制（固定值或学习到的值）
- 预留比例: 由 AdaptiveReservationManager 根据置信度和负载动态计算
- 缓存用户可用: 全部槽位
- 新用户可用: 总槽位 × (1 - 动态预留比例)
Args:
endpoint: ProviderEndpoint对象
key: ProviderAPIKey对象
is_cached_user: 是否是缓存用户
@@ -360,7 +358,7 @@ class CacheAwareScheduler:
(是否可用, 并发快照)
"""
# 获取有效的并发限制
effective_key_limit = self._get_effective_concurrent_limit(key)
effective_key_limit = self._get_effective_rpm_limit(key)
logger.debug(
f" -> 并发检查: _concurrency_manager={self._concurrency_manager is not None}, "
@@ -371,33 +369,23 @@ class CacheAwareScheduler:
# 并发管理器不可用，直接返回True
logger.debug(f" -> 无并发管理器，直接通过")
snapshot = ConcurrencySnapshot(
endpoint_current=0,
endpoint_limit=(
int(endpoint.max_concurrent) if endpoint.max_concurrent is not None else None
),
key_current=0,
key_limit=effective_key_limit,
is_cached_user=is_cached_user,
)
return True, snapshot
# 获取当前并发数
endpoint_count, key_count = await self._concurrency_manager.get_current_concurrency(
endpoint_id=str(endpoint.id),
# 获取当前 RPM 计数
key_count = await self._concurrency_manager.get_key_rpm_count(
key_id=str(key.id),
)
can_use = True
# 检查Endpoint级别限制
if endpoint.max_concurrent is not None:
if endpoint_count >= endpoint.max_concurrent:
can_use = False
# 计算动态预留比例
reservation_result = self._reservation_manager.calculate_reservation(
key=key,
current_concurrent=key_count,
current_usage=key_count,
effective_limit=effective_key_limit,
)
@@ -440,7 +428,8 @@ class CacheAwareScheduler:
# 使用 max 确保至少有 1 个槽位可用
import math
available_for_new = max(1, math.ceil(effective_key_limit * (1 - reservation_ratio)))
# 与 ConcurrencyManager 的 Lua 脚本保持一致：使用 floor 计算新用户可用槽位
available_for_new = max(1, math.floor(effective_key_limit * (1 - reservation_ratio)))
if key_count >= available_for_new:
logger.debug(
f"Key {key.id[:8]}... 新用户配额已满 "
@@ -460,8 +449,6 @@ class CacheAwareScheduler:
key_limit_for_snapshot = None
snapshot = ConcurrencySnapshot(
endpoint_current=endpoint_count,
endpoint_limit=endpoint.max_concurrent,
key_current=key_count,
key_limit=key_limit_for_snapshot,
is_cached_user=is_cached_user,
@@ -475,7 +462,7 @@ class CacheAwareScheduler:
def _get_effective_restrictions(
self,
user_api_key: Optional[ApiKey],
) -> Dict[str, Optional[set]]:
) -> Dict[str, Any]:
"""
获取有效的访问限制（合并 ApiKey 和 User 的限制）
@@ -536,7 +523,10 @@ class CacheAwareScheduler:
)
# 合并 allowed_models
result["allowed_models"] = merge_restrictions(
# allowed_models 支持 list/dict 两种结构，不能转成 set 否则会导致权限校验失效
from src.core.model_permissions import merge_allowed_models
result["allowed_models"] = merge_allowed_models(
user_api_key.allowed_models, user.allowed_models if user else None
)
@@ -617,22 +607,25 @@ class CacheAwareScheduler:
)
return [], global_model_id
# 0.2 检查模型是否被允许
if allowed_models is not None:
if (
requested_model_name not in allowed_models
and resolved_model_name not in allowed_models
):
resolved_note = (
f" (解析为 {resolved_model_name})"
if resolved_model_name != requested_model_name
else ""
)
logger.debug(
f"用户/API Key 不允许使用模型 {requested_model_name}{resolved_note}, "
f"允许的模型: {allowed_models}"
)
return [], global_model_id
# 0.2 检查模型是否被允许（支持简单列表和按格式字典两种模式）
from src.core.model_permissions import check_model_allowed, get_allowed_models_preview
if not check_model_allowed(
model_name=requested_model_name,
allowed_models=allowed_models,
api_format=target_format.value,
resolved_model_name=resolved_model_name,
):
resolved_note = (
f" (解析为 {resolved_model_name})"
if resolved_model_name != requested_model_name
else ""
)
logger.debug(
f"用户/API Key 不允许使用模型 {requested_model_name}{resolved_note}, "
f"允许的模型: {get_allowed_models_preview(allowed_models)}"
)
return [], global_model_id
# 1. 查询 Providers
providers = self._query_providers(
@@ -724,8 +717,11 @@ class CacheAwareScheduler:
provider_query = (
db.query(Provider)
.options(
selectinload(Provider.endpoints).selectinload(ProviderEndpoint.api_keys),
# 同时加载 models 和 global_model 关系，以便 get_effective_* 方法能正确继承默认值
# 预加载 Provider 级别的 api_keys
selectinload(Provider.api_keys),
# 预加载 endpoints（用于按 api_format 选择请求配置）
selectinload(Provider.endpoints),
# 同时加载 models 和 global_model 关系
selectinload(Provider.models).selectinload(Model.global_model),
)
.filter(Provider.is_active == True)
@@ -852,6 +848,7 @@ class CacheAwareScheduler:
def _check_key_availability(
self,
key: ProviderAPIKey,
api_format: Optional[str],
model_name: str,
capability_requirements: Optional[Dict[str, bool]] = None,
resolved_model_name: Optional[str] = None,
@@ -871,20 +868,24 @@ class CacheAwareScheduler:
Returns:
(is_available, skip_reason)
"""
# 检查熔断器状态（使用详细状态方法获取更丰富的跳过原因）
is_available, circuit_reason = health_monitor.get_circuit_breaker_status(key)
# 检查熔断器状态（使用详细状态方法获取更丰富的跳过原因，按 API 格式）
is_available, circuit_reason = health_monitor.get_circuit_breaker_status(
key, api_format=api_format
)
if not is_available:
return False, circuit_reason or "熔断器已打开"
# 模型权限检查：使用 allowed_models 白名单
# 模型权限检查：使用 allowed_models 白名单（支持简单列表和按格式字典两种模式）
# None = 允许所有模型，[] = 拒绝所有模型，["a","b"] = 只允许指定模型
if key.allowed_models is not None and (
model_name not in key.allowed_models
and (not resolved_model_name or resolved_model_name not in key.allowed_models)
from src.core.model_permissions import check_model_allowed, get_allowed_models_preview
if not check_model_allowed(
model_name=model_name,
allowed_models=key.allowed_models,
api_format=api_format,
resolved_model_name=resolved_model_name,
):
allowed_preview = ", ".join(key.allowed_models[:3]) if key.allowed_models else "(无)"
suffix = "..." if len(key.allowed_models) > 3 else ""
return False, f"模型权限不匹配(允许: {allowed_preview}{suffix})"
return False, f"模型权限不匹配(允许: {get_allowed_models_preview(key.allowed_models)})"
# Key 级别的能力匹配检查
# 注意：模型级别的能力检查已在 _check_model_support 中完成
@@ -914,6 +915,8 @@ class CacheAwareScheduler:
"""
构建候选列表
Key 直属 Provider，通过 api_formats 筛选符合目标格式的 Key。
Args:
db: 数据库会话
providers: Provider 列表
@@ -929,10 +932,10 @@ class CacheAwareScheduler:
候选列表
"""
candidates: List[ProviderCandidate] = []
target_format_str = target_format.value
for provider in providers:
# 检查模型支持（同时检查流式支持和模型能力需求）
# 模型能力检查在 Provider 级别进行，如果模型不支持所需能力，整个 Provider 被跳过
supports_model, skip_reason, _model_caps = await self._check_model_support(
db, provider, model_name, is_stream, capability_requirements
)
@@ -940,49 +943,63 @@ class CacheAwareScheduler:
logger.debug(f"Provider {provider.name} 不支持模型 {model_name}: {skip_reason}")
continue
# 查找目标格式对应的 Endpoint（获取请求配置）
target_endpoint = None
for endpoint in provider.endpoints:
# endpoint.api_format 是字符串，target_format 是枚举
endpoint_format_str = (
endpoint.api_format
if isinstance(endpoint.api_format, str)
else endpoint.api_format.value
)
if not endpoint.is_active or endpoint_format_str != target_format.value:
continue
if endpoint.is_active and endpoint_format_str == target_format_str:
target_endpoint = endpoint
break
# 获取活跃的 Key 并按 internal_priority + 负载均衡排序
active_keys = [key for key in endpoint.api_keys if key.is_active]
# 检查是否所有 Key 都是 TTL=0（轮换模式）
# 如果所有 Key 的 cache_ttl_minutes 都是 0 或 None，则使用随机排序
use_random = all(
(key.cache_ttl_minutes or 0) == 0 for key in active_keys
) if active_keys else False
if use_random and len(active_keys) > 1:
logger.debug(
f" Endpoint {endpoint.id[:8]}... 启用 Key 轮换模式 (TTL=0, {len(active_keys)} keys)"
)
keys = self._shuffle_keys_by_internal_priority(active_keys, affinity_key, use_random)
if not target_endpoint:
logger.debug(f"Provider {provider.name} 没有活跃的 {target_format_str} 端点")
continue
for key in keys:
# Key 级别的能力检查（模型级别的能力检查已在上面完成）
is_available, skip_reason = self._check_key_availability(
key,
model_name,
capability_requirements,
resolved_model_name=resolved_model_name,
)
# Key 直属 Provider，通过 api_formats 筛选
active_keys = [
key for key in provider.api_keys
if key.is_active and target_format_str in (key.api_formats or [])
]
candidate = ProviderCandidate(
provider=provider,
endpoint=endpoint,
key=key,
is_skipped=not is_available,
skip_reason=skip_reason,
)
candidates.append(candidate)
if not active_keys:
logger.debug(f"Provider {provider.name} 没有支持 {target_format_str} 的活跃 Key")
continue
if max_candidates and len(candidates) >= max_candidates:
return candidates
# 检查是否所有 Key 都是 TTL=0（轮换模式）
use_random = all(
(key.cache_ttl_minutes or 0) == 0 for key in active_keys
) if active_keys else False
if use_random and len(active_keys) > 1:
logger.debug(
f" Provider {provider.name} 启用 Key 轮换模式 (TTL=0, {len(active_keys)} keys)"
)
keys = self._shuffle_keys_by_internal_priority(active_keys, affinity_key, use_random)
for key in keys:
# Key 级别的能力检查
is_available, skip_reason = self._check_key_availability(
key,
target_format_str,
model_name,
capability_requirements,
resolved_model_name=resolved_model_name,
)
candidate = ProviderCandidate(
provider=provider,
endpoint=target_endpoint,
key=key,
is_skipped=not is_available,
skip_reason=skip_reason,
)
candidates.append(candidate)
if max_candidates and len(candidates) >= max_candidates:
return candidates
return candidates
@@ -1187,7 +1204,6 @@ class CacheAwareScheduler:
from collections import defaultdict
# 使用 tuple 作为统一的 key 类型，兼容两种模式
priority_groups: Dict[tuple, List[ProviderCandidate]] = defaultdict(list)
# 根据优先级模式选择分组方式

34
src/services/cache/provider_cache.py vendored
View File

@@ -27,24 +27,29 @@ class ProviderCacheService:
@staticmethod
async def get_provider_api_key_rate_multiplier(
db: Session, provider_api_key_id: str
db: Session, provider_api_key_id: str, api_format: Optional[str] = None
) -> Optional[float]:
"""
获取 ProviderAPIKey 的 rate_multiplier（带缓存）
优先返回指定 API 格式的倍率，如果没有则返回默认倍率。
Args:
db: 数据库会话
provider_api_key_id: ProviderAPIKey ID
api_format: API 格式（可选），如 "CLAUDE"、"OPENAI"
Returns:
rate_multiplier 或 None（如果找不到）
"""
cache_key = f"provider_api_key:rate_multiplier:{provider_api_key_id}"
# 缓存键包含 api_format
format_suffix = api_format.upper() if api_format else "default"
cache_key = f"provider_api_key:rate_multiplier:{provider_api_key_id}:{format_suffix}"
# 1. 尝试从缓存获取
cached_data = await CacheService.get(cache_key)
if cached_data is not None:
logger.debug(f"ProviderAPIKey rate_multiplier 缓存命中: {provider_api_key_id[:8]}...")
logger.debug(f"ProviderAPIKey rate_multiplier 缓存命中: {provider_api_key_id[:8]}... format={format_suffix}")
# 缓存的 "NOT_FOUND" 表示数据库中不存在
if cached_data == "NOT_FOUND":
return None
@@ -52,18 +57,24 @@ class ProviderCacheService:
# 2. 缓存未命中，查询数据库
provider_key = (
db.query(ProviderAPIKey.rate_multiplier)
db.query(ProviderAPIKey.rate_multiplier, ProviderAPIKey.rate_multipliers)
.filter(ProviderAPIKey.id == provider_api_key_id)
.first()
)
# 3. 写入缓存
# 3. 计算倍率并写入缓存
if provider_key:
# 优先使用 rate_multipliers[api_format]，回退到 rate_multiplier
rate_multiplier = provider_key.rate_multiplier or 1.0
if api_format and provider_key.rate_multipliers:
format_upper = api_format.upper()
if format_upper in provider_key.rate_multipliers:
rate_multiplier = provider_key.rate_multipliers[format_upper]
await CacheService.set(
cache_key, rate_multiplier, ttl_seconds=ProviderCacheService.CACHE_TTL
)
logger.debug(f"ProviderAPIKey rate_multiplier 已缓存: {provider_api_key_id[:8]}...")
logger.debug(f"ProviderAPIKey rate_multiplier 已缓存: {provider_api_key_id[:8]}... format={format_suffix} value={rate_multiplier}")
return rate_multiplier
else:
# 缓存负结果
@@ -125,6 +136,7 @@ class ProviderCacheService:
db: Session,
provider_api_key_id: Optional[str],
provider_id: Optional[str],
api_format: Optional[str] = None,
) -> Tuple[float, bool]:
"""
获取费率倍数和是否免费套餐（带缓存）
@@ -135,6 +147,7 @@ class ProviderCacheService:
db: 数据库会话
provider_api_key_id: ProviderAPIKey ID（可选）
provider_id: Provider ID（可选）
api_format: API 格式（可选），用于获取按格式配置的倍率
Returns:
(rate_multiplier, is_free_tier) 元组
@@ -142,10 +155,10 @@ class ProviderCacheService:
actual_rate_multiplier = 1.0
is_free_tier = False
# 获取费率倍数
# 获取费率倍数（支持按 API 格式查询）
if provider_api_key_id:
rate_multiplier = await ProviderCacheService.get_provider_api_key_rate_multiplier(
db, provider_api_key_id
db, provider_api_key_id, api_format
)
if rate_multiplier is not None:
actual_rate_multiplier = rate_multiplier
@@ -160,8 +173,9 @@ class ProviderCacheService:
@staticmethod
async def invalidate_provider_api_key_cache(provider_api_key_id: str) -> None:
"""清除 ProviderAPIKey 缓存"""
await CacheService.delete(f"provider_api_key:rate_multiplier:{provider_api_key_id}")
"""清除 ProviderAPIKey 缓存（包括所有 API 格式的缓存）"""
# 使用模式匹配删除所有格式的缓存
await CacheService.delete_pattern(f"provider_api_key:rate_multiplier:{provider_api_key_id}:*")
logger.debug(f"ProviderAPIKey 缓存已清除: {provider_api_key_id[:8]}...")
@staticmethod

92
src/services/health/endpoint.py
View File

@@ -70,20 +70,21 @@ class EndpointHealthService:
db.query(ProviderEndpoint).join(Provider).filter(Provider.is_active.is_(True)).all()
)
# 收集所有 endpoint_ids
all_endpoint_ids = [ep.id for ep in endpoints]
# 收集所有 provider_ids
all_provider_ids = list(set(ep.provider_id for ep in endpoints))
# 批量查询所有密钥
# 批量查询所有密钥（通过 provider_id 关联）
all_keys = (
db.query(ProviderAPIKey)
.filter(ProviderAPIKey.endpoint_id.in_(all_endpoint_ids))
.filter(ProviderAPIKey.provider_id.in_(all_provider_ids))
.all()
) if all_endpoint_ids else []
) if all_provider_ids else []
# 按 endpoint_id 分组密钥
keys_by_endpoint: Dict[str, List[ProviderAPIKey]] = defaultdict(list)
# 按 api_format 分组密钥（通过 api_formats 字段）
keys_by_format: Dict[str, List[ProviderAPIKey]] = defaultdict(list)
for key in all_keys:
keys_by_endpoint[key.endpoint_id].append(key)
for fmt in (key.api_formats or []):
keys_by_format[fmt].append(key)
# 按 API 格式聚合
format_stats = defaultdict(
@@ -106,18 +107,36 @@ class EndpointHealthService:
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)
# 统计每个格式的密钥（直接从 keys_by_format 获取）
for api_format, keys in keys_by_format.items():
if api_format not in format_stats:
# 如果有 Key 但没有对应的 Endpoint，跳过
continue
# 统计活跃密钥和健康度
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)
# 去重（同一个 Key 可能支持多个格式）
seen_key_ids = set()
unique_keys = []
for key in keys:
if key.id not in seen_key_ids:
seen_key_ids.add(key.id)
unique_keys.append(key)
format_stats[api_format]["total_keys"] = len(unique_keys)
for key in unique_keys:
format_stats[api_format]["key_ids"].append(key.id)
# 检查该格式的熔断器状态
circuit_by_format = key.circuit_breaker_by_format or {}
format_circuit = circuit_by_format.get(api_format, {})
is_circuit_open = format_circuit.get("open", False)
if key.is_active and not is_circuit_open:
format_stats[api_format]["active_keys"] += 1
# 获取该格式的健康度
health_by_format = key.health_by_format or {}
format_health = health_by_format.get(api_format, {})
health_score = float(format_health.get("health_score") or 1.0)
format_stats[api_format]["health_scores"].append(health_score)
# 批量生成所有格式的时间线数据
all_key_ids = []
@@ -372,7 +391,7 @@ class EndpointHealthService:
segments: int = 100,
) -> Dict[str, Any]:
"""
从真实使用记录生成时间线数据（兼容旧接口，使用批量查询优化）
从真实使用记录生成时间线数据（使用批量查询优化）
Args:
db: 数据库会话
@@ -391,13 +410,34 @@ class EndpointHealthService:
"time_range_end": None,
}
# 先查询该 API 格式下的所有密钥
key_ids = [
k.id
for k in db.query(ProviderAPIKey.id)
.filter(ProviderAPIKey.endpoint_id.in_(endpoint_ids))
# 基于 endpoint_ids 反推 provider_ids 与 api_format，再选出支持该格式的 keys
endpoint_rows = (
db.query(ProviderEndpoint.provider_id, ProviderEndpoint.api_format)
.filter(ProviderEndpoint.id.in_(endpoint_ids))
.all()
]
)
if not endpoint_rows:
return {
"timeline": ["unknown"] * 100,
"time_range_start": None,
"time_range_end": None,
}
provider_ids = {str(pid) for pid, _fmt in endpoint_rows}
# 同一调用中 endpoint_ids 来自同一 api_format（上层已按格式分组）
api_format = (
endpoint_rows[0][1].value
if hasattr(endpoint_rows[0][1], "value")
else str(endpoint_rows[0][1])
)
keys = (
db.query(ProviderAPIKey.id, ProviderAPIKey.api_formats)
.filter(ProviderAPIKey.provider_id.in_(provider_ids))
.all()
)
key_ids = [str(key_id) for key_id, formats in keys if api_format in (formats or [])]
if not key_ids:
return {

749
src/services/health/monitor.py
View File

@@ -1,11 +1,15 @@
"""
健康监控器 - Endpoint 和 Key 的健康度追踪
健康监控器 - Endpoint 和 Key 的健康度追踪（按 API 格式区分）
功能：
1. 基于滑动窗口的错误率计算
2. 三态熔断器：关闭 -> 打开 -> 半开 -> 关闭
1. 基于滑动窗口的错误率计算（按 API 格式独立）
2. 三态熔断器：关闭 -> 打开 -> 半开 -> 关闭（按 API 格式独立）
3. 半开状态允许少量请求验证服务恢复
4. 提供健康度查询和管理 API
数据结构：
- health_by_format: {"CLAUDE": {"health_score": 1.0, "consecutive_failures": 0, ...}, ...}
- circuit_breaker_by_format: {"CLAUDE": {"open": false, "open_at": null, ...}, ...}
"""
import os
@@ -30,8 +34,30 @@ class CircuitState:
HALF_OPEN = "half_open" # 半开（验证恢复）
# 默认健康度数据结构
def _default_health_data() -> Dict[str, Any]:
return {
"health_score": 1.0,
"consecutive_failures": 0,
"last_failure_at": None,
"request_results_window": [],
}
# 默认熔断器数据结构
def _default_circuit_data() -> Dict[str, Any]:
return {
"open": False,
"open_at": None,
"next_probe_at": None,
"half_open_until": None,
"half_open_successes": 0,
"half_open_failures": 0,
}
class HealthMonitor:
"""健康监控器（滑动窗口 + 半开状态模式）"""
"""健康监控器（滑动窗口 + 半开状态模式，按 API 格式区分）"""
# === 滑动窗口配置 ===
WINDOW_SIZE = int(os.getenv("HEALTH_WINDOW_SIZE", str(CircuitBreakerDefaults.WINDOW_SIZE)))
@@ -96,6 +122,38 @@ class HealthMonitor:
_circuit_history: List[Dict[str, Any]] = []
_open_circuit_keys: int = 0
# ==================== 数据访问辅助方法 ====================
@classmethod
def _get_health_data(cls, key: ProviderAPIKey, api_format: str) -> Dict[str, Any]:
"""获取指定格式的健康度数据，不存在则返回默认值"""
health_by_format = key.health_by_format or {}
if api_format not in health_by_format:
return _default_health_data()
return health_by_format[api_format]
@classmethod
def _set_health_data(cls, key: ProviderAPIKey, api_format: str, data: Dict[str, Any]) -> None:
"""设置指定格式的健康度数据"""
health_by_format = dict(key.health_by_format or {})
health_by_format[api_format] = data
key.health_by_format = health_by_format # type: ignore[assignment]
@classmethod
def _get_circuit_data(cls, key: ProviderAPIKey, api_format: str) -> Dict[str, Any]:
"""获取指定格式的熔断器数据，不存在则返回默认值"""
circuit_by_format = key.circuit_breaker_by_format or {}
if api_format not in circuit_by_format:
return _default_circuit_data()
return circuit_by_format[api_format]
@classmethod
def _set_circuit_data(cls, key: ProviderAPIKey, api_format: str, data: Dict[str, Any]) -> None:
"""设置指定格式的熔断器数据"""
circuit_by_format = dict(key.circuit_breaker_by_format or {})
circuit_by_format[api_format] = data
key.circuit_breaker_by_format = circuit_by_format # type: ignore[assignment]
# ==================== 核心方法 ====================
@classmethod
@@ -103,9 +161,21 @@ class HealthMonitor:
cls,
db: Session,
key_id: Optional[str] = None,
api_format: Optional[str] = None,
response_time_ms: Optional[int] = None,
) -> None:
"""记录成功请求"""
"""记录成功请求（按 API 格式）
Args:
db: 数据库会话
key_id: Key ID（必需）
api_format: API 格式（必需，用于区分不同格式的健康度）
response_time_ms: 响应时间（可选）
Note:
api_format 在逻辑上是必需的，但为了向后兼容保持 Optional 签名。
如果未提供，会尝试从 Key 的 api_formats 中获取第一个格式作为 fallback。
"""
try:
if not key_id:
return
@@ -114,39 +184,96 @@ class HealthMonitor:
if not key:
return
# api_format 兼容处理：如果未提供，尝试使用 Key 的第一个格式
effective_api_format = api_format
if not effective_api_format:
if key.api_formats and len(key.api_formats) > 0:
effective_api_format = key.api_formats[0]
logger.debug(
f"record_success: api_format 未提供，使用默认格式 {effective_api_format}"
)
else:
logger.warning(
f"record_success: api_format 未提供且 Key 无可用格式: key_id={key_id[:8]}..."
)
return
now = datetime.now(timezone.utc)
now_ts = now.timestamp()
# 获取当前格式的健康度数据
health_data = cls._get_health_data(key, effective_api_format)
circuit_data = cls._get_circuit_data(key, effective_api_format)
# 1. 更新滑动窗口
cls._add_to_window(key, now_ts, success=True)
window = health_data.get("request_results_window") or []
window.append({"ts": now_ts, "ok": True})
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 :]
health_data["request_results_window"] = window
# 2. 更新健康度（用于展示）
new_score = min(float(key.health_score or 0) + cls.SUCCESS_INCREMENT, 1.0)
key.health_score = new_score # type: ignore[assignment]
current_score = float(health_data.get("health_score") or 0)
new_score = min(current_score + cls.SUCCESS_INCREMENT, 1.0)
health_data["health_score"] = new_score
# 3. 更新统计
key.consecutive_failures = 0 # type: ignore[assignment]
key.last_failure_at = None # type: ignore[assignment]
health_data["consecutive_failures"] = 0
health_data["last_failure_at"] = None
# 4. 处理熔断器状态
state = cls._get_circuit_state_from_data(circuit_data, now)
if state == CircuitState.HALF_OPEN:
# 半开状态：记录成功
circuit_data["half_open_successes"] = int(
circuit_data.get("half_open_successes") or 0
) + 1
if circuit_data["half_open_successes"] >= cls.HALF_OPEN_SUCCESS_THRESHOLD:
# 达到成功阈值，关闭熔断器
cls._close_circuit_data(circuit_data, health_data, reason="半开状态验证成功")
cls._push_circuit_event(
{
"event": "closed",
"key_id": key.id,
"api_format": effective_api_format,
"reason": "半开状态验证成功",
"timestamp": now.isoformat(),
}
)
logger.info(
f"[CLOSED] Key 熔断器关闭: {key.id[:8]}.../{effective_api_format} | 原因: 半开状态验证成功"
)
elif state == CircuitState.OPEN:
# 打开状态下的成功（探测成功），进入半开状态
cls._enter_half_open_data(circuit_data, now)
cls._push_circuit_event(
{
"event": "half_open",
"key_id": key.id,
"api_format": effective_api_format,
"timestamp": now.isoformat(),
}
)
logger.info(
f"[HALF-OPEN] Key 进入半开状态: {key.id[:8]}.../{effective_api_format} | "
f"需要 {cls.HALF_OPEN_SUCCESS_THRESHOLD} 次成功关闭熔断器"
)
# 保存数据
cls._set_health_data(key, effective_api_format, health_data)
cls._set_circuit_data(key, effective_api_format, circuit_data)
# 更新全局统计
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)
@@ -159,9 +286,21 @@ class HealthMonitor:
cls,
db: Session,
key_id: Optional[str] = None,
api_format: Optional[str] = None,
error_type: Optional[str] = None,
) -> None:
"""记录失败请求"""
"""记录失败请求（按 API 格式）
Args:
db: 数据库会话
key_id: Key ID（必需）
api_format: API 格式（必需，用于区分不同格式的健康度）
error_type: 错误类型（可选）
Note:
api_format 在逻辑上是必需的，但为了向后兼容保持 Optional 签名。
如果未提供，会尝试从 Key 的 api_formats 中获取第一个格式作为 fallback。
"""
try:
if not key_id:
return
@@ -170,46 +309,117 @@ class HealthMonitor:
if not key:
return
# api_format 兼容处理：如果未提供，尝试使用 Key 的第一个格式
effective_api_format = api_format
if not effective_api_format:
if key.api_formats and len(key.api_formats) > 0:
effective_api_format = key.api_formats[0]
logger.debug(
f"record_failure: api_format 未提供，使用默认格式 {effective_api_format}"
)
else:
logger.warning(
f"record_failure: api_format 未提供且 Key 无可用格式: key_id={key_id[:8]}..."
)
return
now = datetime.now(timezone.utc)
now_ts = now.timestamp()
# 获取当前格式的健康度数据
health_data = cls._get_health_data(key, effective_api_format)
circuit_data = cls._get_circuit_data(key, effective_api_format)
# 1. 更新滑动窗口
cls._add_to_window(key, now_ts, success=False)
window = health_data.get("request_results_window") or []
window.append({"ts": now_ts, "ok": False})
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 :]
health_data["request_results_window"] = window
# 2. 更新健康度（用于展示）
new_score = max(float(key.health_score or 1) - cls.FAILURE_DECREMENT, 0.0)
key.health_score = new_score # type: ignore[assignment]
current_score = float(health_data.get("health_score") or 1)
new_score = max(current_score - cls.FAILURE_DECREMENT, 0.0)
health_data["health_score"] = new_score
# 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]
health_data["consecutive_failures"] = (
int(health_data.get("consecutive_failures") or 0) + 1
)
health_data["last_failure_at"] = now.isoformat()
# 4. 处理熔断器状态
state = cls._get_circuit_state(key, now)
state = cls._get_circuit_state_from_data(circuit_data, now)
if state == CircuitState.HALF_OPEN:
# 半开状态：记录失败
key.half_open_failures = int(key.half_open_failures or 0) + 1 # type: ignore[assignment]
circuit_data["half_open_failures"] = int(
circuit_data.get("half_open_failures") or 0
) + 1
if int(key.half_open_failures or 0) >= cls.HALF_OPEN_FAILURE_THRESHOLD:
if circuit_data["half_open_failures"] >= cls.HALF_OPEN_FAILURE_THRESHOLD:
# 达到失败阈值，重新打开熔断器
cls._open_circuit(key, now, reason="半开状态验证失败")
# 注意：半开状态本身就是打开状态的子状态，不需要增加计数
consecutive = int(health_data.get("consecutive_failures") or 0)
recovery_seconds = cls._calculate_recovery_seconds(consecutive)
cls._open_circuit_data(
circuit_data, now, recovery_seconds, reason="半开状态验证失败"
)
cls._push_circuit_event(
{
"event": "opened",
"key_id": key.id,
"api_format": effective_api_format,
"reason": "半开状态验证失败",
"recovery_seconds": recovery_seconds,
"timestamp": now.isoformat(),
}
)
logger.warning(
f"[OPEN] Key 熔断器打开: {key.id[:8]}.../{effective_api_format} | 原因: 半开状态验证失败 | "
f"{recovery_seconds}秒后进入半开状态"
)
elif state == CircuitState.CLOSED:
# 关闭状态：检查是否需要打开熔断器
error_rate = cls._calculate_error_rate(key, now_ts)
window = key.request_results_window or []
error_rate = cls._calculate_error_rate_from_window(window, now_ts)
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%}"
consecutive = int(health_data.get("consecutive_failures") or 0)
recovery_seconds = cls._calculate_recovery_seconds(consecutive)
reason = f"错误率 {error_rate:.0%} 超过阈值 {cls.ERROR_RATE_THRESHOLD:.0%}"
cls._open_circuit_data(circuit_data, now, recovery_seconds, reason=reason)
cls._open_circuit_keys += 1
health_open_circuits.set(cls._open_circuit_keys)
cls._push_circuit_event(
{
"event": "opened",
"key_id": key.id,
"api_format": effective_api_format,
"reason": reason,
"recovery_seconds": recovery_seconds,
"timestamp": now.isoformat(),
}
)
logger.warning(
f"[OPEN] Key 熔断器打开: {key.id[:8]}.../{effective_api_format} | 原因: {reason} | "
f"{recovery_seconds}秒后进入半开状态"
)
# 保存数据
cls._set_health_data(key, effective_api_format, health_data)
cls._set_circuit_data(key, effective_api_format, circuit_data)
# 更新全局统计
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]
key.last_error_at = now # type: ignore[assignment]
logger.debug(
f"[WARN] Key 健康度下降: {key_id[:8]}... -> {new_score:.2f} "
f"(连续失败 {key.consecutive_failures} 次, error_type={error_type})"
f"[WARN] Key 健康度下降: {key_id[:8]}.../{effective_api_format} -> {new_score:.2f} "
f"(连续失败 {health_data['consecutive_failures']} 次, error_type={error_type})"
)
db.flush()
@@ -222,31 +432,13 @@ class HealthMonitor:
# ==================== 滑动窗口方法 ====================
@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 []
def _calculate_error_rate_from_window(
cls, window: List[Dict[str, Any]], now_ts: float
) -> float:
"""从窗口数据计算错误率"""
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]
@@ -256,157 +448,158 @@ class HealthMonitor:
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:
def _get_circuit_state_from_data(cls, circuit_data: Dict[str, Any], now: datetime) -> str:
"""从数据字典获取当前熔断器状态"""
if not circuit_data.get("open"):
return CircuitState.CLOSED
# 检查是否在半开状态
if key.half_open_until and now < key.half_open_until:
return CircuitState.HALF_OPEN
half_open_until_str = circuit_data.get("half_open_until")
if half_open_until_str:
half_open_until = datetime.fromisoformat(half_open_until_str)
if now < half_open_until:
return CircuitState.HALF_OPEN
# 检查是否到了探测时间（进入半开）
if key.next_probe_at and now >= key.next_probe_at:
return CircuitState.HALF_OPEN
next_probe_str = circuit_data.get("next_probe_at")
if next_probe_str:
next_probe_at = datetime.fromisoformat(next_probe_str)
if now >= 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(),
}
)
def _open_circuit_data(
cls,
circuit_data: Dict[str, Any],
now: datetime,
recovery_seconds: int,
reason: str,
) -> None:
"""打开熔断器（操作数据字典）"""
circuit_data["open"] = True
circuit_data["open_at"] = now.isoformat()
circuit_data["half_open_until"] = None
circuit_data["half_open_successes"] = 0
circuit_data["half_open_failures"] = 0
circuit_data["next_probe_at"] = (now + timedelta(seconds=recovery_seconds)).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(),
}
)
def _enter_half_open_data(cls, circuit_data: Dict[str, Any], now: datetime) -> None:
"""进入半开状态（操作数据字典）"""
circuit_data["half_open_until"] = (
now + timedelta(seconds=cls.HALF_OPEN_DURATION)
).isoformat()
circuit_data["half_open_successes"] = 0
circuit_data["half_open_failures"] = 0
@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]
def _close_circuit_data(
cls, circuit_data: Dict[str, Any], health_data: Dict[str, Any], reason: str
) -> None:
"""关闭熔断器（操作数据字典）"""
circuit_data["open"] = False
circuit_data["open_at"] = None
circuit_data["next_probe_at"] = None
circuit_data["half_open_until"] = None
circuit_data["half_open_successes"] = 0
circuit_data["half_open_failures"] = 0
# 快速恢复健康度
key.health_score = max(float(key.health_score or 0), cls.PROBE_RECOVERY_SCORE) # type: ignore[assignment]
current_score = float(health_data.get("health_score") or 0)
health_data["health_score"] = max(current_score, cls.PROBE_RECOVERY_SCORE)
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次失败增加一级
exponent = min(consecutive_failures // 5, 4)
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:
def is_circuit_breaker_closed(
cls, resource: ProviderAPIKey, api_format: Optional[str] = None
) -> bool:
"""检查熔断器是否允许请求通过（按 API 格式）"""
if not api_format:
# 兼容旧调用：检查是否有任何格式的熔断器开启
circuit_by_format = resource.circuit_breaker_by_format or {}
for fmt, circuit_data in circuit_by_format.items():
if circuit_data.get("open"):
return False
return True
circuit_data = cls._get_circuit_data(resource, api_format)
if not circuit_data.get("open"):
return True
now = datetime.now(timezone.utc)
state = cls._get_circuit_state(resource, now)
state = cls._get_circuit_state_from_data(circuit_data, 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
next_probe_str = circuit_data.get("next_probe_at")
if next_probe_str:
next_probe_at = datetime.fromisoformat(next_probe_str)
if now >= next_probe_at:
# 自动进入半开状态
cls._enter_half_open_data(circuit_data, now)
cls._set_circuit_data(resource, api_format, circuit_data)
return True
return False
@classmethod
def get_circuit_breaker_status(
cls, resource: ProviderAPIKey
cls, resource: ProviderAPIKey, api_format: Optional[str] = None
) -> Tuple[bool, Optional[str]]:
"""获取熔断器详细状态"""
if not resource.circuit_breaker_open:
"""获取熔断器详细状态（按 API 格式）"""
if not api_format:
# 兼容旧调用：返回第一个开启的熔断器状态
circuit_by_format = resource.circuit_breaker_by_format or {}
for fmt, circuit_data in circuit_by_format.items():
if circuit_data.get("open"):
return cls._get_status_from_circuit_data(circuit_data)
return True, None
circuit_data = cls._get_circuit_data(resource, api_format)
return cls._get_status_from_circuit_data(circuit_data)
@classmethod
def _get_status_from_circuit_data(
cls, circuit_data: Dict[str, Any]
) -> Tuple[bool, Optional[str]]:
"""从熔断器数据获取状态描述"""
if not circuit_data.get("open"):
return True, None
now = datetime.now(timezone.utc)
state = cls._get_circuit_state(resource, now)
state = cls._get_circuit_state_from_data(circuit_data, now)
if state == CircuitState.HALF_OPEN:
successes = int(resource.half_open_successes or 0)
successes = int(circuit_data.get("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:
next_probe_str = circuit_data.get("next_probe_at")
if next_probe_str:
next_probe_at = datetime.fromisoformat(next_probe_str)
if now >= next_probe_at:
return True, None
remaining = resource.next_probe_at - now
remaining = 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"
@@ -417,8 +610,10 @@ class HealthMonitor:
return False, "熔断中"
@classmethod
def get_key_health(cls, db: Session, key_id: str) -> Optional[Dict[str, Any]]:
"""获取 Key 健康状态"""
def get_key_health(
cls, db: Session, key_id: str, api_format: Optional[str] = None
) -> Optional[Dict[str, Any]]:
"""获取 Key 健康状态（支持按格式查询）"""
try:
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
if not key:
@@ -427,24 +622,15 @@ class HealthMonitor:
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 {
# 全局统计
result = {
"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),
@@ -457,25 +643,84 @@ class HealthMonitor:
),
"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),
},
}
# 按格式的健康度数据
health_by_format = key.health_by_format or {}
circuit_by_format = key.circuit_breaker_by_format or {}
if api_format:
# 查询单个格式
health_data = cls._get_health_data(key, api_format)
circuit_data = cls._get_circuit_data(key, api_format)
window = health_data.get("request_results_window") or []
valid_window = [r for r in window if r["ts"] > now_ts - cls.WINDOW_SECONDS]
result["api_format"] = api_format
result["health_score"] = float(health_data.get("health_score") or 1.0)
result["error_rate"] = cls._calculate_error_rate_from_window(window, now_ts)
result["window_size"] = len(valid_window)
result["consecutive_failures"] = int(
health_data.get("consecutive_failures") or 0
)
result["last_failure_at"] = health_data.get("last_failure_at")
result["circuit_breaker"] = {
"state": cls._get_circuit_state_from_data(circuit_data, now),
"open": circuit_data.get("open", False),
"open_at": circuit_data.get("open_at"),
"next_probe_at": circuit_data.get("next_probe_at"),
"half_open_until": circuit_data.get("half_open_until"),
"half_open_successes": int(circuit_data.get("half_open_successes") or 0),
"half_open_failures": int(circuit_data.get("half_open_failures") or 0),
}
else:
# 返回所有格式的健康度数据
formats_health = {}
for fmt in (key.api_formats or []):
health_data = health_by_format.get(fmt, _default_health_data())
circuit_data = circuit_by_format.get(fmt, _default_circuit_data())
window = health_data.get("request_results_window") or []
valid_window = [r for r in window if r["ts"] > now_ts - cls.WINDOW_SECONDS]
formats_health[fmt] = {
"health_score": float(health_data.get("health_score") or 1.0),
"error_rate": cls._calculate_error_rate_from_window(window, now_ts),
"window_size": len(valid_window),
"consecutive_failures": int(
health_data.get("consecutive_failures") or 0
),
"last_failure_at": health_data.get("last_failure_at"),
"circuit_breaker": {
"state": cls._get_circuit_state_from_data(circuit_data, now),
"open": circuit_data.get("open", False),
"open_at": circuit_data.get("open_at"),
"next_probe_at": circuit_data.get("next_probe_at"),
"half_open_until": circuit_data.get("half_open_until"),
"half_open_successes": int(
circuit_data.get("half_open_successes") or 0
),
"half_open_failures": int(
circuit_data.get("half_open_failures") or 0
),
},
}
result["health_by_format"] = formats_health
# 计算整体健康度（取最低值）
if formats_health:
result["health_score"] = min(
h["health_score"] for h in formats_health.values()
)
result["any_circuit_open"] = any(
h["circuit_breaker"]["open"] for h in formats_health.values()
)
else:
result["health_score"] = 1.0
result["any_circuit_open"] = False
return result
except Exception as e:
logger.error(f"获取 Key 健康状态失败: {e}")
return None
@@ -507,23 +752,24 @@ class HealthMonitor:
# ==================== 管理方法 ====================
@classmethod
def reset_health(cls, db: Session, key_id: Optional[str] = None) -> bool:
"""重置健康度"""
def reset_health(
cls, db: Session, key_id: Optional[str] = None, api_format: 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}")
if api_format:
# 重置单个格式
cls._set_health_data(key, api_format, _default_health_data())
cls._set_circuit_data(key, api_format, _default_circuit_data())
logger.info(f"[RESET] 重置 Key 健康度: {key_id}/{api_format}")
else:
# 重置所有格式
key.health_by_format = {} # type: ignore[assignment]
key.circuit_breaker_by_format = {} # type: ignore[assignment]
logger.info(f"[RESET] 重置 Key 所有格式健康度: {key_id}")
db.flush()
get_batch_committer().mark_dirty(db)
@@ -542,7 +788,9 @@ class HealthMonitor:
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]
# 重置所有格式的健康度
key.health_by_format = {} # type: ignore[assignment]
key.circuit_breaker_by_format = {} # type: ignore[assignment]
logger.info(f"[OK] 手动启用 Key: {key_id}")
db.flush()
@@ -566,14 +814,28 @@ class HealthMonitor:
),
).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()
# 统计 Key（需要遍历 JSON 字段计算熔断状态）
keys = db.query(ProviderAPIKey).all()
total_keys = len(keys)
active_keys = sum(1 for k in keys if k.is_active)
unhealthy_keys = 0
circuit_open_keys = 0
for key in keys:
health_by_format = key.health_by_format or {}
circuit_by_format = key.circuit_breaker_by_format or {}
# 检查是否有任何格式健康度低于 0.5
for fmt, health_data in health_by_format.items():
if float(health_data.get("health_score") or 1.0) < 0.5:
unhealthy_keys += 1
break
# 检查是否有任何格式熔断器开启
for fmt, circuit_data in circuit_by_format.items():
if circuit_data.get("open"):
circuit_open_keys += 1
break
return {
"endpoints": {
@@ -582,10 +844,10 @@ class HealthMonitor:
"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,
"total": total_keys,
"active": active_keys,
"unhealthy": unhealthy_keys,
"circuit_open": circuit_open_keys,
},
}
@@ -618,8 +880,9 @@ class HealthMonitor:
db: Session,
endpoint_id: Optional[str] = None,
key_id: Optional[str] = None,
api_format: Optional[str] = None,
) -> bool:
"""检查是否有资格进行探测（兼容旧接口）"""
"""检查是否有资格进行探测（按 API 格式）"""
if not cls.ALLOW_AUTO_RECOVER:
return False
@@ -628,13 +891,53 @@ class HealthMonitor:
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
if key:
if api_format:
circuit_data = cls._get_circuit_data(key, api_format)
if circuit_data.get("open"):
now = datetime.now(timezone.utc)
state = cls._get_circuit_state_from_data(circuit_data, now)
return state == CircuitState.HALF_OPEN
else:
# 兼容旧调用：检查是否有任何格式处于半开状态
circuit_by_format = key.circuit_breaker_by_format or {}
now = datetime.now(timezone.utc)
for fmt, circuit_data in circuit_by_format.items():
if circuit_data.get("open"):
state = cls._get_circuit_state_from_data(circuit_data, now)
if state == CircuitState.HALF_OPEN:
return True
return False
# ==================== 便捷方法 ====================
@classmethod
def get_health_score(
cls, key: ProviderAPIKey, api_format: Optional[str] = None
) -> float:
"""获取指定格式的健康度分数"""
if not api_format:
# 返回所有格式中的最低健康度
health_by_format = key.health_by_format or {}
if not health_by_format:
return 1.0
return min(
float(h.get("health_score") or 1.0) for h in health_by_format.values()
)
health_data = cls._get_health_data(key, api_format)
return float(health_data.get("health_score") or 1.0)
@classmethod
def is_any_circuit_open(cls, key: ProviderAPIKey) -> bool:
"""检查是否有任何格式的熔断器开启"""
circuit_by_format = key.circuit_breaker_by_format or {}
for circuit_data in circuit_by_format.values():
if circuit_data.get("open"):
return True
return False
# 全局健康监控器实例
health_monitor = HealthMonitor()

4
src/services/model/service.py
View File

@@ -216,7 +216,7 @@ class ModelService:
def delete_model(db: Session, model_id: str): # UUID
"""删除模型
新架构删除逻辑：
删除逻辑：
- Model 只是 Provider 对 GlobalModel 的实现，删除不影响 GlobalModel
- 检查是否是该 GlobalModel 的最后一个实现（如果是，警告但允许删除）
"""
@@ -384,7 +384,7 @@ class ModelService:
@staticmethod
def convert_to_response(model: Model) -> ModelResponse:
"""转换为响应模型（新架构：从 GlobalModel 获取显示信息和默认值）"""
"""转换为响应模型（从 GlobalModel 获取显示信息和默认值）"""
return ModelResponse(
id=model.id,
provider_id=model.provider_id,

7
src/services/orchestration/candidate_resolver.py
View File

@@ -171,7 +171,8 @@ class CandidateResolver:
)
candidate_record_map[(candidate_index, 0)] = record_id
else:
max_retries_for_candidate = endpoint.max_retries if candidate.is_cached else 1
# max_retries 已从 Endpoint 迁移到 Provider（Endpoint 仍可能保留旧字段用于兼容）
max_retries_for_candidate = int(provider.max_retries or 2) if candidate.is_cached else 1
for retry_index in range(max_retries_for_candidate):
record_id = str(uuid.uuid4())
@@ -236,7 +237,7 @@ class CandidateResolver:
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
provider = candidate.provider
max_retries = int(provider.max_retries or 2) if candidate.is_cached else 1
total += max_retries
return total

33
src/services/orchestration/error_classifier.py
View File

@@ -26,7 +26,7 @@ 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.adaptive_rpm import get_adaptive_rpm_manager
from src.services.rate_limit.detector import RateLimitType, detect_rate_limit_type
@@ -112,7 +112,7 @@ class ErrorClassifier:
cache_scheduler: 缓存调度器（可选）
"""
self.db = db
self.adaptive_manager = adaptive_manager or get_adaptive_manager()
self.adaptive_manager = adaptive_manager or get_adaptive_rpm_manager()
self.cache_scheduler = cache_scheduler
# 表示客户端错误的 error type（不区分大小写）
@@ -361,7 +361,7 @@ class ErrorClassifier:
self,
key: ProviderAPIKey,
provider_name: str,
current_concurrent: Optional[int],
current_rpm: Optional[int],
exception: ProviderRateLimitException,
request_id: Optional[str] = None,
) -> str:
@@ -371,7 +371,7 @@ class ErrorClassifier:
Args:
key: API Key 对象
provider_name: 提供商名称
current_concurrent: 当前并发数
current_rpm: 当前分钟内的请求数
exception: 速率限制异常
request_id: 请求 ID（用于日志）
@@ -388,27 +388,27 @@ class ErrorClassifier:
rate_limit_info = detect_rate_limit_type(
headers=response_headers,
provider_name=provider_name,
current_concurrent=current_concurrent,
current_usage=current_rpm,
)
logger.info(f" [{request_id}] 429错误分析: "
f"类型={rate_limit_info.limit_type}, "
f"retry_after={rate_limit_info.retry_after}s, "
f"当前并发={current_concurrent}")
f"当前RPM={current_rpm}")
# 调用自适应管理器处理
new_limit = self.adaptive_manager.handle_429_error(
db=self.db,
key=key,
rate_limit_info=rate_limit_info,
current_concurrent=current_concurrent,
current_rpm=current_rpm,
)
if rate_limit_info.limit_type == RateLimitType.CONCURRENT:
logger.warning(f" [{request_id}] 自适应调整: " f"Key {key.id[:8]}... 并发限制 -> {new_limit}")
logger.warning(f" [{request_id}] 并发限制触发（不调整RPM）")
return "concurrent"
elif rate_limit_info.limit_type == RateLimitType.RPM:
logger.info(f" [{request_id}] [RPM] RPM限制，需要切换Provider")
logger.warning(f" [{request_id}] 自适应调整: Key {key.id[:8]}... RPM限制 -> {new_limit}")
return "rpm"
else:
return "unknown"
@@ -439,18 +439,18 @@ class ErrorClassifier:
# 提取可读的错误消息
extracted_message = self._extract_error_message(error_response_text)
# 构建详细错误信息
# 构建详细错误信息（仅用于日志，不暴露给客户端）
if extracted_message:
detailed_message = f"提供商 '{provider_name}' 返回错误 {status}: {extracted_message}"
detailed_message = f"上游服务返回错误 {status}: {extracted_message}"
else:
detailed_message = f"提供商 '{provider_name}' 返回错误: {status}"
detailed_message = f"上游服务返回错误: {status}"
if status == 401:
return ProviderAuthException(provider_name=provider_name)
if status == 429:
return ProviderRateLimitException(
message=error_response_text or f"提供商 '{provider_name}' 速率限制",
message="请求过于频繁，请稍后重试",
provider_name=provider_name,
response_headers=dict(error.response.headers) if error.response else None,
retry_after=(
@@ -583,6 +583,7 @@ class ErrorClassifier:
health_monitor.record_failure(
db=self.db,
key_id=str(key.id),
api_format=api_format_str,
error_type="ProviderAuthException",
)
return extra_data
@@ -592,7 +593,7 @@ class ErrorClassifier:
await self.handle_rate_limit(
key=key,
provider_name=provider_name,
current_concurrent=captured_key_concurrent,
current_rpm=captured_key_concurrent,
exception=converted_error,
request_id=request_id,
)
@@ -620,6 +621,7 @@ class ErrorClassifier:
health_monitor.record_failure(
db=self.db,
key_id=str(key.id),
api_format=api_format_str,
error_type=type(converted_error).__name__,
)
@@ -675,7 +677,7 @@ class ErrorClassifier:
await self.handle_rate_limit(
key=key,
provider_name=provider_name,
current_concurrent=captured_key_concurrent,
current_rpm=captured_key_concurrent,
exception=error,
request_id=request_id,
)
@@ -702,5 +704,6 @@ class ErrorClassifier:
health_monitor.record_failure(
db=self.db,
key_id=str(key.id),
api_format=api_format_str,
error_type=type(error).__name__,
)

20
src/services/orchestration/fallback_orchestrator.py
View File

@@ -44,7 +44,7 @@ from src.services.cache.aware_scheduler import (
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.adaptive_rpm import get_adaptive_rpm_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
@@ -87,7 +87,7 @@ class FallbackOrchestrator:
self.redis = redis_client
self.cache_scheduler: Optional[CacheAwareScheduler] = None
self.concurrency_manager: Any = None
self.adaptive_manager = get_adaptive_manager() # 自适应并发管理器
self.adaptive_manager = get_adaptive_rpm_manager() # 自适应 RPM 管理器
self.request_executor: Optional[RequestExecutor] = None
# 拆分后的组件（延迟初始化）
@@ -558,7 +558,8 @@ class FallbackOrchestrator:
"""尝试单个候选（含重试逻辑），返回执行结果"""
provider = candidate.provider
endpoint = candidate.endpoint
max_retries_for_candidate = int(endpoint.max_retries) if candidate.is_cached else 1
# 从 Provider 读取 max_retries（已从 Endpoint 迁移）
max_retries_for_candidate = int(provider.max_retries or 2) if candidate.is_cached else 1
last_error: Optional[Exception] = None
for retry_index in range(max_retries_for_candidate):
@@ -710,7 +711,7 @@ class FallbackOrchestrator:
upstream_status = getattr(last_error, "upstream_status", None)
upstream_response = getattr(last_error, "upstream_response", None)
# 如果响应为空或无效，使用异常的字符串表示
# 如果响应为空或无效，使用异常的字符串表示作为 upstream_response
if (
not upstream_response
or not upstream_response.strip()
@@ -718,8 +719,17 @@ class FallbackOrchestrator:
):
upstream_response = str(last_error)
# 构建友好的错误消息（用于返回给客户端，不暴露内部信息）
# 如果 last_error 有 message 属性，优先使用（已经是友好提示）
# 否则使用通用提示
friendly_message = "服务暂时不可用，请稍后重试"
if last_error:
last_error_message = getattr(last_error, "message", None)
if last_error_message and isinstance(last_error_message, str):
friendly_message = last_error_message
raise ProviderNotAvailableException(
f"所有Provider均不可用，已尝试{max_attempts}个组合",
friendly_message,
request_metadata=request_metadata,
upstream_status=upstream_status,
upstream_response=upstream_response,

14
src/services/rate_limit/__init__.py
View File

@@ -1,19 +1,23 @@
"""
限流服务模块
包含自适应并发控制、RPM限流、IP限流等功能。
包含自适应 RPM 控制、并发管理、IP限流等功能。
"""
from src.services.rate_limit.adaptive_concurrency import AdaptiveConcurrencyManager
from src.services.rate_limit.adaptive_rpm import (
AdaptiveConcurrencyManager, # 向后兼容别名
AdaptiveRPMManager,
get_adaptive_rpm_manager,
)
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",
"AdaptiveConcurrencyManager", # 向后兼容
"AdaptiveRPMManager",
"ConcurrencyManager",
"IPRateLimiter",
"RPMLimiter",
"RateLimitDetector",
"get_adaptive_rpm_manager",
]

12
src/services/rate_limit/adaptive_reservation.py
View File

@@ -98,7 +98,7 @@ class AdaptiveReservationManager:
def calculate_reservation(
self,
key: "ProviderAPIKey",
current_concurrent: int = 0,
current_usage: int = 0,
effective_limit: Optional[int] = None,
) -> ReservationResult:
"""
@@ -106,8 +106,8 @@ class AdaptiveReservationManager:
Args:
key: ProviderAPIKey 对象
current_concurrent: 当前并发数
effective_limit: 有效并发限制（学习值或配置值）
current_usage: 当前使用量（RPM 计数）
effective_limit: 有效限制（学习值或配置值）
Returns:
ReservationResult 包含预留比例和详细信息
@@ -116,7 +116,7 @@ class AdaptiveReservationManager:
total_requests = self._get_total_requests(key)
# 计算负载率
load_ratio = self._calculate_load_ratio(current_concurrent, effective_limit)
load_ratio = self._calculate_load_ratio(current_usage, effective_limit)
# 阶段1: 探测阶段
if total_requests < self.config.probe_phase_requests:
@@ -165,12 +165,12 @@ class AdaptiveReservationManager:
return request_count
def _calculate_load_ratio(
self, current_concurrent: int, effective_limit: Optional[int]
self, current_usage: 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)
return min(current_usage / effective_limit, 1.0)
def _calculate_confidence(self, key: "ProviderAPIKey") -> float:
"""

232
src/services/rate_limit/adaptive_concurrency.py → src/services/rate_limit/adaptive_rpm.py
View File

@@ -1,16 +1,16 @@
"""
自适应并发调整器 - 基于边界记忆的并发限制调整
自适应 RPM 调整器 - 基于边界记忆的 RPM 限制调整
核心算法：边界记忆 + 渐进探测
- 触发 429 时记录边界（last_concurrent_peak），这就是真实上限
- 缩容策略：新限制 = 边界 - 1，而非乘性减少
- 触发 429 时记录边界（last_rpm_peak），这就是真实上限
- 缩容策略：新限制 = 边界 - 步长，而非乘性减少
- 扩容策略：不超过已知边界，除非是探测性扩容
- 探测性扩容：长时间无 429 时尝试突破边界
设计原则：
1. 快速收敛：一次 429 就能找到接近真实的限制
2. 避免过度保守：不会因为多次 429 而无限下降
3. 安全探测：允许在稳定后尝试更高并发
3. 安全探测：允许在稳定后尝试更高 RPM
"""
from datetime import datetime, timezone
@@ -18,7 +18,7 @@ from typing import Any, Dict, List, Optional, cast
from sqlalchemy.orm import Session
from src.config.constants import ConcurrencyDefaults
from src.config.constants import RPMDefaults
from src.core.batch_committer import get_batch_committer
from src.core.logger import logger
from src.models.database import ProviderAPIKey
@@ -33,14 +33,14 @@ class AdaptiveStrategy:
AGGRESSIVE = "aggressive" # 激进策略（快速探测）
class AdaptiveConcurrencyManager:
class AdaptiveRPMManager:
"""
自适应并发管理器
自适应 RPM 管理器
核心算法：边界记忆 + 渐进探测
- 触发 429 时记录边界（last_concurrent_peak = 触发时的并发数）
- 缩容：新限制 = 边界 - 1（快速收敛到真实限制附近）
- 扩容：不超过边界（即 last_concurrent_peak），允许回到边界值尝试
- 触发 429 时记录边界（last_rpm_peak = 触发时的 RPM）
- 缩容：新限制 = 边界 - 步长（快速收敛到真实限制附近）
- 扩容：不超过边界（即 last_rpm_peak），允许回到边界值尝试
- 探测性扩容：长时间（30分钟）无 429 时，可以尝试 +1 突破边界
扩容条件（满足任一即可）：
@@ -50,35 +50,35 @@ class AdaptiveConcurrencyManager:
关键特性：
1. 快速收敛：一次 429 就能学到接近真实的限制值
2. 边界保护：普通扩容不会超过已知边界
3. 安全探测：长时间稳定后允许尝试更高并发
3. 安全探测：长时间稳定后允许尝试更高 RPM
4. 区分并发限制和 RPM 限制
"""
# 默认配置 - 使用统一常量
DEFAULT_INITIAL_LIMIT = ConcurrencyDefaults.INITIAL_LIMIT
MIN_CONCURRENT_LIMIT = ConcurrencyDefaults.MIN_CONCURRENT_LIMIT
MAX_CONCURRENT_LIMIT = ConcurrencyDefaults.MAX_CONCURRENT_LIMIT
DEFAULT_INITIAL_LIMIT = RPMDefaults.INITIAL_LIMIT
MIN_RPM_LIMIT = RPMDefaults.MIN_RPM_LIMIT
MAX_RPM_LIMIT = RPMDefaults.MAX_RPM_LIMIT
# AIMD 参数
INCREASE_STEP = ConcurrencyDefaults.INCREASE_STEP
INCREASE_STEP = RPMDefaults.INCREASE_STEP
# 滑动窗口参数
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
UTILIZATION_WINDOW_SIZE = RPMDefaults.UTILIZATION_WINDOW_SIZE
UTILIZATION_WINDOW_SECONDS = RPMDefaults.UTILIZATION_WINDOW_SECONDS
UTILIZATION_THRESHOLD = RPMDefaults.UTILIZATION_THRESHOLD
HIGH_UTILIZATION_RATIO = RPMDefaults.HIGH_UTILIZATION_RATIO
MIN_SAMPLES_FOR_DECISION = RPMDefaults.MIN_SAMPLES_FOR_DECISION
# 探测性扩容参数
PROBE_INCREASE_INTERVAL_MINUTES = ConcurrencyDefaults.PROBE_INCREASE_INTERVAL_MINUTES
PROBE_INCREASE_MIN_REQUESTS = ConcurrencyDefaults.PROBE_INCREASE_MIN_REQUESTS
PROBE_INCREASE_INTERVAL_MINUTES = RPMDefaults.PROBE_INCREASE_INTERVAL_MINUTES
PROBE_INCREASE_MIN_REQUESTS = RPMDefaults.PROBE_INCREASE_MIN_REQUESTS
# 记录历史数量
MAX_HISTORY_RECORDS = 20
def __init__(self, strategy: str = AdaptiveStrategy.AIMD):
"""
初始化自适应并发管理器
初始化自适应 RPM 管理器
Args:
strategy: 调整策略
@@ -90,54 +90,54 @@ class AdaptiveConcurrencyManager:
db: Session,
key: ProviderAPIKey,
rate_limit_info: RateLimitInfo,
current_concurrent: Optional[int] = None,
current_rpm: Optional[int] = None,
) -> int:
"""
处理429错误，调整并发限制
处理429错误，调整 RPM 限制
Args:
db: 数据库会话
key: API Key对象
rate_limit_info: 速率限制信息
current_concurrent: 当前并发数
current_rpm: 当前分钟内的请求数
Returns:
调整后的并发限制
调整后的 RPM 限制
"""
# max_concurrent=NULL 表示启用自适应，max_concurrent=数字 表示固定限制
is_adaptive = key.max_concurrent is None
# rpm_limit=NULL 表示启用自适应，rpm_limit=数字 表示固定限制
is_adaptive = key.rpm_limit is None
if not is_adaptive:
logger.debug(
f"Key {key.id} 设置了固定并发限制 ({key.max_concurrent})，跳过自适应调整"
f"Key {key.id} 设置了固定 RPM 限制 ({key.rpm_limit})，跳过自适应调整"
)
return int(key.max_concurrent) # type: ignore[arg-type]
return int(key.rpm_limit) # 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]
# 仅在并发限制且拿到并发数时记录边界（RPM/UNKNOWN 不应覆盖并发边界记忆）
# 仅在 RPM 限制且拿到 RPM 数时记录边界
if (
rate_limit_info.limit_type == RateLimitType.CONCURRENT
and current_concurrent is not None
and current_concurrent > 0
rate_limit_info.limit_type == RateLimitType.RPM
and current_rpm is not None
and current_rpm > 0
):
key.last_concurrent_peak = current_concurrent # type: ignore[assignment]
key.last_rpm_peak = current_rpm # 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]
if rate_limit_info.limit_type == RateLimitType.RPM:
# RPM 限制：减少 RPM 限制
key.rpm_429_count = int(key.rpm_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)
# 获取当前有效限制（自适应模式使用 learned_rpm_limit）
old_limit = int(key.learned_rpm_limit or self.DEFAULT_INITIAL_LIMIT)
new_limit = self._decrease_limit(old_limit, current_rpm)
logger.warning(
f"[CONCURRENT] 并发限制触发: Key {key.id[:8]}... | "
f"当前并发: {current_concurrent} | "
f"[RPM] RPM 限制触发: Key {key.id[:8]}... | "
f"当前 RPM: {current_rpm} | "
f"调整: {old_limit} -> {new_limit}"
)
@@ -146,79 +146,78 @@ class AdaptiveConcurrencyManager:
key,
old_limit=old_limit,
new_limit=new_limit,
reason="concurrent_429",
current_concurrent=current_concurrent,
reason="rpm_429",
current_rpm=current_rpm,
)
# 更新学习到的并发限制
key.learned_max_concurrent = new_limit # type: ignore[assignment]
# 更新学习到的 RPM 限制
key.learned_rpm_limit = 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]
elif rate_limit_info.limit_type == RateLimitType.CONCURRENT:
# 并发限制：不调整 RPM，只记录
key.concurrent_429_count = int(key.concurrent_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"不调整并发限制"
f"[CONCURRENT] 并发限制触发: Key {key.id[:8]}... | "
f"不调整 RPM 限制（这是并发问题，非 RPM 问题）"
)
else:
# 未知类型：保守处理，轻微减少
logger.warning(
f"[UNKNOWN] 未知429类型: Key {key.id[:8]}... | "
f"当前并发: {current_concurrent} | "
f"保守减少并发"
f"当前 RPM: {current_rpm} | "
f"保守减少 RPM"
)
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%
old_limit = int(key.learned_rpm_limit or self.DEFAULT_INITIAL_LIMIT)
new_limit = max(int(old_limit * 0.9), self.MIN_RPM_LIMIT) # 减少10%
self._record_adjustment(
key,
old_limit=old_limit,
new_limit=new_limit,
reason="unknown_429",
current_concurrent=current_concurrent,
current_rpm=current_rpm,
)
key.learned_max_concurrent = new_limit # type: ignore[assignment]
key.learned_rpm_limit = new_limit # type: ignore[assignment]
db.flush()
get_batch_committer().mark_dirty(db)
return int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
return int(key.learned_rpm_limit or self.DEFAULT_INITIAL_LIMIT)
def handle_success(
self,
db: Session,
key: ProviderAPIKey,
current_concurrent: int,
current_rpm: int,
) -> Optional[int]:
"""
处理成功请求，基于滑动窗口利用率考虑增加并发限制
处理成功请求，基于滑动窗口利用率考虑增加 RPM 限制
Args:
db: 数据库会话
key: API Key对象
current_concurrent: 当前并发数（必需，用于计算利用率）
current_rpm: 当前分钟内的请求数（必需，用于计算利用率）
Returns:
调整后的并发限制（如果有调整），否则返回 None
调整后的 RPM 限制（如果有调整），否则返回 None
"""
# max_concurrent=NULL 表示启用自适应
is_adaptive = key.max_concurrent is None
# rpm_limit=NULL 表示启用自适应
is_adaptive = key.rpm_limit is None
if not is_adaptive:
return None
current_limit = int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
current_limit = int(key.learned_rpm_limit or self.DEFAULT_INITIAL_LIMIT)
# 获取已知边界（上次触发 429 时的并发数）
known_boundary = key.last_concurrent_peak
# 获取已知边界（上次触发 429 时的 RPM）
known_boundary = key.last_rpm_peak
# 计算当前利用率
utilization = float(current_concurrent / current_limit) if current_limit > 0 else 0.0
utilization = float(current_rpm / current_limit) if current_limit > 0 else 0.0
now = datetime.now(timezone.utc)
now_ts = now.timestamp()
@@ -229,7 +228,7 @@ class AdaptiveConcurrencyManager:
# 检查是否满足扩容条件
increase_reason = self._check_increase_conditions(key, samples, now, known_boundary)
if increase_reason and current_limit < self.MAX_CONCURRENT_LIMIT:
if increase_reason and current_limit < self.MAX_RPM_LIMIT:
old_limit = current_limit
is_probe = increase_reason == "probe_increase"
new_limit = self._increase_limit(current_limit, known_boundary, is_probe)
@@ -262,12 +261,12 @@ class AdaptiveConcurrencyManager:
avg_utilization=round(avg_util, 2),
high_util_ratio=round(high_util_ratio, 2),
sample_count=len(samples),
current_concurrent=current_concurrent,
current_rpm=current_rpm,
known_boundary=known_boundary,
)
# 更新限制
key.learned_max_concurrent = new_limit # type: ignore[assignment]
key.learned_rpm_limit = new_limit # type: ignore[assignment]
# 如果是探测性扩容，更新探测时间
if is_probe:
@@ -334,7 +333,7 @@ class AdaptiveConcurrencyManager:
key: API Key对象
samples: 利用率采样列表
now: 当前时间
known_boundary: 已知边界（触发 429 时的并发数）
known_boundary: 已知边界（触发 429 时的 RPM）
Returns:
扩容原因（如果满足条件），否则返回 None
@@ -343,7 +342,7 @@ class AdaptiveConcurrencyManager:
if self._is_in_cooldown(key):
return None
current_limit = int(key.learned_max_concurrent or self.DEFAULT_INITIAL_LIMIT)
current_limit = int(key.learned_rpm_limit or self.DEFAULT_INITIAL_LIMIT)
# 条件1：滑动窗口扩容（不超过边界）
if len(samples) >= self.MIN_SAMPLES_FOR_DECISION:
@@ -353,7 +352,7 @@ class AdaptiveConcurrencyManager:
if high_util_ratio >= self.HIGH_UTILIZATION_RATIO:
# 检查是否还有扩容空间（边界保护）
if known_boundary:
# 允许扩容到边界值（而非 boundary - 1），因为缩容时已经 -1 了
# 允许扩容到边界值（而非 boundary - 1），因为缩容时已经 -步长 了
if current_limit < known_boundary:
return "high_utilization"
# 已达边界，不触发普通扩容
@@ -429,34 +428,37 @@ class AdaptiveConcurrencyManager:
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
cooldown_seconds = RPMDefaults.COOLDOWN_AFTER_429_MINUTES * 60
return bool(time_since_429 < cooldown_seconds)
def _decrease_limit(
self,
current_limit: int,
current_concurrent: Optional[int] = None,
current_rpm: Optional[int] = None,
) -> int:
"""
减少并发限制（基于边界记忆策略）
减少 RPM 限制（基于边界记忆策略）
策略：
- 如果知道触发 429 时的并发数，新限制 = 并发数 - 1
- 这样可以快速收敛到真实限制附近，而不会过度保守
- 例如：真实限制 8，触发时并发 8 -> 新限制 7（而非 8*0.85=6）
- 如果知道触发 429 时的 RPM，新限制 = RPM * 0.90（保留 10% 安全边际）
- 10% 的安全边际更保守，考虑到：
1. RPM 报告可能存在延迟，实际触发时的 RPM 可能略高于报告值
2. 上游 API 的限制可能有波动
3. 避免频繁在边界附近触发 429
- 相比固定步长，百分比方式更适应不同量级的限制值
"""
if current_concurrent is not None and current_concurrent > 0:
# 边界记忆策略：新限制 = 触发边界 - 1
candidate = current_concurrent - 1
if current_rpm is not None and current_rpm > 0:
# 边界记忆策略：新限制 = 触发边界 * 0.90（10% 安全边际）
candidate = int(current_rpm * 0.90)
else:
# 没有并发信息时，保守减少 1
candidate = current_limit - 1
# 没有 RPM 信息时，减少 10%
candidate = int(current_limit * 0.9)
# 保证不会“缩容变扩容”（例如 current_concurrent > current_limit 的异常场景）
# 保证不会"缩容变扩容"
candidate = min(candidate, current_limit - 1)
new_limit = max(candidate, self.MIN_CONCURRENT_LIMIT)
new_limit = max(candidate, self.MIN_RPM_LIMIT)
return new_limit
@@ -467,16 +469,15 @@ class AdaptiveConcurrencyManager:
is_probe: bool = False,
) -> int:
"""
增加并发限制（考虑边界保护）
增加 RPM 限制（考虑边界保护）
策略：
- 普通扩容：每次 +INCREASE_STEP，但不超过 known_boundary
（因为缩容时已经 -1 了，这里允许回到边界值尝试）
- 探测性扩容：每次只 +1，可以突破边界，但要谨慎
Args:
current_limit: 当前限制
known_boundary: 已知边界（last_concurrent_peak），即触发 429 时的并发数
known_boundary: 已知边界（last_rpm_peak），即触发 429 时的 RPM
is_probe: 是否是探测性扩容（可以突破边界）
"""
if is_probe:
@@ -486,13 +487,13 @@ class AdaptiveConcurrencyManager:
# 普通模式：每次 +INCREASE_STEP
new_limit = current_limit + self.INCREASE_STEP
# 边界保护：普通扩容不超过 known_boundary（允许回到边界值尝试）
# 边界保护：普通扩容不超过 known_boundary
if known_boundary:
if new_limit > known_boundary:
new_limit = known_boundary
# 全局上限保护
new_limit = min(new_limit, self.MAX_CONCURRENT_LIMIT)
new_limit = min(new_limit, self.MAX_RPM_LIMIT)
# 确保有增长（否则返回原值表示不扩容）
if new_limit <= current_limit:
@@ -509,7 +510,7 @@ class AdaptiveConcurrencyManager:
**extra_data: Any,
) -> None:
"""
记录并发调整历史
记录 RPM 调整历史
Args:
key: API Key对象
@@ -548,10 +549,10 @@ class AdaptiveConcurrencyManager:
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
# rpm_limit=NULL 表示自适应，否则为固定限制
is_adaptive = key.rpm_limit is None
current_limit = int(key.learned_rpm_limit or self.DEFAULT_INITIAL_LIMIT)
effective_limit = current_limit if is_adaptive else int(key.rpm_limit) # type: ignore
# 计算窗口统计
avg_utilization: Optional[float] = None
@@ -570,15 +571,15 @@ class AdaptiveConcurrencyManager:
last_probe_at_str = cast(datetime, key.last_probe_increase_at).isoformat()
# 边界信息
known_boundary = key.last_concurrent_peak
known_boundary = key.last_rpm_peak
return {
"adaptive_mode": is_adaptive,
"max_concurrent": key.max_concurrent, # NULL=自适应，数字=固定限制
"rpm_limit": key.rpm_limit, # NULL=自适应，数字=固定限制
"effective_limit": effective_limit, # 当前有效限制
"learned_limit": key.learned_max_concurrent, # 学习到的限制
"learned_limit": key.learned_rpm_limit, # 学习到的限制
# 边界记忆相关
"known_boundary": known_boundary, # 触发 429 时的并发数（已知上限）
"known_boundary": known_boundary, # 触发 429 时的 RPM（已知上限）
"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,
@@ -607,12 +608,12 @@ class AdaptiveConcurrencyManager:
"""
logger.info(f"[RESET] 重置学习状态: Key {key.id[:8]}...")
key.learned_max_concurrent = None # type: ignore[assignment]
key.learned_rpm_limit = 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.last_rpm_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]
@@ -622,12 +623,17 @@ class AdaptiveConcurrencyManager:
# 全局单例
_adaptive_manager: Optional[AdaptiveConcurrencyManager] = None
_adaptive_rpm_manager: Optional[AdaptiveRPMManager] = None
def get_adaptive_manager() -> AdaptiveConcurrencyManager:
"""获取全局自适应管理器单例"""
global _adaptive_manager
if _adaptive_manager is None:
_adaptive_manager = AdaptiveConcurrencyManager()
return _adaptive_manager
def get_adaptive_rpm_manager() -> AdaptiveRPMManager:
"""获取全局自适应 RPM 管理器单例"""
global _adaptive_rpm_manager
if _adaptive_rpm_manager is None:
_adaptive_rpm_manager = AdaptiveRPMManager()
return _adaptive_rpm_manager
# 向后兼容别名
AdaptiveConcurrencyManager = AdaptiveRPMManager
get_adaptive_manager = get_adaptive_rpm_manager

617
src/services/rate_limit/concurrency_manager.py
View File

@@ -1,29 +1,33 @@
"""
并发管理器 - 支持 Redis 或内存的并发控制
RPM 限制管理器 - 支持 Redis 或内存的 Key 级别 RPM 限制
功能：
1. Endpoint 级别的并发限制
2. ProviderAPIKey 级别的并发限制
3. 分布式环境下优先使用 Redis，多实例共享
4. 在开发/单实例场景下自动降级为内存计数
5. 自动释放和异常处理（Redis 提供 TTL，内存模式请确保手动释放）
1. ProviderAPIKey 级别的 RPM 限制（按分钟窗口计数）
2. 分布式环境下优先使用 Redis，多实例共享
3. 在开发/单实例场景下自动降级为内存计数
4. 支持缓存用户优先级（预留槽位机制）
"""
import asyncio
import math
import os
import time
from contextlib import asynccontextmanager
from datetime import timedelta # noqa: F401 - kept for potential future use
from typing import Optional, Tuple
from typing import Optional
import redis.asyncio as aioredis
from src.config.constants import RPMDefaults
from src.core.logger import logger
class ConcurrencyManager:
"""分布式并发管理器"""
"""Key RPM 限制管理器"""
_instance: Optional["ConcurrencyManager"] = None
_redis: Optional[aioredis.Redis] = None
_key_rpm_bucket_seconds: int = 60
_key_rpm_key_ttl_seconds: int = 120 # 2 分钟，足够覆盖当前分钟与边界
def __new__(cls):
"""单例模式"""
@@ -37,9 +41,22 @@ class ConcurrencyManager:
return
self._memory_lock: asyncio.Lock = asyncio.Lock()
self._memory_endpoint_counts: dict[str, int] = {}
self._memory_key_counts: dict[str, int] = {}
# Key RPM 计数器：{key_id: (bucket, count)}，bucket = floor(now / 60)
self._memory_key_rpm_counts: dict[str, tuple[int, int]] = {}
self._owns_redis: bool = False
self._last_cleanup_bucket: int = 0 # 上次清理时的 bucket，用于定期清理过期数据
self._last_cleanup_time: float = 0 # 上次清理的时间戳，用于强制定期清理
self._cleanup_interval_seconds: int = 300 # 强制清理间隔（5 分钟）
self._cleanup_task: Optional[asyncio.Task] = None # 后台清理任务
# 内存模式下的最大条目限制，防止内存泄漏（支持环境变量覆盖）
self._max_memory_rpm_entries: int = int(
os.getenv("RPM_MAX_MEMORY_ENTRIES", str(RPMDefaults.MAX_MEMORY_RPM_ENTRIES))
)
# 早期告警阈值（达到此比例时记录警告）
self._memory_warning_threshold: float = float(
os.getenv("RPM_MEMORY_WARNING_THRESHOLD", str(RPMDefaults.MEMORY_WARNING_THRESHOLD))
)
self._memory_initialized = True
async def initialize(self) -> None:
@@ -56,220 +73,304 @@ class ConcurrencyManager:
if self._redis:
logger.info("[OK] ConcurrencyManager 已复用全局 Redis 客户端")
else:
logger.warning("[WARN] Redis 不可用，并发控制降级为内存模式（仅在单实例环境下安全）")
logger.warning("[WARN] Redis 不可用，RPM 限制降级为内存模式（仅在单实例环境下安全）")
# 内存模式下启动后台清理任务
self._start_background_cleanup()
except Exception as e:
logger.error(f"[ERROR] 获取全局 Redis 客户端失败: {e}")
logger.warning("[WARN] 并发控制将降级为内存模式（仅在单实例环境下安全）")
logger.warning("[WARN] RPM 限制将降级为内存模式（仅在单实例环境下安全）")
self._redis = None
self._owns_redis = False
# 内存模式下启动后台清理任务
self._start_background_cleanup()
def _start_background_cleanup(self) -> None:
"""启动后台定期清理任务（仅内存模式需要）"""
if self._cleanup_task is not None:
return # 已经启动
async def cleanup_loop():
"""后台清理循环"""
while True:
try:
await asyncio.sleep(60) # 每分钟检查一次
async with self._memory_lock:
current_bucket = self._get_rpm_bucket()
self._cleanup_expired_memory_rpm_counts(current_bucket, force=False)
except asyncio.CancelledError:
break
except Exception as e:
logger.debug(f"后台清理任务异常: {e}")
try:
self._cleanup_task = asyncio.create_task(cleanup_loop())
logger.debug("[OK] 内存模式后台清理任务已启动")
except RuntimeError:
# 没有事件循环时忽略
pass
async def close(self) -> None:
"""关闭 Redis 连接"""
# 停止后台清理任务
if self._cleanup_task is not None:
self._cleanup_task.cancel()
try:
await self._cleanup_task
except asyncio.CancelledError:
pass
self._cleanup_task = None
if self._redis and self._owns_redis:
await self._redis.close()
logger.info("ConcurrencyManager Redis 连接已关闭")
self._redis = None
self._owns_redis = False
def _get_endpoint_key(self, endpoint_id: str) -> str:
"""获取 Endpoint 并发计数的 Redis Key"""
return f"concurrency:endpoint:{endpoint_id}"
@classmethod
def _get_rpm_bucket(cls, now_ts: Optional[float] = None) -> int:
"""获取当前 RPM 计数桶（按分钟）"""
ts = now_ts if now_ts is not None else time.time()
return int(ts // cls._key_rpm_bucket_seconds)
def _get_key_key(self, key_id: str) -> str:
"""获取 ProviderAPIKey 并发计数的 Redis Key"""
return f"concurrency:key:{key_id}"
@classmethod
def _get_key_key(cls, key_id: str, bucket: Optional[int] = None) -> str:
"""获取 ProviderAPIKey RPM 计数的 Redis Key（按分钟桶）"""
b = bucket if bucket is not None else cls._get_rpm_bucket()
return f"rpm:key:{key_id}:{b}"
async def get_current_concurrency(
self, endpoint_id: Optional[str] = None, key_id: Optional[str] = None
) -> Tuple[int, int]:
def _get_memory_key_rpm_count(self, key_id: str, bucket: int) -> int:
"""获取内存模式下 Key 在指定 bucket 的 RPM 计数"""
stored = self._memory_key_rpm_counts.get(key_id)
if not stored:
return 0
stored_bucket, count = stored
if stored_bucket != bucket:
# 旧桶数据已过期，删除以防止内存泄漏
del self._memory_key_rpm_counts[key_id]
return 0
return count
def _set_memory_key_rpm_count(self, key_id: str, bucket: int, count: int) -> None:
"""设置内存模式下 Key 在指定 bucket 的 RPM 计数"""
current_size = len(self._memory_key_rpm_counts)
warning_threshold = int(self._max_memory_rpm_entries * self._memory_warning_threshold)
high_threshold = int(self._max_memory_rpm_entries * 0.8)
critical_threshold = int(self._max_memory_rpm_entries * 0.95)
# 分级告警：根据使用率记录不同级别的日志
if current_size >= critical_threshold and key_id not in self._memory_key_rpm_counts:
logger.critical(
f"[CRITICAL] 内存 RPM 计数器接近上限 ({current_size}/{self._max_memory_rpm_entries})，"
f"强烈建议启用 Redis！继续增长可能导致 RPM 限制失效"
)
elif current_size >= high_threshold and key_id not in self._memory_key_rpm_counts:
# 每 100 个条目告警一次，避免日志过多
if current_size % 100 == 0:
logger.error(
f"[ERROR] 内存 RPM 计数器使用率过高 ({current_size}/{self._max_memory_rpm_entries})，"
f"建议启用 Redis"
)
elif current_size >= warning_threshold and key_id not in self._memory_key_rpm_counts:
if current_size == warning_threshold:
logger.warning(
f"[WARN] 内存 RPM 计数器达到 {self._memory_warning_threshold:.0%} 阈值 "
f"({current_size}/{self._max_memory_rpm_entries})，建议启用 Redis"
)
# 检查是否超过最大条目限制
if (
key_id not in self._memory_key_rpm_counts
and current_size >= self._max_memory_rpm_entries
):
# 触发强制清理
self._cleanup_expired_memory_rpm_counts(bucket, force=True)
# 如果清理后仍然超过限制，执行 LRU 淘汰（删除最旧的 20%）
if len(self._memory_key_rpm_counts) >= self._max_memory_rpm_entries:
evict_count = max(1, self._max_memory_rpm_entries // 5)
# 按 bucket（时间）排序，删除最旧的
sorted_keys = sorted(
self._memory_key_rpm_counts.items(),
key=lambda x: x[1][0] # 按 bucket 排序
)
for k, _ in sorted_keys[:evict_count]:
del self._memory_key_rpm_counts[k]
logger.warning(
f"[WARN] 内存 RPM 计数器达到上限，已淘汰 {evict_count} 个最旧条目"
)
self._memory_key_rpm_counts[key_id] = (bucket, count)
def _cleanup_expired_memory_rpm_counts(self, current_bucket: int, force: bool = False) -> None:
"""
获取当前并发数
清理内存中过期的 RPM 计数（必须在持有 _memory_lock 时调用）
性能优化：使用 MGET 批量获取，减少 Redis 往返次数
清理策略：
- 常规清理：每分钟最多执行一次（当 bucket 变化时）
- 强制清理：每 5 分钟执行一次（防止长时间无请求导致内存泄漏）
"""
now = time.time()
# 检查是否需要清理
should_cleanup = (
current_bucket != self._last_cleanup_bucket # 分钟切换
or force # 强制清理
or (now - self._last_cleanup_time > self._cleanup_interval_seconds) # 超时清理
)
if not should_cleanup:
return
self._last_cleanup_bucket = current_bucket
self._last_cleanup_time = now
expired_keys = []
for key_id, (stored_bucket, _count) in self._memory_key_rpm_counts.items():
if stored_bucket < current_bucket:
expired_keys.append(key_id)
for key_id in expired_keys:
del self._memory_key_rpm_counts[key_id]
if expired_keys:
logger.debug(f"[CLEANUP] 清理了 {len(expired_keys)} 个过期的内存 RPM 计数")
async def get_key_rpm_count(self, key_id: str) -> int:
"""
获取 Key 当前 RPM 计数
Args:
endpoint_id: Endpoint ID（可选）
key_id: ProviderAPIKey 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
bucket = self._get_rpm_bucket()
# 定期清理过期数据，避免内存泄漏
self._cleanup_expired_memory_rpm_counts(bucket)
return self._get_memory_key_rpm_count(key_id, bucket)
try:
# 使用 MGET 批量获取，减少 Redis 往返（2 次 GET -> 1 次 MGET）
keys_to_fetch = []
if endpoint_id:
keys_to_fetch.append(self._get_endpoint_key(endpoint_id))
if key_id:
keys_to_fetch.append(self._get_key_key(key_id))
if keys_to_fetch:
results = await self._redis.mget(keys_to_fetch)
idx = 0
if endpoint_id:
endpoint_count = int(results[idx]) if results[idx] else 0
idx += 1
if key_id:
key_count = int(results[idx]) if results[idx] else 0
key_key = self._get_key_key(key_id)
result = await self._redis.get(key_key)
return int(result) if result else 0
except Exception as e:
logger.error(f"获取并发数失败: {e}")
logger.error(f"获取 RPM 计数失败: {e}")
return 0
return endpoint_count, key_count
async def check_available(
async def check_rpm_available(
self,
endpoint_id: str,
endpoint_max_concurrent: Optional[int],
key_id: str,
key_max_concurrent: Optional[int],
key_rpm_limit: Optional[int],
is_cached_user: bool = False,
cache_reservation_ratio: Optional[float] = None,
) -> bool:
"""
检查是否可以获取并发槽位（不实际获取）
检查是否可以通过 RPM 限制（不实际增加计数）
Args:
endpoint_id: Endpoint ID
endpoint_max_concurrent: Endpoint 最大并发数（None 表示不限制）
key_id: ProviderAPIKey ID
key_max_concurrent: Key 最大并发数（None 表示不限制）
key_rpm_limit: Key RPM 限制（每分钟最大请求数，None 表示不限制）
is_cached_user: 是否是缓存用户
cache_reservation_ratio: 缓存预留比例
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 key_rpm_limit is None:
return True
if (
endpoint_max_concurrent is not None
and endpoint_count >= endpoint_max_concurrent
):
return False
# 从配置读取默认值
from src.config.settings import config
if key_max_concurrent is not None and key_count >= key_max_concurrent:
return False
if cache_reservation_ratio is None:
cache_reservation_ratio = config.cache_reservation_ratio
return True
key_count = await self.get_key_rpm_count(key_id)
endpoint_count, key_count = await self.get_current_concurrency(endpoint_id, key_id)
if is_cached_user:
return key_count < key_rpm_limit
else:
# 新用户只能使用 (1 - cache_reservation_ratio) 的槽位
available_for_new = max(1, math.floor(key_rpm_limit * (1 - cache_reservation_ratio)))
return key_count < available_for_new
# 检查 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(
async def acquire_rpm_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: Optional[float] = None, # 缓存预留比例，None 时从配置读取
ttl_seconds: Optional[int] = None, # TTL 秒数，None 时从配置读取
key_rpm_limit: Optional[int],
is_cached_user: bool = False,
cache_reservation_ratio: Optional[float] = None,
) -> bool:
"""
尝试获取并发槽位（支持缓存用户优先级）
尝试获取 RPM 槽位（支持缓存用户优先级）
Args:
endpoint_id: Endpoint ID
endpoint_max_concurrent: Endpoint 最大并发数（None 表示不限制）
key_id: ProviderAPIKey ID
key_max_concurrent: Key 最大并发数（None 表示不限制）
key_rpm_limit: Key RPM 限制（每分钟最大请求数，None 表示不限制）
is_cached_user: 是否是缓存用户（缓存用户可使用全部槽位）
cache_reservation_ratio: 缓存预留比例，None 时从配置读取
ttl_seconds: TTL 秒数，None 时从配置读取
Returns:
是否成功获取（True/False）
缓存预留机制说明:
- 假设 key_max_concurrent = 10, cache_reservation_ratio = 0.3
- 新用户最多使用: 7个槽位 (10 * (1 - 0.3))
- 缓存用户最多使用: 10个槽位（全部）
- 预留的3个槽位专门给缓存用户，保证他们的请求优先
- 假设 key_rpm_limit = 100, cache_reservation_ratio = 0.3
- 新用户最多使用: 70 RPM (100 * (1 - 0.3))
- 缓存用户最多使用: 100 RPM（全部）
- 预留的 30 RPM 专门给缓存用户，保证他们的请求优先
"""
# 从配置读取默认值
from src.config.settings import config
if cache_reservation_ratio is None:
cache_reservation_ratio = config.cache_reservation_ratio
if ttl_seconds is None:
ttl_seconds = config.concurrency_slot_ttl
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)
bucket = self._get_rpm_bucket()
# 定期清理过期数据，避免内存泄漏
self._cleanup_expired_memory_rpm_counts(bucket)
# Endpoint 限制
if (
endpoint_max_concurrent is not None
and endpoint_count >= endpoint_max_concurrent
):
return False
key_count = self._get_memory_key_rpm_count(key_id, bucket)
# Key 限制，包含缓存预留
if key_max_concurrent is not None:
# Key RPM 限制，包含缓存预留
if key_rpm_limit is not None:
if is_cached_user:
if key_count >= key_max_concurrent:
if key_count >= key_rpm_limit:
return False
else:
# 新用户只能使用 (1 - cache_reservation_ratio) 的槽位
available_for_new = max(
1, math.ceil(key_max_concurrent * (1 - cache_reservation_ratio))
1, math.floor(key_rpm_limit * (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
self._set_memory_key_rpm_count(key_id, bucket, key_count + 1)
return True
endpoint_key = self._get_endpoint_key(endpoint_id)
key_key = self._get_key_key(key_id)
bucket = self._get_rpm_bucket()
key_key = self._get_key_key(key_id, bucket=bucket)
try:
# 使用 Lua 脚本保证原子性（新增缓存预留逻辑）
# 使用 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 key_key = KEYS[1]
local key_max = tonumber(ARGV[1])
local key_ttl = tonumber(ARGV[2])
local is_cached = tonumber(ARGV[3]) -- 0=新用户, 1=缓存用户
local cache_ratio = tonumber(ARGV[4]) -- 缓存预留比例
-- 获取当前值
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))
local available_for_new = math.max(1, math.floor(key_max * (1 - cache_ratio)))
if key_count >= available_for_new then
return 0 -- 失败：新用户配额已满
end
@@ -282,10 +383,8 @@ class ConcurrencyManager:
end
-- 增加计数
redis.call('INCR', endpoint_key)
redis.call('EXPIRE', endpoint_key, ttl)
redis.call('INCR', key_key)
redis.call('EXPIRE', key_key, ttl)
redis.call('EXPIRE', key_key, key_ttl)
return 1 -- 成功
"""
@@ -293,12 +392,10 @@ class ConcurrencyManager:
# 执行脚本
result = await self._redis.eval(
lua_script,
2, # 2 个 KEYS
endpoint_key,
1, # 1 个 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,
key_rpm_limit if key_rpm_limit is not None else -1,
self._key_rpm_key_ttl_seconds,
1 if is_cached_user else 0, # 缓存用户标志
cache_reservation_ratio, # 预留比例
)
@@ -307,143 +404,68 @@ class ConcurrencyManager:
if success:
user_type = "缓存用户" if is_cached_user else "新用户"
logger.debug(
f"[OK] 获取并发槽位成功: endpoint={endpoint_id}, key={key_id}, "
f"类型={user_type}"
)
logger.debug(f"[OK] 获取 RPM 槽位成功: key={key_id}, 类型={user_type}")
else:
endpoint_count, key_count = await self.get_current_concurrency(endpoint_id, key_id)
key_count = await self.get_key_rpm_count(key_id)
# 计算新用户可用槽位
if key_max_concurrent and not is_cached_user:
available_for_new = int(key_max_concurrent * (1 - cache_reservation_ratio))
# 计算新用户可用 RPM
if key_rpm_limit and not is_cached_user:
available_for_new = int(key_rpm_limit * (1 - cache_reservation_ratio))
user_info = f"新用户配额={available_for_new}, 当前={key_count}"
else:
user_info = f"缓存用户, 当前={key_count}/{key_max_concurrent}"
user_info = f"缓存用户, 当前={key_count}/{key_rpm_limit}"
logger.warning(
f"[WARN] 并发槽位已满: endpoint={endpoint_id}({endpoint_count}/{endpoint_max_concurrent}), "
f"key={key_id}({user_info})"
)
logger.warning(f"[WARN] RPM 限制已达上限: key={key_id}({user_info})")
return success
except Exception as e:
logger.error(f"获取并发槽位失败，降级到内存模式: {e}")
logger.error(f"获取 RPM 槽位失败，降级到内存模式: {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)
bucket = self._get_rpm_bucket()
self._cleanup_expired_memory_rpm_counts(bucket)
key_count = self._get_memory_key_rpm_count(key_id, bucket)
# 降级模式下使用更保守的限制（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
fallback_rpm_limit = (
max(1, key_rpm_limit // 2) if key_rpm_limit is not None else None
)
if (
fallback_endpoint_limit is not None
and endpoint_count >= fallback_endpoint_limit
):
if fallback_rpm_limit is not None and key_count >= fallback_rpm_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}"
f"[FALLBACK] Key RPM 达到降级限制: {key_count}/{fallback_rpm_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}"
)
self._set_memory_key_rpm_count(key_id, bucket, key_count + 1)
logger.debug(f"[FALLBACK] 使用内存模式获取 RPM 槽位: 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(
async def rpm_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: Optional[float] = None, # 缓存预留比例，None 时从配置读取
key_rpm_limit: Optional[int],
is_cached_user: bool = False,
cache_reservation_ratio: Optional[float] = None,
):
"""
并发控制上下文管理器（支持缓存用户优先级）
RPM 限制上下文管理器（支持缓存用户优先级）
用法：
async with manager.concurrency_guard(
endpoint_id, endpoint_max, key_id, key_max,
async with manager.rpm_guard(
key_id, key_rpm_limit,
is_cached_user=True # 缓存用户
):
# 执行请求
response = await send_request(...)
如果获取失败，会抛出 ConcurrencyLimitError 异常
注意：RPM 是按分钟窗口计数，不需要在请求结束后释放
"""
# 从配置读取默认值
from src.config.settings import config
@@ -452,11 +474,9 @@ class ConcurrencyManager:
cache_reservation_ratio = config.cache_reservation_ratio
# 尝试获取槽位（传递缓存用户参数）
acquired = await self.acquire_slot(
endpoint_id,
endpoint_max_concurrent,
acquired = await self.acquire_rpm_slot(
key_id,
key_max_concurrent,
key_rpm_limit,
is_cached_user,
cache_reservation_ratio,
)
@@ -466,7 +486,7 @@ class ConcurrencyManager:
user_type = "缓存用户" if is_cached_user else "新用户"
raise ConcurrencyLimitError(
f"并发限制已达上限: endpoint={endpoint_id}, key={key_id}, 类型={user_type}"
f"RPM 限制已达上限: key={key_id}, 类型={user_type}"
)
# 记录开始时间和状态
@@ -477,7 +497,7 @@ class ConcurrencyManager:
try:
yield # 执行请求
except Exception as e:
except Exception:
# 记录异常
exception_occurred = True
raise
@@ -506,7 +526,7 @@ class ConcurrencyManager:
# 告警：槽位占用时间过长（超过 60 秒）
if slot_duration > 60:
logger.warning(
f"[WARN] 并发槽位占用时间过长: "
f"[WARN] 请求耗时过长: "
f"key_id={key_id[:8] if key_id else 'unknown'}..., "
f"duration={slot_duration:.1f}s, "
f"exception={exception_occurred}"
@@ -514,67 +534,64 @@ class ConcurrencyManager:
except Exception as metric_error:
# 指标记录失败不应影响业务逻辑
logger.debug(f"记录并发指标失败: {metric_error}")
logger.debug(f"记录指标失败: {metric_error}")
# 自动释放槽位（即使发生异常）
await self.release_slot(endpoint_id, key_id)
# 注意：RPM 计数不需要在请求结束后释放，它会在分钟窗口过期后自动重置
async def reset_concurrency(
self, endpoint_id: Optional[str] = None, key_id: Optional[str] = None
) -> None:
async def reset_key_rpm(self, key_id: str) -> None:
"""
重置并发计数（管理功能，慎用）
重置 Key RPM 计数（管理功能，慎用）
Args:
endpoint_id: Endpoint ID（可选，None 表示重置所有 endpoint）
key_id: ProviderAPIKey ID（可选，None 表示重置所有 key）
key_id: ProviderAPIKey ID
"""
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} 个")
self._memory_key_rpm_counts.pop(key_id, None)
logger.info(f"[RESET] 重置 Key RPM 计数(内存): {key_id}")
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)} 个")
deleted_count = await self._scan_and_delete(f"rpm:key:{key_id}:*")
logger.info(f"[RESET] 重置 Key RPM 计数: {key_id}, 删除 {deleted_count} 个键")
except Exception as e:
logger.error(f"重置并发计数失败: {e}")
logger.error(f"重置 Key RPM 计数失败: {e}")
async def reset_all_rpm(self) -> None:
"""重置所有 Key RPM 计数（管理功能，慎用）"""
if self._redis is None:
async with self._memory_lock:
count = len(self._memory_key_rpm_counts)
self._memory_key_rpm_counts.clear()
if count:
logger.info(f"[RESET] 重置所有 Key RPM 计数(内存): {count} 个")
return
try:
deleted_count = await self._scan_and_delete("rpm:key:*")
if deleted_count:
logger.info(f"[RESET] 重置所有 Key RPM 计数: {deleted_count} 个")
except Exception as e:
logger.error(f"重置所有 Key RPM 计数失败: {e}")
async def _scan_and_delete(self, pattern: str, batch_size: int = 100) -> int:
"""使用 SCAN 遍历并分批删除匹配的键，避免阻塞 Redis"""
if self._redis is None:
return 0
deleted_count = 0
cursor = 0
while True:
cursor, keys = await self._redis.scan(cursor, match=pattern, count=batch_size)
if keys:
# 分批删除，每批最多 batch_size 个
for i in range(0, len(keys), batch_size):
batch = keys[i : i + batch_size]
await self._redis.delete(*batch)
deleted_count += len(batch)
if cursor == 0:
break
return deleted_count
# 全局单例

190
src/services/rate_limit/detector.py
View File

@@ -3,7 +3,7 @@
"""
from datetime import datetime, timezone
from typing import Any, Dict, Optional, Tuple
from typing import Dict, Optional
from src.core.logger import logger
@@ -62,7 +62,7 @@ class RateLimitDetector:
def detect_from_headers(
headers: Dict[str, str],
provider_name: str = "unknown",
current_concurrent: Optional[int] = None,
current_usage: Optional[int] = None,
) -> RateLimitInfo:
"""
从响应头中检测速率限制类型
@@ -70,7 +70,7 @@ class RateLimitDetector:
Args:
headers: 429响应的HTTP头
provider_name: 提供商名称（用于选择解析策略）
current_concurrent: 当前并发数（用于判断是否为并发限制）
current_usage: 当前使用量（RPM 计数，用于启发式判断是否为并发限制）
Returns:
RateLimitInfo对象
@@ -80,16 +80,16 @@ class RateLimitDetector:
# 根据提供商选择解析策略
if "anthropic" in provider_name.lower() or "claude" in provider_name.lower():
return RateLimitDetector._parse_anthropic_headers(headers_lower, current_concurrent)
return RateLimitDetector._parse_anthropic_headers(headers_lower, current_usage)
elif "openai" in provider_name.lower():
return RateLimitDetector._parse_openai_headers(headers_lower, current_concurrent)
return RateLimitDetector._parse_openai_headers(headers_lower, current_usage)
else:
return RateLimitDetector._parse_generic_headers(headers_lower, current_concurrent)
return RateLimitDetector._parse_generic_headers(headers_lower, current_usage)
@staticmethod
def _parse_anthropic_headers(
headers: Dict[str, str],
current_concurrent: Optional[int] = None,
current_usage: Optional[int] = None,
) -> RateLimitInfo:
"""
解析 Anthropic Claude API 的速率限制头
@@ -127,29 +127,66 @@ class RateLimitDetector:
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) # 短暂等待
)
# 2. 并发限制判断（多条件策略）
# 注意：current_usage 是 RPM 计数（当前分钟请求数），不是真正的并发数
#
# 判断条件（满足任一即可）：
# A. 强判断：remaining > 0 且 retry_after <= 30（Provider 明确告知还有配额但需要等待）
# B. 弱判断：只有 retry_after <= 5 且缺少 remaining 头（短等待时间是并发限制的典型特征）
#
# 选择保守的 retry_after 阈值：
# - 强判断用 30 秒（有 remaining 头时）
# - 弱判断用 5 秒（无 remaining 头时，更保守）
is_likely_concurrent = False
concurrent_reason = ""
# 条件 A：remaining > 0 且 retry_after <= 30
if (
requests_remaining is not None
and requests_remaining > 0
and retry_after is not None
and retry_after <= 30
):
is_likely_concurrent = True
concurrent_reason = f"remaining={requests_remaining} > 0, retry_after={retry_after}s <= 30s"
# 条件 B：无 remaining 头但 retry_after 很短（<= 5 秒）
elif (
requests_remaining is None
and retry_after is not None
and retry_after <= 5
):
is_likely_concurrent = True
concurrent_reason = f"no remaining header, retry_after={retry_after}s <= 5s"
if is_likely_concurrent:
logger.info(
f"检测到可能的并发限制: current_concurrent={current_concurrent}, "
f"remaining={requests_remaining}, retry_after={retry_after}"
)
logger.info(f"检测到并发限制: {concurrent_reason}")
return RateLimitInfo(
limit_type=RateLimitType.CONCURRENT,
retry_after=retry_after,
current_usage=current_concurrent,
current_usage=current_usage,
raw_headers=headers,
)
# 3. 未知类型
# 3. 默认视为 RPM 限制（更保守的处理）
# 无法明确区分时，视为 RPM 限制让系统降低 RPM
# 这比误判为并发限制（不降 RPM）更安全
if retry_after is not None or requests_limit is not None:
logger.info(
f"无法明确区分限制类型，保守视为 RPM 限制: "
f"remaining={requests_remaining}, retry_after={retry_after}"
)
return RateLimitInfo(
limit_type=RateLimitType.RPM,
retry_after=retry_after,
limit_value=requests_limit,
remaining=requests_remaining,
reset_at=requests_reset,
current_usage=current_usage,
raw_headers=headers,
)
# 4. 完全没有信息，标记为未知
return RateLimitInfo(
limit_type=RateLimitType.UNKNOWN,
retry_after=retry_after,
@@ -159,7 +196,7 @@ class RateLimitDetector:
@staticmethod
def _parse_openai_headers(
headers: Dict[str, str],
current_concurrent: Optional[int] = None,
current_usage: Optional[int] = None,
) -> RateLimitInfo:
"""
解析 OpenAI API 的速率限制头
@@ -195,23 +232,55 @@ class RateLimitDetector:
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)
)
# 2. 并发限制判断（多条件策略）
# 判断条件（满足任一即可）：
# A. 强判断：remaining > 0 且 retry_after <= 30
# B. 弱判断：只有 retry_after <= 5 且缺少 remaining 头
is_likely_concurrent = False
concurrent_reason = ""
if (
requests_remaining is not None
and requests_remaining > 0
and retry_after is not None
and retry_after <= 30
):
is_likely_concurrent = True
concurrent_reason = f"remaining={requests_remaining} > 0, retry_after={retry_after}s <= 30s"
elif (
requests_remaining is None
and retry_after is not None
and retry_after <= 5
):
is_likely_concurrent = True
concurrent_reason = f"no remaining header, retry_after={retry_after}s <= 5s"
if is_likely_concurrent:
logger.info(f"检测到并发限制: {concurrent_reason}")
return RateLimitInfo(
limit_type=RateLimitType.CONCURRENT,
retry_after=retry_after,
current_usage=current_concurrent,
current_usage=current_usage,
raw_headers=headers,
)
# 3. 未知类型
# 3. 默认视为 RPM 限制（更保守的处理）
if retry_after is not None or requests_limit is not None:
logger.info(
f"无法明确区分限制类型，保守视为 RPM 限制: "
f"remaining={requests_remaining}, retry_after={retry_after}"
)
return RateLimitInfo(
limit_type=RateLimitType.RPM,
retry_after=retry_after,
limit_value=requests_limit,
remaining=requests_remaining,
reset_at=requests_reset,
current_usage=current_usage,
raw_headers=headers,
)
# 4. 完全没有信息，标记为未知
return RateLimitInfo(
limit_type=RateLimitType.UNKNOWN,
retry_after=retry_after,
@@ -221,7 +290,7 @@ class RateLimitDetector:
@staticmethod
def _parse_generic_headers(
headers: Dict[str, str],
current_concurrent: Optional[int] = None,
current_usage: Optional[int] = None,
) -> RateLimitInfo:
"""
解析通用的速率限制头
@@ -247,23 +316,54 @@ class RateLimitDetector:
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)
)
# 2. 并发限制判断（多条件策略）
# 判断条件（满足任一即可）：
# A. 强判断：remaining > 0 且 retry_after <= 30
# B. 弱判断：只有 retry_after <= 5 且缺少 remaining 头
is_likely_concurrent = False
concurrent_reason = ""
if (
remaining is not None
and remaining > 0
and retry_after is not None
and retry_after <= 30
):
is_likely_concurrent = True
concurrent_reason = f"remaining={remaining} > 0, retry_after={retry_after}s <= 30s"
elif (
remaining is None
and retry_after is not None
and retry_after <= 5
):
is_likely_concurrent = True
concurrent_reason = f"no remaining header, retry_after={retry_after}s <= 5s"
if is_likely_concurrent:
logger.info(f"检测到并发限制: {concurrent_reason}")
return RateLimitInfo(
limit_type=RateLimitType.CONCURRENT,
retry_after=retry_after,
current_usage=current_concurrent,
current_usage=current_usage,
raw_headers=headers,
)
# 3. 未知类型
# 3. 默认视为 RPM 限制（更保守的处理）
if retry_after is not None or limit_value is not None:
logger.info(
f"无法明确区分限制类型，保守视为 RPM 限制: "
f"remaining={remaining}, retry_after={retry_after}"
)
return RateLimitInfo(
limit_type=RateLimitType.RPM,
retry_after=retry_after,
limit_value=limit_value,
remaining=remaining,
current_usage=current_usage,
raw_headers=headers,
)
# 4. 完全没有信息，标记为未知
return RateLimitInfo(
limit_type=RateLimitType.UNKNOWN,
retry_after=retry_after,
@@ -317,7 +417,7 @@ class RateLimitDetector:
def detect_rate_limit_type(
headers: Dict[str, str],
provider_name: str = "unknown",
current_concurrent: Optional[int] = None,
current_usage: Optional[int] = None,
) -> RateLimitInfo:
"""
检测速率限制类型（便捷函数）
@@ -325,9 +425,9 @@ def detect_rate_limit_type(
Args:
headers: 429响应头
provider_name: 提供商名称
current_concurrent: 当前并发数
current_usage: 当前使用量（RPM 计数）
Returns:
RateLimitInfo对象
"""
return RateLimitDetector.detect_from_headers(headers, provider_name, current_concurrent)
return RateLimitDetector.detect_from_headers(headers, provider_name, current_usage)

135
src/services/rate_limit/rpm_limiter.py
View File

@@ -1,135 +0,0 @@
"""
RPM (Requests Per Minute) 限流服务
"""
import time
from datetime import datetime, timedelta, 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 + timedelta(minutes=1)
# 查找或创建追踪记录
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}")

52
src/services/request/executor.py
View File

@@ -89,24 +89,29 @@ class RequestExecutor:
try:
# 计算动态预留比例
reservation_manager = get_adaptive_reservation_manager()
# 获取当前并发数用于计算负载
# 获取当前 RPM 计数用于计算负载
# 注意：key 侧返回的是 RPM 计数（不会在请求结束时减少，靠 TTL 过期）
try:
_, current_key_concurrent = await self.concurrency_manager.get_current_concurrency(
_, current_key_rpm = 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
logger.debug(f"获取 RPM 计数失败（用于预留计算）: {e}")
current_key_rpm = 0
# 获取有效的并发限制（自适应或固定）
effective_key_limit = (
key.learned_max_concurrent if key.max_concurrent is None else key.max_concurrent
)
# 获取有效的 RPM 限制（自适应或固定）
if key.rpm_limit is None:
# 自适应模式：优先使用学习值，否则使用默认初始限制，避免无限制打爆上游
from src.config.constants import RPMDefaults
effective_key_limit = int(key.learned_rpm_limit or RPMDefaults.INITIAL_LIMIT)
else:
effective_key_limit = int(key.rpm_limit)
reservation_result = reservation_manager.calculate_reservation(
key=key,
current_concurrent=current_key_concurrent,
current_usage=current_key_rpm,
effective_limit=effective_key_limit,
)
dynamic_reservation_ratio = reservation_result.ratio
@@ -115,24 +120,22 @@ class RequestExecutor:
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,
async with self.concurrency_manager.rpm_guard(
key_id=key.id,
key_max_concurrent=effective_key_limit,
key_rpm_limit=effective_key_limit,
is_cached_user=is_cached_user,
cache_reservation_ratio=dynamic_reservation_ratio,
):
# 获取当前 RPM 计数（guard 内再次获取以获得最新值）
try:
_, key_concurrent = await self.concurrency_manager.get_current_concurrency(
endpoint_id=endpoint.id,
key_rpm_count = await self.concurrency_manager.get_key_rpm_count(
key_id=key.id,
)
except Exception as e:
logger.debug(f"获取并发数失败（guard 内）: {e}")
key_concurrent = None
logger.debug(f"获取 RPM 计数失败（guard 内）: {e}")
key_rpm_count = None
context.concurrent_requests = key_concurrent
context.concurrent_requests = key_rpm_count # 用于记录，实际是 RPM 计数
context.start_time = time.time()
response = await request_func(provider, endpoint, key)
@@ -142,15 +145,18 @@ class RequestExecutor:
health_monitor.record_success(
db=self.db,
key_id=key.id,
api_format=(
api_format.value if isinstance(api_format, APIFormat) else api_format
),
response_time_ms=context.elapsed_ms,
)
# 自适应模式：max_concurrent = NULL
if key.max_concurrent is None and key_concurrent is not None:
# 自适应模式：rpm_limit = NULL
if key.rpm_limit is None and key_rpm_count is not None:
self.adaptive_manager.handle_success(
db=self.db,
key=key,
current_concurrent=key_concurrent,
current_rpm=key_rpm_count,
)
# 根据是否为流式请求，标记不同状态
@@ -162,7 +168,7 @@ class RequestExecutor:
db=self.db,
candidate_id=candidate_id,
status_code=200,
concurrent_requests=key_concurrent,
concurrent_requests=key_rpm_count,
)
else:
# 非流式请求：标记为 success 状态
@@ -171,7 +177,7 @@ class RequestExecutor:
candidate_id=candidate_id,
status_code=200,
latency_ms=context.elapsed_ms,
concurrent_requests=key_concurrent,
concurrent_requests=key_rpm_count,
extra_data={
"is_cached_user": is_cached_user,
"model_name": model_name,

8
src/services/request/result.py
View File

@@ -289,10 +289,10 @@ class RequestResult:
status_code = 500
error_type = "internal_error"
# 构建错误消息：优先使用上游响应作为主要错误信息
if isinstance(exception, ProviderNotAvailableException) and exception.upstream_response:
error_message = exception.upstream_response
else:
# 构建错误消息：优先使用友好的 message 属性
# upstream_response 仅用于调试/链路追踪，不作为客户端错误消息
error_message = getattr(exception, "message", None)
if not error_message or not isinstance(error_message, str):
error_message = str(exception)
return cls(

6
src/services/usage/quota_scheduler.py
View File

@@ -96,8 +96,6 @@ class QuotaScheduler:
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
@@ -126,8 +124,6 @@ class QuotaScheduler:
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}")
@@ -140,8 +136,6 @@ class QuotaScheduler:
)
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")

7
src/services/usage/recorder.py
View File

@@ -93,6 +93,10 @@ class UsageRecorder:
if metadata.original_model and metadata.original_model != metadata.model:
target_model = metadata.model
# 非流式成功时，返回给客户端的是提供商响应头（透传）+ content-type
client_response_headers = dict(metadata.provider_response_headers) if metadata.provider_response_headers else {}
client_response_headers["content-type"] = "application/json"
await UsageService.record_usage(
db=self.db,
user=self.user,
@@ -115,6 +119,7 @@ class UsageRecorder:
request_body=request_body or result.request_body,
provider_request_headers=metadata.provider_request_headers,
response_headers=metadata.provider_response_headers,
client_response_headers=client_response_headers,
response_body=result.response_data if isinstance(result.response_data, dict) else {},
request_id=self.request_id,
provider_id=metadata.provider_id,
@@ -181,6 +186,8 @@ class UsageRecorder:
request_body=request_body or result.request_body,
provider_request_headers=metadata.provider_request_headers,
response_headers={},
# 失败请求返回给客户端的是 JSON 错误响应
client_response_headers={"content-type": "application/json"},
response_body={"error": result.error_message} if result.error_message else {},
request_id=self.request_id,
provider_id=metadata.provider_id,

27
src/services/usage/service.py
View File

@@ -40,6 +40,7 @@ class UsageRecordParams:
request_body: Optional[Any]
provider_request_headers: Optional[Dict[str, Any]]
response_headers: Optional[Dict[str, Any]]
client_response_headers: Optional[Dict[str, Any]]
response_body: Optional[Any]
request_id: str
provider_id: Optional[str]
@@ -223,6 +224,7 @@ class UsageService:
request_body: Optional[Any],
provider_request_headers: Optional[Dict[str, Any]],
response_headers: Optional[Dict[str, Any]],
client_response_headers: Optional[Dict[str, Any]],
response_body: Optional[Any],
request_id: str,
provider_id: Optional[str],
@@ -288,6 +290,13 @@ class UsageService:
db, response_headers
)
# 处理返回给客户端的响应头
processed_client_response_headers = None
if should_log_headers and client_response_headers:
processed_client_response_headers = SystemConfigService.mask_sensitive_headers(
db, client_response_headers
)
# 计算真实成本（表面成本 * 倍率），免费套餐实际费用为 0
if is_free_tier:
actual_input_cost = 0.0
@@ -351,6 +360,7 @@ class UsageService:
"request_body": processed_request_body,
"provider_request_headers": processed_provider_request_headers,
"response_headers": processed_response_headers,
"client_response_headers": processed_client_response_headers,
"response_body": processed_response_body,
}
@@ -360,12 +370,13 @@ class UsageService:
db: Session,
provider_api_key_id: Optional[str],
provider_id: Optional[str],
api_format: Optional[str] = None,
) -> Tuple[float, bool]:
"""获取费率倍数和是否免费套餐（使用缓存）"""
from src.services.cache.provider_cache import ProviderCacheService
return await ProviderCacheService.get_rate_multiplier_and_free_tier(
db, provider_api_key_id, provider_id
db, provider_api_key_id, provider_id, api_format
)
@classmethod
@@ -484,6 +495,7 @@ class UsageService:
existing_usage.provider_request_headers = usage_params["provider_request_headers"]
existing_usage.response_body = usage_params["response_body"]
existing_usage.response_headers = usage_params["response_headers"]
existing_usage.client_response_headers = usage_params["client_response_headers"]
# 更新 token 和费用信息
existing_usage.input_tokens = usage_params["input_tokens"]
@@ -656,9 +668,9 @@ class UsageService:
Returns:
(usage_params 字典, total_cost 总成本)
"""
# 获取费率倍数和是否免费套餐
# 获取费率倍数和是否免费套餐（传递 api_format 支持按格式配置的倍率）
actual_rate_multiplier, is_free_tier = await cls._get_rate_multiplier_and_free_tier(
params.db, params.provider_api_key_id, params.provider_id
params.db, params.provider_api_key_id, params.provider_id, params.api_format
)
# 计算成本
@@ -704,6 +716,7 @@ class UsageService:
request_body=params.request_body,
provider_request_headers=params.provider_request_headers,
response_headers=params.response_headers,
client_response_headers=params.client_response_headers,
response_body=params.response_body,
request_id=params.request_id,
provider_id=params.provider_id,
@@ -753,6 +766,7 @@ class UsageService:
request_body: Optional[Any] = None,
provider_request_headers: Optional[Dict[str, Any]] = None,
response_headers: Optional[Dict[str, Any]] = None,
client_response_headers: Optional[Dict[str, Any]] = None,
response_body: Optional[Any] = None,
request_id: Optional[str] = None,
provider_id: Optional[str] = None,
@@ -785,7 +799,8 @@ class UsageService:
status_code=status_code, error_message=error_message, metadata=metadata,
request_headers=request_headers, request_body=request_body,
provider_request_headers=provider_request_headers,
response_headers=response_headers, response_body=response_body,
response_headers=response_headers, client_response_headers=client_response_headers,
response_body=response_body,
request_id=request_id, provider_id=provider_id,
provider_endpoint_id=provider_endpoint_id,
provider_api_key_id=provider_api_key_id, status=status,
@@ -844,6 +859,7 @@ class UsageService:
request_body: Optional[Any] = None,
provider_request_headers: Optional[Dict[str, Any]] = None,
response_headers: Optional[Dict[str, Any]] = None,
client_response_headers: Optional[Dict[str, Any]] = None,
response_body: Optional[Any] = None,
request_id: Optional[str] = None,
provider_id: Optional[str] = None,
@@ -878,7 +894,8 @@ class UsageService:
status_code=status_code, error_message=error_message, metadata=metadata,
request_headers=request_headers, request_body=request_body,
provider_request_headers=provider_request_headers,
response_headers=response_headers, response_body=response_body,
response_headers=response_headers, client_response_headers=client_response_headers,
response_body=response_body,
request_id=request_id, provider_id=provider_id,
provider_endpoint_id=provider_endpoint_id,
provider_api_key_id=provider_api_key_id, status=status,

4
src/services/user/service.py
View File

@@ -398,8 +398,8 @@ class UserService:
def get_user_available_models(db: Session, user: User) -> List[Model]:
"""获取用户可用的模型
新架构：通过 GlobalModel + Model 关联查询用户可用模型
逻辑：用户可用提供商 → Provider 的 Model 实现 → 关联的 GlobalModel
通过 GlobalModel + Model 关联查询用户可用模型
逻辑：用户可用提供商 -> Provider 的 Model 实现 -> 关联的 GlobalModel
"""
# 获取用户可用的提供商
if user.role == UserRole.ADMIN: