core/engine/src/runtime/graph_executor.py - main - alexefan136/flowstack

core/engine/src/runtime/graph_executor.py
332 строки12 KB
Alexander Efanov
upd fix
31 июл 2026, 19:17
31 июл 2026, 19:17d146d86
"""Graph Executor — выполняет flows как графы агентов.

Поддерживает два режима:
1. **Runner-based** — flow имеет ``runner`` (async generator событий)
2. **Graph-based** — flow задан графом ``nodes``/``edges`` (обход с ``node_executor``)
"""

from __future__ import annotations

from collections.abc import AsyncIterator, Awaitable, Callable
from typing import Any

import structlog

from src.primitives import (
    Flow,
    FlowNode,
    FlowNodeType,
    Task,
    utc_now,
)

logger = structlog.get_logger()

# Исполнитель узлов графа: (node, context) -> output dict
NodeExecutor = Callable[[FlowNode, dict[str, Any]], Awaitable[dict[str, Any]]]

# События, которые пробрасываются из runner как есть
_PASSTHROUGH_EVENTS = frozenset(
    {
        "agent_start",
        "agent_message",
        "agent_done",
        "reasoning",
        "tool_call",
        "tool_result",
        "flow_done",
    }
)


def _get_flow_attr(flow: Any, attr: str, default: Any = None) -> Any:
    """
    Получить атрибут из flow (поддерживает dict и объекты).

    Args:
        flow: Flow dict или объект с атрибутами
        attr: Имя атрибута
        default: Значение по умолчанию
    """
    if isinstance(flow, dict):
        return flow.get(attr, default)
    return getattr(flow, attr, default)


class GraphExecutor:
    """
    Исполнитель flows.

    Args:
        node_executor: callable для выполнения узлов графа (для graph-based flows).
            Сигнатура: ``(node: FlowNode, context: dict) -> dict``.
    """

    def __init__(self, node_executor: NodeExecutor | None = None) -> None:
        self.node_executor = node_executor
        logger.info("graph_executor.initialized", graph_mode=node_executor is not None)

    async def start(self) -> None:
        """Запустить executor."""
        logger.info("graph_executor.started")

    async def stop(self) -> None:
        """Остановить executor."""
        logger.info("graph_executor.stopped")

    # ========================================================================
    # Flow resolution
    # ========================================================================

    def _ensure_flow(self, flow: Any) -> Flow:
        """Привести flow к объекту Flow (из dict при необходимости)."""
        if isinstance(flow, Flow):
            return flow
        if isinstance(flow, dict):
            return Flow.from_dict(flow)
        raise TypeError(f"Unsupported flow type: {type(flow).__name__}")

    # ========================================================================
    # Non-streaming
    # ========================================================================

    async def run(self, flow: Any, input_data: dict[str, Any]) -> Task:
        """
        Выполнить flow синхронно (собрать все события).

        Выбирает режим автоматически:
        - есть ``runner`` → runner-based
        - есть ``nodes`` → graph-based

        Returns:
            Task с результатом выполнения.
        """
        flow_id = _get_flow_attr(flow, "id", "unknown")
        flow_name = _get_flow_attr(flow, "name", "Unknown Flow")

        task = Task(
            title=flow_name,
            flow_id=flow_id,
            flow_name=flow_name,
            input_data=input_data,
        )
        task.start()

        logger.info("flow.run.start", flow_id=flow_id, flow_name=flow_name)

        try:
            runner = _get_flow_attr(flow, "runner")
            if runner:
                output = await self._run_with_runner(runner, input_data)
            elif _get_flow_attr(flow, "nodes"):
                output = await self._run_graph(flow, input_data)
            else:
                raise ValueError(f"Flow '{flow_id}' has no runner or graph nodes")

            task.complete(output)
            logger.info(
                "flow.run.completed",
                flow_id=flow_id,
                duration_ms=task.duration_ms,
                tokens=output.get("tokens", 0),
            )

        except Exception as e:
            logger.error(
                "flow.run.failed",
                flow_id=flow_id,
                error=str(e),
                duration_ms=task.duration_ms,
            )
            task.fail(str(e))

        return task

    async def _run_with_runner(self, runner: Any, input_data: dict[str, Any]) -> dict[str, Any]:
        """Выполнить flow через runner (async generator), собрать результат."""
        full_output = ""
        reasoning = ""
        agent_outputs: dict[str, str] = {}
        total_tokens = 0

        async for event in runner(input_data):
            event_type = event.get("type")

            if event_type == "agent_message":
                full_output += event.get("content", "")

            elif event_type == "reasoning":
                reasoning += event.get("content", "")

            elif event_type == "agent_done":
                agent = event.get("agent", "unknown")
                agent_outputs[agent] = event.get("output", "")
                total_tokens += event.get("tokens", 0)

            elif event_type == "flow_done":
                full_output = event.get("output", full_output)
                total_tokens = event.get("tokens", total_tokens)

            elif event_type == "error":
                raise RuntimeError(event.get("error", "Unknown error"))

        return {
            "output": full_output,
            "reasoning": reasoning or None,
            "agent_outputs": agent_outputs,
            "tokens": total_tokens,
        }

    async def _run_graph(self, flow: Any, input_data: dict[str, Any]) -> dict[str, Any]:
        """Выполнить flow как граф (обход узлов от START до END)."""
        flow_obj = self._ensure_flow(flow)

        errors = flow_obj.validate()
        if errors:
            raise ValueError(f"Invalid flow: {'; '.join(errors)}")

        start = flow_obj.get_start_node()
        if start is None:
            raise ValueError("Flow has no START node")

        context: dict[str, Any] = dict(input_data)
        visited: set[str] = set()
        frontier: list[FlowNode] = [start]
        steps: list[dict[str, Any]] = []

        while frontier:
            next_frontier: list[FlowNode] = []
            for node in frontier:
                if node.id in visited:
                    continue
                visited.add(node.id)

                if node.type == FlowNodeType.END:
                    continue

                if node.is_executable():
                    output = await self._execute_node(node, context)
                    context[node.id] = output
                    steps.append(
                        {"node_id": node.id, "node_type": node.type.value, "output": output}
                    )

                next_frontier.extend(flow_obj.get_next_nodes(node.id, context))
            frontier = next_frontier

        return {"output": context, "steps": steps}

    async def _execute_node(self, node: FlowNode, context: dict[str, Any]) -> dict[str, Any]:
        """Выполнить узел графа через node_executor."""
        if self.node_executor is None:
            raise ValueError(f"No node executor for node '{node.id}' (type={node.type.value})")
        return await self.node_executor(node, context)

    # ========================================================================
    # Streaming
    # ========================================================================

    async def run_stream(
        self, flow: Any, input_data: dict[str, Any]
    ) -> AsyncIterator[dict[str, Any]]:
        """
        Выполнить flow со стримингом событий.

        Yields:
            События выполнения flow (flow_start, agent_*, reasoning, tool_*,
            node_*, flow_done, error).
        """
        flow_id = _get_flow_attr(flow, "id", "unknown")
        flow_name = _get_flow_attr(flow, "name", "Unknown Flow")

        logger.info("flow.stream.start", flow_id=flow_id, flow_name=flow_name)

        yield {
            "type": "flow_start",
            "flow_id": flow_id,
            "flow_name": flow_name,
            "timestamp": utc_now().isoformat(),
        }

        runner = _get_flow_attr(flow, "runner")
        if runner:
            async for event in self._stream_with_runner(runner, input_data, flow_id):
                yield event
        elif _get_flow_attr(flow, "nodes"):
            async for event in self._stream_graph(flow, input_data, flow_id):
                yield event
        else:
            yield {"type": "error", "error": f"Flow '{flow_id}' has no runner or graph nodes"}

    async def _stream_with_runner(
        self, runner: Any, input_data: dict[str, Any], flow_id: str
    ) -> AsyncIterator[dict[str, Any]]:
        """Стриминг через runner (пробрасывает события)."""
        try:
            async for event in runner(input_data):
                event_type = event.get("type")

                if event_type == "error":
                    yield {"type": "error", "error": event.get("error", "Unknown error")}
                    return

                # Известные события пробрасываем, неизвестные — как есть
                yield event

        except Exception as e:
            logger.error("flow.stream.error", flow_id=flow_id, error=str(e))
            yield {"type": "error", "error": str(e)}

    async def _stream_graph(
        self, flow: Any, input_data: dict[str, Any], flow_id: str
    ) -> AsyncIterator[dict[str, Any]]:
        """Стриминг обхода графа (node_start / node_done)."""
        try:
            flow_obj = self._ensure_flow(flow)

            errors = flow_obj.validate()
            if errors:
                yield {"type": "error", "error": f"Invalid flow: {'; '.join(errors)}"}
                return

            start = flow_obj.get_start_node()
            if start is None:
                yield {"type": "error", "error": "Flow has no START node"}
                return

            context: dict[str, Any] = dict(input_data)
            visited: set[str] = set()
            frontier: list[FlowNode] = [start]

            while frontier:
                next_frontier: list[FlowNode] = []
                for node in frontier:
                    if node.id in visited:
                        continue
                    visited.add(node.id)

                    if node.type == FlowNodeType.END:
                        continue

                    if node.is_executable():
                        yield {
                            "type": "node_start",
                            "node_id": node.id,
                            "node_type": node.type.value,
                            "label": node.label,
                        }
                        output = await self._execute_node(node, context)
                        context[node.id] = output
                        yield {"type": "node_done", "node_id": node.id, "output": output}

                    next_frontier.extend(flow_obj.get_next_nodes(node.id, context))
                frontier = next_frontier

            yield {"type": "flow_done", "flow_id": flow_id, "output": context}

        except Exception as e:
            logger.error("flow.stream.graph_error", flow_id=flow_id, error=str(e))
            yield {"type": "error", "error": str(e)}


__all__ = ["GraphExecutor", "NodeExecutor"]