Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -19,16 +19,15 @@ state_schema = {
     "chat_log": list,
 }
 
-# Wrap each agent
-pm_node = ToolNode.from_callable(product_manager_agent.run)
-proj_node = ToolNode.from_callable(project_manager_agent.run)
-arch_node = ToolNode.from_callable(software_architect_agent.run)
-dev_node = ToolNode.from_callable(software_engineer_agent.run)
-qa_node = ToolNode.from_callable(quality_assurance_agent.run)
-
-# Define the graph
+# Define the ToolNodes directly
+pm_node = ToolNode(product_manager_agent.run)
+proj_node = ToolNode(project_manager_agent.run)
+arch_node = ToolNode(software_architect_agent.run)
+dev_node = ToolNode(software_engineer_agent.run)
+qa_node = ToolNode(quality_assurance_agent.run)
+
+# Build the graph
 graph = StateGraph(state_schema)
-
 graph.add_node("ProductManager", pm_node)
 graph.add_node("ProjectManager", proj_node)
 graph.add_node("SoftwareArchitect", arch_node)
@@ -42,15 +41,12 @@ graph.add_edge("SoftwareArchitect", "SoftwareEngineer")
 graph.add_edge("SoftwareEngineer", "QualityAssurance")
 graph.add_edge("QualityAssurance", END)
 
-# Compile it
 compiled_graph = graph.compile()
 
-# Define the wrapper
 def run_pipeline_and_save(prompt: str):
     initial_state = {
         "input": prompt,
         "chat_log": [],
     }
     final_state = compiled_graph.invoke(initial_state)
-
     return final_state["chat_log"], final_state["qa_output"]