Aether/_deprecated_py_src/api/admin/endpoints/concurrency.py

"""
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 计数已重置"}