Delete tools/multi_agent_workflow_for_research.py
Browse files
tools/multi_agent_workflow_for_research.py
DELETED
@@ -1,237 +0,0 @@
|
|
1 |
-
"""
|
2 |
-
This file contains the multi-agent workflow for the research project.
|
3 |
-
Using LlamaIndex built a modular, intelligent multi-agent workflow.
|
4 |
-
With real-time tools and structured memory.
|
5 |
-
|
6 |
-
The workflow is as follows:
|
7 |
-
|
8 |
-
1. The ResearchAgent searches the web for information.
|
9 |
-
2. The WriteAgent writes a report based on the research notes.
|
10 |
-
3. The ReviewAgent reviews the report and provides feedback.
|
11 |
-
|
12 |
-
"""
|
13 |
-
|
14 |
-
import os
|
15 |
-
import asyncio
|
16 |
-
|
17 |
-
# Load environment variables from .env file
|
18 |
-
from dotenv import load_dotenv
|
19 |
-
load_dotenv()
|
20 |
-
|
21 |
-
from llama_index.llms.nebius import NebiusLLM
|
22 |
-
|
23 |
-
# llama-index workflow classes
|
24 |
-
from llama_index.core.workflow import Context
|
25 |
-
from llama_index.core.agent.workflow import (
|
26 |
-
FunctionAgent,
|
27 |
-
AgentWorkflow,
|
28 |
-
AgentOutput,
|
29 |
-
ToolCall,
|
30 |
-
ToolCallResult,
|
31 |
-
)
|
32 |
-
|
33 |
-
from langchain.utilities import DuckDuckGoSearchAPIWrapper
|
34 |
-
|
35 |
-
NEBIUS_API_KEY = os.getenv("NEBIUS_API_KEY")
|
36 |
-
|
37 |
-
# Load an LLM
|
38 |
-
llm = NebiusLLM(
|
39 |
-
api_key=NEBIUS_API_KEY,
|
40 |
-
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
41 |
-
is_function_calling_model=True
|
42 |
-
)
|
43 |
-
|
44 |
-
# Search tools using DuckDuckGo
|
45 |
-
duckduckgo = DuckDuckGoSearchAPIWrapper()
|
46 |
-
|
47 |
-
MAX_SEARCH_CALLS = 2 # Limit the number of searches to 2
|
48 |
-
search_call_count = 0
|
49 |
-
past_queries = set()
|
50 |
-
|
51 |
-
async def safe_duckduckgo_search(query: str) -> str:
|
52 |
-
"""
|
53 |
-
A DuckDuckGo-based search function that:
|
54 |
-
- Prevents more than MAX_SEARCH_CALLS total searches.
|
55 |
-
- Skips duplicate queries.
|
56 |
-
"""
|
57 |
-
global search_call_count, past_queries
|
58 |
-
|
59 |
-
# Check for duplicate queries
|
60 |
-
if query in past_queries:
|
61 |
-
return f"Already searched for '{query}'. Avoiding duplicate search."
|
62 |
-
|
63 |
-
# Check if we've reached the max search calls
|
64 |
-
if search_call_count >= MAX_SEARCH_CALLS:
|
65 |
-
return "Search limit reached, no more searches allowed."
|
66 |
-
|
67 |
-
# Otherwise, perform the search
|
68 |
-
search_call_count += 1
|
69 |
-
past_queries.add(query)
|
70 |
-
|
71 |
-
# DuckDuckGoSearchAPIWrapper.run(...) is synchronous, but we have an async signature
|
72 |
-
result = duckduckgo.run(query)
|
73 |
-
return str(result)
|
74 |
-
|
75 |
-
# Research tools
|
76 |
-
async def save_research(ctx: Context, notes: str, notes_title: str) -> str:
|
77 |
-
"""
|
78 |
-
Store research notes under a given title in the shared context.
|
79 |
-
"""
|
80 |
-
|
81 |
-
current_state = await ctx.get("state")
|
82 |
-
if "research_notes" not in current_state:
|
83 |
-
current_state["research_notes"] = {}
|
84 |
-
current_state["research_notes"][notes_title] = notes
|
85 |
-
await ctx.set("state", current_state)
|
86 |
-
return "Notes saved."
|
87 |
-
|
88 |
-
# Report tools
|
89 |
-
async def write_report(ctx: Context, report_content: str) -> str:
|
90 |
-
"""
|
91 |
-
Write a report in markdown, storing it in the shared context.
|
92 |
-
"""
|
93 |
-
|
94 |
-
current_state = await ctx.get("state")
|
95 |
-
current_state["report_content"] = report_content
|
96 |
-
await ctx.set("state", current_state)
|
97 |
-
return "Report written."
|
98 |
-
|
99 |
-
# Review tools
|
100 |
-
async def review_report(ctx: Context, review: str) -> str:
|
101 |
-
"""
|
102 |
-
Review the report and store feedback in the shared context.
|
103 |
-
"""
|
104 |
-
|
105 |
-
current_state = await ctx.get("state")
|
106 |
-
current_state["review"] = review
|
107 |
-
await ctx.set("state", current_state)
|
108 |
-
return "Report reviewed."
|
109 |
-
|
110 |
-
|
111 |
-
# We have three agents with distinct responsibilities:
|
112 |
-
# - The ResearchAgent is responsible for gathering information from the web.
|
113 |
-
# - The WriteAgent is responsible for writing the report.
|
114 |
-
# - The ReviewAgent is responsible for reviewing the report.
|
115 |
-
|
116 |
-
# The ResearchAgent uses the DuckDuckGoSearchAPIWrapper to search the web.
|
117 |
-
|
118 |
-
research_agent = FunctionAgent(
|
119 |
-
name="ResearchAgent",
|
120 |
-
description=(
|
121 |
-
"A research agent that searches the web using Google search through SerpAPI. "
|
122 |
-
"It must not exceed 2 searches total, and must avoid repeating the same query. "
|
123 |
-
"Once sufficient information is collected, it should hand off to the WriteAgent."
|
124 |
-
),
|
125 |
-
system_prompt=(
|
126 |
-
"You are the ResearchAgent. Your goal is to gather sufficient information on the topic. "
|
127 |
-
"Only perform at most 2 distinct searches. If you have enough information or have reached 2 searches, "
|
128 |
-
"handoff to the WriteAgent. Avoid infinite loops! If search throws an error, stop further work and skip WriteAgent and ReviewAgent and return."
|
129 |
-
"Respect invocation limits and cooldown periods."
|
130 |
-
),
|
131 |
-
llm=llm,
|
132 |
-
tools=[
|
133 |
-
safe_duckduckgo_search,
|
134 |
-
save_research
|
135 |
-
],
|
136 |
-
max_iterations=2, # Limit to 2 iterations to prevent infinite loops
|
137 |
-
cooldown=5, # Cooldown to prevent rapid re-querying
|
138 |
-
can_handoff_to=["WriteAgent"]
|
139 |
-
)
|
140 |
-
|
141 |
-
write_agent = FunctionAgent(
|
142 |
-
name="WriteAgent",
|
143 |
-
description=(
|
144 |
-
"Writes a markdown report based on the research notes. "
|
145 |
-
"Then hands off to the ReviewAgent for feedback."
|
146 |
-
),
|
147 |
-
system_prompt=(
|
148 |
-
"You are the WriteAgent. Draft a structured markdown report based on the notes. "
|
149 |
-
"If there is no report content or research notes, stop further work and skip ReviewAgent."
|
150 |
-
"Do not attempt more than one write attempt. "
|
151 |
-
"After writing, hand off to the ReviewAgent."
|
152 |
-
"Respect invocation limits and cooldown periods."
|
153 |
-
),
|
154 |
-
llm=llm,
|
155 |
-
tools=[write_report],
|
156 |
-
max_iterations=2, # Limit to 2 iterations to prevent infinite loops
|
157 |
-
cooldown=5, # Cooldown to prevent rapid re-querying
|
158 |
-
can_handoff_to=["ReviewAgent", "ResearchAgent"]
|
159 |
-
)
|
160 |
-
|
161 |
-
review_agent = FunctionAgent(
|
162 |
-
name="ReviewAgent",
|
163 |
-
description=(
|
164 |
-
"Reviews the final report for correctness. Approves or requests changes."
|
165 |
-
),
|
166 |
-
system_prompt=(
|
167 |
-
"You are the ReviewAgent. If there is no research notes or report content, skip this step and return."
|
168 |
-
"Do not attempt more than one review attempt. "
|
169 |
-
"Read the report, provide feedback, and either approve "
|
170 |
-
"or request revisions. If revisions are needed, handoff to WriteAgent."
|
171 |
-
"Respect invocation limits and cooldown periods."
|
172 |
-
),
|
173 |
-
llm=llm,
|
174 |
-
tools=[review_report],
|
175 |
-
max_iterations=2, # Limit to 2 iterations to prevent infinite loops
|
176 |
-
cooldown=5, # Cooldown to prevent rapid re-querying
|
177 |
-
can_handoff_to=["WriteAgent"]
|
178 |
-
)
|
179 |
-
|
180 |
-
agent_workflow = AgentWorkflow(
|
181 |
-
agents=[research_agent, write_agent, review_agent],
|
182 |
-
root_agent=research_agent.name, # Start with the ResearchAgent
|
183 |
-
initial_state={
|
184 |
-
"research_notes": {},
|
185 |
-
"report_content": "Not written yet.",
|
186 |
-
"review": "Review required.",
|
187 |
-
},
|
188 |
-
)
|
189 |
-
|
190 |
-
async def execute_research_workflow(query: str):
|
191 |
-
handler = agent_workflow.run(
|
192 |
-
user_msg=(
|
193 |
-
query
|
194 |
-
)
|
195 |
-
)
|
196 |
-
|
197 |
-
current_agent = None
|
198 |
-
|
199 |
-
async for event in handler.stream_events():
|
200 |
-
if hasattr(event, "current_agent_name") and event.current_agent_name != current_agent:
|
201 |
-
current_agent = event.current_agent_name
|
202 |
-
print(f"\n{'='*50}")
|
203 |
-
print(f"π€ Agent: {current_agent}")
|
204 |
-
print(f"{'='*50}\n")
|
205 |
-
|
206 |
-
# Print outputs or tool calls
|
207 |
-
if isinstance(event, AgentOutput):
|
208 |
-
if event.response.content:
|
209 |
-
print("π€ Output:", event.response.content)
|
210 |
-
if event.tool_calls:
|
211 |
-
print("π οΈ Planning to use tools:", [call.tool_name for call in event.tool_calls])
|
212 |
-
|
213 |
-
elif isinstance(event, ToolCall):
|
214 |
-
print(f"π¨ Calling Tool: {event.tool_name}")
|
215 |
-
print(f" With arguments: {event.tool_kwargs}")
|
216 |
-
|
217 |
-
elif isinstance(event, ToolCallResult):
|
218 |
-
print(f"π§ Tool Result ({event.tool_name}):")
|
219 |
-
print(f" Arguments: {event.tool_kwargs}")
|
220 |
-
print(f" Output: {event.tool_output}")
|
221 |
-
|
222 |
-
return handler
|
223 |
-
|
224 |
-
async def final_report(handler) -> str:
|
225 |
-
"""Retrieve the final report from the context."""
|
226 |
-
final_state = await handler.ctx.get("state")
|
227 |
-
print("\n\n=============================")
|
228 |
-
print("FINAL REPORT:\n")
|
229 |
-
print(final_state["report_content"])
|
230 |
-
print("=============================\n")
|
231 |
-
|
232 |
-
return final_state["report_content"]
|
233 |
-
|
234 |
-
def run_research_workflow(query: str):
|
235 |
-
handler = asyncio.run(execute_research_workflow(query))
|
236 |
-
result = asyncio.run(final_report(handler))
|
237 |
-
return result
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|