Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -1,57 +1,56 @@
-import os, json
-from langgraph.graph import StateGraph
-from langgraph.prebuilt.tool_node import ToolNode
-
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode
 from agents import (
     product_manager_agent,
     project_manager_agent,
     software_architect_agent,
     software_engineer_agent,
-    quality_assurance_agent
+    quality_assurance_agent,
 )
-from utils.zip_output import zip_output
-
-def run_pipeline_and_save(prompt):
-    conversation = []
-
-    # Initialize LangGraph state graph
-    graph = StateGraph()
-    graph.set_input_type(str)
-    graph.set_output_type(dict)
-
-    # Add tool-wrapped nodes
-    pm_node = ToolNode(product_manager_agent.run)
-    proj_node = ToolNode(project_manager_agent.run)
-    arch_node = ToolNode(software_architect_agent.run)
-    eng_node = ToolNode(software_engineer_agent.run)
-    qa_node = ToolNode(quality_assurance_agent.run)
-
-    # Add nodes to graph
-    graph.add_node("product_manager", pm_node)
-    graph.add_node("project_manager", proj_node)
-    graph.add_node("software_architect", arch_node)
-    graph.add_node("software_engineer", eng_node)
-    graph.add_node("quality_assurance", qa_node)
-
-    # Define execution flow
-    graph.set_entry_point("product_manager")
-    graph.add_edge("product_manager", "project_manager")
-    graph.add_edge("project_manager", "software_architect")
-    graph.add_edge("software_architect", "software_engineer")
-    graph.add_edge("software_engineer", "quality_assurance")
-
-    # Compile and execute
-    app = graph.compile()
-    outputs = app.invoke(prompt)
-
-    # Capture outputs
-    for k, v in outputs.items():
-        conversation.append({"agent": k, "message": v})
-
-    # Save to log and zip output
-    os.makedirs("output", exist_ok=True)
-    with open("output/agent_log.json", "w") as f:
-        json.dump(conversation, f, indent=2)
-    
-    zip_path = zip_output()
-    return [(x["agent"], x["message"]) for x in conversation], zip_path
+
+# Define the state schema
+state_schema = {
+    "input": str,
+    "pm_output": str,
+    "proj_output": str,
+    "arch_output": str,
+    "dev_output": str,
+    "qa_output": str,
+    "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
+graph = StateGraph(state_schema)
+
+graph.add_node("ProductManager", pm_node)
+graph.add_node("ProjectManager", proj_node)
+graph.add_node("SoftwareArchitect", arch_node)
+graph.add_node("SoftwareEngineer", dev_node)
+graph.add_node("QualityAssurance", qa_node)
+
+graph.set_entry_point("ProductManager")
+graph.add_edge("ProductManager", "ProjectManager")
+graph.add_edge("ProjectManager", "SoftwareArchitect")
+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"]