Update utils/langgraph_pipeline.py

utils/langgraph_pipeline.py

@@ -2,13 +2,12 @@
 
 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 langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages.base import BaseMessage
+
 from agents import (
     product_manager_agent,
     project_manager_agent,
@@ -17,9 +16,8 @@ from agents import (
     quality_assurance_agent,
 )
 
-# ——————————————
-# 1) State types
-# ——————————————
+
+# 1) Define your state shapes
 class InputState(TypedDict):
     messages: List[BaseMessage]
     chat_log: list
@@ -32,86 +30,51 @@ 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")
-
-    user_text = msgs[-1].content
-    spec = f"""# Stakeholder Prompt
-
-A new product request has been submitted:
-
-"{user_text}"
-
-Please convert this into a structured product specification including:
-- Goals
-- Key Features
-- User Stories
-- Success Metrics
-"""
-    ai = AIMessage(content=spec)
-    return {
-        "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])
-
-# ——————————————
-# 3) Wrap your LLM agents
-# ——————————————
-pm_node   = ToolNode([product_manager_agent.run])
+# 2) Wrap ONLY the downstream agents in ToolNodes
 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])
 
-# ——————————————
-# 4) Build & compile the graph
-# ——————————————
+# 3) Build the LangGraph starting at ProjectManager
 graph = StateGraph(input=InputState, output=OutputState)
 
-graph.add_node("Bridge", bridge_node)
-graph.add_node("ProductManager", pm_node)
-graph.add_node("ProjectManager", proj_node)
+graph.add_node("ProjectManager",   proj_node)
 graph.add_node("SoftwareArchitect", arch_node)
-graph.add_node("SoftwareEngineer", dev_node)
-graph.add_node("QualityAssurance", qa_node)
-
-graph.set_entry_point("Bridge")
-graph.add_edge("Bridge", "ProductManager")
-graph.add_edge("ProductManager", "ProjectManager")
-graph.add_edge("ProjectManager", "SoftwareArchitect")
-graph.add_edge("SoftwareArchitect", "SoftwareEngineer")
+graph.add_node("SoftwareEngineer",  dev_node)
+graph.add_node("QualityAssurance",  qa_node)
+
+graph.set_entry_point("ProjectManager")
+graph.add_edge("ProjectManager",   "SoftwareArchitect")
+graph.add_edge("SoftwareArchitect","SoftwareEngineer")
 graph.add_edge("SoftwareEngineer", "QualityAssurance")
 graph.add_edge("QualityAssurance", END)
 
 compiled_graph = graph.compile()
 
-# ——————————————
-# 5) Pipeline entrypoint
-# ——————————————
+
+# 4) The one‐and‐only run function
 def run_pipeline_and_save(prompt: str):
+    # a) Start with the raw HumanMessage
+    messages = [HumanMessage(content=prompt)]
+    chat_log = []
+
+    # b) Run your ProductManager agent manually
+    pm_state = product_manager_agent.run({
+        "messages": messages,
+        "chat_log": chat_log,
+    })
+
+    # Sanity check: ensure we got back an AIMessage
+    assert isinstance(pm_state["messages"][-1], AIMessage), "ProductManager must produce an AIMessage"
+
+    # c) Feed that into the graph
     initial_state = {
-        "messages": [HumanMessage(content=prompt)],
-        "chat_log": [],
+        "messages": pm_state["messages"],
+        "chat_log": pm_state["chat_log"],
     }
-    # 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"]