ui/agent/graph/workflow.py

# ui/agent/graph/workflow.py
from __future__ import annotations

from functools import lru_cache

from langgraph.graph import END, START, StateGraph

from .nodes import (
    consolidator_node,
    discovery_node,
    fan_out_research,
    planner_node,
    researcher_node,
)
from .state import AgentState


@lru_cache(maxsize=1)
def build_workflow():
    builder = StateGraph(AgentState)
    builder.add_node("discovery", discovery_node)
    builder.add_node("planner", planner_node)
    builder.add_node("researcher", researcher_node)
    builder.add_node("consolidator", consolidator_node)

    builder.add_edge(START, "discovery")
    builder.add_edge("discovery", "planner")
    builder.add_conditional_edges("planner", fan_out_research, ["researcher"])
    builder.add_edge("researcher", "consolidator")
    builder.add_edge("consolidator", END)
    return builder.compile()