Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -1,6 +1,7 @@
+import os, json
 from langgraph.graph import StateGraph
-from langgraph.prebuilt import ToolNode
-from typing import Annotated, TypedDict
+from langgraph.prebuilt.tool_node import ToolNode
+
 from agents import (
     product_manager_agent,
     project_manager_agent,
@@ -8,65 +9,49 @@ from agents import (
     software_engineer_agent,
     quality_assurance_agent
 )
-import os, json
-
-# 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)
-
-# 3. Build graph with defined input/output
-def run_pipeline_and_save(prompt: str):
-    graph = StateGraph(State)
-
-    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("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_graph = graph.compile()
-
-    # 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"]}
-    ]
-
+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(chat_log, f, indent=2)
-
-    # Write HTML file if generated
-    with open("output/index.html", "w") as f:
-        f.write(results["qa_output"])
-
-    return chat_log, results["qa_output"]
+        json.dump(conversation, f, indent=2)
+    
+    zip_path = zip_output()
+    return [(x["agent"], x["message"]) for x in conversation], zip_path