Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -12,9 +12,6 @@ from langchain_core.messages.base import BaseMessage
 from typing import TypedDict, List
 
 
-# --------------------------
-# State Definitions
-# --------------------------
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: list
@@ -29,58 +26,49 @@ class OutputState(TypedDict):
     chat_log: list
 
 
-# --------------------------
-# Correctly Implemented Bridge Node
-# --------------------------
-def bridge_to_product_manager(state: dict) -> dict:
-    """Ensure AIMessage is added after HumanMessage."""
+# ✅ Fixed Bridge node (must return a proper state with messages and chat_log)
+def bridge_to_pm(state: dict) -> dict:
     messages = state.get("messages", [])
     chat_log = state.get("chat_log", [])
 
     if not messages or not isinstance(messages[-1], HumanMessage):
         raise ValueError("Expected last message to be a HumanMessage")
 
-    prompt = messages[-1].content
+    user_prompt = messages[-1].content
 
-    structured = f"""# Stakeholder Prompt
+    structured_context = f"""# Stakeholder Prompt
 
-A new product request has been submitted:
+"{user_prompt}"
 
-"{prompt}"
-
-Please convert it into a structured product specification including:
+Please convert this into a structured product specification including:
 - Goals
-- Features
+- Key Features
 - User Stories
-- KPIs
+- Success Metrics
 """
 
-    ai_msg = AIMessage(content=structured)
+    new_ai_msg = AIMessage(content=structured_context)
 
     return {
-        "messages": messages + [ai_msg],
-        "chat_log": chat_log + [{"role": "System", "content": structured}],
+        "messages": messages + [new_ai_msg],  # Important: last item must be AIMessage
+        "chat_log": chat_log + [{"role": "System", "content": structured_context}],
     }
 
 
-# --------------------------
-# ToolNodes for All Agents
-# --------------------------
+# ✅ Wrap bridge node in a ToolNode
+bridge_node = ToolNode([bridge_to_pm])
+
+# ✅ Other agents as ToolNodes
 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])
 
-
-# --------------------------
-# Graph Definition
-# --------------------------
+# ✅ Build LangGraph
 graph = StateGraph(input=InputState, output=OutputState)
 
-# ✅ Native function for bridge node (no ToolNode needed)
-graph.add_node("Bridge", bridge_to_product_manager)
-
+graph.add_node("Bridge", bridge_node)
 graph.add_node("ProductManager", pm_node)
 graph.add_node("ProjectManager", proj_node)
 graph.add_node("SoftwareArchitect", arch_node)
@@ -98,14 +86,15 @@ graph.add_edge("QualityAssurance", END)
 compiled_graph = graph.compile()
 
 
-# --------------------------
-# Pipeline Entry Function
-# --------------------------
+# ✅ Run function
 def run_pipeline_and_save(prompt: str):
     initial_state = {
         "messages": [HumanMessage(content=prompt)],
         "chat_log": [],
     }
 
+    # Sanity check
+    assert isinstance(initial_state["messages"][-1], HumanMessage)
+
     final_state = compiled_graph.invoke(initial_state)
-    return final_state["chat_log"], final_state["qa_output"]
+    return final_state["chat_log"], final_state["qa_output"]