Upload 2 files

utils/langgraph_pipeline.py

@@ -21,6 +21,8 @@ from agents import (
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: List[Dict[str, Any]]
+    iteration: int
+    feedback: str
 
 class OutputState(TypedDict):
     pm_output: str
@@ -30,6 +32,8 @@ class OutputState(TypedDict):
     dev_output: str
     qa_output: str
     chat_log: List[Dict[str, Any]]
+    iteration: int
+    feedback: str
 
 # ——————————————
 # 2) Wrap agents so they see full history
@@ -37,12 +41,21 @@ class OutputState(TypedDict):
 def wrap_agent(agent_run, output_key: str):
     def node(state: Dict[str, Any]) -> Dict[str, Any]:
         history = state["messages"]
-        log     = state["chat_log"]
-        result  = agent_run({"messages": history, "chat_log": log})
+        log = state["chat_log"]
+        iteration = state.get("iteration", 0)
+        feedback = state.get("feedback", "")
+        
+        # Add feedback to the prompt if it exists
+        if feedback:
+            history = history + [AIMessage(content=f"Previous feedback: {feedback}")]
+        
+        result = agent_run({"messages": history, "chat_log": log})
         return {
             "messages": history + result["messages"],
-            "chat_log":  result["chat_log"],
-            output_key:  result[output_key],
+            "chat_log": result["chat_log"],
+            output_key: result[output_key],
+            "iteration": iteration,
+            "feedback": feedback
         }
     return node
 
@@ -51,50 +64,84 @@ def wrap_agent(agent_run, output_key: str):
 # ——————————————
 def bridge_to_pm(state: Dict[str, Any]) -> Dict[str, Any]:
     history = state["messages"]
-    log     = state["chat_log"]
+    log = state["chat_log"]
+    iteration = state.get("iteration", 0)
+    feedback = state.get("feedback", "")
+    
     if not history or not isinstance(history[-1], HumanMessage):
         raise ValueError("bridge_to_pm expected a HumanMessage at history end")
+    
     prompt = history[-1].content
     spec_prompt = (
-        f"# Stakeholder Prompt\n\n"
+        f"# Stakeholder Prompt (Iteration {iteration})\n\n"
         f"\"{prompt}\"\n\n"
+    )
+    
+    if feedback:
+        spec_prompt += f"Previous feedback to consider:\n{feedback}\n\n"
+    
+    spec_prompt += (
         "Generate a structured product specification including:\n"
         "- Goals\n"
         "- Key features\n"
         "- User stories\n"
         "- Success metrics\n"
     )
+    
     return {
         "messages": [AIMessage(content=spec_prompt)],
         "chat_log": log + [{"role": "System", "content": spec_prompt}],
+        "iteration": iteration,
+        "feedback": feedback
     }
 
 # ——————————————
-# 4) Build & compile the LangGraph
+# 4) Feedback Loop Handler
 # ——————————————
-graph = StateGraph(input=InputState, output=OutputState)
+def handle_feedback(state: Dict[str, Any]) -> Dict[str, Any]:
+    qa_output = state["qa_output"]
+    iteration = state.get("iteration", 0)
+    
+    # Check if we need another iteration
+    if iteration < 3:  # Maximum 3 iterations
+        return {
+            "messages": state["messages"],
+            "chat_log": state["chat_log"],
+            "iteration": iteration + 1,
+            "feedback": f"Iteration {iteration + 1} feedback: {qa_output}"
+        }
+    return END
 
-graph.add_node("BridgePM",         bridge_to_pm)
-graph.add_node("ProductManager",   wrap_agent(product_manager_agent.run,   "pm_output"))
-graph.add_node("ProjectManager",   wrap_agent(project_manager_agent.run,   "proj_output"))
-graph.add_node("SoftwareArchitect",wrap_agent(software_architect_agent.run, "arch_output"))
-graph.add_node("UIDesigner",       wrap_agent(ui_designer_agent.run,       "ui_design_output"))
-graph.add_node("SoftwareEngineer", wrap_agent(software_engineer_agent.run,  "dev_output"))
-graph.add_node("QualityAssurance", wrap_agent(quality_assurance_agent.run,  "qa_output"))
+# ——————————————
+# 5) Build & compile the LangGraph
+# ——————————————
+graph = StateGraph(input=InputState, output=OutputState)
 
+# Add nodes
+graph.add_node("BridgePM", bridge_to_pm)
+graph.add_node("ProductManager", wrap_agent(product_manager_agent.run, "pm_output"))
+graph.add_node("ProjectManager", wrap_agent(project_manager_agent.run, "proj_output"))
+graph.add_node("SoftwareArchitect", wrap_agent(software_architect_agent.run, "arch_output"))
+graph.add_node("UIDesigner", wrap_agent(ui_designer_agent.run, "ui_design_output"))
+graph.add_node("SoftwareEngineer", wrap_agent(software_engineer_agent.run, "dev_output"))
+graph.add_node("QualityAssurance", wrap_agent(quality_assurance_agent.run, "qa_output"))
+graph.add_node("FeedbackHandler", handle_feedback)
+
+# Add edges with feedback loop
 graph.set_entry_point("BridgePM")
-graph.add_edge("BridgePM",         "ProductManager")
-graph.add_edge("ProductManager",   "ProjectManager")
-graph.add_edge("ProjectManager",   "SoftwareArchitect")
-graph.add_edge("SoftwareArchitect","UIDesigner")
-graph.add_edge("UIDesigner",       "SoftwareEngineer")
+graph.add_edge("BridgePM", "ProductManager")
+graph.add_edge("ProductManager", "ProjectManager")
+graph.add_edge("ProjectManager", "SoftwareArchitect")
+graph.add_edge("SoftwareArchitect", "UIDesigner")
+graph.add_edge("UIDesigner", "SoftwareEngineer")
 graph.add_edge("SoftwareEngineer", "QualityAssurance")
-graph.add_edge("QualityAssurance", END)
+graph.add_edge("QualityAssurance", "FeedbackHandler")
+graph.add_edge("FeedbackHandler", "BridgePM")  # Feedback loop back to start
 
 compiled_graph = graph.compile()
 
 # ——————————————
-# 5) Parse spec into sections
+# 6) Parse spec into sections
 # ——————————————
 def parse_spec(spec: str) -> Dict[str, List[str]]:
     sections: Dict[str, List[str]] = {}
@@ -105,11 +152,11 @@ def parse_spec(spec: str) -> Dict[str, List[str]]:
     return sections
 
 # ——————————————
-# 6) Run pipeline, generate site, zip, return (chat_log, zip_path)
+# 7) Run pipeline, generate site, zip, return (chat_log, zip_path)
 # ——————————————
 def run_pipeline_and_save(prompt: str) -> Tuple[List[Dict[str, Any]], str]:
     # a) invoke agents
-    initial_state = {"messages": [HumanMessage(content=prompt)], "chat_log": []}
+    initial_state = {"messages": [HumanMessage(content=prompt)], "chat_log": [], "iteration": 0, "feedback": ""}
     final_state   = compiled_graph.invoke(initial_state)
 
     chat_log  = final_state["chat_log"]