1. fix bug at getting id and filename from modified simulations. 2. …

…Change stucture in make_tools so modifyscript is with other llm tools. 3. Fix bug at modifyscriputils during its init (added the llm when called)

mdagent/tools/base_tools/simulation_tools/create_simulation.py

@@ -8,12 +8,14 @@
 from langchain.tools import BaseTool
 from pydantic import BaseModel, Field
 
-from mdagent.utils import PathRegistry
+from mdagent.utils import FileType, PathRegistry
 
 
 class ModifyScriptUtils:
+    llm: Optional[BaseLanguageModel]
+
     def __init__(self, llm):
-        llm = llm
+        self.llm = llm
 
     Examples = [
         """
@@ -153,12 +155,16 @@ class ModifyBaseSimulationScriptTool(BaseTool):
     llm: Optional[BaseLanguageModel]
     path_registry: Optional[PathRegistry]
 
-    def __init__(self, path_registry: Optional[PathRegistry], llm: BaseLanguageModel):
+    def __init__(self, path_registry: Optional[PathRegistry], llm):
         super().__init__()
         self.path_registry = path_registry
         self.llm = llm
+        print(f"fModifyScriptTool initialized, llm is {llm}")
 
     def _run(self, *args, **input):
+        if self.llm is None:  # this should not happen
+            print("No language model provided at ModifyScriptTool")
+            return "llm not initialized"
         if len(args) > 0:
             return (
                 "This tool expects you to provide the input as a "
@@ -178,7 +184,7 @@ def _run(self, *args, **input):
         with open(base_script_path, "r") as file:
             base_script = file.read()
         base_script = "".join(base_script)
-        utils = ModifyScriptUtils()
+        utils = ModifyScriptUtils(self.llm)
 
         description = input.get("query")
         answer = utils._prompt_summary(
@@ -194,9 +200,9 @@ def _run(self, *args, **input):
         script_content = textwrap.dedent(script_content).strip()
         # Write to file
         filename = self.path_registry.write_file_name(
-            type="SIMULATION", Sim_id=base_script_id, modified=True
+            type=FileType.SIMULATION, Sim_id=base_script_id, modified=True
         )
-        file_id = self.path_registry.get_fileid(filename, type="SIMULATION")
+        file_id = self.path_registry.get_fileid(filename, type=FileType.SIMULATION)
         directory = "files/simulations"
         if not os.path.exists(directory):
             os.makedirs(directory)

mdagent/tools/maketools.py

@@ -24,7 +24,6 @@
     PPIDistance,
     RMSDCalculator,
     Scholar2ResultLLM,
-    SerpGitTool,
     SetUpandRunFunction,
     SimulationOutputFigures,
     VisualizeProtein,
@@ -63,15 +62,17 @@ def make_all_tools(
 ):
     load_dotenv()
     all_tools = []
-
+    path_instance = PathRegistry.get_instance()  # get instance first
     if llm:
         all_tools += agents.load_tools(["llm-math"], llm)
         all_tools += [PythonREPLTool()]  # or PythonREPLTool(llm=llm)?
+        all_tools += [
+            ModifyBaseSimulationScriptTool(path_registry=path_instance, llm=llm)
+        ]
         if human:
             all_tools += [agents.load_tools(["human"], llm)[0]]
 
     # get path registry
-    path_instance = PathRegistry.get_instance()  # get instance first
 
     # add base tools
     base_tools = [
@@ -110,10 +111,10 @@ def make_all_tools(
     all_tools += base_tools + subagents_tools + learned_tools
 
     # add other tools depending on api keys
-    serp_key = os.getenv("SERP_API_KEY")
+    os.getenv("SERP_API_KEY")
     pqa_key = os.getenv("PQA_API_KEY")
-    if serp_key:
-        all_tools.append(SerpGitTool(serp_key))  # github issues search
+    # if serp_key:
+    #    all_tools.append(SerpGitTool(serp_key))  # github issues search
     if pqa_key:
         all_tools.append(Scholar2ResultLLM(pqa_key))  # literature search
     return all_tools

mdagent/utils/path_registry.py

@@ -159,17 +159,20 @@ def write_file_name(self, type: FileType, **kwargs):
         conditions = kwargs.get("conditions", None)
         Sim_id = kwargs.get("Sim_id", None)
         modified = kwargs.get("modified", False)
-
+        file_name = ""
         if type == FileType.PROTEIN:
-            file_name = f"{protein_name}_{description}_{time_stamp}.{file_format}"
+            file_name += f"{protein_name}_{description}_{time_stamp}.{file_format}"
         if type == FileType.SIMULATION:
+            print("im here inside")
             if conditions:
-                file_name = f"{type_of_sim}_{protein_file_id}_{conditions}_{time_stamp}"
+                file_name += (
+                    f"{type_of_sim}_{protein_file_id}_{conditions}_{time_stamp}.py"
+                )
             elif modified:
-                file_name = f"{Sim_id}_MOD_{time_stamp}"
+                print("I got here!!!!")
+                file_name += f"{Sim_id}_MOD_{time_stamp}.py"
             else:
-                file_name = f"{type_of_sim}_{protein_file_id}_{time_stamp}"
-        if type == FileType.RECORD:
-            file_name = f"{protein_file_id}_{Sim_id}_{time_stamp}"
-
+                file_name += f"{type_of_sim}_{protein_file_id}_{time_stamp}.py"
+        if file_name == "":
+            file_name += "ErrorDuringNaming_error.py"
         return file_name