refactor agent system (#84)

Author: TataKKKL

examples/atomic/IOAgent.py

@@ -1,8 +1,6 @@
-from litemultiagent.core.agent_manager import AgentManager
-
 import logging
-
 from litemultiagent.core.agent_system import AgentSystem
+import uuid
 
 # Configure logging
 logging.basicConfig(
@@ -26,8 +24,10 @@ def main():
         "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."
     }
+
     system_config = {
-        "meta_task_id": "io_subtask",
+        "system_name": "io_agent_system",
+        "system_runtime_id": str(uuid.uuid4()),
         "save_to": "csv",
         "log_dir": "log",
         "model_name": "gpt-4o-mini",

examples/composite/MasterAgent.py

@@ -1,6 +1,7 @@
 from litemultiagent.core.agent_manager import AgentManager
 from litemultiagent.core.agent_system import AgentSystem
 from litemultiagent.tools.registry import ToolRegistry, Tool
+import uuid
 import logging
 
 # Configure logging
@@ -101,7 +102,8 @@ def main():
     }
 
     system_config = {
-        "meta_task_id": "master_agent_task",
+        "system_name": "master_agent_system",
+        "system_runtime_id": str(uuid.uuid4()),
         "save_to": "csv",
         "log_dir": "log",
         "model_name": "gpt-4o-mini",

examples/new_tool/add_function_example.py

@@ -1,4 +1,5 @@
-from litemultiagent.core.agent_manager import AgentManager
+from litemultiagent.core.agent_system import AgentSystem
+import uuid
 
 import logging
 
@@ -15,7 +16,6 @@
 # Create a logger
 logger = logging.getLogger(__name__)
 def main():
-    agent_manager = AgentManager()
     from litemultiagent.tools.registry import Tool
     from dotenv import load_dotenv
     _ = load_dotenv()
@@ -48,29 +48,30 @@ def calculate(operation, num1, num2):
         "name": "test_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"],
         "self_defined_tools": [new_tool],
         "agent_description": "test ai agent",
         "parameter_description": "test ai agent"
     }
-    test_agent = agent_manager.get_agent(test_agent_config)
+
+    system_config = {
+        "system_name": "test_agent_system",
+        "system_runtime_id": str(uuid.uuid4()),
+        "save_to": "csv",
+        "log_dir": "log",
+        "model_name": "gpt-4o-mini",
+        "tool_choice": "auto"
+    }
+    agent_system = AgentSystem(test_agent_config, system_config)
 
     # Example usage
     task = "calculate 3+4"
-    result = test_agent.execute(task)
+    result = agent_system.execute(task)
     print("Test Agent Result:", result)
 
     task = "calculate 3 times 4"
-    result = test_agent.execute(task)
+    result = agent_system.execute(task)
     print("Test Agent Result:", result)
 
 

examples/new_tool/add_llm_generated_function_example.py

@@ -3,6 +3,9 @@
 import traceback
 from dotenv import load_dotenv
 from litemultiagent.tools.tool_creation_agent import ToolCreationAgent
+from litemultiagent.core.agent_system import AgentSystem
+from litemultiagent.tools.registry import Tool
+import uuid
 
 # Configure logging
 logging.basicConfig(
@@ -48,38 +51,35 @@ def main():
     # Cleanup resources
     tool_creation_agent.cleanup_cache()
 
-    # Set up agent manager and tools
-    from litemultiagent.core.agent_manager import AgentManager
-    from litemultiagent.tools.registry import Tool
 
-    agent_manager = AgentManager()
     new_tools = [
         Tool(name, func, description, parameters)
         for spec in mapping
         for name, func, description, parameters in [mapping[spec]]
     ]
 
-    # Configure test agent
+
     test_agent_config = {
         "name": "test_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"],
         "self_defined_tools": new_tools,
         "agent_description": "test ai agent",
         "parameter_description": "test ai agent"
     }
 
-    # Initialize and test the agent
-    test_agent = agent_manager.get_agent(test_agent_config)
+    system_config = {
+        "system_name": "test_agent_system",
+        "system_runtime_id": str(uuid.uuid4()),
+        "save_to": "csv",
+        "log_dir": "log",
+        "model_name": "gpt-4o-mini",
+        "tool_choice": "auto"
+    }
+    agent_system = AgentSystem(test_agent_config, system_config)
+
 
     # Test calculator functionality
     test_cases = [
@@ -88,7 +88,7 @@ def main():
     ]
 
     for task, description in test_cases:
-        result = test_agent.execute(task)
+        result = agent_system.execute(task)
         print(f"{description} - Task: {task}")
         print(f"Result: {result}")
 

litemultiagent/agents/agent_class/base.py

@@ -7,8 +7,9 @@
 from dotenv import load_dotenv
 from litellm import completion
 from datetime import datetime
-
-from litemultiagent.core.agent_system import AgentSystem
+from supabase import create_client, Client
+import os
+import csv
 
 _ = load_dotenv()
 
@@ -39,6 +40,16 @@
     },
 }
 
+# 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],
@@ -82,10 +93,10 @@ 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 set_shared_config(self, shared_config):
+        self.shared_config = shared_config
+        self.model_name = self.model_name or self.shared_config["model_name"]
+        self.tool_choice = self.tool_choice or self.shared_config["tool_choice"]
 
     def _send_completion_request(self, plan, depth: int = 0) -> str:
         pass
@@ -129,16 +140,17 @@ def _log_response(self, response, depth):
         logger.info(f'Agent: {self.agent_name}, depth: {depth}, response: {response}')
 
     def _save_response(self, response, depth):
-        if self.system.save_to == "supabase":
+        if self.shared_config["save_to"] == "supabase":
             self._save_to_supabase(response, depth)
-        if self.system.save_to == "csv":
+        if self.shared_config["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.system.meta_task_id,
-            "task_id": self.system.task_id,
+            "system_name": self.shared_config["system_name"],
+            "system_runtime_id": self.shared_config["system_runtime_id"],
+            "task_id": self.shared_config["task_id"],
             "agent": self.agent_name,
             "depth": depth,
             "role": "assistant",
@@ -151,13 +163,34 @@ def _save_to_csv(self, response, depth):
             "model_name": self.model_name,
             "timestamp": datetime.now().isoformat()
         }
-        self.system.save_to_csv(data)
+        filename = os.path.join(self.shared_config["log_dir"], 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}")
 
     def _save_to_supabase(self, response, depth):
         usage_dict = self._extract_cost(response)
         data = {
-            "meta_task_id": self.system.meta_task_id,
-            "task_id": self.system.task_id,
+            "system_name": self.shared_config["system_name"],
+            "system_runtime_id": self.shared_config["system_runtime_id"],
+            "task_id": self.shared_config["task_id"],
             "agent": self.agent_name,
             "depth": depth,
             "role": "assistant",
@@ -169,7 +202,13 @@ def _save_to_supabase(self, response, depth):
             "total_cost": usage_dict["total_cost"],
             "model_name": self.model_name,
         }
-        self.system.save_to_csv(data)
+        if supabase is None:
+            logger.warning("Supabase client is not initialized. Skipping database save.")
+            return
+        try:
+            supabase.table("multiagent").insert(data).execute()
+        except Exception as e:
+            logger.error(f"Failed to save data to Supabase: {e}")
 
     def _extract_cost(self, response):
         prompt_tokens = response.usage.prompt_tokens

litemultiagent/agents/agent_type/atomic.py

@@ -1,6 +1,5 @@
 from typing import List, Dict, Any, Type
 from litemultiagent.agents.agent_class.base import BaseAgent
-from litemultiagent.core.agent_system import AgentSystem
 from litemultiagent.tools.registry import ToolRegistry
 
 
@@ -24,8 +23,8 @@ def __init__(self, agent_name: str, agent_description: str, parameter_descriptio
     def execute(self, task: str) -> str:
         return self.agent.send_prompt(task)
 
-    def set_system(self, system: AgentSystem):
-        self.agent.set_system(system)
+    def set_shared_config(self, shared_config):
+        self.agent.set_shared_config(shared_config)
 
     def __getattr__(self, name):
         return getattr(self.agent, name)

litemultiagent/agents/agent_type/composite.py

@@ -1,6 +1,5 @@
 from typing import List, Dict, Any, Type
 from litemultiagent.agents.agent_class.base import BaseAgent
-from litemultiagent.core.agent_system import AgentSystem
 from litemultiagent.tools.registry import ToolRegistry, Tool
 
 class CompositeAgent:
@@ -24,10 +23,10 @@ def __init__(self, agent_name: str, agent_description: str, parameter_descriptio
         self.agent = agent_class(agent_name, agent_description, parameter_description,
                                  self.tools, self.available_tools, meta_data)
 
-    def set_system(self, system: AgentSystem):
-        self.agent.set_system(system)
+    def set_shared_config(self, shared_config):
+        self.agent.set_shared_config(shared_config)
         for sub_agent in self.sub_agents:
-            sub_agent.set_system(system)
+            sub_agent.set_shared_config(shared_config)
 
     def _build_sub_agents(self, sub_agent_configs: List[Dict[str, Any]]) -> List[BaseAgent]:
         from litemultiagent.core.agent_factory import AgentFactory