docs/examples/agent/custom_multi_agent.ipynb
<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/custom_multi_agent.ipynb" target="_parent"></a>
In this notebook, we will explore how to prompt an LLM to write, refine, and follow a plan to generate a report using multiple agents.
This is not meant to be a comprehensive guide to creating a report generation system, but rather, giving you the knowledge and tools to build your own robust systems that can plan and orchestrate multiple agents to achieve a goal.
This notebook will assume that you have already either read the basic agent workflow notebook or the agent workflow documentation, as well as the workflow documentation.
In this example, we will use OpenAI as our LLM. For all LLMs, check out the examples documentation or LlamaHub for a list of all supported LLMs and how to install/use them.
If we wanted, each agent could have a different LLM, but for this example, we will use the same LLM for all agents.
%pip install llama-index
from llama_index.llms.openai import OpenAI
sub_agent_llm = OpenAI(model="gpt-4.1-mini", api_key="sk-...")
Our system will have three agents:
ResearchAgent that will search the web for information on the given topic.WriteAgent that will write the report using the information found by the ResearchAgent.ReviewAgent that will review the report and provide feedback.We will then use a top-level LLM to manually orchestrate and plan around the other agents to write our report.
While there are many ways to implement this system, in this case, we will use a single web_search tool to search the web for information on the given topic.
%pip install tavily-python
from tavily import AsyncTavilyClient
async def search_web(query: str) -> str:
"""Useful for using the web to answer questions."""
client = AsyncTavilyClient(api_key="tvly-...")
return str(await client.search(query))
With our tool defined, we can now create our sub-agents.
If the LLM you are using supports tool calling, you can use the FunctionAgent class. Otherwise, you can use the ReActAgent class.
from llama_index.core.agent.workflow import FunctionAgent, ReActAgent
research_agent = FunctionAgent(
name="ResearchAgent",
description="Useful for recording research notes based on a specific prompt.",
system_prompt=(
"You are the ResearchAgent that can search the web for information on a given topic and record notes on the topic. "
"You should output notes on the topic in a structured format."
),
llm=sub_agent_llm,
tools=[search_web],
)
write_agent = FunctionAgent(
name="WriteAgent",
description="Useful for writing a report based on the research notes or revising the report based on feedback.",
system_prompt=(
"You are the WriteAgent that can write a report on a given topic. "
"Your report should be in a markdown format. The content should be grounded in the research notes. "
"Return your markdown report surrounded by <report>...</report> tags."
),
llm=sub_agent_llm,
tools=[],
)
review_agent = FunctionAgent(
name="ReviewAgent",
description="Useful for reviewing a report and providing feedback.",
system_prompt=(
"You are the ReviewAgent that can review the write report and provide feedback. "
"Your review should either approve the current report or request changes to be implemented."
),
llm=sub_agent_llm,
tools=[],
)
With each agent defined, we can also write helper functions to help execute each agent.
import re
from llama_index.core.workflow import Context
async def call_research_agent(ctx: Context, prompt: str) -> str:
"""Useful for recording research notes based on a specific prompt."""
result = await research_agent.run(
user_msg=f"Write some notes about the following: {prompt}"
)
async with ctx.store.edit_state() as ctx_state:
ctx_state["state"]["research_notes"].append(str(result))
return str(result)
async def call_write_agent(ctx: Context) -> str:
"""Useful for writing a report based on the research notes or revising the report based on feedback."""
async with ctx.store.edit_state() as ctx_state:
notes = ctx_state["state"].get("research_notes", None)
if not notes:
return "No research notes to write from."
user_msg = f"Write a markdown report from the following notes. Be sure to output the report in the following format: <report>...</report>:\n\n"
# Add the feedback to the user message if it exists
feedback = ctx_state["state"].get("review", None)
if feedback:
user_msg += f"<feedback>{feedback}</feedback>\n\n"
# Add the research notes to the user message
notes = "\n\n".join(notes)
user_msg += f"<research_notes>{notes}</research_notes>\n\n"
# Run the write agent
result = await write_agent.run(user_msg=user_msg)
report = re.search(
r"<report>(.*)</report>", str(result), re.DOTALL
).group(1)
ctx_state["state"]["report_content"] = str(report)
return str(report)
async def call_review_agent(ctx: Context) -> str:
"""Useful for reviewing the report and providing feedback."""
async with ctx.store.edit_state() as ctx_state:
report = ctx_state["state"].get("report_content", None)
if not report:
return "No report content to review."
result = await review_agent.run(
user_msg=f"Review the following report: {report}"
)
ctx_state["state"]["review"] = result
return result
In order to plan around the other agents, we will write a custom workflow that will explicitly orchestrate and plan the other agents.
Here our prompt assumes a sequential plan, but we can expand it in the future to support parallel steps. (This just involves more complex parsing and prompting, which is left as an exercise for the reader.)
import re
import xml.etree.ElementTree as ET
from pydantic import BaseModel, Field
from typing import Any, Optional
from llama_index.core.llms import ChatMessage
from llama_index.core.workflow import (
Context,
Event,
StartEvent,
StopEvent,
Workflow,
step,
)
PLANNER_PROMPT = """You are a planner chatbot.
Given a user request and the current state, break the solution into ordered <step> blocks. Each step must specify the agent to call and the message to send, e.g.
<plan>
<step agent=\"ResearchAgent\">search for …</step>
<step agent=\"WriteAgent\">draft a report …</step>
...
</plan>
<state>
{state}
</state>
<available_agents>
{available_agents}
</available_agents>
The general flow should be:
- Record research notes
- Write a report
- Review the report
- Write the report again if the review is not positive enough
If the user request does not require any steps, you can skip the <plan> block and respond directly.
"""
class InputEvent(StartEvent):
user_msg: Optional[str] = Field(default=None)
chat_history: list[ChatMessage]
state: Optional[dict[str, Any]] = Field(default=None)
class OutputEvent(StopEvent):
response: str
chat_history: list[ChatMessage]
state: dict[str, Any]
class StreamEvent(Event):
delta: str
class PlanEvent(Event):
step_info: str
# Modelling the plan
class PlanStep(BaseModel):
agent_name: str
agent_input: str
class Plan(BaseModel):
steps: list[PlanStep]
class ExecuteEvent(Event):
plan: Plan
chat_history: list[ChatMessage]
class PlannerWorkflow(Workflow):
llm: OpenAI = OpenAI(
model="o3-mini",
api_key="sk-...",
)
agents: dict[str, FunctionAgent] = {
"ResearchAgent": research_agent,
"WriteAgent": write_agent,
"ReviewAgent": review_agent,
}
@step
async def plan(
self, ctx: Context, ev: InputEvent
) -> ExecuteEvent | OutputEvent:
# Set initial state if it exists
if ev.state:
await ctx.store.set("state", ev.state)
chat_history = ev.chat_history
if ev.user_msg:
user_msg = ChatMessage(
role="user",
content=ev.user_msg,
)
chat_history.append(user_msg)
# Inject the system prompt with state and available agents
state = await ctx.store.get("state")
available_agents_str = "\n".join(
[
f'<agent name="{agent.name}">{agent.description}</agent>'
for agent in self.agents.values()
]
)
system_prompt = ChatMessage(
role="system",
content=PLANNER_PROMPT.format(
state=str(state),
available_agents=available_agents_str,
),
)
# Stream the response from the llm
response = await self.llm.astream_chat(
messages=[system_prompt] + chat_history,
)
full_response = ""
async for chunk in response:
full_response += chunk.delta or ""
if chunk.delta:
ctx.write_event_to_stream(
StreamEvent(delta=chunk.delta),
)
# Parse the response into a plan and decide whether to execute or output
xml_match = re.search(r"(<plan>.*</plan>)", full_response, re.DOTALL)
if not xml_match:
chat_history.append(
ChatMessage(
role="assistant",
content=full_response,
)
)
return OutputEvent(
response=full_response,
chat_history=chat_history,
state=state,
)
else:
xml_str = xml_match.group(1)
root = ET.fromstring(xml_str)
plan = Plan(steps=[])
for step in root.findall("step"):
plan.steps.append(
PlanStep(
agent_name=step.attrib["agent"],
agent_input=step.text.strip() if step.text else "",
)
)
return ExecuteEvent(plan=plan, chat_history=chat_history)
@step
async def execute(self, ctx: Context, ev: ExecuteEvent) -> InputEvent:
chat_history = ev.chat_history
plan = ev.plan
for step in plan.steps:
agent = self.agents[step.agent_name]
agent_input = step.agent_input
ctx.write_event_to_stream(
PlanEvent(
step_info=f'<step agent="{step.agent_name}">{step.agent_input}</step>'
),
)
if step.agent_name == "ResearchAgent":
await call_research_agent(ctx, agent_input)
elif step.agent_name == "WriteAgent":
# Note: we aren't passing the input from the plan since
# we're using the state to drive the write agent
await call_write_agent(ctx)
elif step.agent_name == "ReviewAgent":
await call_review_agent(ctx)
state = await ctx.store.get("state")
chat_history.append(
ChatMessage(
role="user",
content=f"I've completed the previous steps, here's the updated state:\n\n<state>\n{state}\n</state>\n\nDo you need to continue and plan more steps?, If not, write a final response.",
)
)
return InputEvent(
chat_history=chat_history,
)
With our custom planner defined, we can now run the workflow and see it in action!
As the workflow is running, we will stream the events to get an idea of what is happening under the hood.
planner_workflow = PlannerWorkflow(timeout=None)
handler = planner_workflow.run(
user_msg=(
"Write me a report on the history of the internet. "
"Briefly describe the history of the internet, including the development of the internet, the development of the web, "
"and the development of the internet in the 21st century."
),
chat_history=[],
state={
"research_notes": [],
"report_content": "Not written yet.",
"review": "Review required.",
},
)
current_agent = None
current_tool_calls = ""
async for event in handler.stream_events():
if isinstance(event, PlanEvent):
print("Executing plan step: ", event.step_info)
elif isinstance(event, ExecuteEvent):
print("Executing plan: ", event.plan)
result = await handler
print(result.response)
Now, we can retrieve the final report in the system for ourselves.
state = await handler.ctx.store.get("state")
print(state["report_content"])
print(state["review"])