Day 17/30: Orchestrator-Worker with LangGraph

I recently spent a few late nights debugging an issue with our support bot, which was supposed to handle multiple conversations simultaneously. The problem was that it would often drop context or respond to the wrong user, even though we were using LangGraph's StateGraph to manage the conversation state. It turned out that the issue was due to the way we were running our agents - we were using a single agent to handle all conversations, which was leading to a lot of conflicts and inconsistencies.

To fix this, we decided to switch to an orchestrator-worker model, where a single orchestrator agent would manage the overall workflow and assign tasks to multiple worker agents. This way, each conversation would have its own dedicated worker agent, which would prevent the context switching issues we were seeing. But to make this work, we needed a way for the orchestrator to send tasks to the worker agents, and for the workers to send results back to the orchestrator.

That's where LangGraph's Send API comes in. The Send API allows you to send messages between agents, which is exactly what we needed to implement our orchestrator-worker model. Here's an example of how we used the Send API to send a task from the orchestrator to a worker:

import langgraph

# Create the orchestrator agent
orchestrator = langgraph.Agent("orchestrator")

# Create a worker agent
worker = langgraph.Agent("worker")

# Define a task for the worker to complete
task = langgraph.Task("respond_to_user", user_id=123, message="Hello!")

# Send the task to the worker using the Send API
orchestrator.send(worker, task)

# Define a callback function to handle the response from the worker
def handle_response(response):
    print(f"Received response from worker: {response}")

# Send a message to the worker and wait for a response
response = orchestrator.send_and_wait(worker, task, callback=handle_response)

In this example, the orchestrator agent sends a task to the worker agent using the send method, and then waits for a response using the send_and_wait method. The worker agent receives the task and completes it, and then sends a response back to the orchestrator using the send method.

To run multiple workers in parallel, we can use LangGraph's add_node method to create a new node for each worker, and then use the add_conditional_edges method to define the workflow between the orchestrator and the workers. Here's an example:

# Create a new node for each worker
workers = [langgraph.Agent(f"worker_{i}") for i in range(5)]

# Define the workflow between the orchestrator and the workers
orchestrator.add_node("dispatch_task", lambda task: task.user_id % len(workers))
for i, worker in enumerate(workers):
    orchestrator.add_conditional_edge("dispatch_task", worker, lambda task: task.user_id % len(workers) == i)

In this example, we create a new node for each worker, and then define a workflow that dispatches tasks to the workers based on the user ID.

One practical gotcha to watch out for when using the Send API is that it can be slow if you're sending a lot of messages between agents. To mitigate this, you can use LangGraph's checkpointers to cache the results of expensive computations and reduce the number of messages that need to be sent. For example:

# Create a checkpoint to cache the results of expensive computations
checkpoint = langgraph.Checkpoint("expensive_computation")

# Define a function to compute the result
def compute_result(task):
    # Simulate an expensive computation
    import time
    time.sleep(1)
    return task.user_id * 2

# Use the checkpoint to cache the result
@langgraph.checkpointed(compute_result, checkpoint)
def cached_compute_result(task):
    return compute_result(task)

By using the checkpointed decorator, we can cache the result of the expensive computation and avoid recomputing it every time the worker receives a task.

As we continue to build out our agentic AI system, we'll need to tackle even more complex challenges, such as integrating with external systems and handling failures and errors. Tomorrow, we'll explore some of the tools and techniques we can use to build more robust and resilient agents.