Kye Gomez commited on
Commit
89c7bc8
1 Parent(s): a7db6ca

Create agents.py

Browse files
Files changed (1) hide show
  1. agents.py +337 -0
agents.py ADDED
@@ -0,0 +1,337 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from typing import Any, Dict
3
+
4
+ import requests
5
+ import tavily
6
+ from dotenv import load_dotenv
7
+
8
+ from swarms import Agent, OpenAIChat
9
+ from swarms.tools.prebuilt.bing_api import fetch_web_articles_bing_api
10
+ from swarms.utils.loguru_logger import logger
11
+
12
+ load_dotenv()
13
+
14
+ try:
15
+ from openai import OpenAI
16
+
17
+ from swarms import BaseLLM
18
+ except ImportError as e:
19
+ raise ImportError(f"Required modules are not available: {e}")
20
+
21
+
22
+ def perplexity_api_key():
23
+ try:
24
+ api_key = os.getenv("PPLX_API_KEY")
25
+ return api_key
26
+ except Exception as e:
27
+ print(f"Error: {e}")
28
+
29
+
30
+ class Perplexity(BaseLLM):
31
+ """
32
+ A class to interact with the Perplexity API using OpenAI's interface.
33
+ """
34
+
35
+ def __init__(self, api_key: str = perplexity_api_key(), *args, **kwargs):
36
+ """
37
+ Initialize the Perplexity class with an API key.
38
+
39
+ Args:
40
+ api_key (str): The API key for authenticating with the OpenAI client.
41
+ """
42
+ super().__init__(*args, **kwargs)
43
+ self.client = OpenAI(
44
+ api_key=api_key,
45
+ base_url="https://api.perplexity.ai",
46
+ *args,
47
+ **kwargs,
48
+ )
49
+
50
+ def run(self, task: str, *args, **kwargs):
51
+ """
52
+ Run the model to process the given task.
53
+
54
+ Args:
55
+ task (str): The task to be performed.
56
+
57
+ Returns:
58
+ dict: The processed output from the model.
59
+ """
60
+ messages = [
61
+ {
62
+ "role": "system",
63
+ "content": (
64
+ "You are an artificial intelligence assistant and you need to "
65
+ "engage in a helpful, detailed, polite conversation with a user."
66
+ ),
67
+ },
68
+ {
69
+ "role": "user",
70
+ "content": task,
71
+ },
72
+ ]
73
+ try:
74
+ response = self.client.chat.completions.create(
75
+ model="llama-3-sonar-large-32k-online",
76
+ messages=messages,
77
+ )
78
+ return response
79
+ except Exception as e:
80
+ raise RuntimeError(f"Error running the model: {e}")
81
+
82
+
83
+ def check_exa_api():
84
+ try:
85
+ api_key = os.getenv("EXA_API_KEY")
86
+ return api_key
87
+ except Exception as e:
88
+ print(f"Error: {e}")
89
+
90
+
91
+ class ExaAgent(BaseLLM):
92
+ """
93
+ A class to interact with the Exa API.
94
+ """
95
+
96
+ def __init__(self, api_key: str = check_exa_api(), *args, **kwargs):
97
+ """
98
+ Initialize the ExaAgent class with an API key.
99
+
100
+ Args:
101
+ api_key (str): The API key for authenticating with the Exa client.
102
+ """
103
+ super().__init__(*args, **kwargs)
104
+ try:
105
+ from exa_py import Exa
106
+
107
+ self.exa = Exa(api_key=api_key)
108
+ except ImportError as e:
109
+ raise ImportError(f"Failed to import Exa: {e}")
110
+
111
+ def run(self, task: str, *args, **kwargs):
112
+ """
113
+ Run a search query using the Exa API.
114
+
115
+ Args:
116
+ task (str): The search query.
117
+
118
+ Returns:
119
+ dict: The search results from the Exa API.
120
+ """
121
+ try:
122
+ results = self.exa.search(
123
+ task, use_autoprompt=True, *args, **kwargs
124
+ )
125
+ return results
126
+ except Exception as e:
127
+ raise RuntimeError(f"Error running the search query: {e}")
128
+
129
+
130
+ class ResearchAgent:
131
+ """
132
+ A class to represent a research agent that uses an LLM to summarize content from various sources.
133
+ """
134
+
135
+ def __init__(
136
+ self,
137
+ api_key: str,
138
+ output_dir: str = "research_base",
139
+ n_results: int = 2,
140
+ temperature: float = 0.2,
141
+ max_tokens: int = 3500,
142
+ ):
143
+ """
144
+ Initialize the ResearchAgent class with necessary parameters.
145
+
146
+ Args:
147
+ api_key (str): The API key for the Bing API.
148
+ output_dir (str): The directory for storing memory outputs. Default is "research_base".
149
+ n_results (int): Number of results to return from the memory. Default is 2.
150
+ temperature (float): The temperature setting for the LLM. Default is 0.2.
151
+ max_tokens (int): The maximum number of tokens for the LLM. Default is 3500.
152
+ """
153
+ self.api_key = api_key
154
+
155
+ self.llm = OpenAIChat(
156
+ temperature=temperature, max_tokens=max_tokens
157
+ )
158
+ self.agent = self._initialize_agent()
159
+
160
+ def _initialize_agent(self):
161
+ """
162
+ Initialize the agent with the provided parameters and system prompt.
163
+
164
+ Returns:
165
+ Agent: An initialized Agent instance.
166
+ """
167
+ research_system_prompt = """
168
+ Research Agent LLM Prompt: Summarizing Sources and Content
169
+ Objective: Your task is to summarize the provided sources and the content within those sources. The goal is to create concise, accurate, and informative summaries that capture the key points of the original content.
170
+ Instructions:
171
+ 1. Identify Key Information: ...
172
+ 2. Summarize Clearly and Concisely: ...
173
+ 3. Preserve Original Meaning: ...
174
+ 4. Include Relevant Details: ...
175
+ 5. Structure: ...
176
+ """
177
+
178
+ return Agent(
179
+ agent_name="Research Agent",
180
+ system_prompt=research_system_prompt,
181
+ llm=self.llm,
182
+ max_loops=1,
183
+ autosave=True,
184
+ dashboard=False,
185
+ # tools=[fetch_web_articles_bing_api],
186
+ verbose=True,
187
+ )
188
+
189
+ def run(self, task: str, *args, **kwargs):
190
+ """
191
+ Run the research agent to fetch and summarize web articles related to the task.
192
+
193
+ Args:
194
+ task (str): The task or query for the agent to process.
195
+
196
+ Returns:
197
+ str: The agent's response after processing the task.
198
+ """
199
+ articles = fetch_web_articles_bing_api(
200
+ task, subscription_key=self.api_key
201
+ )
202
+ sources_prompts = "".join([task, articles])
203
+ agent_response = self.agent.run(sources_prompts)
204
+ return agent_response
205
+
206
+
207
+ def check_tavily_api():
208
+ try:
209
+ api_key = os.getenv("TAVILY_API_KEY")
210
+ return api_key
211
+ except Exception as e:
212
+ print(f"Error: {e}")
213
+
214
+
215
+ class TavilyWrapper:
216
+ """
217
+ A wrapper class for the Tavily API to facilitate searches and retrieve relevant information.
218
+ """
219
+
220
+ def __init__(self, api_key: str = check_tavily_api()):
221
+ """
222
+ Initialize the TavilyWrapper with the provided API key.
223
+
224
+ Args:
225
+ api_key (str): The API key for authenticating with the Tavily API.
226
+ """
227
+ if not isinstance(api_key, str):
228
+ raise TypeError("API key must be a string")
229
+
230
+ self.api_key = api_key
231
+ self.client = self._initialize_client(api_key)
232
+
233
+ def _initialize_client(self, api_key: str) -> Any:
234
+ """
235
+ Initialize the Tavily client with the provided API key.
236
+
237
+ Args:
238
+ api_key (str): The API key for authenticating with the Tavily API.
239
+
240
+ Returns:
241
+ TavilyClient: An initialized Tavily client instance.
242
+ """
243
+ try:
244
+ return tavily.TavilyClient(api_key=api_key)
245
+ except Exception as e:
246
+ raise RuntimeError(f"Error initializing Tavily client: {e}")
247
+
248
+ def run(self, task: str) -> Dict[str, Any]:
249
+ """
250
+ Perform a search query using the Tavily API.
251
+
252
+ Args:
253
+ task (str): The search query.
254
+
255
+ Returns:
256
+ dict: The search results from the Tavily API.
257
+ """
258
+ if not isinstance(task, str):
259
+ raise TypeError("Task must be a string")
260
+
261
+ try:
262
+ response = self.client.search(
263
+ query=task, search_depth="advanced"
264
+ )
265
+ return response
266
+ except Exception as e:
267
+ raise RuntimeError(f"Error performing search: {e}")
268
+
269
+
270
+ def you_search_api_key():
271
+ try:
272
+ api_key = os.getenv("YOU_API_KEY")
273
+ return api_key
274
+ except Exception as e:
275
+ print(f"Error: {e}")
276
+
277
+
278
+ class YouSearchAgent:
279
+ """
280
+ A wrapper class for the YDC Index API to facilitate fetching AI snippets based on a query.
281
+ """
282
+
283
+ def __init__(self, api_key: str = you_search_api_key()):
284
+ """
285
+ Initialize the AISnippetsWrapper with the provided API key.
286
+
287
+ Args:
288
+ api_key (str): The API key for authenticating with the YDC Index API.
289
+ """
290
+
291
+ self.api_key = api_key
292
+
293
+ def run(self, task: str) -> Dict[str, Any]:
294
+ """
295
+ Fetch AI snippets for the given query using the YDC Index API.
296
+
297
+ Args:
298
+ task (str): The search query.
299
+
300
+ Returns:
301
+ dict: The search results from the YDC Index API.
302
+ """
303
+ if not isinstance(task, str):
304
+ raise TypeError("Task must be a string")
305
+
306
+ headers = {"X-API-Key": self.api_key}
307
+ params = {"query": task}
308
+
309
+ try:
310
+ response = requests.get(
311
+ "https://api.ydc-index.io/search",
312
+ params=params,
313
+ headers=headers,
314
+ )
315
+ response.raise_for_status() # Raise an error for bad status codes
316
+ return response.json()
317
+ except requests.RequestException as e:
318
+ raise RuntimeError(f"Error fetching AI snippets: {e}")
319
+
320
+
321
+ task = "What is the swarmms framework"
322
+
323
+ # Run all of the agents
324
+ agents = [
325
+ Perplexity,
326
+ ExaAgent,
327
+ # ResearchAgent,
328
+ TavilyWrapper,
329
+ # YouSearchAgent,
330
+ ]
331
+
332
+ # Run each agent with the given task
333
+ for agent_class in agents:
334
+ logger.info(f"Running agent: {agent_class.__name__}")
335
+ agent = agent_class()
336
+ response = agent.run(task)
337
+ print(response)