How to Use Sub-Agents in Your Chat Agent

In this guide, we will explore how to implement and use sub-agents within a main chat agent. Sub-agents allow you to modularize your agent's functionality, making your code more organized and maintainable. It also enables you to create more complex conversational flows.

Prerequisites

Before you begin, make sure you have completed the Getting Started with the Agents SDK tutorial and have a basic understanding of how to create agents.

Simple Example

Let's start by scaffolding the code for the agents. Open your project directory and run the following commands to create the main agent and the sub-agent:

za agents add --custom "main_agent"
za agents add --custom "sub_agent"

Open the file named sub_agent.py and paste the following code:

from typing import List, Optional
from zav.agents_sdk import ChatAgent, ChatAgentClassRegistry, ChatMessage


@ChatAgentClassRegistry.register()
class SubAgent(ChatAgent):
    agent_name = "sub_agent"

    async def execute(self, conversation: List[ChatMessage]) -> Optional[ChatMessage]:
        return ChatMessage(
            sender="bot",
            content="Hello from the sub-agent!",
        )

Next, open the file named main_agent.py and paste the following code:

from typing import List, Optional
from zav.agents_sdk import ChatAgent, ChatAgentClassRegistry, ChatMessage
from sub_agent import SubAgent


@ChatAgentClassRegistry.register()
class MainAgent(ChatAgent):
    agent_name = "main_agent"

    def __init__(self, sub_agent: SubAgent):
        self.sub_agent = sub_agent

    async def execute(self, conversation: List[ChatMessage]) -> Optional[ChatMessage]:
        # Call the sub-agent
        sub_agent_response = await self.sub_agent.execute(conversation)
        assert sub_agent_response is not None

        return ChatMessage(
            sender="bot",
            content=f"Main agent received: {sub_agent_response.content}",
        )

In this code, the main agent calls the sub-agent and includes its response in the final message.