cognitrix 0.2.1__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. {cognitrix-0.2.1 → cognitrix-0.2.2}/PKG-INFO +11 -6
  2. {cognitrix-0.2.1 → cognitrix-0.2.2}/README.md +5 -5
  3. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/__init__.py +1 -1
  4. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/agents/assistant.py +2 -2
  5. cognitrix-0.2.2/cognitrix/agents/base.py +417 -0
  6. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/agents/templates.py +152 -3
  7. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/app.py +2 -6
  8. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/cli.py +69 -49
  9. cognitrix-0.2.2/cognitrix/config.py +30 -0
  10. cognitrix-0.2.2/cognitrix/llms/__init__.py +10 -0
  11. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/llms/anthropic_llm.py +4 -4
  12. cognitrix-0.2.2/cognitrix/llms/base.py +90 -0
  13. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/llms/cohere_llm.py +2 -7
  14. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/llms/google_llm.py +15 -9
  15. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/llms/groq_llm.py +33 -4
  16. cognitrix-0.2.2/cognitrix/llms/local_llm.py +123 -0
  17. cognitrix-0.2.2/cognitrix/llms/mindsdb_llm.py +40 -0
  18. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/llms/openai_llm.py +16 -6
  19. cognitrix-0.2.2/cognitrix/llms/session.py +69 -0
  20. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/llms/together_llm.py +26 -36
  21. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/tasks/base.py +6 -1
  22. cognitrix-0.2.2/cognitrix/tools/__init__.py +12 -0
  23. cognitrix-0.2.2/cognitrix/tools/base.py +60 -0
  24. cognitrix-0.2.2/cognitrix/tools/misc.py +654 -0
  25. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/tools/python.py +1 -1
  26. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/tools/tool.py +1 -1
  27. cognitrix-0.2.2/cognitrix/transcriber.py +95 -0
  28. cognitrix-0.2.2/cognitrix/utils/__init__.py +151 -0
  29. {cognitrix-0.2.1 → cognitrix-0.2.2}/pyproject.toml +6 -1
  30. cognitrix-0.2.1/cognitrix/agents/base.py +0 -363
  31. cognitrix-0.2.1/cognitrix/config.py +0 -19
  32. cognitrix-0.2.1/cognitrix/llms/__init__.py +0 -8
  33. cognitrix-0.2.1/cognitrix/llms/base.py +0 -145
  34. cognitrix-0.2.1/cognitrix/tools/__init__.py +0 -12
  35. cognitrix-0.2.1/cognitrix/tools/base.py +0 -27
  36. cognitrix-0.2.1/cognitrix/tools/misc.py +0 -535
  37. cognitrix-0.2.1/cognitrix/tools/serpapi.py +0 -134
  38. cognitrix-0.2.1/cognitrix/utils/__init__.py +0 -12
  39. {cognitrix-0.2.1 → cognitrix-0.2.2}/LICENSE +0 -0
  40. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/agents/__init__.py +0 -0
  41. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/agents/core.py +0 -0
  42. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/llms/clarifai_llm.py +0 -0
  43. {cognitrix-0.2.1 → cognitrix-0.2.2}/cognitrix/tasks/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: cognitrix
3
- Version: 0.2.1
3
+ Version: 0.2.2
4
4
  Summary: Package for creating AI Agents using llms
5
5
  Home-page: https://github.com/theonlyamos/cognitrix
6
6
  License: Apache2
@@ -23,10 +23,13 @@ Classifier: Programming Language :: Python :: 3.8
23
23
  Classifier: Programming Language :: Python :: 3.9
24
24
  Classifier: Topic :: Software Development :: Build Tools
25
25
  Classifier: Topic :: Software Development :: Libraries :: Python Modules
26
+ Requires-Dist: aiofiles (>=23.2.1,<24.0.0)
26
27
  Requires-Dist: aiohttp (>=3.9.3,<4.0.0)
27
28
  Requires-Dist: anthropic (>=0.17.0,<0.18.0)
29
+ Requires-Dist: beautifulsoup4 (>=4.12.3,<5.0.0)
28
30
  Requires-Dist: clarifai (>=10.1.0,<11.0.0)
29
31
  Requires-Dist: cohere (>=4.50,<5.0)
32
+ Requires-Dist: deepgram-sdk (==3.*)
30
33
  Requires-Dist: google-generativeai (>=0.4.0,<0.5.0)
31
34
  Requires-Dist: google-search-results (>=2.4.2,<3.0.0)
32
35
  Requires-Dist: groq (>=0.4.1,<0.5.0)
@@ -35,10 +38,12 @@ Requires-Dist: pillow (>=10.2.0,<11.0.0)
35
38
  Requires-Dist: pyautogui (>=0.9.54,<0.10.0)
36
39
  Requires-Dist: pydantic (>=2.6.2,<3.0.0)
37
40
  Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
41
+ Requires-Dist: pyttsx3 (>=2.90,<3.0)
38
42
  Requires-Dist: requests (>=2.31.0,<3.0.0)
39
43
  Requires-Dist: serpapi (>=0.1.5,<0.2.0)
40
44
  Requires-Dist: together (>=0.2.11,<0.3.0)
41
45
  Requires-Dist: uuid (>=1.30,<2.0)
46
+ Requires-Dist: xmltodict (>=0.13.0,<0.14.0)
42
47
  Project-URL: Bug Tracker, https://github.com/theonlyamos/cognitrix/issues
43
48
  Project-URL: Repository, https://github.com/theonlyamos/cognitrix
44
49
  Description-Content-Type: text/markdown
@@ -82,14 +87,14 @@ Run with default settings
82
87
  cognitrix
83
88
  ```
84
89
 
85
- List supported platforms
90
+ List supported providers
86
91
  ```bash
87
- cognitrix --platforms
92
+ cognitrix --providers
88
93
  ```
89
94
 
90
- Run with specific platform
95
+ Run with specific provider
91
96
  ```bash
92
- cognitrix --platform <platform_name>
97
+ cognitrix --provider <provider_name>
93
98
  ```
94
99
 
95
100
  Create a new agent
@@ -119,4 +124,4 @@ Cognitrix is open source and contributions are welcome. See [CONTRIBUTING.md](CO
119
124
  This project is licensed under the MIT license. See [LICENSE.md](LICENSE.md) for more details.
120
125
 
121
126
  ### Acknowledgments
122
- Cognitrix was created by [Amos Amissah](https://github.com/theonlyamos) and is heavily inspired by projects like AutoGPT and EngineerGPT. Special thanks to the open-source community for their contributions and the AI companies providing LLM APIs.
127
+ Cognitrix was created by [Amos Amissah](https://github.com/theonlyamos) and is heavily inspired by projects like AutoGPT and GPT Engineer. Special thanks to the open-source community for their contributions and the AI companies providing LLM APIs.
@@ -37,14 +37,14 @@ Run with default settings
37
37
  cognitrix
38
38
  ```
39
39
 
40
- List supported platforms
40
+ List supported providers
41
41
  ```bash
42
- cognitrix --platforms
42
+ cognitrix --providers
43
43
  ```
44
44
 
45
- Run with specific platform
45
+ Run with specific provider
46
46
  ```bash
47
- cognitrix --platform <platform_name>
47
+ cognitrix --provider <provider_name>
48
48
  ```
49
49
 
50
50
  Create a new agent
@@ -74,4 +74,4 @@ Cognitrix is open source and contributions are welcome. See [CONTRIBUTING.md](CO
74
74
  This project is licensed under the MIT license. See [LICENSE.md](LICENSE.md) for more details.
75
75
 
76
76
  ### Acknowledgments
77
- Cognitrix was created by [Amos Amissah](https://github.com/theonlyamos) and is heavily inspired by projects like AutoGPT and EngineerGPT. Special thanks to the open-source community for their contributions and the AI companies providing LLM APIs.
77
+ Cognitrix was created by [Amos Amissah](https://github.com/theonlyamos) and is heavily inspired by projects like AutoGPT and GPT Engineer. Special thanks to the open-source community for their contributions and the AI companies providing LLM APIs.
@@ -1,3 +1,3 @@
1
1
  from .tools import Tool
2
2
  from .agents import Agent, AIAssistant
3
- from .llms import LLM, Clarifai, Cohere, Groq, OpenAI, Together, Gemini
3
+ from .llms import LLM, Clarifai, Cohere, Groq, OpenAI, Together, Google
@@ -1,6 +1,6 @@
1
1
  from pydantic import Field
2
2
  from ..agents.base import Agent
3
- from ..agents.templates import ASSISTANT_TEMPLATE
3
+ from ..agents.templates import AUTONOMOUSE_AGENT_2
4
4
 
5
5
  class AIAssistant(Agent):
6
6
  """_summary_
@@ -8,4 +8,4 @@ class AIAssistant(Agent):
8
8
  Args:
9
9
  Agent (_type_): _description_
10
10
  """
11
- prompt_template: str = Field(default=ASSISTANT_TEMPLATE)
11
+ prompt_template: str = Field(default=AUTONOMOUSE_AGENT_2)
@@ -0,0 +1,417 @@
1
+ import inspect
2
+ import json
3
+ import sys
4
+ import uuid
5
+ import asyncio
6
+ import logging
7
+ import aiofiles
8
+ from pathlib import Path
9
+ from threading import Thread
10
+ from typing import List, Optional, Self, TypeAlias, Union, Type, Any
11
+ from dataclasses import dataclass
12
+
13
+ from pydantic import BaseModel, Field
14
+
15
+ from cognitrix.tasks import Task
16
+ from cognitrix.llms.base import LLM
17
+ from cognitrix.tools.base import Tool
18
+ from cognitrix.utils import extract_json, json_return_format
19
+ from cognitrix.agents.templates import AUTONOMOUSE_AGENT_2
20
+ from cognitrix.config import AGENTS_FILE, SESSIONS_FILE
21
+ from cognitrix.llms.session import Session
22
+ from cognitrix.transcriber import Transcriber
23
+
24
+ logger = logging.getLogger('cognitrix.log')
25
+
26
+ AgentList: TypeAlias = List['Agent']
27
+
28
+ class Agent(BaseModel):
29
+ name: str = Field(default='Agent')
30
+ llm: LLM
31
+ tools: List[Tool] = Field(default_factory=list)
32
+ prompt_template: str = Field(default=AUTONOMOUSE_AGENT_2)
33
+ verbose: bool = Field(default=False)
34
+ sub_agents: AgentList = Field(default_factory=list)
35
+ task: Optional[Task] = None
36
+ is_sub_agent: bool = Field(default=False)
37
+ id: str = Field(default_factory=lambda: uuid.uuid4().hex)
38
+ parent_id: Optional[str] = None
39
+ autostart: bool = False
40
+
41
+ @property
42
+ def available_tools(self) -> List[str]:
43
+ return [tool.name for tool in self.tools]
44
+
45
+ def format_system_prompt(self):
46
+ tools_str = self._format_tools_string()
47
+ subagents_str = self._format_subagents_string()
48
+ llms_str = self._format_llms_string()
49
+
50
+ prompt = self.prompt_template
51
+ prompt = prompt.replace("{name}", self.name)
52
+ prompt = prompt.replace("{tools}", tools_str)
53
+ prompt = prompt.replace("{subagents}", subagents_str)
54
+ prompt = prompt.replace("{available_tools}", json.dumps(self.available_tools))
55
+ prompt = prompt.replace("{llms}", llms_str)
56
+ prompt = prompt.replace("{return_format}", json_return_format)
57
+
58
+ if 'json' not in prompt.lower():
59
+ prompt += f"\n{json_return_format}"
60
+
61
+ self.llm.system_prompt = prompt
62
+
63
+ def _format_tools_string(self) -> str:
64
+ return "You have access to the following Tools:\n" + "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
65
+
66
+ def _format_subagents_string(self) -> str:
67
+ subagents_str = "Available Subagents:\n"
68
+ subagents_str += "\n".join([f"-- {agent.name}: {agent.task.description}" for agent in self.sub_agents if agent.task])
69
+ subagents_str += "\nYou should always use a subagent for a task if there is one specifically created for that task."
70
+ subagents_str += "\nWhen creating a sub agent, it's description should be a comprehensive prompt of the agent's behavior, capabilities and example tasks."
71
+ return subagents_str
72
+
73
+ def _format_llms_string(self) -> str:
74
+ llms = LLM.list_llms()
75
+ llms_str = "Available LLM Providers:\n" + ", ".join(llms) + "\nChoose one for each subagent."
76
+ return llms_str
77
+
78
+ def generate_prompt(self, query: str | dict, role: str = 'User') -> dict:
79
+ processed_query = self._process_query(query)
80
+ prompt: dict[str, Any] = {'role': role, 'type': 'text'}
81
+
82
+ if isinstance(processed_query, dict):
83
+ if self.is_sub_agent:
84
+ print("=======is sub agent===========")
85
+ print(processed_query)
86
+ result = processed_query['result']
87
+ if isinstance(result, list):
88
+ if result[0] == 'image':
89
+ prompt['type'] = 'image'
90
+ prompt['image'] = result[1]
91
+ elif result[0] == 'agent':
92
+ new_agent: Agent = result[1]
93
+ new_agent.parent_id = self.id
94
+ self.add_sub_agent(new_agent)
95
+
96
+ prompt['message'] = result[2]
97
+ else:
98
+ prompt['message'] = result
99
+ else:
100
+ prompt['message'] = result
101
+ else:
102
+ prompt['message'] = processed_query
103
+
104
+ return prompt
105
+
106
+ def _process_query(self, query: str | dict) -> str | dict:
107
+ return extract_json(query) if isinstance(query, str) else query
108
+
109
+ def add_sub_agent(self, agent: 'Agent'):
110
+ self.sub_agents.append(agent)
111
+
112
+ def get_sub_agent_by_name(self, name: str) -> Optional['Agent']:
113
+ return next((agent for agent in self.sub_agents if agent.name.lower() == name.lower()), None)
114
+
115
+ def get_tool_by_name(self, name: str) -> Optional[Tool]:
116
+ return next((tool for tool in self.tools if tool.name.lower() == name.lower()), None)
117
+
118
+ async def process_response(self, response: str) -> Union[dict, str]:
119
+ # response = response.replace("'", '"')
120
+ response = response.replace('\\n', '')
121
+ response = response.replace("'", "\"")
122
+ # response = response.replace('"', '\\"')
123
+ response_data = extract_json(response)
124
+
125
+
126
+ try:
127
+ if isinstance(response_data, dict):
128
+ final_result_keys = ['final_answer', 'function_call_result', 'respons']
129
+
130
+ if response_data['type'].replace('\\', '') in final_result_keys:
131
+ return response_data['result']
132
+
133
+ tool = self.get_tool_by_name(response_data['function'])
134
+ if isinstance(response_data['arguments'], dict):
135
+ response_data['arguments'] = list(response_data['arguments'].values())
136
+
137
+ if response_data['function'].lower() == 'create agents':
138
+ response_data['arguments'] = [*response_data['arguments'], self.id]
139
+
140
+ if not tool:
141
+ raise Exception(f"Tool {response_data['function']} not found")
142
+
143
+ print(f"\nRunning tool '{tool.name.title()}' with parameters: {response_data['arguments']}")
144
+
145
+ if 'sub agent' in tool.name.lower():
146
+ response_data['arguments'].append(self)
147
+
148
+ if tool.name.lower() == 'create sub agent':
149
+ result = await tool.arun(*response_data['arguments'])
150
+ else:
151
+ result = tool.run(*response_data['arguments'])
152
+
153
+ response_json = {
154
+ 'type': 'function_call_result',
155
+ 'result': result
156
+ }
157
+
158
+ return response_json
159
+ else:
160
+ raise Exception('Not a json object')
161
+ except Exception as e:
162
+ # logger.exception(e)
163
+ return response_data
164
+
165
+ def add_tool(self, tool: Tool):
166
+ self.tools.append(tool)
167
+
168
+ def initialize(self, session_id: Optional[str] = None):
169
+ session: Session = asyncio.run(Session.load(session_id)) if session_id else Session(chat=self.llm.chat_history, agent_id=self.id)
170
+ self.llm.chat_history = session.chat
171
+
172
+ query: str | dict = input("\nUser (q to quit): ")
173
+ while True:
174
+ try:
175
+ if not query:
176
+ query: str | dict = input("\nUser (q to quit): ")
177
+ continue
178
+ if isinstance(query, str):
179
+ if query.lower() in ['q', 'quit', 'exit']:
180
+ print('Exiting...')
181
+ break
182
+
183
+ elif query.lower() == 'add agent':
184
+ new_agent = asyncio.run(self.create_agent(is_sub_agent=True, parent_id=self.id))
185
+ if new_agent:
186
+ self.add_sub_agent(new_agent)
187
+ print(f"\nAgent {new_agent.name} added successfully!")
188
+ else:
189
+ print("\nError creating agent")
190
+
191
+ query = input("\nUser (q to quit): ")
192
+ continue
193
+
194
+ elif query.lower() == 'list agents':
195
+ agents_str = "\nAvailable Agents:"
196
+ sub_agents = [agent for agent in asyncio.run(self.list_agents()) if agent.parent_id == self.id]
197
+ for index, agent in enumerate(sub_agents):
198
+ agents_str += (f"\n[{index}] {agent.name}")
199
+ print(agents_str)
200
+ query = input("\nUser (q to quit): ")
201
+ continue
202
+
203
+ self.format_system_prompt()
204
+
205
+ full_prompt = self.generate_prompt(query)
206
+ response: Any = self.llm(full_prompt)
207
+
208
+ self.llm.chat_history.append(full_prompt)
209
+ self.llm.chat_history.append({'role': self.name, 'type': 'text', 'message': response})
210
+
211
+ if self.verbose:
212
+ print(response)
213
+
214
+ result: dict[Any, Any] | str = asyncio.run(self.process_response(response))
215
+
216
+ if isinstance(result, dict) and result['type'] == 'function_call_result':
217
+ query = result
218
+ else:
219
+ print(f"\n{self.name}: {result}")
220
+ query = input("\nUser (q to quit): ")
221
+
222
+ self.save_session(session)
223
+
224
+ except KeyboardInterrupt:
225
+ print('Exiting...')
226
+ break
227
+ except Exception as e:
228
+ logger.exception(e)
229
+ break
230
+
231
+ def start(self, session_id: Optional[str] = None):
232
+ self.initialize(session_id)
233
+
234
+ def handle_transcription(self, sentence: str, transcriber: Transcriber):
235
+ if sentence:
236
+ self.format_system_prompt()
237
+
238
+ full_prompt = self.generate_prompt(sentence)
239
+ response: Any = self.llm(full_prompt)
240
+ self.llm.chat_history.append(full_prompt)
241
+ self.llm.chat_history.append({'role': self.name, 'type': 'text', 'message': response})
242
+
243
+ if self.verbose:
244
+ print(response)
245
+
246
+ processsed_response: dict[Any, Any] | str = asyncio.run(self.process_response(response))
247
+
248
+ if isinstance(processsed_response, dict) and processsed_response['type'] == 'function_call_result':
249
+ query = processsed_response
250
+ else:
251
+ print(f"\n{self.name}: {processsed_response}")
252
+ if isinstance(processsed_response, str):
253
+ transcriber.text_to_speech(processsed_response)
254
+
255
+
256
+ def start_audio(self, **kwargs):
257
+ transcriber = Transcriber(on_message_callback=self.handle_transcription) # Pass callback to Transcriber
258
+ if transcriber.start_transcription():
259
+ print("Audio transcription started.")
260
+ input("\n\nPress Enter to stop recording...\n\n")
261
+ transcriber.stop_transcription()
262
+ else:
263
+ print("Failed to start audio transcription.")
264
+
265
+
266
+ @staticmethod
267
+ def start_task_thread(agent: 'Agent', parent: 'Agent'):
268
+ agent_thread = Thread(target=agent.run_task, args=(parent,))
269
+ agent_thread.name = agent.name.lower()
270
+ agent_thread.start()
271
+
272
+ def run_task(self, parent: Self):
273
+ self.format_system_prompt()
274
+ query = self.task.description if self.task else None
275
+
276
+ while query:
277
+ full_prompt = self.generate_prompt(query)
278
+ response: Any = self.llm(full_prompt)
279
+ self.llm.chat_history.append(full_prompt)
280
+ self.llm.chat_history.append({'role': 'user', 'type': 'text', 'message': response})
281
+
282
+ if parent.verbose:
283
+ print(response)
284
+
285
+ agent_result = asyncio.run(self.process_response(response))
286
+ if isinstance(agent_result, dict) and agent_result['type'] == 'function_call_result':
287
+ query = agent_result
288
+ else:
289
+ print(f"\n--{self.name}: {agent_result}")
290
+ parent_prompt = parent.generate_prompt(response, 'user')
291
+ parent_prompt['message'] = self.name + ": " + response
292
+ parent_response: Any = parent.llm(parent_prompt)
293
+ parent.llm.chat_history.append(parent_prompt)
294
+ parent.llm.chat_history.append({'role': 'assistant', 'type': 'text', 'message': parent_response})
295
+ parent_result = asyncio.run(parent.process_response(parent_response))
296
+ print(f"\n\n{parent.name}: {parent_result}")
297
+ query = ""
298
+
299
+ def call_sub_agent(self, agent_name: str, task_description: str):
300
+ sub_agent = self.get_sub_agent_by_name(agent_name)
301
+ if sub_agent:
302
+ sub_agent.task = Task(description=task_description)
303
+ if sub_agent.task:
304
+ self.start_task_thread(sub_agent, self)
305
+ else:
306
+ full_prompt = self.generate_prompt(f'Sub-agent with name {agent_name} was not found.')
307
+ self.llm(full_prompt)
308
+
309
+ @classmethod
310
+ async def create_agent(cls, name: str = '', description: str = '', task_description: str = '', tools: List[Tool] = [],
311
+ llm: Optional[LLM] = None, is_sub_agent: bool = False, parent_id=None) -> Optional[Self]:
312
+ try:
313
+ name = name or input("\n[Enter agent name]: ")
314
+
315
+ while not llm:
316
+ llms = LLM.list_llms()
317
+ llms_str = "\nAvailable LLMs:"
318
+ for index, llm_l in enumerate(llms):
319
+ llms_str += (f"\n[{index}] {llm_l}")
320
+ print(llms_str)
321
+ agent_llm = int(input("\n[Select LLM]: "))
322
+ selected_llm = llms[agent_llm]
323
+ loaded_llm = LLM.load_llm(selected_llm)
324
+
325
+ if loaded_llm:
326
+ llm = loaded_llm()
327
+ if llm:
328
+ llm.model = input(f"\nEnter model name [{llm.model}]: ") or llm.model
329
+ temp = input(f"\nEnter model temperature [{llm.temperature}]: ")
330
+ llm.temperature = float(temp) if temp else llm.temperature
331
+
332
+ if llm:
333
+ task = Task(description=task_description)
334
+ new_agent = cls(name=name, llm=llm, task=task, tools=tools, is_sub_agent=is_sub_agent, parent_id=parent_id)
335
+
336
+ if description:
337
+ new_agent.prompt_template = description
338
+
339
+ agents = []
340
+ async with aiofiles.open(AGENTS_FILE, 'r') as file:
341
+ content = await file.read()
342
+ agents = json.loads(content) if content else []
343
+ agents.append(new_agent.model_dump())
344
+
345
+ async with aiofiles.open(AGENTS_FILE, 'w') as file:
346
+ await file.write(json.dumps(agents, indent=4))
347
+
348
+ return new_agent
349
+
350
+ except Exception as e:
351
+ logger.error(str(e))
352
+
353
+ @classmethod
354
+ async def list_agents(cls, parent_id: Optional[str] = None) -> AgentList:
355
+ try:
356
+ agents: AgentList = []
357
+ async with aiofiles.open(AGENTS_FILE, 'r') as file:
358
+ content = await file.read()
359
+ loaded_agents: list[dict] = json.loads(content) if content else []
360
+ for agent in loaded_agents:
361
+ llm = LLM.load_llm(agent["llm"]["provider"])
362
+ loaded_agent = Agent(**agent)
363
+ if llm:
364
+ loaded_agent.llm = llm(**agent["llm"])
365
+ agents.append(loaded_agent)
366
+
367
+ if parent_id:
368
+ agents = [agent for agent in agents if agent.parent_id == parent_id]
369
+
370
+ return agents
371
+
372
+ except Exception as e:
373
+ logger.exception(e)
374
+ return []
375
+
376
+ @classmethod
377
+ async def get(cls, id) -> Optional['Agent']:
378
+ try:
379
+ agents = await cls.list_agents()
380
+ loaded_agents: list[Agent] = [agent for agent in agents if agent.id == id]
381
+ if len(loaded_agents):
382
+ return loaded_agents[0]
383
+ except Exception as e:
384
+ logger.exception(e)
385
+ return None
386
+
387
+ @classmethod
388
+ def load_agent(cls, agent_name: str) -> Optional['Agent']:
389
+ try:
390
+ agent_name = agent_name.lower()
391
+ agents = asyncio.run(cls.list_agents())
392
+ loaded_agents: list[Agent] = [agent for agent in agents if agent.name.lower() == agent_name]
393
+ if len(loaded_agents):
394
+ agent = loaded_agents[0]
395
+ agent.sub_agents = asyncio.run(cls.list_agents(agent.id))
396
+ return agent
397
+ except Exception as e:
398
+ logger.exception(e)
399
+ return None
400
+
401
+ async def save(self):
402
+ """Save current agent"""
403
+ agents = await Agent.list_agents()
404
+ updated_agents = []
405
+ for index, agent in enumerate(agents):
406
+ if agent.id == self.id:
407
+ agents[index] = self
408
+
409
+ updated_agents.append(agent.model_dump())
410
+
411
+ async with aiofiles.open(AGENTS_FILE, 'w') as file:
412
+ await file.write(json.dumps(updated_agents, indent=4))
413
+
414
+ def save_session(self, session: Session):
415
+ save_thread = Thread(target=session.save, args=(self.llm.chat_history,))
416
+ save_thread.daemon = True
417
+ save_thread.start()