update CreateNewTool to execute as well

mdagent/subagents/agents/skill.py

@@ -149,7 +149,7 @@ def update_skill_library(self, function, code_script, description, arguments):
         )
         self.vectordb.persist()
 
-    def execute_skill_function(self, tool_name, path_registry, **kwargs):
+    def execute_skill_function(self, tool_name, **kwargs):
         code = self.skills.get(tool_name, {}).get("code", None)
         if not code:
             raise ValueError(
@@ -158,7 +158,7 @@ def execute_skill_function(self, tool_name, path_registry, **kwargs):
             )
         # capture initial state
         initial_files = set(os.listdir("."))
-        initial_registry = path_registry.list_path_names()
+        initial_registry = self.path_registry.list_path_names()
 
         try:
             self._check_arguments(tool_name, **kwargs)
@@ -172,7 +172,7 @@ def execute_skill_function(self, tool_name, path_registry, **kwargs):
         # capture final state
         new_files = list(set(os.listdir(".")) - initial_files)
         new_registry = list(
-            set(path_registry.list_path_names()) - set(initial_registry)
+            set(self.path_registry.list_path_names()) - set(initial_registry)
         )
         message = "Successfully executed code."
         if new_files:

mdagent/subagents/subagent_fxns.py

@@ -2,22 +2,19 @@
 import os
 from typing import Optional
 
-from mdagent.utils import PathRegistry
-
 from .subagent_setup import SubAgentInitializer, SubAgentSettings
 
 
 class Iterator:
     def __init__(
         self,
-        path_registry: Optional[PathRegistry],
         subagent_settings: Optional[SubAgentSettings],
         all_tools_string: Optional[str] = None,
         current_tools: Optional[dict] = None,
     ):
-        self.path_registry = path_registry
         if subagent_settings is None:
             raise ValueError("Subagent settings cannot be None")  # shouldn't happen
+        self.path_registry = subagent_settings.path_registry
         self.ckpt_dir = subagent_settings.ckpt_dir
         self.all_tools_string = all_tools_string
         self.current_tools = current_tools

mdagent/tools/maketools.py

@@ -11,7 +11,7 @@
 from langchain_experimental.tools import PythonREPLTool
 from pydantic import BaseModel, Field
 
-from mdagent.subagents import Iterator, SubAgentSettings
+from mdagent.subagents import Iterator, SubAgentInitializer, SubAgentSettings
 from mdagent.utils import PathRegistry, _make_llm
 
 from .base_tools import (
@@ -29,7 +29,7 @@
     SimulationOutputFigures,
     VisualizeProtein,
 )
-from .subagent_tools import ExecuteSkill, SkillRetrieval, WorkflowPlan
+from .subagent_tools import RetryExecuteSkill, SkillRetrieval, WorkflowPlan
 
 
 def get_learned_tools(ckpt_dir="ckpt"):
@@ -77,7 +77,6 @@ def make_all_tools(
     base_tools = [
         CleaningToolFunction(path_registry=path_instance),
         CheckDirectoryFiles(),
-        #    InstructionSummary(path_registry=path_instance),
         ListRegistryPaths(path_registry=path_instance),
         #    MapPath2Name(path_registry=path_instance),
         Name2PDBTool(path_registry=path_instance),
@@ -97,7 +96,7 @@ def make_all_tools(
     if not skip_subagents:
         subagents_tools = [
             CreateNewTool(subagent_settings=subagent_settings),
-            ExecuteSkill(subagent_settings=subagent_settings),
+            RetryExecuteSkill(subagent_settings=subagent_settings),
             SkillRetrieval(subagent_settings=subagent_settings),
             WorkflowPlan(subagent_settings=subagent_settings),
         ]
@@ -135,7 +134,7 @@ def get_tools(
         PathRegistry.get_instance()
         retrieved_tools = [
             CreateNewTool(subagent_settings=subagent_settings),
-            ExecuteSkill(subagent_settings=subagent_settings),
+            RetryExecuteSkill(subagent_settings=subagent_settings),
             SkillRetrieval(subagent_settings=subagent_settings),
             WorkflowPlan(subagent_settings=subagent_settings),
         ]
@@ -148,13 +147,13 @@ def get_tools(
             llm, subagent_settings, skip_subagents=False, human=human
         )
 
-    # create vector DB for all tools
+    # set vector DB for all tools
     vectordb = Chroma(
         collection_name="all_tools_vectordb",
         embedding_function=OpenAIEmbeddings(),
         persist_directory=f"{ckpt_dir}/all_tools_vectordb",
     )
-    # vectordb.delete_collection() # to clear vectordb directory
+    # vectordb.delete_collection()      #<--- to clear previous vectordb directory
     for i, tool in enumerate(all_tools):
         vectordb.add_texts(
             texts=[tool.description],
@@ -185,15 +184,23 @@ class CreateNewToolInputSchema(BaseModel):
         description="""List of all tools you have access to. Such as
         this tool, 'ExecuteSkill', 'SkillRetrieval', and maybe `Name2PDBTool`, etc."""
     )
+    execute: Optional[bool] = Field(
+        True,
+        description="Whether to execute the new tool or not.",
+    )
+    args: Optional[dict] = Field(
+        description="Input variables as a dictionary to pass to the skill"
+    )
 
 
-# move this here to avoid circular import error (since it gets a list of all tools)
 class CreateNewTool(BaseTool):
     name: str = "CreateNewTool"
     description: str = """
         Only use if you don't have right tools for sure and need a different tool.
-        If succeeded, it will return the name of the tool.
-        You can then use the tool in subsequent steps.
+        If succeeded, it will return the name of the tool. Unless you set
+        'execute' to False, it will also execute the tool and return the result.
+        Make sure to provide any necessary input arguments for the tool.
+        If execution fails, move on to ReTryExecuteSkill.
     """
     args_schema: Type[BaseModel] = CreateNewToolInputSchema
     subagent_settings: Optional[SubAgentSettings]
@@ -210,30 +217,44 @@ def get_all_tools_string(self):
             all_tools_string += f"{tool.name}: {tool.description}\n"
         return all_tools_string
 
-    def _run(self, task, orig_prompt, curr_tools):
-        # def _run(self, task, orig_prompt):
+    def _run(self, task, orig_prompt, curr_tools, execute, args=None):
+        # run iterator
         try:
-            # run iterator
-            path_registry = self.subagent_settings.path_registry
-            print("getting all tools info")
             all_tools_string = self.get_all_tools_string()
-            print("setting up iterator")
             newcode_iterator = Iterator(
-                path_registry,
                 self.subagent_settings,
                 all_tools_string=all_tools_string,
                 current_tools=curr_tools,
             )
-            print("running iterator")
+            print("running iterator to draft a new tool")
             tool_name = newcode_iterator.run(task, orig_prompt)
-            # tool_name = newcode_iterator.run(task, task)
-            if tool_name:
-                return f"""Tool created successfully: {tool_name}
-                You can now use the tool in subsequent steps."""
-            else:
+            if not tool_name:
                 return "The 'CreateNewTool' tool failed to build a new tool."
         except Exception as e:
-            return f"Something went wrong. {type(e).__name__}: {e}"
+            return f"Something went wrong while creating tool. {type(e).__name__}: {e}"
+
+        # execute the new tool code
+        if execute:
+            try:
+                print("\nexecuting tool")
+                agent_initializer = SubAgentInitializer(self.subagent_settings)
+                skill = agent_initializer.create_skill_manager(resume=True)
+                if skill is None:
+                    return "SubAgent for this tool not initialized"
+                if args is not None:
+                    print("input args: ", args)
+                    return skill.execute_skill_function(tool_name, **args)
+                else:
+                    return skill.execute_skill_function(tool_name)
+            except TypeError as e:
+                return f"""{type(e).__name__}: {e}. Executing new tool failed.
+                    Please check your inputs and make sure to use a dictionary.\n"""
+            except ValueError as e:
+                return f"{type(e).__name__}: {e}. Provide correct arguments for tool.\n"
+            except Exception as e:
+                return f"Something went wrong while executing. {type(e).__name__}:{e}\n"
+        else:
+            return f"A new tool is created: {tool_name}. You can use it in next prompt."
 
     async def _arun(self, query) -> str:
         """Use the tool asynchronously."""

mdagent/tools/subagent_tools.py

@@ -15,11 +15,11 @@ class ExecuteSkillInputSchema(BaseModel):
     )
 
 
-class ExecuteSkill(BaseTool):
-    name = "ExecuteSkill"
-    description = """Executes the code for a new tool or skill that has
-    been recently made during the current iteration. Make sure to include
-    function name and inputs arguments.
+class RetryExecuteSkill(BaseTool):
+    name = "RetryExecuteSkill"
+    description = """Only use this tool to execute a tool recently created
+    that previously failed to execute or hasn't been executed yet.
+    Make sure to include tool name and inputs arguments.
     """
     subagent_settings: Optional[SubAgentSettings]
     args_schema: Optional[Type[BaseModel]] = ExecuteSkillInputSchema
@@ -30,18 +30,15 @@ def __init__(self, subagent_settings: Optional[SubAgentSettings] = None):
 
     def _run(self, skill_name, args=None):
         try:
-            path_registry = self.subagent_settings.path_registry
             agent_initializer = SubAgentInitializer(self.subagent_settings)
             skill = agent_initializer.create_skill_manager(resume=True)
             if skill is None:
                 return "SubAgent for this tool not initialized"
             if args is not None:
                 print("args: ", args)
-                code_result = skill.execute_skill_function(
-                    skill_name, path_registry, **args
-                )
+                code_result = skill.execute_skill_function(skill_name, **args)
             else:
-                code_result = skill.execute_skill_function(skill_name, path_registry)
+                code_result = skill.execute_skill_function(skill_name)
             return code_result
         except TypeError as e:
             return f"""{type(e).__name__}: {e}. Please check your inputs