from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select from app.database import get_db from app.models.agent import Agent from app.models.user import User from app.routers.auth import get_current_user from app.schemas.agent import AgentCreate, AgentOut, AgentStats, AgentConfigUpdate, AgentConfigOut router = APIRouter(prefix="/api/agents", tags=["Agent"]) # 运行时调用统计（内存中，非持久化） _agent_call_counts: dict[str, int] = {} _agent_current_tasks: dict[str, str | None] = {} _agent_statuses: dict[str, str] = {} # 默认 Agent 角色列表 DEFAULT_AGENT_ROLES = ["master", "planner", "executor", "librarian", "analyst"] def record_agent_call(agent_id: str): _agent_call_counts[agent_id] = _agent_call_counts.get(agent_id, 0) + 1 def set_agent_task(agent_id: str, task: str | None): _agent_current_tasks[agent_id] = task _agent_statuses[agent_id] = "active" if task else "idle" def set_agent_status(agent_id: str, status: str): _agent_statuses[agent_id] = status @router.get("", response_model=list[AgentOut]) async def list_agents( current_user: User = Depends(get_current_user), db: AsyncSession = Depends(get_db), ): result = await db.execute( select(Agent).where(Agent.is_active == True).order_by(Agent.role) ) return result.scalars().all() # ———— 运行时统计（必须在 /{agent_id} 之前）———— @router.get("/stats", response_model=list[AgentStats]) async def get_agent_stats( current_user: User = Depends(get_current_user), ): """ 获取各 Agent 的运行时统计（调用次数、当前任务、状态） """ stats = [] for role in DEFAULT_AGENT_ROLES: stats.append(AgentStats( agent_id=role, call_count=_agent_call_counts.get(role, 0), current_task=_agent_current_tasks.get(role), status=_agent_statuses.get(role, "idle"), )) return stats # ———— 配置管理（必须在 /{agent_id} 之前）———— @router.get("/config/{agent_id}", response_model=AgentConfigOut) async def get_agent_config( agent_id: str, db: AsyncSession = Depends(get_db), ): """获取单个 Agent 完整配置""" result = await db.execute(select(Agent).where(Agent.role == agent_id)) agent = result.scalar_one_or_none() if not agent: from app.agents.prompts import MASTER_SYSTEM_PROMPT, PLANNER_SYSTEM_PROMPT, EXECUTOR_SYSTEM_PROMPT, LIBRARIAN_SYSTEM_PROMPT, ANALYST_SYSTEM_PROMPT defaults = { "master": ("JARVIS", "主控制核心", MASTER_SYSTEM_PROMPT), "planner": ("PLANNER", "规划专家", PLANNER_SYSTEM_PROMPT), "executor": ("EXECUTOR", "执行专家", EXECUTOR_SYSTEM_PROMPT), "librarian": ("LIBRARIAN", "知识管理员", LIBRARIAN_SYSTEM_PROMPT), "analyst": ("ANALYST", "数据分析师", ANALYST_SYSTEM_PROMPT), } if agent_id not in defaults: raise HTTPException(status_code=404, detail="Agent 不存在") name, desc, prompt = defaults[agent_id] return AgentConfigOut( id=agent_id, name=name, role=agent_id, description=desc, system_prompt=prompt, enabled=True, is_active=True, ) return AgentConfigOut( id=agent.role, name=agent.name, role=agent.role, description=agent.description, system_prompt=agent.system_prompt, enabled=agent.is_active, is_active=agent.is_active, ) @router.put("/config/{agent_id}", response_model=AgentConfigOut) async def update_agent_config( agent_id: str, data: AgentConfigUpdate, current_user: User = Depends(get_current_user), db: AsyncSession = Depends(get_db), ): """更新 Agent 配置（名称、描述、提示词、启用状态）""" result = await db.execute(select(Agent).where(Agent.role == agent_id)) agent = result.scalar_one_or_none() if not agent: raise HTTPException(status_code=404, detail="Agent 不存在") if data.name is not None: agent.name = data.name if data.description is not None: agent.description = data.description if data.system_prompt is not None: agent.system_prompt = data.system_prompt if data.enabled is not None: agent.is_active = data.enabled _agent_statuses[agent_id] = "disabled" if not data.enabled else "idle" await db.commit() await db.refresh(agent) return AgentConfigOut( id=agent.role, name=agent.name, role=agent.role, description=agent.description, system_prompt=agent.system_prompt, enabled=agent.is_active, is_active=agent.is_active, ) @router.post("", response_model=AgentOut, status_code=201) async def create_agent( data: AgentCreate, current_user: User = Depends(get_current_user), db: AsyncSession = Depends(get_db), ): agent = Agent( name=data.name, role=data.role, description=data.description, system_prompt=data.system_prompt, ) db.add(agent) await db.commit() await db.refresh(agent) return agent @router.get("/{agent_id}", response_model=AgentOut) async def get_agent( agent_id: str, db: AsyncSession = Depends(get_db), ): result = await db.execute(select(Agent).where(Agent.id == agent_id)) agent = result.scalar_one_or_none() if not agent: raise HTTPException(status_code=404, detail="Agent 不存在") return agent # ———— 配置管理 ———— @router.get("/config/{agent_id}", response_model=AgentConfigOut) async def get_agent_config( agent_id: str, db: AsyncSession = Depends(get_db), ): """获取单个 Agent 完整配置""" result = await db.execute(select(Agent).where(Agent.role == agent_id)) agent = result.scalar_one_or_none() if not agent: # 如果数据库中没有，返回默认配置 from app.agents.prompts import MASTER_SYSTEM_PROMPT, PLANNER_SYSTEM_PROMPT, EXECUTOR_SYSTEM_PROMPT, LIBRARIAN_SYSTEM_PROMPT, ANALYST_SYSTEM_PROMPT defaults = { "master": ("JARVIS", "主控制核心", MASTER_SYSTEM_PROMPT), "planner": ("PLANNER", "规划专家", PLANNER_SYSTEM_PROMPT), "executor": ("EXECUTOR", "执行专家", EXECUTOR_SYSTEM_PROMPT), "librarian": ("LIBRARIAN", "知识管理员", LIBRARIAN_SYSTEM_PROMPT), "analyst": ("ANALYST", "数据分析师", ANALYST_SYSTEM_PROMPT), } if agent_id not in defaults: raise HTTPException(status_code=404, detail="Agent 不存在") name, desc, prompt = defaults[agent_id] return AgentConfigOut( id=agent_id, name=name, role=agent_id, description=desc, system_prompt=prompt, enabled=True, is_active=True, ) return AgentConfigOut( id=agent.role, name=agent.name, role=agent.role, description=agent.description, system_prompt=agent.system_prompt, enabled=agent.is_active, is_active=agent.is_active, ) @router.put("/config/{agent_id}", response_model=AgentConfigOut) async def update_agent_config( agent_id: str, data: AgentConfigUpdate, current_user: User = Depends(get_current_user), db: AsyncSession = Depends(get_db), ): """更新 Agent 配置（名称、描述、提示词、启用状态）""" result = await db.execute(select(Agent).where(Agent.role == agent_id)) agent = result.scalar_one_or_none() if not agent: raise HTTPException(status_code=404, detail="Agent 不存在") if data.name is not None: agent.name = data.name if data.description is not None: agent.description = data.description if data.system_prompt is not None: agent.system_prompt = data.system_prompt if data.enabled is not None: agent.is_active = data.enabled _agent_statuses[agent_id] = "disabled" if not data.enabled else "idle" await db.commit() await db.refresh(agent) return AgentConfigOut( id=agent.role, name=agent.name, role=agent.role, description=agent.description, system_prompt=agent.system_prompt, enabled=agent.is_active, is_active=agent.is_active, )