Design/langgraph_workflow_orchestration.md

LangGraph Workflow Orchestration - Design Document

Date: January 2025 Author: BeatDebate Team Status: Draft Review Status: Pending

---

0. References

• Overall Phase 2 Design: phase2_planner_agent_design.md
• Main Project Design: Plans/beatdebate-design-doc.md
• PlannerAgent Design: phase2_planner_agent_design.md (contains PlannerAgent details)
• JudgeAgent Design: phase2_judge_agent_design.md
• Agent Implementations: src/agents/
• State Model: src/models/agent_models.py (expecting MusicRecommenderState or similar)

---

1. Problem Statement

The primary objective is to orchestrate the four core agents of the BeatDebate system (PlannerAgent, GenreMoodAgent, DiscoveryAgent, and JudgeAgent) into a cohesive, sequential, and partially parallel workflow using LangGraph. This orchestration is essential for processing a user's music query from initial planning through to final recommendation selection and explanation.

The current state is that individual agents (PlannerAgent, GenreMoodAgent, DiscoveryAgent, JudgeAgent) have been designed and implemented (as per src/agents/). The API clients (src/api/) and data models (src/models/) are also in place. What's missing is the central nervous system that connects these components, manages the flow of data (state), and handles the execution sequence.

The value of this orchestration includes:

• Demonstrating a functional multi-agent system.
• Enabling end-to-end processing of user queries.
• Providing a clear structure for managing complex agent interactions.
• Facilitating transparency and debugging through LangGraph's inherent state tracking.

---

2. Goals & Non-Goals

✅ In Scope

• Defining the LangGraph graph structure, including nodes and edges, to represent the 4-agent workflow.
• Implementing the graph in src/services/recommendation_engine.py.
• Utilizing the MusicRecommenderState (likely defined in src/models/agent_models.py or a dedicated state management file like src/services/state_manager.py as hinted in Design/phase2_planner_agent_design.md) as the shared state object that flows through the graph.
• Defining nodes for each agent:
  • PlannerAgentNode: Executes the PlannerAgent's create_music_discovery_strategy (or similar method as per src/agents/planner_agent.py).
  • GenreMoodAgentNode: Executes the GenreMoodAgent's strategy execution method.
  • DiscoveryAgentNode: Executes the DiscoveryAgent's strategy execution method.
  • JudgeAgentNode: Executes the JudgeAgent's evaluate_and_select method.
• Implementing parallel execution for GenreMoodAgentNode and DiscoveryAgentNode.
• Defining conditional edges if necessary (e.g., handling cases where the Planner fails or no candidates are found by advocates).
• Ensuring the final output is a populated MusicRecommenderState with final_recommendations and reasoning_log.
• Basic error handling within graph nodes to update the state with error information.
• Logging of graph execution steps and state transitions.

❌ Out of Scope

• Advanced error recovery strategies beyond logging and basic state updates (e.g., automated retries with modified parameters).
• Complex dynamic routing or agent selection beyond the defined 4-agent sequence.
• User interface integration for visualizing the graph execution (this is a UI concern).
• Detailed performance optimization of individual agent calls (focus is on orchestration).
• Implementation of the MusicRecommenderState itself if not already fully defined (this design assumes it exists or will be created based on Design/phase2_planner_agent_design.md).

---

3. Technical Architecture

3.1 System Overview

The LangGraph workflow will be the core of the recommendation_engine.py service. It will take an initial MusicRecommenderState (primarily with user_query populated) and pass it through a sequence of agent nodes.

3.2 LangGraph State

The shared state for the graph will be an instance of MusicRecommenderState. Based on Design/phase2_planner_agent_design.md, this state should include at least:

# Expected structure from src/models/agent_models.py or similar
from typing import Dict, List, Any, Optional
from pydantic import BaseModel, Field # Assuming Pydantic is used

class MusicRecommenderState(BaseModel):
    user_query: str
    user_profile: Optional[Dict[str, Any]] = None # Not used in MVP workflow directly but good to have
    
    # Planning phase output
    planning_strategy: Optional[Dict[str, Any]] = Field(default_factory=dict)
    
    # Advocate phase outputs (as lists of dictionaries, to be parsed by Judge or if models are used directly)
    genre_mood_recommendations: List[Dict] = Field(default_factory=list) # Or List[TrackRecommendationDict]
    discovery_recommendations: List[Dict] = Field(default_factory=list) # Or List[TrackRecommendationDict]
    
    # Judge phase output
    final_recommendations: List[Dict] = Field(default_factory=list) # Or List[TrackRecommendationDict]
    
    # Reasoning & Logging
    reasoning_log: List[str] = Field(default_factory=list)
    agent_deliberations: List[Dict] = Field(default_factory=list) # For more structured logs
    
    # Error Handling
    error_info: Optional[Dict[str, str]] = None # e.g., {"agent": "PlannerAgent", "message": "API timeout"}

    # Metadata
    processing_start_time: Optional[float] = None
    total_processing_time: Optional[float] = None

Actual Pydantic model for TrackRecommendation would be used if advocates return model instances, but design docs lean towards List[Dict] passed between agents in the state.

3.3 Graph Definition

The graph will be constructed using langgraph.StateGraph.

Nodes:

1. planner_node:

   • Function: Wraps PlannerAgent.create_music_discovery_strategy() (or its equivalent async method from src/agents/planner_agent.py).
   • Input: MusicRecommenderState (uses user_query).
   • Output: Updates state.planning_strategy and state.reasoning_log.
   • Error Handling: If planning fails, updates state.error_info and state.reasoning_log.

2. genre_mood_advocate_node:

   • Function: Wraps GenreMoodAgent.execute_strategy() (or its equivalent from src/agents/genre_mood_agent.py).
   • Input: MusicRecommenderState (uses planning_strategy).
   • Output: Updates state.genre_mood_recommendations and state.reasoning_log.
   • Error Handling: Updates state.error_info if it fails.

3. discovery_advocate_node:

   • Function: Wraps DiscoveryAgent.execute_strategy() (or its equivalent from src/agents/discovery_agent.py).
   • Input: MusicRecommenderState (uses planning_strategy).
   • Output: Updates state.discovery_recommendations and state.reasoning_log.
   • Error Handling: Updates state.error_info if it fails.

4. judge_node:

   • Function: Wraps JudgeAgent.evaluate_and_select().
   • Input: MusicRecommenderState (uses planning_strategy, genre_mood_recommendations, discovery_recommendations).
   • Output: Updates state.final_recommendations and state.reasoning_log.
   • Error Handling: Updates state.error_info if it fails.

Edges:

• START → planner_node

• planner_node → should_proceed_after_planning (Conditional Edge)

  • If state.error_info is set by planner_node OR state.planning_strategy is empty/invalid: → END
  • Else: → execute_advocates_parallel (a conceptual branch point)

• execute_advocates_parallel (This isn't a node itself, but LangGraph supports fanning out to multiple nodes that can run in parallel if their inputs are ready. We'll add edges from planner_node (or the conditional node success branch) to both advocate nodes.)

  • planner_node (on success) → genre_mood_advocate_node
  • planner_node (on success) → discovery_advocate_node (LangGraph will execute these concurrently as they don't depend on each other's direct output, only on the planner_node's output in the state.)

• After both genre_mood_advocate_node AND discovery_advocate_node complete, they will converge to the judge_node. LangGraph handles this implicitly by defining judge_node as the next step after both. We need a mechanism to gather results, which LangGraph handles via the shared state object. The judge_node will be added after both advocate nodes.

  • genre_mood_advocate_node → judge_node
  • discovery_advocate_node → judge_node (LangGraph's compilation will create a join point before the judge if these are set as direct next steps for the parallel branches.)

• judge_node → END

Workflow Diagram:

[START]
   ↓
┌────────────────┐
│  PlannerAgent  │
│ (planner_node) │
└────────────────┘
   ↓
┌───────────────────────────┐
│ should_proceed_after_plan │ (Conditional)
└───────────────────────────┘
   ├── (Error/No Plan) → [END]
   ↓ (Success)
┌────────────────┐   ┌──────────────────┐
│ GenreMoodAgent │   │ DiscoveryAgent   │
│ (gm_node)      │   │ (discovery_node) │
└────────────────┘   └──────────────────┘
   │                 │
   └───────▼─────────┘
           ↓ (Implicit Join)
┌────────────────┐
│   JudgeAgent   │
│  (judge_node)  │
└────────────────┘
   ↓
[END]

3.4 RecommendationEngine Service

A class or a set of functions in src/services/recommendation_engine.py will be responsible for:

1. Initializing agent instances.
2. Defining the LangGraph StateGraph.
3. Adding nodes and edges as described above.
4. Compiling the graph.
5. Providing a main method to invoke the graph with an initial user_query.

# Example structure in src/services/recommendation_engine.py
from langgraph.graph import StateGraph, END
from src.models.agent_models import MusicRecommenderState # Assuming this path
from src.agents.planner_agent import PlannerAgent
from src.agents.genre_mood_agent import GenreMoodAgent
from src.agents.discovery_agent import DiscoveryAgent
from src.agents.judge_agent import JudgeAgent
import time # For processing_start_time

class RecommendationEngine:
    def __init__(self):
        self.planner_agent = PlannerAgent()
        self.genre_mood_agent = GenreMoodAgent()
        self.discovery_agent = DiscoveryAgent()
        self.judge_agent = JudgeAgent()
        self.graph = self._build_graph()

    async def _planner_node_func(self, state: MusicRecommenderState):
        # self.logger.info("Executing Planner Node")
        try:
            # Assuming PlannerAgent has an async method like 'process' or 'run'
            # that takes the state and updates it.
            # For example, if it's 'create_music_discovery_strategy'
            updated_planning_strategy = await self.planner_agent.create_music_discovery_strategy(state) # Or equivalent
            state.planning_strategy = updated_planning_strategy
            state.reasoning_log.append("PlannerAgent: Strategy created.")
        except Exception as e:
            # self.logger.error(f"Planner Node Error: {e}")
            state.error_info = {"agent": "PlannerAgent", "message": str(e)}
            state.reasoning_log.append(f"PlannerAgent: Error - {str(e)}")
        return state

    async def _genre_mood_advocate_node_func(self, state: MusicRecommenderState):
        # self.logger.info("Executing GenreMood Advocate Node")
        try:
            # Assuming agent has a method that takes state (esp. planning_strategy)
            # and returns recommendations or updates state directly
            # Example: recommendations = await self.genre_mood_agent.execute_strategy(state.planning_strategy)
            # For now, let's assume agents modify state directly for simplicity in graph nodes
            await self.genre_mood_agent.execute_strategy(state) # Modifies state.genre_mood_recommendations
            state.reasoning_log.append("GenreMoodAgent: Recommendations generated.")
        except Exception as e:
            # self.logger.error(f"GenreMoodAdvocate Node Error: {e}")
            state.error_info = {"agent": "GenreMoodAgent", "message": str(e)} # Potentially overwrite, needs careful thought
            state.reasoning_log.append(f"GenreMoodAgent: Error - {str(e)}")
        return state
        
    async def _discovery_advocate_node_func(self, state: MusicRecommenderState):
        # self.logger.info("Executing Discovery Advocate Node")
        try:
            await self.discovery_agent.execute_strategy(state) # Modifies state.discovery_recommendations
            state.reasoning_log.append("DiscoveryAgent: Recommendations generated.")
        except Exception as e:
            # self.logger.error(f"DiscoveryAdvocate Node Error: {e}")
            state.error_info = {"agent": "DiscoveryAgent", "message": str(e)}
            state.reasoning_log.append(f"DiscoveryAgent: Error - {str(e)}")
        return state

    async def _judge_node_func(self, state: MusicRecommenderState):
        # self.logger.info("Executing Judge Node")
        try:
            # JudgeAgent's evaluate_and_select updates state.final_recommendations
            await self.judge_agent.evaluate_and_select(state)
            state.reasoning_log.append("JudgeAgent: Final selection complete.")
        except Exception as e:
            # self.logger.error(f"Judge Node Error: {e}")
            state.error_info = {"agent": "JudgeAgent", "message": str(e)}
            state.reasoning_log.append(f"JudgeAgent: Error - {str(e)}")
        return state

    def _should_proceed_after_planning(self, state: MusicRecommenderState):
        # self.logger.debug(f"Conditional Edge: Checking planner output. Error: {state.error_info}")
        if state.error_info and state.error_info.get("agent") == "PlannerAgent":
            return "end_workflow"
        if not state.planning_strategy or not state.planning_strategy.get("evaluation_framework"): # Basic check
            # self.logger.warning("Planner did not produce a valid strategy. Ending workflow.")
            state.reasoning_log.append("Workflow: Planner failed to produce strategy, ending.")
            return "end_workflow"
        return "execute_advocates"

    def _build_graph(self):
        workflow = StateGraph(MusicRecommenderState)

        workflow.add_node("planner", self._planner_node_func)
        workflow.add_node("genre_mood_advocate", self._genre_mood_advocate_node_func)
        workflow.add_node("discovery_advocate", self._discovery_advocate_node_func)
        workflow.add_node("judge", self._judge_node_func)

        workflow.set_entry_point("planner")

        workflow.add_conditional_edges(
            "planner",
            self._should_proceed_after_planning,
            {
                "execute_advocates": "genre_mood_advocate", # Start of one parallel branch
                "end_workflow": END
            }
        )
        # The other parallel branch also starts after successful planning
        workflow.add_edge("planner", "discovery_advocate") # This needs careful handling with conditional edges.
                                                           # Langchain typically routes from conditional output.

        # Revised approach for parallel execution after conditional check:
        # The conditional edge from 'planner' will route to a dummy "fan_out_advocates" node if successful,
        # or directly to END. This dummy node then has unconditional edges to both advocates.
        # However, LangGraph allows multiple targets from a conditional branch.
        # A simpler way is for `_should_proceed_after_planning` to return the name of the
        # *first* node in the parallel execution if successful, or END.
        # Let's assume `genre_mood_advocate` is one, and we add another edge for `discovery_advocate`.
        # For true parallelism, they should not depend on each other.

        # If planning succeeds, it goes to 'genre_mood_advocate'.
        # We also need an edge from 'planner' (on success) to 'discovery_advocate'.
        # This is where conditional logic becomes tricky for multiple next steps.
        # A better pattern for conditional branching to parallel execution:
        # planner -> conditional_router
        # conditional_router:
        #   if success: "fan_out_node" (dummy node)
        #   if failure: END
        # fan_out_node -> genre_mood_advocate
        # fan_out_node -> discovery_advocate
        #
        # Or, if the conditional function can return a list of next nodes, that's ideal.
        # LangGraph's `add_conditional_edges` maps a string return to a single next node.
        # For fan-out after condition:
        # 1. `planner`
        # 2. `conditional_planner_gate` (routes to `run_advocates` or `END`)
        # 3. `run_advocates` (dummy node, or just ensure edges are set)
        # From `run_advocates` (or directly from `conditional_planner_gate`'s "success" branch):
        # Edge to `genre_mood_advocate`
        # Edge to `discovery_advocate`

        # For now, let's try to make the conditional router point to one, and add another edge.
        # If `_should_proceed_after_planning` returns "execute_advocates", we go to `genre_mood_advocate`.
        # We also need `discovery_advocate` to run.
        # LangGraph requires that if a node is specified as a target in `add_conditional_edges`,
        # other direct `add_edge` calls from the source of the conditional edge might be ignored or cause conflict.

        # A cleaner pattern:
        # workflow.add_node("fan_out_advocates_anchor", lambda x: x) # Dummy node
        # workflow.add_conditional_edges(
        #     "planner",
        #     self._should_proceed_after_planning,
        #     {
        #         "execute_advocates": "fan_out_advocates_anchor",
        #         "end_workflow": END
        #     }
        # )
        # workflow.add_edge("fan_out_advocates_anchor", "genre_mood_advocate")
        # workflow.add_edge("fan_out_advocates_anchor", "discovery_advocate")

        # Simpler: The conditional edge path for "execute_advocates" will go to "genre_mood_advocate".
        # An additional edge will also be triggered to "discovery_advocate" if the condition isn't "end_workflow".
        # This needs to be expressed carefully.
        # The most straightforward way for parallel branches after a condition is to route to *one* of the parallel
        # tasks and ensure the other parallel task is also set to start *after the same condition point*.
        # This is usually handled by LangGraph if the tasks truly can run in parallel.

        # Edges for advocate to judge (these create an implicit join)
        workflow.add_edge("genre_mood_advocate", "judge")
        workflow.add_edge("discovery_advocate", "judge")
        
        workflow.add_edge("judge", END)
        
        return workflow.compile()

    async def process_query(self, user_query: str) -> MusicRecommenderState:
        initial_state = MusicRecommenderState(user_query=user_query)
        initial_state.processing_start_time = time.time()
        
        final_state = await self.graph.ainvoke(initial_state)
        
        if final_state.processing_start_time: # Should always be true
            final_state.total_processing_time = time.time() - final_state.processing_start_time
        
        return final_state

The parallel execution setup with conditional edges needs careful handling in LangGraph. The above _build_graph provides a conceptual sketch. It might require a dummy node or ensuring the conditional routing correctly triggers all parallel starts.

---

4. Implementation Plan

1. Verify/Define MusicRecommenderState: Ensure MusicRecommenderState is fully defined in src/models/agent_models.py (or src/services/state_manager.py) covering all necessary fields for the 4-agent workflow, including error handling fields.
2. Implement Agent Invocation Wrappers: Create the async def _<agent_name>_node_func(self, state: MusicRecommenderState) methods within RecommendationEngine. These methods will:
   • Call the respective agent's main processing method (e.g., planner_agent.create_music_discovery_strategy(state)).
   • Ensure agents update the passed-in state object or return the necessary parts to update the state.
   • Implement try-except blocks to catch errors, update state.error_info and state.reasoning_log.
3. Implement Conditional Logic: Implement _should_proceed_after_planning to check state.error_info and the validity of state.planning_strategy.
4. Build Graph in RecommendationEngine:
   • Instantiate StateGraph with MusicRecommenderState.
   • Add all agent nodes and the conditional logic node.
   • Define the entry point.
   • Add edges, carefully configuring the parallel execution of advocate agents and the subsequent join at the judge agent. This might involve an anchor/dummy node if direct conditional fanning out is complex.
   • Compile the graph.
5. Implement process_query method: This public method on RecommendationEngine will initialize the state with the user query, invoke the compiled graph, and return the final state.
6. Basic Logging: Add logging calls within node functions and the RecommendationEngine for key events and state changes.
7. Initial Unit Tests: Write basic tests for:
   • RecommendationEngine initialization.
   • Each node function wrapper (mocking the underlying agent) to ensure it correctly updates state and handles errors.
   • The conditional logic function _should_proceed_after_planning.

---

5. Testing Strategy

• Unit Tests (tests/services/test_recommendation_engine.py):
  • Test RecommendationEngine initialization.
  • Test each node function (_planner_node_func, _genre_mood_advocate_node_func, etc.) by mocking the actual agent calls:
    • Verify correct state updates on successful agent execution.
    • Verify correct state.error_info and reasoning_log updates on agent failure.
  • Test the _should_proceed_after_planning conditional logic with various states.
  • Test the process_query method by mocking the compiled graph's ainvoke to ensure state initialization and timing logic works.
• Integration Tests (Covered in overall Phase 2.3):
  • Once implemented, this engine will be the core of end-to-end tests.
  • Test the full graph execution with mocked agents to verify flow and parallel execution.
  • Then, test with actual agent instances (potentially mocking their external API calls) to ensure data compatibility and correct state flow through the real agents via the graph. Scenarios:
    • Successful end-to-end run.
    • Planner failure.
    • One advocate agent fails, the other succeeds.
    • Both advocate agents fail or return no candidates.
    • Judge agent failure.

---

6. Success Criteria

• The LangGraph workflow is successfully implemented in src/services/recommendation_engine.py.
• The system can process a user query from start to finish using the 4-agent pipeline.
• GenreMoodAgent and DiscoveryAgent execute in parallel after the PlannerAgent.
• The MusicRecommenderState is correctly passed and updated between agents.
• Basic error conditions (e.g., planner failure) are handled by routing to END or updating state appropriately.
• The final state contains final_recommendations and a reasoning_log detailing the process.
• Unit tests for the orchestration logic pass.

---

7. Risk Mitigation

• Risk: Complexity in configuring parallel execution with conditional logic in LangGraph.
  • Mitigation: Refer to LangGraph documentation for best practices on parallel execution and conditional branching. Start with a simpler sequential flow if necessary and incrementally add parallelism. Consider using an explicit "fan-out" dummy node if direct conditional branching to multiple parallel starts is problematic.
• Risk: Agents might not correctly update the shared MusicRecommenderState object, leading to data inconsistencies.
  • Mitigation: Clearly define the responsibility of each agent node function to update specific parts of the state. Implement thorough unit tests for node functions verifying state modifications. Ensure agents either modify the state object by reference or return specific Pydantic models that the node function then uses to update the state.
• Risk: Error handling logic in one node might incorrectly halt the entire graph or mask errors.
  • Mitigation: Design error handling in nodes to update state.error_info clearly, allowing downstream conditional logic or the final result to reflect partial failures. Ensure reasoning_log captures all errors.
• Risk: Difficulty in debugging the state flow through the compiled graph.
  • Mitigation: Utilize LangGraph's built-in debugging capabilities (e.g., get_graph().print_ascii()). Implement comprehensive logging within each node and in the RecommendationEngine to trace state changes.

---

8. Open Questions & Discussion Points

• State Modification: Should agent methods directly modify the passed-in state object, or should they return their results, with the node wrapper function responsible for updating the state?
  • Proposed Resolution: For cleaner separation, agent methods (e.g., planner_agent.create_music_discovery_strategy) should ideally return their primary output (e.g., a PlanningStrategy object or dict). The LangGraph node wrapper function (_planner_node_func) will then be responsible for updating the MusicRecommenderState based on this return value and handling any errors. This keeps agent logic more self-contained. However, some agents like the JudgeAgent are designed to take the full state and modify it. This needs to be consistent or clearly documented for each agent node. For MVP, direct state modification within agent methods (if they are designed that way) is acceptable for the node functions to call, as long as it's clear. The current agent designs seem to imply they update the state given to them.
• Detailed Parallelism Configuration: The exact LangGraph syntax for ensuring the advocate agents run in parallel after a conditional check on the planner's output needs to be finalized.
  • Proposed Resolution: Consult LangGraph examples for conditional fan-out. If add_conditional_edges path mapping only supports a single string target, an anchor/dummy node post-condition for fanning out is a standard pattern. E.g., planner -> planner_gate (conditional) -> (if success) advocate_fanout_anchor -> (parallel edges to) genre_mood_advocate AND discovery_advocate.
• Error Overwriting: If multiple parallel nodes fail, how should state.error_info be handled? Last error wins, or a list of errors?
  • Proposed Resolution: For MVP, state.error_info can store the first critical error that halts a main path or the last one encountered. More sophisticated error aggregation can be a future enhancement. The reasoning_log should capture all errors encountered.

---