allow adding self defined tool (#76)

TataKKKL · web-flow · commit 6257709a58c5 · 2024-10-28T18:23:40.000-07:00
diff --git a/examples/atomic/DBRetrievalAgent.py b/examples/atomic/DBRetrievalAgent.py
@@ -28,7 +28,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["retrieve_db"], # Changed from "write_file" to "write_to_file"
+        "tool_names": ["retrieve_db"], # Changed from "write_file" to "write_to_file"
         "agent_description": "Use a database retrieval agent to fetch information based on a given query.",
         "parameter_description": "The query to be processed by the database retrieval agent."
     }
diff --git a/examples/atomic/ExecAgent.py b/examples/atomic/ExecAgent.py
@@ -28,7 +28,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Execute some script in a subprocess, either run a bash script, or run a python script ",
         "parameter_description": "The task description describing what to execute in the subprocess."
     }
diff --git a/examples/atomic/FileRetrievalAgent.py b/examples/atomic/FileRetrievalAgent.py
@@ -28,7 +28,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["retrieve_file"],
+        "tool_names": ["retrieve_file"],
         "agent_description": "Retrieve information from local documents to answer questions or perform tasks.",
         "parameter_description": "The task description specifying the local file and the question to be answered. specify this in natural language"
     }
diff --git a/examples/atomic/IOAgent.py b/examples/atomic/IOAgent.py
@@ -29,7 +29,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["read_file", "write_to_file", "generate_and_download_image"],  
+        "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],
         "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."
     }
diff --git a/examples/atomic/WebAgent.py b/examples/atomic/WebAgent.py
@@ -27,7 +27,7 @@ def main():
             "model_name": "gpt-4o-mini",
             "tool_choice": "auto"
         },
-        "tools": ["call_webagent"],
+        "tool_names": ["call_webagent"],
         "agent_description": "Use a web agent to perform tasks and fetch information from web pages based on a given instruction.",
         "parameter_description": "A natural language instruction describing the task to be performed by the web agent, including the starting URL and the goal."
     }
diff --git a/examples/atomic/WebRetrievalAgent.py b/examples/atomic/WebRetrievalAgent.py
@@ -28,7 +28,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Perform a search using API and return the searched results.",
         "parameter_description": "The task description describing what to read or write."
     }
diff --git a/examples/composite/MasterAgent.py b/examples/composite/MasterAgent.py
@@ -31,7 +31,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
         "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."
     }
@@ -49,7 +49,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["retrieve_db"], # Changed from "write_file" to "write_to_file"
+        "tool_names": ["retrieve_db"], # Changed from "write_file" to "write_to_file"
         "agent_description": "Use a database retrieval agent to fetch information based on a given query.",
         "parameter_description": "The query to be processed by the database retrieval agent."
     }
@@ -67,7 +67,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["retrieve_file"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["retrieve_file"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Retrieve information from local documents to answer questions or perform tasks.",
         "parameter_description": "The task description specifying the local file and the question to be answered. specify this in natural language"
     }
@@ -85,7 +85,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Perform a search using API and return the searched results.",
         "parameter_description": "The task description describing what to read or write."
     }
@@ -103,7 +103,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Execute some script in a subprocess, either run a bash script, or run a python script ",
         "parameter_description": "The task description describing what to execute in the subprocess."
     }
@@ -121,7 +121,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": [],
+        "tool_names": [],
         "sub_agents": [
             web_retrieval_agent_config,
             file_retrieval_agent_config,
@@ -144,7 +144,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["scan_folder"],
+        "tool_names": ["scan_folder"],
         "sub_agents": [
             retrieval_agent_config,
             exec_agent_config,
diff --git a/examples/composite/RetrievalAgent.py b/examples/composite/RetrievalAgent.py
@@ -31,7 +31,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["retrieve_db"], # Changed from "write_file" to "write_to_file"
+        "tool_names": ["retrieve_db"], # Changed from "write_file" to "write_to_file"
         "agent_description": "Use a database retrieval agent to fetch information based on a given query.",
         "parameter_description": "The query to be processed by the database retrieval agent."
     }
@@ -49,7 +49,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["retrieve_file"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["retrieve_file"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Retrieve information from local documents to answer questions or perform tasks.",
         "parameter_description": "The task description specifying the local file and the question to be answered. specify this in natural language"
     }
@@ -67,7 +67,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Perform a search using API and return the searched results.",
         "parameter_description": "The task description describing what to read or write."
     }
@@ -85,7 +85,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": [],
+        "tool_names": [],
         "sub_agents": [
             web_retrieval_agent_config,
             file_retrieval_agent_config,
diff --git a/examples/composite/coding/CodingAgent.py b/examples/composite/coding/CodingAgent.py
@@ -32,7 +32,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
         "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."
     }
@@ -51,7 +51,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Execute some script in a subprocess, either run a bash script, or run a python script ",
         "parameter_description": "The task description describing what to execute in the subprocess."
     }
@@ -70,7 +70,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["scan_folder"],
+        "tool_names": ["scan_folder"],
         "sub_agents": [
             exec_agent_config,
             io_agent_config,
diff --git a/examples/composite/research/ResearchAgent.py b/examples/composite/research/ResearchAgent.py
@@ -31,7 +31,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
         "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."
     }
@@ -49,7 +49,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["bing_search", "scrape"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Perform a search using API and return the searched results.",
         "parameter_description": "The task description describing what to read or write."
     }
@@ -67,7 +67,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["execute_shell_command", "run_python_script"],  # Changed from "write_file" to "write_to_file"
         "agent_description": "Execute some script in a subprocess, either run a bash script, or run a python script ",
         "parameter_description": "The task description describing what to execute in the subprocess."
     }
@@ -86,7 +86,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["scan_folder"],
+        "tool_names": ["scan_folder"],
         "sub_agents": [
             web_retrieval_agent_config,
             exec_agent_config,
diff --git a/examples/composite/research/WebBrowsingResearchAgent.py b/examples/composite/research/WebBrowsingResearchAgent.py
@@ -31,7 +31,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
+        "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],  # Changed from "write_file" to "write_to_file"
         "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."
     }
@@ -48,7 +48,7 @@ def main():
             "model_name": "gpt-4o-mini",
             "tool_choice": "auto"
         },
-        "tools": ["call_webagent"],
+        "tool_names": ["call_webagent"],
         "agent_description": "Use a web agent to perform tasks and fetch information from web pages based on a given instruction.",
         "parameter_description": "A natural language instruction describing the task to be performed by the web agent, including the starting URL and the goal."
     }
@@ -66,7 +66,7 @@ def main():
                 "model_name": "gpt-4o-mini",
                 "tool_choice": "auto"
             },
-        "tools": ["scan_folder"],
+        "tool_names": ["scan_folder"],
         "sub_agents": [
             web_agent_config,
             io_agent_config,
diff --git a/examples/new_tool/add_function_example.py b/examples/new_tool/add_function_example.py
@@ -0,0 +1,79 @@
+from litemultiagent.core.agent_manager import AgentManager
+
+import logging
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("log.txt", mode="w"),
+        logging.StreamHandler()
+    ]
+)
+
+# 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()
+
+    def calculate(operation, num1, num2):
+        if operation == 'add':
+            return num1 + num2
+        elif operation == 'subtract':
+            return num1 - num2
+        elif operation == 'multiply':
+            return num1 * num2
+        elif operation == 'divide':
+            if num2 != 0:
+                return num1 / num2
+            else:
+                raise ValueError("Cannot divide by zero")
+        else:
+            raise ValueError("Invalid operation. Choices are 'add', 'subtract', 'multiply', 'divide'")
+
+    name = 'calculate'
+    func = calculate
+    description = 'This is a function that takes two numbers and an operation (add, subtract, multiply, divide) as inputs and returns the result of the operation.'
+    parameters = {'operation': {'type': 'string',
+                                'description': 'Operation to be performed. Allowed operations are add, subtract, multiply and divide.'},
+                  'num1': {'type': 'number', 'description': 'First number.'},
+                  'num2': {'type': 'number', 'description': 'Second number.'}}
+
+    new_tool = Tool(name, func, description, parameters)
+    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"
+            },
+        "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)
+
+    # Example usage
+    task = "calculate 3+4"
+    result = test_agent.execute(task)
+    print("Test Agent Result:", result)
+
+    task = "calculate 3 times 4"
+    result = test_agent.execute(task)
+    print("Test Agent Result:", result)
+
+
+
+if __name__ == "__main__":
+    main()
diff --git a/litemultiagent/agents/agent_type/atomic.py b/litemultiagent/agents/agent_type/atomic.py
@@ -5,11 +5,14 @@
 
 class AtomicAgent:
     def __init__(self, agent_name: str, agent_description: str, parameter_description: str,
-                 tool_names: List[str], meta_data: Dict[str, Any],
+                 tool_names: List[str], self_defined_tools, meta_data: Dict[str, Any],
                  agent_class: Type[BaseAgent]):
         tool_registry = ToolRegistry()
         available_tools = {}
         tools = []
+        for self_defined_tool in self_defined_tools:
+            tool_registry.register(self_defined_tool)
+            tool_names.append(self_defined_tool.name)
         for tool_name in tool_names:
             available_tools[tool_name] = tool_registry.get_tool(tool_name).func
             tools.append(tool_registry.get_tool_description(tool_name))
diff --git a/litemultiagent/agents/agent_type/composite.py b/litemultiagent/agents/agent_type/composite.py
@@ -4,11 +4,15 @@
 
 class CompositeAgent:
     def __init__(self, agent_name: str, agent_description: str, parameter_description: str,
-                 sub_agent_configs: List[Dict[str, Any]], tool_names: List[str], meta_data: Dict[str, Any],
+                 sub_agent_configs: List[Dict[str, Any]], tool_names: List[str], self_defined_tools, meta_data: Dict[str, Any],
                  agent_class: Type[BaseAgent]):
         self.available_tools = {}
         self.tools = []
 
+        for self_defined_tool in self_defined_tools:
+            tool_registry.register(self_defined_tool)
+            tool_names.append(self_defined_tool.name)
+
         for tool_name in tool_names:
             self.available_tools[tool_name] = ToolRegistry.get_tool(tool_name).func
             self.tools.append(ToolRegistry.get_tool_description(tool_name))
diff --git a/litemultiagent/core/agent_factory.py b/litemultiagent/core/agent_factory.py
@@ -12,7 +12,8 @@ def create_agent(config: Dict[str, Any]) -> Union[AtomicAgent, CompositeAgent]:
         agent_class = config["agent_class"]
         agent_description = config["agent_description"]
         parameter_description = config["parameter_description"]
-        tools = config.get("tools", [])
+        tool_names = config.get("tool_names", [])
+        self_defined_tools = config.get("self_defined_tools", [])
         meta_data = config.get("meta_data", {})
 
         agent_class_map = {
@@ -28,10 +29,10 @@ def create_agent(config: Dict[str, Any]) -> Union[AtomicAgent, CompositeAgent]:
 
         if agent_type == "atomic":
             return AtomicAgent(agent_name, agent_description, parameter_description,
-                               tools, meta_data, selected_agent_class)
+                               tool_names, self_defined_tools, meta_data, selected_agent_class)
         elif agent_type == "composite":
             sub_agent_configs = config.get("sub_agents", [])
             return CompositeAgent(agent_name, agent_description, parameter_description,
-                                  sub_agent_configs, tools, meta_data, selected_agent_class)
+                                  sub_agent_configs, tool_names, self_defined_tools, meta_data, selected_agent_class)
         else:
             raise ValueError(f"Unknown agent type: {agent_type}")

Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@ def main():`
`28`	`28`	`"model_name": "gpt-4o-mini",`
`29`	`29`	`"tool_choice": "auto"`
`30`	`30`	`},`
`31`		`- "tools": ["retrieve_db"], # Changed from "write_file" to "write_to_file"`
	`31`	`+ "tool_names": ["retrieve_db"], # Changed from "write_file" to "write_to_file"`
`32`	`32`	`"agent_description": "Use a database retrieval agent to fetch information based on a given query.",`
`33`	`33`	`"parameter_description": "The query to be processed by the database retrieval agent."`
`34`	`34`	`}`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ def main():`
`29`	`29`	`"model_name": "gpt-4o-mini",`
`30`	`30`	`"tool_choice": "auto"`
`31`	`31`	`},`
`32`		`- "tools": ["read_file", "write_to_file", "generate_and_download_image"],`
	`32`	`+ "tool_names": ["read_file", "write_to_file", "generate_and_download_image"],`
`33`	`33`	`"agent_description": "Read or write content from/to a file, or generate and save an image using text input",`
`34`	`34`	`"parameter_description": "The task description detailing what to read, write, or generate. This can include file operations or image generation requests."`
`35`	`35`	`}`
Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@ def main():`
`27`	`27`	`"model_name": "gpt-4o-mini",`
`28`	`28`	`"tool_choice": "auto"`
`29`	`29`	`},`
`30`		`- "tools": ["call_webagent"],`
	`30`	`+ "tool_names": ["call_webagent"],`
`31`	`31`	`"agent_description": "Use a web agent to perform tasks and fetch information from web pages based on a given instruction.",`
`32`	`32`	`"parameter_description": "A natural language instruction describing the task to be performed by the web agent, including the starting URL and the goal."`
`33`	`33`	`}`