File size: 1,745 Bytes
bf39d5e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
# graph/research_graph.py

from langgraph.graph import StateGraph, END
from state.research_state import ResearchState
from agents.research_agent import research_agent
from agents.analyst_agent import analyst_agent
from agents.critic_agent import critic_agent
from agents.writer_agent import writer_agent

def should_revise(state: ResearchState) -> str:
    """

    Router function — reads state, decides next node.

    """
    needs_revision = state["confidence_scores"].get("needs_revision", False)
    iteration_count = state["iteration_count"]

    if needs_revision and iteration_count < 2:
        return "analyst"
    return "writer"

def increment_iteration(state: ResearchState) -> ResearchState:
    """

    Orchestrator updates iteration count after each critic cycle.

    """
    return {
        **state,
        "iteration_count": state["iteration_count"] + 1
    }

def build_graph():
    graph = StateGraph(ResearchState)

    # Register all nodes
    graph.add_node("researcher", research_agent)
    graph.add_node("analyst", analyst_agent)
    graph.add_node("critic", critic_agent)
    graph.add_node("iterate", increment_iteration)
    graph.add_node("writer", writer_agent)

    # Define edges — flow between nodes
    graph.set_entry_point("researcher")
    graph.add_edge("researcher", "analyst")
    graph.add_edge("analyst", "critic")
    graph.add_edge("critic", "iterate")

    # Conditional routing after iteration count updates
    graph.add_conditional_edges(
        "iterate",
        should_revise,
        {
            "analyst": "analyst",
            "writer": "writer"
        }
    )

    graph.add_edge("writer", END)

    return graph.compile()