Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -9,11 +9,12 @@ from agents import (
 )
 from langchain_core.messages import HumanMessage, AIMessage
 from langchain_core.messages.base import BaseMessage
-from langchain_core.tools import StructuredTool
 from typing import TypedDict, List
 
 
-# Define LangGraph input/output state types
+# --------------------------
+# State Definitions
+# --------------------------
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: list
@@ -28,53 +29,43 @@ class OutputState(TypedDict):
     chat_log: list
 
 
-# -----------------------------
-# Bridge Node Logic (Fixed ✅)
-# -----------------------------
-def bridge_to_product_manager_fn(state: dict) -> dict:
-    """
-    Converts HumanMessage into structured AIMessage so the Product Manager agent can consume it.
-    """
-    if "messages" not in state or not isinstance(state["messages"], list):
-        raise ValueError("Input must contain a 'messages' list.")
+# --------------------------
+# Correctly Implemented Bridge Node
+# --------------------------
+def bridge_to_product_manager(state: dict) -> dict:
+    """Ensure AIMessage is added after HumanMessage."""
+    messages = state.get("messages", [])
+    chat_log = state.get("chat_log", [])
 
-    last_msg = state["messages"][-1]
-    if not isinstance(last_msg, HumanMessage):
-        raise ValueError("Last message must be a HumanMessage.")
+    if not messages or not isinstance(messages[-1], HumanMessage):
+        raise ValueError("Expected last message to be a HumanMessage")
 
-    # Structured prompt injected as system-level AIMessage
-    structured_prompt = f"""# Stakeholder Prompt
+    prompt = messages[-1].content
+
+    structured = f"""# Stakeholder Prompt
 
 A new product request has been submitted:
 
-"{last_msg.content}"
+"{prompt}"
 
-Please analyze and convert this into a structured product specification including:
+Please convert it into a structured product specification including:
 - Goals
 - Features
-- User stories
+- User Stories
 - KPIs
 """
 
+    ai_msg = AIMessage(content=structured)
+
     return {
-        "messages": state["messages"] + [AIMessage(content=structured_prompt)],
-        "chat_log": state["chat_log"] + [{"role": "System", "content": structured_prompt}],
+        "messages": messages + [ai_msg],
+        "chat_log": chat_log + [{"role": "System", "content": structured}],
     }
 
 
-# ✅ Manually wrap bridge function as StructuredTool
-bridge_tool = StructuredTool.from_function(
-    func=bridge_to_product_manager_fn,
-    name="bridge_to_product_manager",
-    description="Converts HumanMessage into structured system prompt for Product Manager."
-)
-
-bridge_node = ToolNode([bridge_tool])
-
-
-# -----------------------------
-# Define Other Agent ToolNodes
-# -----------------------------
+# --------------------------
+# ToolNodes for All Agents
+# --------------------------
 pm_node = ToolNode([product_manager_agent.run])
 proj_node = ToolNode([project_manager_agent.run])
 arch_node = ToolNode([software_architect_agent.run])
@@ -82,12 +73,14 @@ dev_node = ToolNode([software_engineer_agent.run])
 qa_node = ToolNode([quality_assurance_agent.run])
 
 
-# -----------------------------
-# Build LangGraph
-# -----------------------------
+# --------------------------
+# Graph Definition
+# --------------------------
 graph = StateGraph(input=InputState, output=OutputState)
 
-graph.add_node("Bridge", bridge_node)
+# ✅ Native function for bridge node (no ToolNode needed)
+graph.add_node("Bridge", bridge_to_product_manager)
+
 graph.add_node("ProductManager", pm_node)
 graph.add_node("ProjectManager", proj_node)
 graph.add_node("SoftwareArchitect", arch_node)
@@ -105,9 +98,9 @@ graph.add_edge("QualityAssurance", END)
 compiled_graph = graph.compile()
 
 
-# -----------------------------
-# Run the full pipeline
-# -----------------------------
+# --------------------------
+# Pipeline Entry Function
+# --------------------------
 def run_pipeline_and_save(prompt: str):
     initial_state = {
         "messages": [HumanMessage(content=prompt)],
@@ -115,5 +108,4 @@ def run_pipeline_and_save(prompt: str):
     }
 
     final_state = compiled_graph.invoke(initial_state)
-
     return final_state["chat_log"], final_state["qa_output"]