diff --git a/examples/atomic/IOAgent.py b/examples/atomic/IOAgent.py index 22935f0..1b231a6 100644 --- a/examples/atomic/IOAgent.py +++ b/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", diff --git a/examples/composite/MasterAgent.py b/examples/composite/MasterAgent.py index 367f33d..264a5d8 100644 --- a/examples/composite/MasterAgent.py +++ b/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", diff --git a/examples/new_tool/add_function_example.py b/examples/new_tool/add_function_example.py index 9a025c8..0455e2c 100644 --- a/examples/new_tool/add_function_example.py +++ b/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) diff --git a/examples/new_tool/add_llm_generated_function_example.py b/examples/new_tool/add_llm_generated_function_example.py index 75eb5ae..6270606 100644 --- a/examples/new_tool/add_llm_generated_function_example.py +++ b/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}") diff --git a/litemultiagent/agents/agent_class/base.py b/litemultiagent/agents/agent_class/base.py index ce2feb9..be829b5 100644 --- a/litemultiagent/agents/agent_class/base.py +++ b/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 diff --git a/litemultiagent/agents/agent_type/atomic.py b/litemultiagent/agents/agent_type/atomic.py index fe5aaa3..8a58631 100644 --- a/litemultiagent/agents/agent_type/atomic.py +++ b/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) diff --git a/litemultiagent/agents/agent_type/composite.py b/litemultiagent/agents/agent_type/composite.py index a6075cd..5429b09 100644 --- a/litemultiagent/agents/agent_type/composite.py +++ b/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 diff --git a/litemultiagent/core/agent_system.py b/litemultiagent/core/agent_system.py index df86180..0171715 100644 --- a/litemultiagent/core/agent_system.py +++ b/litemultiagent/core/agent_system.py @@ -5,6 +5,7 @@ from typing import List, Dict, Any, Optional from supabase import create_client, Client from dotenv import load_dotenv +from litemultiagent.core.agent_manager import AgentManager _ = load_dotenv() logger = logging.getLogger(__name__) @@ -21,51 +22,26 @@ class AgentSystem: def __init__(self, main_agent_config, system_config: Dict[str, Any]) -> None: - self.task_id = 0 - self.log_dir = system_config["log_dir"] - self.save_to = system_config["save_to"] - self.meta_task_id = system_config["meta_task_id"] - self.model_name = system_config["model_name"] - self.tool_choice = system_config["tool_choice"] - # Move import here to avoid cirecular import - from litemultiagent.core.agent_manager import AgentManager - agent_manager = AgentManager() - self.main_agent = agent_manager.get_agent(main_agent_config) - self.main_agent.set_system(self) - - def execute(self, task: str): - self.task_id += 1 - self.main_agent.execute(task) - - def save_to_csv(self, data): - filename = os.path.join(self.log_dir, f"multiagent_data_{datetime.now().strftime('%Y%m%d')}.csv") - file_exists = os.path.isfile(filename) + # TODO: add shared_memory + # TODO: add shared_context - # Ensure the directory exists - os.makedirs(os.path.dirname(filename), exist_ok=True) + # create shared_config for agents to share + self.shared_config = { + "task_id": 0, + "system_name": system_config["system_name"], + "log_dir": system_config["log_dir"], + "save_to": system_config["save_to"], + "system_runtime_id": system_config["system_runtime_id"], + "model_name": system_config["model_name"], + "tool_choice": system_config["tool_choice"] + } - # 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}") + agent_manager = AgentManager() + self.main_agent = agent_manager.get_agent(main_agent_config) + self.main_agent.set_shared_config(self.shared_config) - def save_to_supabase(self, 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}") \ No newline at end of file + def execute(self, task: str): + self.shared_config["task_id"] += 1 + return self.main_agent.execute(task) diff --git a/supabase_db_setup.py b/supabase_db_setup.py index 01f6750..b50cf28 100644 --- a/supabase_db_setup.py +++ b/supabase_db_setup.py @@ -47,7 +47,8 @@ class DynamicTable(Base): output_cost = Column(Float) total_cost = Column(Float) model_name = Column(Text) - meta_task_id = Column(Text) + system_name = Column(Text) + system_runtime_id = Column(Text) task_id = Column(BigInteger) # Create the table @@ -69,7 +70,8 @@ class DynamicTable(Base): output_cost=0.3, total_cost=0.8, model_name="Test Model", - meta_task_id="12345", + system_name="ai system", + system_runtime_id="12345", task_id=1 )