Updated & Revamped Tool Registry

.gitignore

@@ -29,3 +29,7 @@ process_message.py
 # database
 ./code/files/chroma/*
 
+1.txt
+2.txt
+3.txt
+ginger_cat.png

examples/atomic/DBRetrievalAgent.py

@@ -1,5 +1,4 @@
 from litemultiagent.core.agent_manager import AgentManager
-from litemultiagent.tools.registry import ToolRegistry, Tool
 import logging
 
 # Configure logging

examples/atomic/ExecAgent.py

@@ -1,5 +1,4 @@
 from litemultiagent.core.agent_manager import AgentManager
-from litemultiagent.tools.registry import ToolRegistry, Tool
 import logging
 
 # Configure logging

examples/atomic/FileRetrievalAgent.py

@@ -1,5 +1,4 @@
 from litemultiagent.core.agent_manager import AgentManager
-from litemultiagent.tools.registry import ToolRegistry, Tool
 import logging
 
 # Configure logging

examples/atomic/IOAgent.py

@@ -1,5 +1,5 @@
 from litemultiagent.core.agent_manager import AgentManager
-from litemultiagent.tools.registry import ToolRegistry, Tool
+
 import logging
 
 # Configure logging

examples/atomic/WebAgent.py

@@ -1,5 +1,4 @@
 from litemultiagent.core.agent_manager import AgentManager
-from litemultiagent.tools.registry import ToolRegistry, Tool
 import logging
 
 # Configure logging

examples/atomic/WebRetrievalAgent.py

@@ -1,5 +1,4 @@
 from litemultiagent.core.agent_manager import AgentManager
-from litemultiagent.tools.registry import ToolRegistry, Tool
 import logging
 
 # Configure logging

litemultiagent/agents/atomic.py

@@ -1,9 +1,7 @@
 from litemultiagent.agents.base import BaseAgent
-from typing import List, Optional
-from litemultiagent.tools.registry import ToolRegistry, Tool
+from litemultiagent.tools.registry import ToolRegistry
 class AtomicAgent(BaseAgent):
-    def __init__(self, agent_name: str, agent_description, parameter_description, tool_names: List[str], meta_data):
-        print(tool_names)
+    def __init__(self, agent_name: str, agent_description, parameter_description, tool_names: list[str], meta_data):
         tool_registry = ToolRegistry()
         available_tools = {}
         tools = []

litemultiagent/agents/base.py

@@ -16,7 +16,7 @@
 logger = logging.getLogger(__name__)
 
 
-model_cost = {
+MODEL_COST = {
     "gpt-4o-mini": {
         "input_price_per_1m": 0.15,
         "output_price_per_1m": 0.6,
@@ -274,8 +274,8 @@ def _extract_cost(self, response):
         total_tokens = response.usage.total_tokens
 
         # Pricing for gpt-4-0314
-        input_price_per_1m = model_cost[self.model_name]["input_price_per_1m"]
-        output_price_per_1m = model_cost[self.model_name]["output_price_per_1m"]
+        input_price_per_1m = MODEL_COST[self.model_name]["input_price_per_1m"]
+        output_price_per_1m = MODEL_COST[self.model_name]["output_price_per_1m"]
 
         input_cost = (prompt_tokens / 1000) * input_price_per_1m
         output_cost = (completion_tokens / 1000) * output_price_per_1m

litemultiagent/agents/composite.py

@@ -4,14 +4,18 @@
 
 class CompositeAgent(BaseAgent):
     def __init__(self, agent_name: str, agent_description, parameter_description, sub_agent_configs: List[Dict[str, Any]], tool_names: List[str], meta_data):
-        self.tool_registry = ToolRegistry()
         self.available_tools = {}
+
         self.tools = []
+
         for tool_name in tool_names:
-            self.available_tools[tool_name] = self.tool_registry.get_tool(tool_name).func
-            self.tools.append(self.tool_registry.get_tool_description(tool_name))
+            self.available_tools[tool_name] = ToolRegistry.get_tool(tool_name).func
+            self.tools.append(ToolRegistry.get_tool_description(tool_name))
+
         self.sub_agents = self._build_sub_agents(sub_agent_configs)
+
         self._register_sub_agents_as_tools()
+
         super().__init__(agent_name, agent_description, parameter_description, self.tools, self.available_tools, meta_data)
 
 
@@ -34,7 +38,7 @@ def _register_sub_agents_as_tools(self):
                 }
             ))
         # Update the tools and available_tools after registering sub-agents
-        self.tools.extend([self.tool_registry.get_tool_description(sub_agent.agent_name) for sub_agent in self.sub_agents])
+        self.tools.extend([ToolRegistry.get_tool_description(sub_agent.agent_name) for sub_agent in self.sub_agents])
         self.available_tools.update({sub_agent.agent_name: sub_agent for sub_agent in self.sub_agents})
 
     def execute(self, task: str) -> str:

litemultiagent/tools/db_retrieval.py

@@ -0,0 +1,61 @@
+import os
+
+from supabase import create_client, Client
+
+from litemultiagent.tools.registry import Tool
+
+def retrieve_db(client, db, input_column, output_column, input_value):
+    if client == "SUPABASE":
+        url: str = os.getenv("SUPABASE_URL")
+        key: str = os.getenv("SUPABASE_ANON_KEY")
+
+        if not url or not key:
+            return "Error: SUPABASE_URL or SUPABASE_ANON_KEY environment variables are not set."
+
+        try:
+            supabase: Client = create_client(url, key)
+            data = supabase.table(db).select(output_column).eq(
+                input_column, input_value).execute()
+
+            if data.data:
+                return data.data[0][output_column]
+            else:
+                return f"No data found for {input_value} in {db}"
+
+        except Exception as e:
+            return f"Failed to retrieve data from Supabase: {str(e)}"
+    else:
+        return "not defined clients"
+
+retrieve_db_tool = Tool(
+    "retrieve_db",
+    retrieve_db,
+    "Retrieve data from a specified database (currently supports Supabase) based on input parameters.",
+    {
+        "client": {
+            "type": "string",
+            "description": "The database client to use. Currently supports 'SUPABASE'.",
+            "required": True
+        },
+        "db": {
+            "type": "string",
+            "description": "The name of the database table to query.",
+            "required": True
+        },
+        "input_column": {
+            "type": "string",
+            "description": "The column name to search in.",
+            "required": True
+        },
+        "output_column": {
+            "type": "string",
+            "description": "The column name to retrieve data from.",
+            "required": True
+        },
+        "input_value": {
+            "type": "string",
+            "description": "The value to search for in the input column.",
+            "required": True
+        }
+    }
+)

litemultiagent/tools/exec.py

@@ -0,0 +1,60 @@
+import subprocess
+
+from litemultiagent.tools.registry import Tool
+
+
+def execute_shell_command(command, wait=True):
+    try:
+        if wait:
+            result = subprocess.run(
+                command, shell=True, capture_output=True, text=True, check=True)
+            return result.stdout.strip() if result.stdout else result.stderr.strip()
+        else:
+            subprocess.Popen(command, shell=True)
+            return "Command executed in non-blocking mode."
+    except subprocess.CalledProcessError as e:
+        return f"Error executing command '{command}': {e.stderr.strip()}"
+
+
+def run_python_script(script_name):
+    try:
+        result = subprocess.run(["python", script_name],
+                                capture_output=True, text=True, check=True)
+        res = f"stdout:{result.stdout}"
+        if result.stderr:
+            res += f"stderr:{result.stderr}"
+        return res
+    except subprocess.CalledProcessError as e:
+        return f"Error:{e}"
+
+
+run_python_script_tool = Tool(
+    "run_python_script",
+    run_python_script,
+    "Execute a Python script in a subprocess.",
+    {
+        "script_name": {
+            "type": "string",
+            "description": "The name with path of the script to be executed.",
+            "required": True
+        }
+    }
+)
+
+execute_shell_command_tool = Tool(
+    "execute_shell_command",
+    execute_shell_command,
+    "Execute a shell command in a subprocess.",
+    {
+        "command": {
+            "type": "string",
+            "description": "The shell command to be executed.",
+            "required": True
+        },
+        "wait": {
+            "type": "boolean",
+            "description": "Wait for the command to complete. Set to true for blocking execution and false for non-blocking.",
+            "required": True
+        }
+    }
+)