refactor: introduce AgentSystem class (#83)

Author: zzfoo

examples/atomic/IOAgent.py

@@ -2,6 +2,8 @@
 
 import logging
 
+from litemultiagent.core.agent_system import AgentSystem
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -15,34 +17,32 @@
 # Create a logger
 logger = logging.getLogger(__name__)
 def main():
-    agent_manager = AgentManager()
     io_agent_config = {
         "name": "io_agent",
         "type": "atomic",
         "agent_class": "FunctionCallingAgent",
-        "meta_data":
-            {
-                "meta_task_id": "io_subtask",
-                "task_id": 1,
-                "save_to": "supabase",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        "meta_data":{},
         "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],
         "agent_description": "Read or write content from/to a file, or generate and save an image using text input",
         "parameter_description": "The task description detailing what to read, write, or generate. This can include file operations or image generation requests."
     }
-    io_agent = agent_manager.get_agent(io_agent_config)
+    system_config = {
+        "meta_task_id": "io_subtask",
+        "save_to": "csv",
+        "log_dir": "log",
+        "model_name": "gpt-4o-mini",
+        "tool_choice": "auto"
+    }
+    agent_system = AgentSystem(io_agent_config, system_config)
 
     # Example usage
     task = "write aaa to 1.txt, bbb to 2.txt, ccc to 3.txt"
-    result = io_agent.execute(task)
+    result = agent_system.execute(task)
     print("IO Agent Result:", result)
 
 
     task = "generate a image of a ginger cat and save it as ginger_cat.png"
-    result = io_agent.execute(task)
+    result = agent_system.execute(task)
     print("IO Agent Result:", result)
 
 

examples/composite/MasterAgent.py

@@ -1,4 +1,5 @@
 from litemultiagent.core.agent_manager import AgentManager
+from litemultiagent.core.agent_system import AgentSystem
 from litemultiagent.tools.registry import ToolRegistry, Tool
 import logging
 
@@ -15,22 +16,15 @@
 # Create a logger
 logger = logging.getLogger(__name__)
 def main():
-    agent_manager = AgentManager()
-    meta_task_id = "master_agent_task"
-    task_id = 1
     io_agent_config = {
         "name": "io_agent",
         "type": "atomic",
         "agent_class": "FunctionCallingAgent",
         "meta_data":
-            {
-                "meta_task_id": meta_task_id,
-                "task_id": task_id,
-                "save_to": "csv",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        {
+            "model_name": "gpt-4o-mini",
+            "tool_choice": "auto"
+        },
         "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Read or write content from/to a file, or generate and save an image using text input",
         "parameter_description": "The task description detailing what to read, write, or generate. This can include file operations or image generation requests."
@@ -40,15 +34,7 @@ def main():
         "name": "db_retrieval_agent",
         "type": "atomic",
         "agent_class": "FunctionCallingAgent",
-        "meta_data":
-            {
-                "meta_task_id": meta_task_id,
-                "task_id": task_id,
-                "save_to": "csv",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        "meta_data": {},
         "tool_names": ["retrieve_db"], # Changed from "write_file" to "write_to_file"
         "agent_description": "Use a database retrieval agent to fetch information based on a given query.",
         "parameter_description": "The query to be processed by the database retrieval agent."
@@ -58,15 +44,7 @@ def main():
         "name": "file_retrieval_agent",
         "type": "atomic",
         "agent_class": "FunctionCallingAgent",
-        "meta_data":
-            {
-                "meta_task_id": meta_task_id,
-                "task_id": task_id,
-                "save_to": "csv",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        "meta_data": {},
         "tool_names": ["retrieve_file"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Retrieve information from local documents to answer questions or perform tasks.",
         "parameter_description": "The task description specifying the local file and the question to be answered. specify this in natural language"
@@ -76,15 +54,7 @@ def main():
         "name": "web_retrieval_agent",
         "type": "atomic",
         "agent_class": "FunctionCallingAgent",
-        "meta_data":
-            {
-                "meta_task_id": meta_task_id,
-                "task_id": task_id,
-                "save_to": "csv",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        "meta_data": {},
         "tool_names": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Perform a search using API and return the searched results.",
         "parameter_description": "The task description describing what to read or write."
@@ -94,15 +64,7 @@ def main():
         "name": "exec_agent",
         "type": "atomic",
         "agent_class": "FunctionCallingAgent",
-        "meta_data":
-            {
-                "meta_task_id": meta_task_id,
-                "task_id": task_id,
-                "save_to": "csv",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        "meta_data": {},
         "tool_names": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Execute some script in a subprocess, either run a bash script, or run a python script ",
         "parameter_description": "The task description describing what to execute in the subprocess."
@@ -112,15 +74,7 @@ def main():
         "name": "retrieval_agent",
         "type": "composite",
         "agent_class": "FunctionCallingAgent",
-        "meta_data":
-            {
-                "meta_task_id": meta_task_id,
-                "task_id": task_id,
-                "save_to": "csv",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        "meta_data": {},
         "tool_names": [],
         "sub_agents": [
             web_retrieval_agent_config,
@@ -132,18 +86,10 @@ def main():
     }
 
     main_agent_config = {
-        "name": "retrieval_agent",
+        "name": "main_agent",
         "type": "composite",
         "agent_class": "FunctionCallingAgent",
-        "meta_data":
-            {
-                "meta_task_id": meta_task_id,
-                "task_id": task_id,
-                "save_to": "csv",
-                "log": "log",
-                "model_name": "gpt-4o-mini",
-                "tool_choice": "auto"
-            },
+        "meta_data": {},
         "tool_names": ["scan_folder"],
         "sub_agents": [
             retrieval_agent_config,
@@ -153,20 +99,27 @@ def main():
         "agent_description": None,
         "parameter_description": None
     }
-    # Create the master agent
-    main_agent = agent_manager.get_agent(main_agent_config)
 
-    # # # Example usage
+    system_config = {
+        "meta_task_id": "master_agent_task",
+        "save_to": "csv",
+        "log_dir": "log",
+        "model_name": "gpt-4o-mini",
+        "tool_choice": "auto"
+    }
+    agent_system = AgentSystem(main_agent_config, system_config)
+
+    # # Example usage
     task = "generate a image of a ginger cat and save it as ginger_cat.png"
-    result = main_agent.execute(task)
+    result = agent_system.execute(task)
     print("IO Agent Result:", result)
 
     task = "write python script to calculate the sum from 1 to 10, and run the python script to get result"
-    result = main_agent.execute(task)
-    print("IO Agent Result:", result)
+    result = agent_system.execute(task)
+    print("Exec Agent Result:", result)
 
     task = "browse web to search and check the brands of dining table, and summarize the results in a table, save the table into a markdown file called summary.md"
-    result = main_agent.execute(task)
+    result = agent_system.execute(task)
     print("Agent Result:", result)
 
 

litemultiagent/agents/agent_class/base.py

@@ -3,13 +3,12 @@
 from openai import OpenAI
 import json
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from supabase import create_client, Client
 from litellm import completion
-import os
 from dotenv import load_dotenv
 from litellm import completion
 from datetime import datetime
-import csv
+
+from litemultiagent.core.agent_system import AgentSystem
 
 _ = load_dotenv()
 
@@ -40,17 +39,6 @@
     },
 }
 
-# Initialize Supabase client
-url = os.getenv("SUPABASE_URL")
-key = os.getenv("SUPABASE_ANON_KEY")
-supabase: Optional[Client] = None
-if url and key:
-    try:
-        supabase = create_client(url, key)
-    except Exception as e:
-        logger.error(f"Failed to initialize Supabase client: {e}")
-
-
 class BaseAgent:
     def __init__(self, agent_name: str, agent_description, parameter_description, tools: List[Dict[str, Any]],
                  available_tools: Dict[str, callable],
@@ -59,12 +47,8 @@ def __init__(self, agent_name: str, agent_description, parameter_description, to
         self.tools = tools
         self.available_tools = available_tools
         self.messages = []
-        self.model_name = meta_data["model_name"]
-        self.tool_choice = meta_data["tool_choice"]
-        self.meta_task_id = meta_data["meta_task_id"]
-        self.task_id = meta_data["task_id"]
-        self.save_to = meta_data.get("save_to", "csv")
-        self.log = meta_data["log"]
+        self.model_name = meta_data.get("model_name", None)
+        self.tool_choice = meta_data.get("tool_choice", None)
         self.agent_description = agent_description
         self.parameter_description = parameter_description
         self.goal = None
@@ -98,6 +82,11 @@ def send_prompt(self, goal: str) -> str:
         self.goal = goal
         return self._send_completion_request(plan=goal, depth=0)
 
+    def set_system(self, system: AgentSystem):
+        self.system = system
+        self.model_name = self.model_name or self.system.model_name
+        self.tool_choice = self.tool_choice or self.system.tool_choice
+
     def _send_completion_request(self, plan, depth: int = 0) -> str:
         pass
 
@@ -139,11 +128,17 @@ def _log_response(self, response, depth):
             f'Agent: {self.agent_name}, prompt tokens: {response.usage.prompt_tokens}, completion tokens: {response.usage.completion_tokens}')
         logger.info(f'Agent: {self.agent_name}, depth: {depth}, response: {response}')
 
+    def _save_response(self, response, depth):
+        if self.system.save_to == "supabase":
+            self._save_to_supabase(response, depth)
+        if self.system.save_to == "csv":
+            self._save_to_csv(response, depth)
+
     def _save_to_csv(self, response, depth):
         usage_dict = self._extract_cost(response)
         data = {
-            "meta_task_id": self.meta_task_id,
-            "task_id": self.task_id,
+            "meta_task_id": self.system.meta_task_id,
+            "task_id": self.system.task_id,
             "agent": self.agent_name,
             "depth": depth,
             "role": "assistant",
@@ -156,38 +151,13 @@ def _save_to_csv(self, response, depth):
             "model_name": self.model_name,
             "timestamp": datetime.now().isoformat()
         }
-
-        filename = os.path.join(self.log, f"multiagent_data_{datetime.now().strftime('%Y%m%d')}.csv")
-        file_exists = os.path.isfile(filename)
-
-        # Ensure the directory exists
-        os.makedirs(os.path.dirname(filename), exist_ok=True)
-
-        # If file doesn't exist, create it with header
-        if not file_exists:
-            with open(filename, 'w', newline='') as csvfile:
-                fieldnames = list(data.keys())
-                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-                writer.writeheader()
-            logger.info(f"Created new CSV file with header: {filename}")
-
-        # Append data to the file
-        with open(filename, 'a', newline='') as csvfile:
-            fieldnames = list(data.keys())
-            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-            writer.writerow(data)
-
-        logger.info(f"Data saved to CSV: {filename}")
+        self.system.save_to_csv(data)
 
     def _save_to_supabase(self, response, depth):
-        if supabase is None:
-            logger.warning("Supabase client is not initialized. Skipping database save.")
-            return
-
         usage_dict = self._extract_cost(response)
         data = {
-            "meta_task_id": self.meta_task_id,
-            "task_id": self.task_id,
+            "meta_task_id": self.system.meta_task_id,
+            "task_id": self.system.task_id,
             "agent": self.agent_name,
             "depth": depth,
             "role": "assistant",
@@ -199,11 +169,7 @@ def _save_to_supabase(self, response, depth):
             "total_cost": usage_dict["total_cost"],
             "model_name": self.model_name,
         }
-
-        try:
-            supabase.table("multiagent").insert(data).execute()
-        except Exception as e:
-            logger.error(f"Failed to save data to Supabase: {e}")
+        self.system.save_to_csv(data)
 
     def _extract_cost(self, response):
         prompt_tokens = response.usage.prompt_tokens