""" 分组-Key号池管理路由 支持一个分组绑定多个key,形成号池 """ from __future__ import annotations from datetime import datetime from typing import List, Optional from fastapi import APIRouter, HTTPException from pydantic import BaseModel from app.db import get_pool router = APIRouter(tags=["group_keys"]) class GroupKeyIn(BaseModel): """添加key到分组""" group_id: int api_key_id: int priority: Optional[int] = 0 weight: Optional[int] = 1 class GroupKeyUpdate(BaseModel): """更新分组-key关系""" priority: Optional[int] = None weight: Optional[int] = None is_active: Optional[bool] = None status: Optional[str] = None status_msg: Optional[str] = None class GroupKeyOut(BaseModel): id: int group_id: int group_name: str api_key_id: int api_key_name: str api_key_masked: str priority: int weight: int is_active: bool status: str status_msg: Optional[str] last_used_at: Optional[datetime] total_tokens: int request_count: int error_count: int created_at: datetime updated_at: datetime class BatchAddKeysIn(BaseModel): """批量添加key到分组""" group_id: int api_key_ids: List[int] priority: Optional[int] = 0 weight: Optional[int] = 1 def _mask_key(v: str) -> str: if len(v) <= 10: return '••••••••' return v[:6] + '••••••••' + v[-4:] @router.get("/group-keys", response_model=List[GroupKeyOut]) async def list_group_keys(group_id: Optional[int] = None) -> List[GroupKeyOut]: """获取分组-key列表""" pool = get_pool() query = """ SELECT gk.id, gk.group_id, mg.name AS group_name, gk.api_key_id, ak.name AS api_key_name, ak.key_value, gk.priority, gk.weight, gk.is_active, gk.status, gk.status_msg, gk.last_used_at, gk.total_tokens, gk.request_count, gk.error_count, gk.created_at, gk.updated_at FROM group_keys gk JOIN model_groups mg ON mg.id = gk.group_id JOIN api_keys ak ON ak.id = gk.api_key_id """ if group_id is not None: query += " WHERE gk.group_id = $1" rows = await pool.fetch(query, group_id) else: rows = await pool.fetch(query) return [ GroupKeyOut(**{ **dict(r), "api_key_masked": _mask_key(r["key_value"]) }) for r in rows ] @router.post("/group-keys", response_model=GroupKeyOut, status_code=201) async def add_key_to_group(body: GroupKeyIn) -> GroupKeyOut: """添加key到分组(号池)""" pool = get_pool() # 检查分组是否存在 group = await pool.fetchrow("SELECT id, name FROM model_groups WHERE id = $1", body.group_id) if not group: raise HTTPException(status_code=404, detail="分组不存在") # 检查key是否存在 key = await pool.fetchrow("SELECT id, name, key_value FROM api_keys WHERE id = $1", body.api_key_id) if not key: raise HTTPException(status_code=404, detail="API Key不存在") # 检查是否已存在 existing = await pool.fetchrow( "SELECT id FROM group_keys WHERE group_id = $1 AND api_key_id = $2", body.group_id, body.api_key_id ) if existing: raise HTTPException(status_code=400, detail="该key已在此分组中") row = await pool.fetchrow( """ INSERT INTO group_keys (group_id, api_key_id, priority, weight) VALUES ($1, $2, $3, $4) RETURNING id, group_id, api_key_id, priority, weight, is_active, status, status_msg, last_used_at, total_tokens, request_count, error_count, created_at, updated_at """, body.group_id, body.api_key_id, body.priority or 0, body.weight or 1 ) return GroupKeyOut(**{ **dict(row), "group_name": group["name"], "api_key_name": key["name"], "api_key_masked": _mask_key(key["key_value"]), }) @router.post("/group-keys/batch-add", response_model=List[GroupKeyOut], status_code=201) async def batch_add_keys_to_group(body: BatchAddKeysIn) -> List[GroupKeyOut]: """批量添加key到分组""" pool = get_pool() # 检查分组 group = await pool.fetchrow("SELECT id, name FROM model_groups WHERE id = $1", body.group_id) if not group: raise HTTPException(status_code=404, detail="分组不存在") results = [] for key_id in body.api_key_ids: key = await pool.fetchrow("SELECT id, name, key_value FROM api_keys WHERE id = $1", key_id) if not key: continue # 检查是否已存在 existing = await pool.fetchrow( "SELECT id FROM group_keys WHERE group_id = $1 AND api_key_id = $2", body.group_id, key_id ) if existing: continue row = await pool.fetchrow( """ INSERT INTO group_keys (group_id, api_key_id, priority, weight) VALUES ($1, $2, $3, $4) RETURNING id, group_id, api_key_id, priority, weight, is_active, status, status_msg, last_used_at, total_tokens, request_count, error_count, created_at, updated_at """, body.group_id, key_id, body.priority or 0, body.weight or 1 ) results.append(GroupKeyOut(**{ **dict(row), "group_name": group["name"], "api_key_name": key["name"], "api_key_masked": _mask_key(key["key_value"]), })) return results @router.put("/group-keys/{id}", response_model=GroupKeyOut) async def update_group_key(id: int, body: GroupKeyUpdate) -> GroupKeyOut: """更新分组-key关系""" pool = get_pool() existing = await pool.fetchrow( """ SELECT gk.*, mg.name AS group_name, ak.name AS api_key_name, ak.key_value FROM group_keys gk JOIN model_groups mg ON mg.id = gk.group_id JOIN api_keys ak ON ak.id = gk.api_key_id WHERE gk.id = $1 """, id ) if not existing: raise HTTPException(status_code=404, detail="记录不存在") updates = [] values = [] param_idx = 1 if body.priority is not None: updates.append(f"priority = ${param_idx}") values.append(body.priority) param_idx += 1 if body.weight is not None: updates.append(f"weight = ${param_idx}") values.append(body.weight) param_idx += 1 if body.is_active is not None: updates.append(f"is_active = ${param_idx}") values.append(body.is_active) param_idx += 1 if body.status is not None: updates.append(f"status = ${param_idx}") values.append(body.status) param_idx += 1 if body.status_msg is not None: updates.append(f"status_msg = ${param_idx}") values.append(body.status_msg) param_idx += 1 if updates: updates.append("updated_at = NOW()") values.append(id) await pool.execute( f"UPDATE group_keys SET {', '.join(updates)} WHERE id = ${param_idx}", *values ) # 重新获取 row = await pool.fetchrow( """ SELECT gk.id, gk.group_id, mg.name AS group_name, gk.api_key_id, ak.name AS api_key_name, ak.key_value, gk.priority, gk.weight, gk.is_active, gk.status, gk.status_msg, gk.last_used_at, gk.total_tokens, gk.request_count, gk.error_count, gk.created_at, gk.updated_at FROM group_keys gk JOIN model_groups mg ON mg.id = gk.group_id JOIN api_keys ak ON ak.id = gk.api_key_id WHERE gk.id = $1 """, id ) return GroupKeyOut(**{ **dict(row), "api_key_masked": _mask_key(row["key_value"]) }) @router.delete("/group-keys/{id}", status_code=204) async def remove_key_from_group(id: int): """从分组中移除key""" pool = get_pool() result = await pool.execute("DELETE FROM group_keys WHERE id = $1", id) if result == "DELETE 0": raise HTTPException(status_code=404, detail="记录不存在") @router.get("/groups/{group_id}/pool-summary") async def get_group_pool_summary(group_id: int): """获取分组号池摘要(健康度统计)""" pool = get_pool() # 检查分组 group = await pool.fetchrow("SELECT id, name FROM model_groups WHERE id = $1", group_id) if not group: raise HTTPException(status_code=404, detail="分组不存在") # 统计 stats = await pool.fetchrow( """ SELECT COUNT(*) AS total_keys, COUNT(*) FILTER (WHERE status = 'active' AND is_active = TRUE) AS active_keys, COUNT(*) FILTER (WHERE status = 'rate_limited') AS rate_limited_keys, COUNT(*) FILTER (WHERE status = 'banned') AS banned_keys, COUNT(*) FILTER (WHERE status = 'failed') AS failed_keys, SUM(total_tokens) AS total_tokens, SUM(request_count) AS total_requests, SUM(error_count) AS total_errors FROM group_keys WHERE group_id = $1 """, group_id ) # 健康度计算 total = stats["total_keys"] or 0 active = stats["active_keys"] or 0 health_rate = (active / total * 100) if total > 0 else 0 return { "group_id": group_id, "group_name": group["name"], "total_keys": total, "active_keys": active, "rate_limited_keys": stats["rate_limited_keys"] or 0, "banned_keys": stats["banned_keys"] or 0, "failed_keys": stats["failed_keys"] or 0, "health_rate": round(health_rate, 2), "total_tokens": stats["total_tokens"] or 0, "total_requests": stats["total_requests"] or 0, "total_errors": stats["total_errors"] or 0, } @router.get("/groups/{group_id}/keys") async def get_group_keys_detail(group_id: int): """获取分组的key详情列表""" pool = get_pool() rows = await pool.fetch( """ SELECT gk.id, gk.api_key_id, ak.name AS api_key_name, ak.key_value, gk.priority, gk.weight, gk.is_active, gk.status, gk.status_msg, gk.last_used_at, gk.total_tokens, gk.request_count, gk.error_count, gk.created_at, gk.updated_at FROM group_keys gk JOIN api_keys ak ON ak.id = gk.api_key_id WHERE gk.group_id = $1 ORDER BY gk.priority DESC, gk.weight DESC """, group_id ) return [ { "id": r["id"], "api_key_id": r["api_key_id"], "api_key_name": r["api_key_name"], "api_key_masked": _mask_key(r["key_value"]), "priority": r["priority"], "weight": r["weight"], "is_active": r["is_active"], "status": r["status"], "status_msg": r["status_msg"], "last_used_at": r["last_used_at"].isoformat() if r["last_used_at"] else None, "total_tokens": r["total_tokens"], "request_count": r["request_count"], "error_count": r["error_count"], "created_at": r["created_at"].isoformat() if r["created_at"] else None, "updated_at": r["updated_at"].isoformat() if r["updated_at"] else None, } for r in rows ] # ============================================ # 推送到company_base_pool # ============================================ @router.post("/groups/{group_id}/push") async def push_group(group_id: int): """手动推送单个分组到company_base_pool""" pool = get_pool() try: await push_group_keys_to_pool(pool, group_id, []) return {"success": True, "message": "推送成功"} except Exception as e: raise HTTPException(status_code=500, detail=f"推送失败: {str(e)}") @router.post("/groups/push-all") async def push_all_groups(): """手动推送所有分组到company_base_pool""" pool = get_pool() try: await push_all_groups_to_pool(pool) return {"success": True, "message": "推送完成"} except Exception as e: raise HTTPException(status_code=500, detail=f"推送失败: {str(e)}") @router.get("/internal/group-keys-raw") async def get_group_keys_raw(): """内部接口:返回分组名和原始key值(用于company_base_pool同步)""" pool = get_pool() rows = await pool.fetch(""" SELECT mg.name AS group_name, ak.key_value FROM group_keys gk JOIN model_groups mg ON mg.id = gk.group_id JOIN api_keys ak ON ak.id = gk.api_key_id WHERE gk.is_active = TRUE """) return [{"group_name": r["group_name"], "key_value": r["key_value"]} for r in rows]