vent_ai_platform/app/agents/workflows.py

from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import Any

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.agents.context import build_knowledge_context, build_mine_context
from app.agents.prompts import (
    DIAGNOSIS_PROMPT,
    DOCUMENT_ANALYSIS_PROMPT,
    EVALUATION_PROMPT,
    KNOWLEDGE_QA_PROMPT,
)
from app.agents.runtime import get_agent_runtime
from app.core.errors import APIError
from app.db.base import new_id
from app.db.models.agent import AgentRun, AgentSession
from app.db.repositories import document as document_repo
from app.schemas.diagnosis import DiagnosisRunRequest
from app.schemas.evaluation import EvaluationRunRequest
from app.schemas.work_order import WorkOrderCreate
from app.services.diagnosis_service import run_diagnosis
from app.services.evaluation_service import run_evaluation
from app.services.work_order_service import create_work_order


def _json_context(payload: dict[str, Any], limit: int = 14000) -> str:
    text = json.dumps(payload, ensure_ascii=False, default=str, indent=2)
    return text[:limit]


def _fallback_output(workflow: str, context: dict[str, Any], error: str = "") -> str:
    if workflow == "diagnosis":
        high = context.get("result", {}).get("high_risk_count", 0)
        return f"规则引擎已完成诊断，识别高风险缺陷 {high} 项。请按证据链逐项复核并生成整改任务。"
    if workflow == "evaluation":
        result = context.get("result", {})
        return (
            f"规则引擎已完成评价，综合得分 {result.get('total_score', '--')}，"
            f"等级 {result.get('grade', '--')}。建议按评价结果中的短板指标开展专项治理。"
        )
    if workflow == "knowledge_qa":
        count = context.get("knowledge", {}).get("total", 0)
        return f"检索到 {count} 条相关知识，请结合引用结果进行人工复核。"
    if workflow == "document_analysis":
        return "资料解析结果已生成，请对关键字段和缺失项进行人工确认。"
    return "工作流已执行。"


def _get_or_create_session(
    db: Session,
    *,
    mine_id: str | None,
    thread_id: str | None,
    workflow: str,
    user_name: str,
) -> AgentSession:
    if thread_id:
        session = db.scalar(select(AgentSession).where(AgentSession.thread_id == thread_id))
        if session is not None:
            return session
    session = AgentSession(
        id=new_id(),
        mine_id=mine_id,
        user_name=user_name,
        thread_id=thread_id or new_id(),
        title=f"{workflow} 工作流",
        workflow=workflow,
        status="active",
        context={},
        last_message_at=datetime.now(timezone.utc),
    )
    db.add(session)
    db.commit()
    db.refresh(session)
    return session


async def run_workflow(
    db: Session,
    *,
    mine_id: str,
    workflow: str,
    payload: dict[str, Any],
    requested_by: str = "",
) -> dict[str, Any]:
    session = _get_or_create_session(
        db,
        mine_id=mine_id,
        thread_id=payload.get("thread_id"),
        workflow=workflow,
        user_name=requested_by or "default",
    )
    run = AgentRun(
        id=new_id(),
        session_id=session.id,
        run_id=new_id(),
        workflow=workflow,
        status="running",
        input_payload=payload,
        output_payload={},
        model_name="",
        started_at=datetime.now(timezone.utc),
    )
    db.add(run)
    db.commit()
    db.refresh(run)

    context: dict[str, Any]
    prompt: str

    if workflow == "diagnosis":
        task = run_diagnosis(
            db,
            DiagnosisRunRequest(
                mine_id=mine_id,
                scene=payload.get("scene", "all"),
                context=payload.get("context", {}),
                requested_by=requested_by,
            ),
        )
        from app.db.repositories import diagnosis as diagnosis_repo

        results = diagnosis_repo.list_results(db, task.id)
        context = {
            "mine": build_mine_context(db, mine_id),
            "result": {
                "task_no": task.task_no,
                "scene": task.scene,
                "defects": [
                    {
                        "name": item.defect_name,
                        "scene": item.scene,
                        "severity": item.severity,
                        "confidence": item.confidence,
                        "conclusion": item.conclusion,
                        "recommendation": item.recommendation,
                    }
                    for item in results
                ],
                "high_risk_count": sum(1 for item in results if item.severity == "high"),
            },
        }
        prompt = f"{DIAGNOSIS_PROMPT}\n\n业务上下文：\n{_json_context(context)}"
    elif workflow == "evaluation":
        evaluation = run_evaluation(
            db,
            EvaluationRunRequest(
                mine_id=mine_id,
                context=payload.get("context", {}),
                requested_by=requested_by,
            ),
        )
        context = {
            "mine": build_mine_context(db, mine_id),
            "result": {
                "run_no": evaluation.run_no,
                "total_score": evaluation.total_score,
                "grade": evaluation.grade,
                "summary": evaluation.summary,
                "recommendations": evaluation.recommendations,
                "model_output": evaluation.model_output,
            },
        }
        prompt = f"{EVALUATION_PROMPT}\n\n业务上下文：\n{_json_context(context)}"
    elif workflow == "knowledge_qa":
        question = str(payload.get("question") or payload.get("message") or "").strip()
        if not question:
            raise APIError(422, "question_required", "知识问答必须提供 question")
        knowledge = build_knowledge_context(db, mine_id, question)
        context = {"question": question, "knowledge": knowledge}
        prompt = f"{KNOWLEDGE_QA_PROMPT}\n\n业务上下文：\n{_json_context(context)}"
    elif workflow == "document_analysis":
        document_id = str(payload.get("document_id") or "")
        if not document_id:
            raise APIError(422, "document_id_required", "资料分析必须提供 document_id")
        document = document_repo.get_document(db, document_id)
        if document is None:
            raise APIError(404, "document_not_found", "业务资料不存在")
        extractions = document_repo.list_extractions(db, document_id)
        context = {
            "document": {
                "id": document.id,
                "title": document.title,
                "category": document.category,
                "parse_status": document.parse_status,
                "summary": document.summary,
            },
            "fields": [
                {
                    "name": item.field_name,
                    "value": item.field_value,
                    "normalized_value": item.normalized_value,
                    "verified": item.verified,
                }
                for item in extractions
            ],
        }
        prompt = f"{DOCUMENT_ANALYSIS_PROMPT}\n\n业务上下文：\n{_json_context(context)}"
    else:
        raise APIError(422, "unsupported_workflow", f"不支持的工作流：{workflow}")

    runtime = get_agent_runtime()
    agent_result = await runtime.invoke(
        message=prompt,
        thread_id=session.thread_id,
        context=context,
    )
    output = agent_result.output or _fallback_output(workflow, context, agent_result.error)

    if workflow == "diagnosis" and context["result"]["defects"]:
        first = context["result"]["defects"][0]
        create_work_order(
            db,
            WorkOrderCreate(
                mine_id=mine_id,
                title=f"处置：{first['name']}",
                content=output,
                priority="high" if first["severity"] == "high" else "medium",
                assignee_name="",
                department="通防科",
                metadata={"source": "diagnosis_workflow"},
            ),
        )

    run.status = "success" if agent_result.status != "failed" else "degraded"
    run.output_payload = {
        "output": output,
        "agent_status": agent_result.status,
        "agent_error": agent_result.error,
        "context": context,
    }
    run.model_name = agent_result.model
    run.finished_at = datetime.now(timezone.utc)
    session.last_message_at = datetime.now(timezone.utc)
    db.add_all([run, session])
    db.commit()
    db.refresh(run)

    return {
        "run_id": run.run_id,
        "session_id": session.id,
        "workflow": workflow,
        "status": run.status,
        "output": output,
        "agent": runtime.describe(),
        "result": context.get("result", {}),
    }