Skip to content

Commit

Permalink
refactor agent system (#84)
Browse files Browse the repository at this point in the history
  • Loading branch information
TataKKKL authored Oct 30, 2024
1 parent 3785a85 commit c29b871
Show file tree
Hide file tree
Showing 9 changed files with 119 additions and 101 deletions.
8 changes: 4 additions & 4 deletions examples/atomic/IOAgent.py
Original file line number Diff line number Diff line change
@@ -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(
Expand All @@ -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",
Expand Down
4 changes: 3 additions & 1 deletion examples/composite/MasterAgent.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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",
Expand Down
29 changes: 15 additions & 14 deletions examples/new_tool/add_function_example.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from litemultiagent.core.agent_manager import AgentManager
from litemultiagent.core.agent_system import AgentSystem
import uuid

import logging

Expand All @@ -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()
Expand Down Expand Up @@ -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)


Expand Down
32 changes: 16 additions & 16 deletions examples/new_tool/add_llm_generated_function_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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 = [
Expand All @@ -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}")

Expand Down
67 changes: 53 additions & 14 deletions litemultiagent/agents/agent_class/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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
Expand Down
5 changes: 2 additions & 3 deletions litemultiagent/agents/agent_type/atomic.py
Original file line number Diff line number Diff line change
@@ -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


Expand All @@ -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)
Expand Down
7 changes: 3 additions & 4 deletions litemultiagent/agents/agent_type/composite.py
Original file line number Diff line number Diff line change
@@ -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:
Expand All @@ -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
Expand Down
Loading

0 comments on commit c29b871

Please sign in to comment.