betki commited on
Commit
192f945
Β·
verified Β·
1 Parent(s): cd28f35

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