docs/examples/agent/agent_workflow_multi.ipynb
<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/agent_workflow_multi.ipynb" target="_parent"></a>
In this notebook, we will explore how to use the AgentWorkflow class to create multi-agent systems. Specifically, we will create a system that can generate a report on a given topic.
This notebook will assume that you have already either read the basic agent workflow notebook or the agent 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
llm = OpenAI(model="gpt-4o", 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 use the AgentWorkflow class to create a multi-agent system that will execute these agents in order.
While there are many ways to implement this system, in this case, we will use a few tools to help with the research and writing processes.
web_search tool to search the web for information on the given topic.record_notes tool to record notes on the given topic.write_report tool to write the report using the information found by the ResearchAgent.review_report tool to review the report and provide feedback.Utilizing the Context class, we can pass state between agents, and each agent will have access to the current state of the system.
%pip install tavily-python
from tavily import AsyncTavilyClient
from llama_index.core.workflow import Context
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))
async def record_notes(ctx: Context, notes: str, notes_title: str) -> str:
"""Useful for recording notes on a given topic. Your input should be notes with a title to save the notes under."""
async with ctx.store.edit_state() as ctx_state:
if "research_notes" not in ctx_state["state"]:
ctx_state["state"]["research_notes"] = {}
ctx_state["state"]["research_notes"][notes_title] = notes
return "Notes recorded."
async def write_report(ctx: Context, report_content: str) -> str:
"""Useful for writing a report on a given topic. Your input should be a markdown formatted report."""
async with ctx.store.edit_state() as ctx_state:
ctx_state["state"]["report_content"] = report_content
return "Report written."
async def review_report(ctx: Context, review: str) -> str:
"""Useful for reviewing a report and providing feedback. Your input should be a review of the report."""
async with ctx.store.edit_state() as ctx_state:
ctx_state["state"]["review"] = review
return "Report reviewed."
With our tools defined, we can now create our agents.
If the LLM you are using supports tool calling, you can use the FunctionAgent class. Otherwise, you can use the ReActAgent class.
Here, the name and description of each agent is used so that the system knows what each agent is responsible for and when to hand off control to the next agent.
from llama_index.core.agent.workflow import FunctionAgent, ReActAgent
research_agent = FunctionAgent(
name="ResearchAgent",
description="Useful for searching the web for information on a given topic and recording notes on the topic.",
system_prompt=(
"You are the ResearchAgent that can search the web for information on a given topic and record notes on the topic. "
"Once notes are recorded and you are satisfied, you should hand off control to the WriteAgent to write a report on the topic. "
"You should have at least some notes on a topic before handing off control to the WriteAgent."
),
llm=llm,
tools=[search_web, record_notes],
can_handoff_to=["WriteAgent"],
)
write_agent = FunctionAgent(
name="WriteAgent",
description="Useful for writing a report on a given topic.",
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. "
"Once the report is written, you should get feedback at least once from the ReviewAgent."
),
llm=llm,
tools=[write_report],
can_handoff_to=["ReviewAgent", "ResearchAgent"],
)
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 for the WriteAgent to implement. "
"If you have feedback that requires changes, you should hand off control to the WriteAgent to implement the changes after submitting the review."
),
llm=llm,
tools=[review_report],
can_handoff_to=["WriteAgent"],
)
With our agents defined, we can create our AgentWorkflow and run it.
from llama_index.core.agent.workflow import AgentWorkflow
agent_workflow = AgentWorkflow(
agents=[research_agent, write_agent, review_agent],
root_agent=research_agent.name,
initial_state={
"research_notes": {},
"report_content": "Not written yet.",
"review": "Review required.",
},
)
As the workflow is running, we will stream the events to get an idea of what is happening under the hood.
from llama_index.core.agent.workflow import (
AgentInput,
AgentOutput,
ToolCall,
ToolCallResult,
AgentStream,
)
handler = agent_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."
)
)
current_agent = None
current_tool_calls = ""
async for event in handler.stream_events():
if (
hasattr(event, "current_agent_name")
and event.current_agent_name != current_agent
):
current_agent = event.current_agent_name
print(f"\n{'='*50}")
print(f"🤖 Agent: {current_agent}")
print(f"{'='*50}\n")
# if isinstance(event, AgentStream):
# if event.delta:
# print(event.delta, end="", flush=True)
# elif isinstance(event, AgentInput):
# print("📥 Input:", event.input)
elif isinstance(event, AgentOutput):
if event.response.content:
print("📤 Output:", event.response.content)
if event.tool_calls:
print(
"🛠️ Planning to use tools:",
[call.tool_name for call in event.tool_calls],
)
elif isinstance(event, ToolCallResult):
print(f"🔧 Tool Result ({event.tool_name}):")
print(f" Arguments: {event.tool_kwargs}")
print(f" Output: {event.tool_output}")
elif isinstance(event, ToolCall):
print(f"🔨 Calling Tool: {event.tool_name}")
print(f" With arguments: {event.tool_kwargs}")
Now, we can retrieve the final report in the system for ourselves.
state = await handler.ctx.store.get("state")
print(state["report_content"])