api/src/workflows/engine.py - master - ncit/multiagentsystem

api/src/workflows/engine.py
144 строки4 KB
Nikita
feat: initial multi-agent platform scaffold
09 июн 2026, 00:48
09 июн 2026, 00:4819cef10
"""DAG workflow executor — runs agent pipelines defined as directed acyclic graphs.

Workflow DAGs are stored as JSONB in Supabase and executed via Trigger.dev.
This engine handles the local execution logic for each DAG node.
"""
from __future__ import annotations

import asyncio
from dataclasses import dataclass
from enum import Enum
from typing import Any

from src.agents.factory import create_agent


class NodeType(str, Enum):
    AGENT = "agent"
    CONDITION = "condition"
    MERGE = "merge"
    PARALLEL = "parallel"


@dataclass
class DAGNode:
    id: str
    type: NodeType
    agent_type: str | None = None
    prompt_id: str | None = None
    config: dict[str, Any] | None = None


@dataclass
class DAGEdge:
    from_node: str
    to_node: str
    condition: str | None = None  # for conditional edges


@dataclass
class WorkflowDAG:
    nodes: list[DAGNode]
    edges: list[DAGEdge]

    def get_node(self, node_id: str) -> DAGNode | None:
        return next((n for n in self.nodes if n.id == node_id), None)

    def get_children(self, node_id: str) -> list[str]:
        return [e.to_node for e in self.edges if e.from_node == node_id]

    def get_parents(self, node_id: str) -> list[str]:
        return [e.from_node for e in self.edges if e.to_node == node_id]

    def get_roots(self) -> list[str]:
        child_nodes = {e.to_node for e in self.edges}
        return [n.id for n in self.nodes if n.id not in child_nodes]


async def execute_workflow(
    dag: WorkflowDAG,
    user_id: str,
    input_data: dict[str, Any],
) -> dict[str, Any]:
    """Execute a DAG workflow, running agent nodes in topological order.

    Returns a dict mapping node_id → output for each node.
    """
    results: dict[str, Any] = {}
    roots = dag.get_roots()

    if not roots:
        raise ValueError("DAG has no root nodes (no entry point)")

    # BFS topological execution
    queue = list(roots)
    visited: set[str] = set()

    while queue:
        # Find nodes whose parents have all completed
        ready = [
            nid for nid in queue
            if all(p in visited for p in dag.get_parents(nid))
        ]

        if not ready:
            break

        # Run ready nodes in parallel if multiple
        tasks = []
        for node_id in ready:
            queue.remove(node_id)
            visited.add(node_id)
            node = dag.get_node(node_id)
            if node:
                parent_outputs = {
                    p: results.get(p) for p in dag.get_parents(node_id)
                }
                tasks.append(_execute_node(node, user_id, input_data, parent_outputs))

        node_results = await asyncio.gather(*tasks)

        for node_id, result in zip(ready, node_results):
            results[node_id] = result
            # Add children to queue
            for child in dag.get_children(node_id):
                if child not in queue:
                    queue.append(child)

    return results


async def _execute_node(
    node: DAGNode,
    user_id: str,
    input_data: dict[str, Any],
    parent_outputs: dict[str, Any],
) -> Any:
    """Execute a single DAG node."""
    if node.type == NodeType.AGENT:
        # Build message from parent outputs + input
        parent_text = "\n".join(
            str(out) for out in parent_outputs.values() if out
        )
        message = parent_text or input_data.get("message", "")

        agent = create_agent(
            agent_type=node.agent_type or "general",
            user_id=user_id,
            instructions_override=None,  # could load from prompt_id
        )
        response = await agent.arun(message)
        return response.content

    elif node.type == NodeType.MERGE:
        # Merge all parent outputs into one string
        merged = "\n\n---\n\n".join(
            str(out) for out in parent_outputs.values() if out
        )
        return merged

    elif node.type == NodeType.CONDITION:
        # TODO: evaluate condition expression
        return parent_outputs

    return None