vent_ai_platform/app/api/v1/endpoints/ai.py

from __future__ import annotations

import json

from fastapi import APIRouter, Depends, Query
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.agents.runtime import get_agent_runtime
from app.agents.workflows import run_workflow
from app.db.session import get_db
from app.schemas.ai import (
    AgentCapability,
    AgentChatRequest,
    AgentRunResponse,
    WorkflowRunRequest,
)


router = APIRouter(prefix="/ai", tags=["DeepAgent 智能体"])


@router.get(
    "/capabilities",
    response_model=list[AgentCapability],
    summary="智能体与工作流能力",
)
def capabilities() -> list[AgentCapability]:
    runtime = get_agent_runtime()
    enabled = runtime.enabled
    return [
        AgentCapability(
            name="knowledge_qa",
            label="通防知识问答",
            enabled=enabled,
            description="基于知识库检索结果回答问题并返回依据。",
        ),
        AgentCapability(
            name="diagnosis",
            label="缺陷根因分析",
            enabled=enabled,
            description="增强诊断解释、耦合分析和处置建议。",
        ),
        AgentCapability(
            name="evaluation",
            label="评价解释与优化",
            enabled=enabled,
            description="解释评价结果并生成优化建议。",
        ),
        AgentCapability(
            name="document_analysis",
            label="业务资料分析",
            enabled=enabled,
            description="分析资料字段、缺失项和可沉淀知识。",
        ),
    ]


@router.post(
    "/workflows/run",
    response_model=AgentRunResponse,
    summary="运行 DeepAgent 业务工作流",
)
async def run_business_workflow(
    payload: WorkflowRunRequest,
    db: Session = Depends(get_db),
) -> AgentRunResponse:
    result = await run_workflow(
        db,
        mine_id=payload.mine_id,
        workflow=payload.workflow,
        payload=payload.payload,
        requested_by=payload.requested_by,
    )
    return AgentRunResponse(**result)


@router.post("/chat", summary="智能问答，支持 SSE")
async def chat(
    payload: AgentChatRequest,
    stream: bool = Query(True),
    db: Session = Depends(get_db),
):
    workflow_payload = {
        "question": payload.message,
        "thread_id": payload.thread_id,
        **payload.context,
    }
    if not stream:
        result = await run_workflow(
            db,
            mine_id=payload.mine_id or "",
            workflow="knowledge_qa",
            payload=workflow_payload,
            requested_by="default",
        )
        return AgentRunResponse(**result)

    async def event_stream():
        yield f"data: {json.dumps({'event': 'start', 'workflow': 'knowledge_qa'}, ensure_ascii=False)}\n\n"
        result = await run_workflow(
            db,
            mine_id=payload.mine_id or "",
            workflow="knowledge_qa",
            payload=workflow_payload,
            requested_by="default",
        )
        yield f"data: {json.dumps({'event': 'message', 'content': result['output']}, ensure_ascii=False)}\n\n"
        yield f"data: {json.dumps({'event': 'done', 'run_id': result['run_id'], 'session_id': result['session_id']}, ensure_ascii=False)}\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )