DEMO-AGENT/review_agent/application_form_fill/workflow.py

from __future__ import annotations

import logging
from threading import Thread
from uuid import uuid4

from django.conf import settings
from django.db import transaction
from django.utils import timezone

from review_agent.application_form_fill.constants import DEFAULT_OUTPUT_TYPES, FORM_FILL_NODE_DEFINITIONS, WORKFLOW_TYPE
from review_agent.application_form_fill.events import record_event
from review_agent.application_form_fill.storage import build_batch_work_dir
from review_agent.models import ApplicationFormFillBatch, Conversation, FileSummaryBatch, Message, WorkflowNodeRun


logger = logging.getLogger("review_agent.application_form_fill.workflow")


def build_batch_no() -> str:
    return f"AFF-{timezone.localtime().strftime('%Y%m%d%H%M%S')}-{uuid4().hex[:6]}"


def find_latest_successful_summary_batch(conversation: Conversation) -> FileSummaryBatch | None:
    return (
        FileSummaryBatch.objects.filter(
            conversation=conversation,
            status=FileSummaryBatch.Status.SUCCESS,
        )
        .order_by("-finished_at", "-created_at", "-id")
        .first()
    )


@transaction.atomic
def create_application_form_fill_batch(
    *,
    conversation: Conversation,
    user,
    source_summary_batch: FileSummaryBatch,
    trigger_message: Message | None = None,
    requested_templates: list[str] | None = None,
    output_types: list[str] | None = None,
) -> ApplicationFormFillBatch:
    batch_no = build_batch_no()
    work_dir = build_batch_work_dir(batch_no=batch_no)
    work_dir.mkdir(parents=True, exist_ok=True)
    batch = ApplicationFormFillBatch.objects.create(
        conversation=conversation,
        user=user,
        trigger_message=trigger_message,
        source_summary_batch=source_summary_batch,
        batch_no=batch_no,
        requested_templates=requested_templates or [],
        output_types=output_types or DEFAULT_OUTPUT_TYPES,
        work_dir=str(work_dir),
    )
    for code, name, group in FORM_FILL_NODE_DEFINITIONS:
        WorkflowNodeRun.objects.create(
            workflow_type=WORKFLOW_TYPE,
            workflow_batch_id=batch.pk,
            node_group=group,
            node_code=code,
            node_name=name,
        )
    record_event(batch, "workflow_created", {"batch_id": batch.pk, "batch_no": batch.batch_no})
    return batch


class FormFillWorkflowExecutor:
    """Runs the auto-fill workflow skeleton; later stages fill node bodies."""

    def __init__(self, batch: ApplicationFormFillBatch):
        self.batch = batch

    def run(self) -> None:
        logger.info("自动填表工作流开始 batch_no=%s batch_id=%s", self.batch.batch_no, self.batch.pk)
        self.batch.status = ApplicationFormFillBatch.Status.RUNNING
        self.batch.started_at = timezone.now()
        self.batch.save(update_fields=["status", "started_at"])
        record_event(self.batch, "workflow_started", {"batch_id": self.batch.pk})

        try:
            for node in self._nodes():
                if node.status in {WorkflowNodeRun.Status.SUCCESS, WorkflowNodeRun.Status.SKIPPED}:
                    continue
                self._run_node(node)
        except Exception as exc:
            logger.exception("Application form fill workflow failed", extra={"batch_id": self.batch.pk})
            self.batch.status = ApplicationFormFillBatch.Status.FAILED
            self.batch.error_message = str(exc)
            self.batch.finished_at = timezone.now()
            self.batch.save(update_fields=["status", "error_message", "finished_at"])
            record_event(self.batch, "workflow_failed", {"message": str(exc)})
            return

        self.batch.status = ApplicationFormFillBatch.Status.SUCCESS
        self.batch.finished_at = timezone.now()
        self.batch.save(update_fields=["status", "finished_at"])
        record_event(self.batch, "workflow_completed", {"batch_id": self.batch.pk})
        logger.info("自动填表工作流完成 batch_no=%s", self.batch.batch_no)

    def _nodes(self):
        return WorkflowNodeRun.objects.filter(
            workflow_type=WORKFLOW_TYPE,
            workflow_batch_id=self.batch.pk,
        ).order_by("id")

    def _run_node(self, node: WorkflowNodeRun) -> None:
        node.status = WorkflowNodeRun.Status.RUNNING
        node.progress = 10
        node.started_at = timezone.now()
        node.message = f"{node.node_name}处理中"
        node.save(update_fields=["status", "progress", "started_at", "message"])
        record_event(
            self.batch,
            "node_progress",
            {"node_code": node.node_code, "status": node.status, "progress": node.progress, "message": node.message},
        )

        if node.node_code == "pdf_convert":
            node.status = WorkflowNodeRun.Status.SKIPPED
            node.progress = 100
            node.finished_at = timezone.now()
            node.message = "PDF 转换为后续增强项，本次跳过"
            node.save(update_fields=["status", "progress", "finished_at", "message"])
            record_event(
                self.batch,
                "node_progress",
                {"node_code": node.node_code, "status": node.status, "progress": node.progress, "message": node.message},
            )
            return

        node.status = WorkflowNodeRun.Status.SUCCESS
        node.progress = 100
        node.finished_at = timezone.now()
        node.message = f"{node.node_name}完成"
        node.save(update_fields=["status", "progress", "finished_at", "message"])
        record_event(
            self.batch,
            "node_progress",
            {"node_code": node.node_code, "status": node.status, "progress": node.progress, "message": node.message},
        )


def start_application_form_fill_workflow(batch: ApplicationFormFillBatch, *, async_run: bool = True) -> None:
    executor = FormFillWorkflowExecutor(batch)
    if not async_run:
        executor.run()
        return
    Thread(target=executor.run, daemon=True).start()