Files
Aether/_deprecated_py_src/api/admin/adaptive.py

412 lines
14 KiB
Python
Raw Normal View History

2025-12-10 20:52:44 +08:00
"""
自适应 RPM 管理 API 端点
2025-12-10 20:52:44 +08:00
设计原则
- 自适应模式由 rpm_limit 字段决定
- rpm_limit = NULL启用自适应模式系统自动学习并调整 RPM 限制
- rpm_limit = 数字固定限制模式使用用户指定的 RPM 限制
- learned_rpm_limit自适应模式下学习到的 RPM 限制值
- adaptive_mode 是计算字段基于 rpm_limit 是否为 NULL
2025-12-10 20:52:44 +08:00
"""
from __future__ import annotations
2025-12-10 20:52:44 +08:00
from dataclasses import dataclass
from typing import Any
2025-12-10 20:52:44 +08:00
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from pydantic import BaseModel, Field, ValidationError
from sqlalchemy import func
from sqlalchemy.orm import Session, load_only
2025-12-10 20:52:44 +08:00
from src.api.base.admin_adapter import AdminApiAdapter
from src.api.base.context import ApiRequestContext
from src.api.base.pipeline import get_pipeline
2025-12-10 20:52:44 +08:00
from src.core.exceptions import InvalidRequestException, translate_pydantic_error
from src.database import get_db
from src.models.database import ProviderAPIKey
from src.services.rate_limit.adaptive_rpm import get_adaptive_rpm_manager
2025-12-10 20:52:44 +08:00
router = APIRouter(prefix="/api/admin/adaptive", tags=["Adaptive RPM"])
pipeline = get_pipeline()
2025-12-10 20:52:44 +08:00
# ==================== Pydantic Models ====================
class EnableAdaptiveRequest(BaseModel):
"""启用自适应模式请求"""
enabled: bool = Field(..., description="是否启用自适应模式true=自适应false=固定限制)")
fixed_limit: int | None = Field(
None, ge=1, le=100, description="固定 RPM 限制(仅当 enabled=false 时生效1-100"
2025-12-10 20:52:44 +08:00
)
class AdaptiveStatsResponse(BaseModel):
"""自适应统计响应"""
adaptive_mode: bool = Field(..., description="是否为自适应模式rpm_limit=NULL")
rpm_limit: int | None = Field(None, description="用户配置的固定限制NULL=自适应)")
effective_limit: int | None = Field(
2025-12-10 20:52:44 +08:00
None, description="当前有效限制(自适应使用学习值,固定使用配置值)"
)
learned_limit: int | None = Field(None, description="学习到的 RPM 限制")
2025-12-10 20:52:44 +08:00
concurrent_429_count: int
rpm_429_count: int
last_429_at: str | None
last_429_type: str | None
2025-12-10 20:52:44 +08:00
adjustment_count: int
recent_adjustments: list[dict]
# 置信度相关
learning_confidence: float | None = Field(None, description="学习置信度 (0.0-1.0)")
enforcement_active: bool | None = Field(None, description="是否正在执行本地 RPM 限制")
observation_count: int = Field(0, description="429 观察记录数")
header_observation_count: int = Field(0, description="带 header 的观察记录数")
latest_upstream_limit: int | None = Field(None, description="最近一次上游 header 限制值")
2025-12-10 20:52:44 +08:00
class KeyListItem(BaseModel):
"""Key 列表项"""
id: str
name: str | None
provider_id: str
api_formats: list[str] = Field(default_factory=list)
is_adaptive: bool = Field(..., description="是否为自适应模式rpm_limit=NULL")
rpm_limit: int | None = Field(None, description="固定 RPM 限制NULL=自适应)")
effective_limit: int | None = Field(None, description="当前有效限制")
learned_rpm_limit: int | None = Field(None, description="学习到的 RPM 限制")
2025-12-10 20:52:44 +08:00
concurrent_429_count: int
rpm_429_count: int
# ==================== API Endpoints ====================
@router.get(
"/keys",
response_model=list[KeyListItem],
2025-12-10 20:52:44 +08:00
summary="获取所有启用自适应模式的Key",
)
async def list_adaptive_keys(
request: Request,
provider_id: str | None = Query(None, description="按 Provider 过滤"),
2025-12-10 20:52:44 +08:00
db: Session = Depends(get_db),
) -> Any:
2025-12-10 20:52:44 +08:00
"""
获取所有启用自适应模式的Key列表
可选参数
- provider_id: Provider 过滤
2025-12-10 20:52:44 +08:00
"""
adapter = ListAdaptiveKeysAdapter(provider_id=provider_id)
2025-12-10 20:52:44 +08:00
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.patch(
"/keys/{key_id}/mode",
summary="Toggle key's RPM control mode",
2025-12-10 20:52:44 +08:00
)
async def toggle_adaptive_mode(
key_id: str,
request: Request,
db: Session = Depends(get_db),
) -> Any:
2025-12-10 20:52:44 +08:00
"""
Toggle the RPM control mode for a specific key
2025-12-10 20:52:44 +08:00
Parameters:
- enabled: true=adaptive mode (rpm_limit=NULL), false=fixed limit mode
2025-12-10 20:52:44 +08:00
- fixed_limit: fixed limit value (required when enabled=false)
"""
adapter = ToggleAdaptiveModeAdapter(key_id=key_id)
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.get(
"/keys/{key_id}/stats",
response_model=AdaptiveStatsResponse,
summary="获取Key的自适应统计",
)
async def get_adaptive_stats(
key_id: str,
request: Request,
db: Session = Depends(get_db),
) -> Any:
2025-12-10 20:52:44 +08:00
"""
获取指定Key的自适应 RPM 统计信息
2025-12-10 20:52:44 +08:00
包括
- 当前配置
- 学习到的限制
- 429错误统计
- 调整历史
"""
adapter = GetAdaptiveStatsAdapter(key_id=key_id)
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.delete(
"/keys/{key_id}/learning",
summary="Reset key's learning state",
)
async def reset_adaptive_learning(
key_id: str,
request: Request,
db: Session = Depends(get_db),
) -> Any:
2025-12-10 20:52:44 +08:00
"""
Reset the adaptive learning state for a specific key
Clears:
- Learned RPM limit (learned_rpm_limit)
2025-12-10 20:52:44 +08:00
- 429 error counts
- Adjustment history
Does not change:
- rpm_limit config (determines adaptive mode)
2025-12-10 20:52:44 +08:00
"""
adapter = ResetAdaptiveLearningAdapter(key_id=key_id)
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.patch(
"/keys/{key_id}/limit",
summary="Set key to fixed RPM limit mode",
2025-12-10 20:52:44 +08:00
)
async def set_rpm_limit(
2025-12-10 20:52:44 +08:00
key_id: str,
request: Request,
limit: int = Query(..., ge=1, le=100, description="RPM limit value (1-100)"),
2025-12-10 20:52:44 +08:00
db: Session = Depends(get_db),
) -> Any:
2025-12-10 20:52:44 +08:00
"""
Set key to fixed RPM limit mode
2025-12-10 20:52:44 +08:00
Note:
- After setting this value, key switches to fixed limit mode and won't auto-adjust
- To restore adaptive mode, use PATCH /keys/{key_id}/mode
"""
adapter = SetRPMLimitAdapter(key_id=key_id, limit=limit)
2025-12-10 20:52:44 +08:00
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.get(
"/summary",
summary="获取自适应 RPM 的全局统计",
2025-12-10 20:52:44 +08:00
)
async def get_adaptive_summary(
request: Request,
db: Session = Depends(get_db),
) -> Any:
2025-12-10 20:52:44 +08:00
"""
获取自适应 RPM 的全局统计摘要
2025-12-10 20:52:44 +08:00
包括
- 启用自适应模式的Key数量
- 总429错误数
- RPM 限制调整次数
2025-12-10 20:52:44 +08:00
"""
adapter = AdaptiveSummaryAdapter()
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
# ==================== Pipeline 适配器 ====================
@dataclass
class ListAdaptiveKeysAdapter(AdminApiAdapter):
provider_id: str | None = None
2025-12-10 20:52:44 +08:00
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
# 自适应模式rpm_limit = NULL
query = context.db.query(ProviderAPIKey).filter(ProviderAPIKey.rpm_limit.is_(None))
if self.provider_id:
query = query.filter(ProviderAPIKey.provider_id == self.provider_id)
2025-12-10 20:52:44 +08:00
keys = query.all()
adaptive_manager = get_adaptive_rpm_manager()
2025-12-10 20:52:44 +08:00
return [
KeyListItem(
id=key.id,
name=key.name,
provider_id=key.provider_id,
api_formats=key.api_formats or [],
is_adaptive=key.rpm_limit is None,
rpm_limit=key.rpm_limit,
effective_limit=adaptive_manager.get_effective_limit(key),
learned_rpm_limit=key.learned_rpm_limit,
2025-12-10 20:52:44 +08:00
concurrent_429_count=key.concurrent_429_count or 0,
rpm_429_count=key.rpm_429_count or 0,
)
for key in keys
]
@dataclass
class ToggleAdaptiveModeAdapter(AdminApiAdapter):
key_id: str
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
2025-12-10 20:52:44 +08:00
key = context.db.query(ProviderAPIKey).filter(ProviderAPIKey.id == self.key_id).first()
if not key:
raise HTTPException(status_code=404, detail="Key not found")
payload = context.ensure_json_body()
try:
body = EnableAdaptiveRequest.model_validate(payload)
except ValidationError as e:
errors = e.errors()
if errors:
raise InvalidRequestException(translate_pydantic_error(errors[0]))
raise InvalidRequestException("请求数据验证失败")
if body.enabled:
# 启用自适应模式:将 rpm_limit 设为 NULL
key.rpm_limit = None
message = "已切换为自适应模式,系统将自动学习并调整 RPM 限制"
2025-12-10 20:52:44 +08:00
else:
# 禁用自适应模式:设置固定限制
if body.fixed_limit is None:
raise HTTPException(
status_code=400, detail="禁用自适应模式时必须提供 fixed_limit 参数"
)
key.rpm_limit = body.fixed_limit
message = f"已切换为固定限制模式RPM 限制设为 {body.fixed_limit}"
2025-12-10 20:52:44 +08:00
context.db.commit()
context.db.refresh(key)
is_adaptive = key.rpm_limit is None
adaptive_manager = get_adaptive_rpm_manager()
2025-12-10 20:52:44 +08:00
return {
"message": message,
"key_id": key.id,
"is_adaptive": is_adaptive,
"rpm_limit": key.rpm_limit,
"effective_limit": adaptive_manager.get_effective_limit(key),
2025-12-10 20:52:44 +08:00
}
@dataclass
class GetAdaptiveStatsAdapter(AdminApiAdapter):
key_id: str
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
2025-12-10 20:52:44 +08:00
key = context.db.query(ProviderAPIKey).filter(ProviderAPIKey.id == self.key_id).first()
if not key:
raise HTTPException(status_code=404, detail="Key not found")
adaptive_manager = get_adaptive_rpm_manager()
2025-12-10 20:52:44 +08:00
stats = adaptive_manager.get_adjustment_stats(key)
# 转换字段名以匹配响应模型
return AdaptiveStatsResponse(
adaptive_mode=stats["adaptive_mode"],
rpm_limit=stats["rpm_limit"],
2025-12-10 20:52:44 +08:00
effective_limit=stats["effective_limit"],
learned_limit=stats["learned_limit"],
concurrent_429_count=stats["concurrent_429_count"],
rpm_429_count=stats["rpm_429_count"],
last_429_at=stats["last_429_at"],
last_429_type=stats["last_429_type"],
adjustment_count=stats["adjustment_count"],
recent_adjustments=stats["recent_adjustments"],
learning_confidence=stats.get("learning_confidence"),
enforcement_active=stats.get("enforcement_active"),
observation_count=stats.get("observation_count", 0),
header_observation_count=stats.get("header_observation_count", 0),
latest_upstream_limit=stats.get("latest_upstream_limit"),
2025-12-10 20:52:44 +08:00
)
@dataclass
class ResetAdaptiveLearningAdapter(AdminApiAdapter):
key_id: str
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
2025-12-10 20:52:44 +08:00
key = context.db.query(ProviderAPIKey).filter(ProviderAPIKey.id == self.key_id).first()
if not key:
raise HTTPException(status_code=404, detail="Key not found")
adaptive_manager = get_adaptive_rpm_manager()
2025-12-10 20:52:44 +08:00
adaptive_manager.reset_learning(context.db, key)
return {"message": "学习状态已重置", "key_id": key.id}
@dataclass
class SetRPMLimitAdapter(AdminApiAdapter):
2025-12-10 20:52:44 +08:00
key_id: str
limit: int
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
2025-12-10 20:52:44 +08:00
key = context.db.query(ProviderAPIKey).filter(ProviderAPIKey.id == self.key_id).first()
if not key:
raise HTTPException(status_code=404, detail="Key not found")
was_adaptive = key.rpm_limit is None
key.rpm_limit = self.limit
2025-12-10 20:52:44 +08:00
context.db.commit()
context.db.refresh(key)
return {
"message": f"已设置为固定限制模式RPM 限制为 {self.limit}",
2025-12-10 20:52:44 +08:00
"key_id": key.id,
"is_adaptive": False,
"rpm_limit": key.rpm_limit,
2025-12-10 20:52:44 +08:00
"previous_mode": "adaptive" if was_adaptive else "fixed",
}
class AdaptiveSummaryAdapter(AdminApiAdapter):
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
db = context.db
is_adaptive = ProviderAPIKey.rpm_limit.is_(None)
# SQL 聚合获取 count / sum避免全表 ORM 加载
total_keys, total_concurrent_429, total_rpm_429 = (
db.query(
func.count(ProviderAPIKey.id),
func.coalesce(func.sum(ProviderAPIKey.concurrent_429_count), 0),
func.coalesce(func.sum(ProviderAPIKey.rpm_429_count), 0),
)
.filter(is_adaptive)
.one()
)
# adjustment_history 是 JSON 列,长度只能在 Python 侧统计;
# 只加载有历史记录的 key 的必要列
keys_with_history = (
db.query(ProviderAPIKey)
.options(
load_only(ProviderAPIKey.id, ProviderAPIKey.name, ProviderAPIKey.adjustment_history)
)
.filter(is_adaptive, ProviderAPIKey.adjustment_history.isnot(None))
.all()
2025-12-10 20:52:44 +08:00
)
total_adjustments = sum(len(key.adjustment_history or []) for key in keys_with_history)
2025-12-10 20:52:44 +08:00
recent_adjustments = []
for key in keys_with_history:
2025-12-10 20:52:44 +08:00
if key.adjustment_history:
for adj in key.adjustment_history[-3:]:
recent_adjustments.append(
{
"key_id": key.id,
"key_name": key.name,
**adj,
}
)
recent_adjustments.sort(key=lambda item: item.get("timestamp", ""), reverse=True)
return {
"total_adaptive_keys": total_keys,
"total_concurrent_429_errors": int(total_concurrent_429),
"total_rpm_429_errors": int(total_rpm_429),
2025-12-10 20:52:44 +08:00
"total_adjustments": total_adjustments,
"recent_adjustments": recent_adjustments[:10],
}