refactor: restructure graph creation to incorporate DocumentState and simplify node management

planning_ai/graph.py

@@ -1,4 +1,4 @@
-from langgraph.constants import Send
+from langgraph.constants import START, Send
 from langgraph.graph import END, StateGraph
 
 from planning_ai.nodes.hallucination_node import (
@@ -8,45 +8,44 @@ from planning_ai.nodes.hallucination_node import (
     map_hallucinations,
 )
 from planning_ai.nodes.map_node import (
-    all_summaries,
     collect_summaries,
     generate_summary,
     map_summaries,
 )
 from planning_ai.nodes.reduce_node import generate_final_summary
-from planning_ai.states import OverallState
+from planning_ai.states import DocumentState, OverallState
 
 
 def create_graph():
+    subgraph = StateGraph(DocumentState)
+    subgraph.add_node("generate_summary", generate_summary)
+    # subgraph.add_node("check_hallucination", check_hallucination)
+    # subgraph.add_node("fix_hallucination", fix_hallucination)
+
+    subgraph.add_conditional_edges(START, map_summaries, ["generate_summary"])
+    # subgraph.add_conditional_edges(
+    #     "generate_summary", map_hallucinations, ["check_hallucination"]
+    # )
+    # subgraph.add_conditional_edges(
+    #     "check_hallucination", map_fix_hallucinations, ["fix_hallucination"]
+    # )
+    # subgraph.add_conditional_edges(
+    #     "fix_hallucination", map_hallucinations, ["check_hallucination"]
+    # )
+    subgraph = subgraph.compile()
+
     graph = StateGraph(OverallState)
-    graph.add_node("generate_summary", generate_summary)
-    graph.add_node("check_hallucination", check_hallucination)
-    graph.add_node("fix_hallucination", fix_hallucination)
-    graph.add_node("collect_summaries", collect_summaries)
-    graph.add_node("generate_final_summary", generate_final_summary)
-
-    graph.set_conditional_entry_point(map_summaries, ["generate_summary"])
-    graph.add_conditional_edges(
-        "generate_summary", map_hallucinations, ["check_hallucination"]
-    )
-    graph.add_conditional_edges(
-        "check_hallucination",
-        map_fix_hallucinations,
-        ["fix_hallucination", "collect_summaries"],
-    )
-    graph.add_conditional_edges(
-        "fix_hallucination",
-        map_hallucinations,
-        ["check_hallucination", "collect_summaries"],
-    )
-    graph.add_conditional_edges(
-        "collect_summaries",
-        all_summaries,
-        {True: "generate_final_summary", False: "collect_summaries"},
-    )
-    graph.add_edge("generate_final_summary", END)
+    graph.add_node("summary_graph", subgraph)
+    # graph.add_node("collect_summaries", collect_summaries)
+
+    graph.add_edge(START, "summary_graph")
+    # graph.add_conditional_edges(
+    #     "summary_graph", map_hallucinations, ["collect_summaries"]
+    # )
+    # graph.add_edge("generate_final_summary", END)
 
     return graph.compile()
 
 
-print(create_graph().get_graph().draw_ascii())
+# graph = create_graph()
+# graph.get_graph().draw_png("test.png")