DocsGPT/application/agents/workflow_agent.py

import logging
from datetime import datetime, timezone
from typing import Any, Dict, Generator, List, Optional

from application.agents.base import BaseAgent
from application.agents.workflows.schemas import (
    ExecutionStatus,
    Workflow,
    WorkflowEdge,
    WorkflowGraph,
    WorkflowNode,
    WorkflowRun,
)
from application.agents.workflows.workflow_engine import WorkflowEngine
from application.logging import log_activity, LogContext
from application.storage.db.base_repository import looks_like_uuid
from application.storage.db.repositories.workflow_edges import WorkflowEdgesRepository
from application.storage.db.repositories.workflow_nodes import WorkflowNodesRepository
from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository
from application.storage.db.repositories.workflows import WorkflowsRepository
from application.storage.db.session import db_readonly, db_session

logger = logging.getLogger(__name__)

# Per-run cap on attachments staged as run-scoped artifacts; the remainder is
# dropped (the per-user artifact quota is only a best-effort soft cap).
_MAX_INPUT_DOCUMENTS = 25


class WorkflowAgent(BaseAgent):
    """A specialized agent that executes predefined workflows."""

    def __init__(
        self,
        *args,
        workflow_id: Optional[str] = None,
        workflow: Optional[Dict[str, Any]] = None,
        workflow_owner: Optional[str] = None,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.workflow_id = workflow_id
        self.workflow_owner = workflow_owner
        self._workflow_data = workflow
        self._engine: Optional[WorkflowEngine] = None
        self._run_persisted = False

    @log_activity()
    def gen(
        self, query: str, log_context: LogContext = None
    ) -> Generator[Dict[str, str], None, None]:
        yield from self._gen_inner(query, log_context)

    def _gen_inner(
        self, query: str, log_context: LogContext
    ) -> Generator[Dict[str, str], None, None]:
        graph = self._load_workflow_graph()
        if not graph:
            yield {"type": "error", "error": "Failed to load workflow configuration."}
            return
        self._engine = WorkflowEngine(graph, self)

        owner_id = self._resolve_owner_id()
        pg_workflow_id = self._precreate_workflow_run(owner_id, query)
        self._run_persisted = pg_workflow_id is not None

        input_documents = self._bridge_attachments(owner_id, persisted=self._run_persisted)

        yield from self._engine.execute({"input_documents": input_documents}, query)
        self._finalize_workflow_run(owner_id, pg_workflow_id, query)

    def _load_workflow_graph(self) -> Optional[WorkflowGraph]:
        if self._workflow_data:
            return self._parse_embedded_workflow()
        if self.workflow_id:
            return self._load_from_database()
        return None

    def _parse_embedded_workflow(self) -> Optional[WorkflowGraph]:
        try:
            nodes_data = self._workflow_data.get("nodes", [])
            edges_data = self._workflow_data.get("edges", [])

            workflow = Workflow(
                name=self._workflow_data.get("name", "Embedded Workflow"),
                description=self._workflow_data.get("description"),
            )

            nodes = []
            for n in nodes_data:
                node_config = n.get("data", {})
                nodes.append(
                    WorkflowNode(
                        id=n["id"],
                        workflow_id=self.workflow_id or "embedded",
                        type=n["type"],
                        title=n.get("title", "Node"),
                        description=n.get("description"),
                        position=n.get("position", {"x": 0, "y": 0}),
                        config=node_config,
                    )
                )
            edges = []
            for e in edges_data:
                edges.append(
                    WorkflowEdge(
                        id=e["id"],
                        workflow_id=self.workflow_id or "embedded",
                        source=e.get("source") or e.get("source_id"),
                        target=e.get("target") or e.get("target_id"),
                        sourceHandle=e.get("sourceHandle") or e.get("source_handle"),
                        targetHandle=e.get("targetHandle") or e.get("target_handle"),
                    )
                )
            return WorkflowGraph(workflow=workflow, nodes=nodes, edges=edges)
        except Exception as e:
            logger.error(f"Invalid embedded workflow: {e}")
            return None

    def _load_from_database(self) -> Optional[WorkflowGraph]:
        try:
            if not self.workflow_id:
                logger.error("Missing workflow ID for load")
                return None
            owner_id = self.workflow_owner
            if not owner_id and isinstance(self.decoded_token, dict):
                owner_id = self.decoded_token.get("sub")
            if not owner_id:
                logger.error(
                    f"Workflow owner not available for workflow load: {self.workflow_id}"
                )
                return None

            with db_readonly() as conn:
                wf_repo = WorkflowsRepository(conn)
                if looks_like_uuid(self.workflow_id):
                    workflow_row = wf_repo.get(self.workflow_id, owner_id)
                else:
                    workflow_row = wf_repo.get_by_legacy_id(self.workflow_id, owner_id)
                if workflow_row is None:
                    logger.error(
                        f"Workflow {self.workflow_id} not found or inaccessible "
                        f"for user {owner_id}"
                    )
                    return None
                pg_workflow_id = str(workflow_row["id"])
                graph_version = workflow_row.get("current_graph_version", 1)
                try:
                    graph_version = int(graph_version)
                    if graph_version <= 0:
                        graph_version = 1
                except (ValueError, TypeError):
                    graph_version = 1

                node_rows = WorkflowNodesRepository(conn).find_by_version(
                    pg_workflow_id, graph_version,
                )
                edge_rows = WorkflowEdgesRepository(conn).find_by_version(
                    pg_workflow_id, graph_version,
                )

            workflow = Workflow(
                name=workflow_row.get("name"),
                description=workflow_row.get("description"),
            )
            nodes = [
                WorkflowNode(
                    id=n["node_id"],
                    workflow_id=pg_workflow_id,
                    type=n["node_type"],
                    title=n.get("title") or "Node",
                    description=n.get("description"),
                    position=n.get("position") or {"x": 0, "y": 0},
                    config=n.get("config") or {},
                )
                for n in node_rows
            ]
            edges = [
                WorkflowEdge(
                    id=e["edge_id"],
                    workflow_id=pg_workflow_id,
                    source=e.get("source_id"),
                    target=e.get("target_id"),
                    sourceHandle=e.get("source_handle"),
                    targetHandle=e.get("target_handle"),
                )
                for e in edge_rows
            ]

            return WorkflowGraph(workflow=workflow, nodes=nodes, edges=edges)
        except Exception as e:
            logger.error(f"Failed to load workflow from database: {e}")
            return None

    def _resolve_owner_id(self) -> Optional[str]:
        """Resolve the run owner from the explicit workflow owner or the token ``sub``."""
        owner_id = self.workflow_owner
        if not owner_id and isinstance(self.decoded_token, dict):
            owner_id = self.decoded_token.get("sub")
        return owner_id

    def _resolve_owned_workflow_pg_id(
        self, conn: Any, owner_id: Optional[str]
    ) -> Optional[str]:
        """Return the owned workflow's PG id, or None for an unowned/draft id."""
        if not self.workflow_id or not owner_id:
            return None
        wf_repo = WorkflowsRepository(conn)
        if looks_like_uuid(self.workflow_id):
            workflow_row = wf_repo.get(self.workflow_id, owner_id)
        else:
            workflow_row = wf_repo.get_by_legacy_id(self.workflow_id, owner_id)
        return str(workflow_row["id"]) if workflow_row is not None else None

    def _precreate_workflow_run(self, owner_id: Optional[str], query: str) -> Optional[str]:
        """Insert the run row up front so run-scoped artifacts are authz-reachable mid-run."""
        if not self._engine or not self.workflow_id or not owner_id:
            return None
        try:
            with db_session() as conn:
                pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, owner_id)
                if pg_workflow_id is None:
                    return None
                WorkflowRunsRepository(conn).create(
                    pg_workflow_id,
                    owner_id,
                    ExecutionStatus.RUNNING.value,
                    run_id=self._engine.workflow_run_id,
                    inputs={"query": query},
                    started_at=datetime.now(timezone.utc),
                )
            return pg_workflow_id
        except Exception as e:
            logger.error(f"Failed to pre-create workflow run: {e}")
            return None

    def _bridge_attachments(
        self, owner_id: Optional[str], *, persisted: bool
    ) -> List[Dict[str, Any]]:
        """Stage uploaded attachments as run-scoped artifacts the nodes can read.

        Bytes are read server-side from each attachment's ``upload_path`` and
        re-persisted through ``persist_new_artifact`` (size/sha256/storage key all
        derived server-side); only the resulting references enter the run state.
        """
        if not self._engine or not self.attachments or not owner_id:
            return []
        # Without a persisted run row the artifacts would be orphaned (no authz
        # parent), so skip the bridge for unowned/draft ids.
        if not persisted:
            return []
        from application.sandbox.artifacts_capture import persist_new_artifact
        from application.storage.storage_creator import StorageCreator

        storage = StorageCreator.get_storage()
        if len(self.attachments) > _MAX_INPUT_DOCUMENTS:
            dropped = len(self.attachments) - _MAX_INPUT_DOCUMENTS
            logger.warning(
                "Workflow run input documents exceed cap (%d); dropping %d attachment(s)",
                _MAX_INPUT_DOCUMENTS,
                dropped,
            )
        refs: List[Dict[str, Any]] = []
        for index, attachment in enumerate(self.attachments[:_MAX_INPUT_DOCUMENTS]):
            upload_path = attachment.get("upload_path") or attachment.get("path")
            if not upload_path:
                continue
            filename = attachment.get("filename") or "attachment"
            mime_type = attachment.get("mime_type") or "application/octet-stream"
            attachment_id = attachment.get("id", index)
            try:
                data = storage.get_file(upload_path).read()
            except Exception as exc:
                logger.error(
                    "Failed to read attachment %s for workflow run: %s",
                    attachment_id,
                    type(exc).__name__,
                )
                continue
            try:
                ref = persist_new_artifact(
                    user_id=owner_id,
                    kind="file",
                    data=data,
                    filename=filename,
                    mime_type=mime_type,
                    title=filename,
                    workflow_run_id=self._engine.workflow_run_id,
                )
            except Exception as exc:
                logger.error(
                    "Failed to persist attachment %s artifact: %s",
                    attachment_id,
                    type(exc).__name__,
                )
                continue
            if ref is None:
                continue
            refs.append(
                {
                    "artifact_id": ref["artifact_id"],
                    "ref": ref.get("ref"),
                    "filename": ref["filename"],
                    "mime_type": ref["mime_type"],
                }
            )
        return refs

    def _finalize_workflow_run(
        self, owner_id: Optional[str], pg_workflow_id: Optional[str], query: str
    ) -> None:
        """Write the run's terminal status/result; upsert the row if pre-creation was skipped."""
        if not self._engine:
            return
        try:
            run = WorkflowRun(
                workflow_id=self.workflow_id or "unknown",
                user=owner_id,
                status=self._determine_run_status(),
                inputs={"query": query},
                outputs=self._serialize_state(self._engine.state),
                steps=self._engine.get_execution_summary(),
                created_at=datetime.now(timezone.utc),
                completed_at=datetime.now(timezone.utc),
            )
            steps_json = [step.model_dump(mode="json") for step in run.steps]

            if not self.workflow_id or not owner_id:
                return
            with db_session() as conn:
                if pg_workflow_id is None:
                    pg_workflow_id = self._resolve_owned_workflow_pg_id(conn, owner_id)
                    if pg_workflow_id is None:
                        return
                runs_repo = WorkflowRunsRepository(conn)
                updated = False
                if self._run_persisted:
                    updated = runs_repo.finalize(
                        self._engine.workflow_run_id,
                        owner_id,
                        run.status.value,
                        result=run.outputs,
                        steps=steps_json,
                        ended_at=run.completed_at,
                    )
                    if not updated:
                        logger.warning(
                            "Workflow run %s finalize matched no row; "
                            "recovering via insert so terminal data is not lost",
                            self._engine.workflow_run_id,
                        )
                if not self._run_persisted or not updated:
                    runs_repo.create(
                        pg_workflow_id,
                        owner_id,
                        run.status.value,
                        run_id=self._engine.workflow_run_id,
                        inputs=run.inputs,
                        result=run.outputs,
                        steps=steps_json,
                        started_at=run.created_at,
                        ended_at=run.completed_at,
                    )
        except Exception as e:
            logger.error(f"Failed to save workflow run: {e}")

    def _determine_run_status(self) -> ExecutionStatus:
        if not self._engine or not self._engine.execution_log:
            return ExecutionStatus.COMPLETED
        for log in self._engine.execution_log:
            if log.get("status") == ExecutionStatus.FAILED.value:
                return ExecutionStatus.FAILED
        return ExecutionStatus.COMPLETED

    def _serialize_state(self, state: Dict[str, Any]) -> Dict[str, Any]:
        serialized: Dict[str, Any] = {}
        for key, value in state.items():
            serialized[key] = self._serialize_state_value(value)
        return serialized

    def _serialize_state_value(self, value: Any) -> Any:
        if isinstance(value, dict):
            return {
                str(dict_key): self._serialize_state_value(dict_value)
                for dict_key, dict_value in value.items()
            }
        if isinstance(value, list):
            return [self._serialize_state_value(item) for item in value]
        if isinstance(value, tuple):
            return [self._serialize_state_value(item) for item in value]
        if isinstance(value, datetime):
            return value.isoformat()
        if isinstance(value, (str, int, float, bool, type(None))):
            return value
        return str(value)