Merge branch 'main' of github.com:ur-whitelab/md-agent into r_gyration

merge main

mdagent/mainagent/agent.py

@@ -45,8 +45,10 @@ def __init__(
         subagents_model="gpt-4-1106-preview",
         ckpt_dir="ckpt",
         resume=False,
+        learn=True,
         top_k_tools=20,  # set "all" if you want to use all tools (& skills if resume)
         use_human_tool=False,
+        curriculum=True,
         uploaded_files=[],  # user input files to add to path registry
     ):
         if path_registry is None:
@@ -69,7 +71,11 @@ def __init__(
             callbacks=[StreamingStdOutCallbackHandler()],
         )
 
-        # assign prompt
+        if learn:
+            self.skip_subagents = False
+        else:
+            self.skip_subagents = True
+
         if agent_type == "Structured":
             self.prompt = structured_prompt
         elif agent_type == "OpenAIFunctionsAgent":
@@ -83,6 +89,7 @@ def __init__(
             verbose=verbose,
             ckpt_dir=ckpt_dir,
             resume=resume,
+            curriculum=curriculum,
         )
 
     def _initialize_tools_and_agent(self, user_input=None):
@@ -97,13 +104,15 @@ def _initialize_tools_and_agent(self, user_input=None):
                     llm=self.tools_llm,
                     subagent_settings=self.subagents_settings,
                     human=self.use_human_tool,
+                    skip_subagents=self.skip_subagents,
                 )
             else:
                 # retrieve all tools, including new tools if any
                 self.tools = make_all_tools(
                     self.tools_llm,
                     subagent_settings=self.subagents_settings,
                     human=self.use_human_tool,
+                    skip_subagents=self.skip_subagents,
                 )
         return AgentExecutor.from_agent_and_tools(
             tools=self.tools,

mdagent/subagents/subagent_setup.py

@@ -15,6 +15,7 @@ def __init__(
         ckpt_dir="ckpt",
         resume=False,
         retrieval_top_k=5,
+        curriculum=True,
     ):
         self.path_registry = path_registry
         self.subagents_model = subagents_model
@@ -24,6 +25,7 @@ def __init__(
         self.ckpt_dir = ckpt_dir
         self.resume = resume
         self.retrieval_top_k = retrieval_top_k
+        self.curriculum = curriculum
 
 
 class SubAgentInitializer:
@@ -40,6 +42,7 @@ def __init__(self, settings: Optional[SubAgentSettings] = None):
         self.ckpt_dir = settings.ckpt_dir
         self.resume = settings.resume
         self.retrieval_top_k = settings.retrieval_top_k
+        self.curriculum = settings.curriculum
 
     def create_action(self, **overrides):
         params = {
@@ -61,6 +64,8 @@ def create_critic(self, **overrides):
         return Critic(**params)
 
     def create_curriculum(self, **overrides):
+        if not self.curriculum:
+            return None
         params = {
             "model": self.subagents_model,
             "temp": self.temp,

mdagent/tools/base_tools/__init__.py

@@ -13,7 +13,12 @@
     RemoveWaterCleaningTool,
     SpecializedCleanTool,
 )
-from .preprocess_tools.pdb_tools import Name2PDBTool, PackMolTool, get_pdb
+from .preprocess_tools.pdb_tools import (
+    PackMolTool,
+    ProteinName2PDBTool,
+    SmallMolPDB,
+    get_pdb,
+)
 from .simulation_tools.create_simulation import ModifyBaseSimulationScriptTool
 from .simulation_tools.setup_and_run import (
     InstructionSummary,
@@ -32,9 +37,10 @@
     "InstructionSummary",
     "ListRegistryPaths",
     "MapPath2Name",
-    "Name2PDBTool",
+    "ProteinName2PDBTool",
     "PackMolTool",
     "PPIDistance",
+    "SmallMolPDB",
     "VisualizeProtein",
     "RMSDCalculator",
     "RemoveWaterCleaningTool",

mdagent/tools/base_tools/preprocess_tools/__init__.py

@@ -5,15 +5,16 @@
     RemoveWaterCleaningTool,
     SpecializedCleanTool,
 )
-from .pdb_tools import Name2PDBTool, PackMolTool, get_pdb
+from .pdb_tools import PackMolTool, ProteinName2PDBTool, SmallMolPDB, get_pdb
 
 __all__ = [
     "AddHydrogensCleaningTool",
     "CleaningTools",
-    "Name2PDBTool",
+    "ProteinName2PDBTool",
     "PackMolTool",
     "RemoveWaterCleaningTool",
     "SpecializedCleanTool",
     "get_pdb",
     "CleaningToolFunction",
+    "SmallMolPDB",
 ]

mdagent/tools/base_tools/preprocess_tools/clean_tools.py

@@ -296,17 +296,17 @@ def _run(self, **input_args) -> str:
             file_description = "Cleaned File: "
             CleaningTools()
             try:
-                pdbfile = self.path_registry.get_mapped_path(pdbfile_id)
-                if "/" in pdbfile:
-                    pdbfile = pdbfile.split("/")[-1]
-
-                name = pdbfile.split("_")[0]
-                end = pdbfile.split(".")[1]
+                pdbfile_path = self.path_registry.get_mapped_path(pdbfile_id)
+                if "/" in pdbfile_path:
+                    pdbfile = pdbfile_path.split("/")[-1]
+                else:
+                    pdbfile = pdbfile_path
+                name, end = pdbfile.split(".")
 
             except Exception as e:
                 print(f"error retrieving from path_registry, trying to read file {e}")
                 return "File not found in path registry. "
-            fixer = PDBFixer(filename=pdbfile)
+            fixer = PDBFixer(filename=pdbfile_path)
             try:
                 fixer.findMissingResidues()
             except Exception:
@@ -353,7 +353,7 @@ def _run(self, **input_args) -> str:
             file_mode = "w" if add_hydrogens else "a"
             file_name = self.path_registry.write_file_name(
                 type=FileType.PROTEIN,
-                protein_name=name,
+                protein_name=name.split("_")[0],
                 description="Clean",
                 file_format=end,
             )