Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -1,39 +1,72 @@
 from langgraph.graph import StateGraph
+from langgraph.prebuilt import ToolNode
+from typing import Annotated, TypedDict
 from agents import (
     product_manager_agent,
     project_manager_agent,
     software_architect_agent,
     software_engineer_agent,
-    quality_assurance_agent,
+    quality_assurance_agent
 )
+import os, json
 
-def run_pipeline_and_save(prompt: str):
-    state = {"prompt": prompt, "chat_log": []}
+# 1. Define your State
+class State(TypedDict):
+    prompt: str
+    pm_output: str
+    proj_output: str
+    arch_output: str
+    eng_output: str
+    qa_output: str
+
+# 2. Define each agent as a ToolNode
+pm_node = ToolNode.from_function(product_manager_agent.run)
+proj_node = ToolNode.from_function(project_manager_agent.run)
+arch_node = ToolNode.from_function(software_architect_agent.run)
+eng_node = ToolNode.from_function(software_engineer_agent.run)
+qa_node = ToolNode.from_function(quality_assurance_agent.run)
 
-    graph = StateGraph()
+# 3. Build graph with defined input/output
+def run_pipeline_and_save(prompt: str):
+    graph = StateGraph(State)
 
-    graph.add_node("PM", lambda s: product_manager_agent.run(s))
-    graph.add_node("ProjM", lambda s: project_manager_agent.run(s))
-    graph.add_node("Architect", lambda s: software_architect_agent.run(s))
-    graph.add_node("Engineer", lambda s: software_engineer_agent.run(s))
-    graph.add_node("QA", lambda s: quality_assurance_agent.run(s))
+    graph.add_node("ProductManager", pm_node)
+    graph.add_node("ProjectManager", proj_node)
+    graph.add_node("Architect", arch_node)
+    graph.add_node("Engineer", eng_node)
+    graph.add_node("QA", qa_node)
 
-    graph.set_entry_point("PM")
-    graph.add_edge("PM", "ProjM")
-    graph.add_edge("ProjM", "Architect")
+    graph.set_entry_point("ProductManager")
+    graph.add_edge("ProductManager", "ProjectManager")
+    graph.add_edge("ProjectManager", "Architect")
     graph.add_edge("Architect", "Engineer")
     graph.add_edge("Engineer", "QA")
-
     graph.set_finish_point("QA")
 
-    final_state = graph.execute(state)
+    final_graph = graph.compile()
 
-    # Save outputs
-    with open("output/index.html", "w") as f:
-        f.write(final_state.get("final_output", ""))
+    # Initial input
+    inputs = {"prompt": prompt}
+
+    # Run full chain
+    results = final_graph.invoke(inputs)
+
+    # Save log
+    chat_log = [
+        {"Product Manager": results["pm_output"]},
+        {"Project Manager": results["proj_output"]},
+        {"Software Architect": results["arch_output"]},
+        {"Software Engineer": results["eng_output"]},
+        {"QA Engineer": results["qa_output"]}
+    ]
+
+    os.makedirs("output", exist_ok=True)
 
     with open("output/agent_log.json", "w") as f:
-        import json
-        json.dump(final_state["chat_log"], f, indent=2)
+        json.dump(chat_log, f, indent=2)
+
+    # Write HTML file if generated
+    with open("output/index.html", "w") as f:
+        f.write(results["qa_output"])
 
-    return final_state["chat_log"], final_state.get("final_output", "")
+    return chat_log, results["qa_output"]