Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -8,13 +8,17 @@ from agents import (
     quality_assurance_agent,
 )
 from langchain_core.messages import HumanMessage, AIMessage
-from typing import TypedDict, List
 from langchain_core.messages.base import BaseMessage
+from langchain_core.tools import StructuredTool
+from typing import TypedDict, List
 
+
+# Define LangGraph input/output state types
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: list
 
+
 class OutputState(TypedDict):
     pm_output: str
     proj_output: str
@@ -23,36 +27,64 @@ class OutputState(TypedDict):
     qa_output: str
     chat_log: list
 
-def bridge_to_product_manager(state: dict) -> dict:
+
+# -----------------------------
+# Bridge Node Logic (Fixed ✅)
+# -----------------------------
+def bridge_to_product_manager_fn(state: dict) -> dict:
     """
-    Convert HumanMessage into structured AIMessage for Product Manager agent.
+    Converts HumanMessage into structured AIMessage so the Product Manager agent can consume it.
     """
-    user_prompt = state["messages"][-1].content
+    if "messages" not in state or not isinstance(state["messages"], list):
+        raise ValueError("Input must contain a 'messages' list.")
+
+    last_msg = state["messages"][-1]
+    if not isinstance(last_msg, HumanMessage):
+        raise ValueError("Last message must be a HumanMessage.")
 
-    system_spec = f"""# Stakeholder Prompt
+    # Structured prompt injected as system-level AIMessage
+    structured_prompt = f"""# Stakeholder Prompt
 
 A new product request has been submitted:
 
-"{user_prompt}"
+"{last_msg.content}"
 
-Please analyze and generate a structured product specification including goals, features, user needs, and KPIs.
+Please analyze and convert this into a structured product specification including:
+- Goals
+- Features
+- User stories
+- KPIs
 """
 
     return {
-        "messages": state["messages"] + [AIMessage(content=system_spec)],
-        "chat_log": state["chat_log"] + [{"role": "System", "content": system_spec}],
+        "messages": state["messages"] + [AIMessage(content=structured_prompt)],
+        "chat_log": state["chat_log"] + [{"role": "System", "content": structured_prompt}],
     }
 
-# Define the bridge node as a simple function node
-def bridge_node(state: dict) -> dict:
-    return bridge_to_product_manager(state)
 
+# ✅ 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
+# -----------------------------
 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])
 
+
+# -----------------------------
+# Build LangGraph
+# -----------------------------
 graph = StateGraph(input=InputState, output=OutputState)
 
 graph.add_node("Bridge", bridge_node)
@@ -72,10 +104,16 @@ graph.add_edge("QualityAssurance", END)
 
 compiled_graph = graph.compile()
 
+
+# -----------------------------
+# Run the full pipeline
+# -----------------------------
 def run_pipeline_and_save(prompt: str):
     initial_state = {
         "messages": [HumanMessage(content=prompt)],
         "chat_log": [],
     }
+
     final_state = compiled_graph.invoke(initial_state)
+
     return final_state["chat_log"], final_state["qa_output"]