Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor agent system #84

Merged
merged 1 commit into from
Oct 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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