Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -1,5 +1,14 @@
+# MAC/utils/langgraph_pipeline.py
+
+from typing import TypedDict, List
+
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages.base import BaseMessage
+from langchain_core.tools.structured import StructuredTool
+
 from langgraph.graph import StateGraph, END
 from langgraph.prebuilt import ToolNode
+
 from agents import (
     product_manager_agent,
     project_manager_agent,
@@ -7,16 +16,14 @@ from agents import (
     software_engineer_agent,
     quality_assurance_agent,
 )
-from langchain_core.messages import HumanMessage, AIMessage
-from langchain_core.messages.base import BaseMessage
-from typing import TypedDict, List
-
 
+# ——————————————
+# 1) State types
+# ——————————————
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: list
 
-
 class OutputState(TypedDict):
     pm_output: str
     proj_output: str
@@ -25,20 +32,25 @@ class OutputState(TypedDict):
     qa_output: str
     chat_log: list
 
+# ——————————————
+# 2) Bridge → ProductManager
+# ——————————————
+def bridge_to_product_manager(state: dict) -> dict:
+    """
+    Convert the last HumanMessage into a structured system AIMessage
+    that the Product Manager agent can consume.
+    """
+    msgs = state.get("messages")
+    logs = state.get("chat_log", [])
+    if not isinstance(msgs, list) or not msgs or not isinstance(msgs[-1], HumanMessage):
+        raise ValueError("Expected state['messages'] to be a non-empty list ending in a 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")
-
-    user_prompt = messages[-1].content
+    user_text = msgs[-1].content
+    spec = f"""# Stakeholder Prompt
 
-    structured_context = f"""# Stakeholder Prompt
+A new product request has been submitted:
 
-"{user_prompt}"
+"{user_text}"
 
 Please convert this into a structured product specification including:
 - Goals
@@ -46,26 +58,32 @@ Please convert this into a structured product specification including:
 - User Stories
 - Success Metrics
 """
-
-    new_ai_msg = AIMessage(content=structured_context)
-
+    ai = AIMessage(content=spec)
     return {
-        "messages": messages + [new_ai_msg],  # Important: last item must be AIMessage
-        "chat_log": chat_log + [{"role": "System", "content": structured_context}],
+        "messages": msgs + [ai],
+        "chat_log": logs + [{"role": "System", "content": spec}],
     }
 
+# Create a StructuredTool (with explicit description) for the bridge
+bridge_tool = StructuredTool.from_function(
+    func=bridge_to_product_manager,
+    name="bridge_to_product_manager",
+    description="Generate a structured AIMessage from a HumanMessage for the Product Manager agent."
+)
+bridge_node = ToolNode([bridge_tool])
 
-# ✅ Wrap bridge node in a ToolNode
-bridge_node = ToolNode([bridge_to_pm])
-
-# ✅ Other agents as ToolNodes
-pm_node = ToolNode([product_manager_agent.run])
+# ——————————————
+# 3) Wrap your LLM agents
+# ——————————————
+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])
+dev_node  = ToolNode([software_engineer_agent.run])
+qa_node   = ToolNode([quality_assurance_agent.run])
 
-# ✅ Build LangGraph
+# ——————————————
+# 4) Build & compile the graph
+# ——————————————
 graph = StateGraph(input=InputState, output=OutputState)
 
 graph.add_node("Bridge", bridge_node)
@@ -85,16 +103,15 @@ graph.add_edge("QualityAssurance", END)
 
 compiled_graph = graph.compile()
 
-
-# ✅ Run function
+# ——————————————
+# 5) Pipeline entrypoint
+# ——————————————
 def run_pipeline_and_save(prompt: str):
     initial_state = {
         "messages": [HumanMessage(content=prompt)],
         "chat_log": [],
     }
-
-    # Sanity check
-    assert isinstance(initial_state["messages"][-1], HumanMessage)
-
+    # this invoke will now see:
+    # Bridge → PM → ProjectManager → Architect → Engineer → QA → END
     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"]