mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
"""
|
|
Key RPM 限制管理 API
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, Depends, Request
|
|
from sqlalchemy.orm import Session
|
|
|
|
from src.api.base.admin_adapter import AdminApiAdapter
|
|
from src.api.base.context import ApiRequestContext
|
|
from src.api.base.pipeline import get_pipeline
|
|
from src.core.exceptions import NotFoundException
|
|
from src.database import get_db
|
|
from src.models.database import ProviderAPIKey
|
|
from src.models.endpoint_models import KeyRpmStatusResponse
|
|
from src.services.rate_limit.concurrency_manager import get_concurrency_manager
|
|
|
|
router = APIRouter(tags=["RPM Control"])
|
|
pipeline = get_pipeline()
|
|
|
|
|
|
@router.get("/rpm/key/{key_id}", response_model=KeyRpmStatusResponse)
|
|
async def get_key_rpm(
|
|
key_id: str,
|
|
request: Request,
|
|
db: Session = Depends(get_db),
|
|
) -> KeyRpmStatusResponse:
|
|
"""
|
|
获取 Key 当前 RPM 状态
|
|
|
|
查询指定 API Key 的实时 RPM 使用情况，包括当前 RPM 计数和最大 RPM 限制。
|
|
|
|
**路径参数**:
|
|
- `key_id`: API Key ID
|
|
|
|
**返回字段**:
|
|
- `key_id`: API Key ID
|
|
- `current_rpm`: 当前 RPM 计数
|
|
- `rpm_limit`: RPM 限制
|
|
"""
|
|
adapter = AdminKeyRpmAdapter(key_id=key_id)
|
|
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
|
|
|
|
|
@router.delete("/rpm/key/{key_id}")
|
|
async def reset_key_rpm(
|
|
key_id: str,
|
|
http_request: Request,
|
|
db: Session = Depends(get_db),
|
|
) -> dict:
|
|
"""
|
|
重置 Key RPM 计数器
|
|
|
|
重置指定 API Key 的 RPM 计数器，用于解决计数不准确的问题。
|
|
管理员功能，请谨慎使用。
|
|
|
|
**路径参数**:
|
|
- `key_id`: API Key ID
|
|
|
|
**返回字段**:
|
|
- `message`: 操作结果消息
|
|
"""
|
|
adapter = AdminResetKeyRpmAdapter(key_id=key_id)
|
|
return await pipeline.run(adapter=adapter, http_request=http_request, db=db, mode=adapter.mode)
|
|
|
|
|
|
# -------- Adapters --------
|
|
|
|
|
|
@dataclass
|
|
class AdminKeyRpmAdapter(AdminApiAdapter):
|
|
key_id: str
|
|
|
|
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
|
db = context.db
|
|
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == self.key_id).first()
|
|
if not key:
|
|
raise NotFoundException(f"Key {self.key_id} 不存在")
|
|
|
|
concurrency_manager = await get_concurrency_manager()
|
|
key_count = await concurrency_manager.get_key_rpm_count(key_id=self.key_id)
|
|
|
|
return KeyRpmStatusResponse(
|
|
key_id=self.key_id,
|
|
current_rpm=key_count,
|
|
rpm_limit=key.rpm_limit,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class AdminResetKeyRpmAdapter(AdminApiAdapter):
|
|
key_id: str
|
|
|
|
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
|
concurrency_manager = await get_concurrency_manager()
|
|
await concurrency_manager.reset_key_rpm(key_id=self.key_id)
|
|
return {"message": "RPM 计数已重置"}
|