merge main into test_rmsd

.env.example

@@ -4,8 +4,5 @@
 # OpenAI API Key
 OPENAI_API_KEY=YOUR_OPENAI_API_KEY_GOES_HERE  # pragma: allowlist secret
 
-# PQA API Key
-PQA_API_KEY=YOUR_PQA_API_KEY_GOES_HERE  # pragma: allowlist secret
-
 # Serp API key
 SERP_API_KEY=YOUR_SERP_API_KEY_GOES_HERE  # pragma: allowlist secret

.github/workflows/build.yml

@@ -13,10 +13,10 @@ jobs:
 
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python "3.9"
+    - name: Set up Python "3.11"
       uses: actions/setup-python@v2
       with:
-        python-version: "3.9"
+        python-version: "3.11"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

.github/workflows/tests.yml

@@ -23,10 +23,10 @@ jobs:
         environment-file: environment.yaml
         python-version: ${{ matrix.python-version }}
         auto-activate-base: true
-    - name: Install openmm pdbfixer mdanalysis with conda
+    - name: Install pdbfixer with conda
       shell: bash -l {0}
       run: |
-        conda install -c conda-forge openmm pdbfixer mdanalysis
+        conda install -c conda-forge pdbfixer
     - name: Install dependencies
       shell: bash -l {0}
       run: |
@@ -45,6 +45,5 @@ jobs:
       env:
         OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
         SEMANTIC_SCHOLAR_API_KEY: ${{ secrets.SEMANTIC_SCHOLAR_API_KEY }}
-        PQA_API_KEY : ${{ secrets.PQA_API_TOKEN }}
       run: |
          pytest -m "not skip" tests

.secrets.baseline

@@ -3,8 +3,5 @@
 # Rule for detecting OpenAI API keys
 OpenAI API Key: \b[secrets]{3}_[a-zA-Z0-9]{32}\b
 
-# Rule for detecting pqa API keys
-PQA API Key: "pqa[a-zA-Z0-9-._]+"
-
 # Rule for detecting serp API keys
 # Serp API Key: "[a-zA-Z0-9]{64}"

README.md

@@ -7,8 +7,8 @@ To use the OpenMM features in the agent, please set up a conda environment, foll
 - Create conda environment: `conda env create -n mdagent -f environment.yaml`
 - Activate your environment: `conda activate mdagent`
 
-If you already have a conda environment, you can install the necessary dependencies with the following steps.
-- Install the necessary conda dependencies: `conda install -c conda-forge openmm pdbfixer mdanalysis`
+If you already have a conda environment, you can install, pdbfixer, a necessary dependency with the following steps.
+- Install the necessary conda dependencies: `conda install -c conda-forge pdbfixer`
 
 
 ## Installation

dev-requirements.txt

@@ -1,2 +1,3 @@
 pre-commit
 pytest
+pytest-mock

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

@@ -1,27 +1,23 @@
 from .analysis_tools.plot_tools import SimulationOutputFigures
 from .analysis_tools.ppi_tools import PPIDistance
-from .analysis_tools.rmsd_tools import RMSDCalculator
-from .analysis_tools.vis_tools import (
-    CheckDirectoryFiles,
-    VisFunctions,
-    VisualizeProtein,
+from .analysis_tools.rgy import (
+    RadiusofGyrationAverage,
+    RadiusofGyrationPerFrame,
+    RadiusofGyrationPlot,
 )
+from .analysis_tools.rmsd_tools import RMSDCalculator
+from .analysis_tools.vis_tools import VisFunctions, VisualizeProtein
 from .preprocess_tools.clean_tools import (
     AddHydrogensCleaningTool,
     CleaningToolFunction,
     CleaningTools,
     RemoveWaterCleaningTool,
     SpecializedCleanTool,
 )
-from .preprocess_tools.pdb_tools import (
-    PackMolTool,
-    ProteinName2PDBTool,
-    SmallMolPDB,
-    get_pdb,
-)
+from .preprocess_tools.packing import PackMolTool
+from .preprocess_tools.pdb_get import ProteinName2PDBTool, SmallMolPDB, get_pdb
 from .simulation_tools.create_simulation import ModifyBaseSimulationScriptTool
 from .simulation_tools.setup_and_run import (
-    InstructionSummary,
     SetUpandRunFunction,
     SetUpAndRunTool,
     SimulationFunctions,
@@ -32,9 +28,7 @@
 
 __all__ = [
     "AddHydrogensCleaningTool",
-    "CheckDirectoryFiles",
     "CleaningTools",
-    "InstructionSummary",
     "ListRegistryPaths",
     "MapPath2Name",
     "ProteinName2PDBTool",
@@ -44,6 +38,9 @@
     "VisualizeProtein",
     "RMSDCalculator",
     "RemoveWaterCleaningTool",
+    "RadiusofGyrationAverage",
+    "RadiusofGyrationPerFrame",
+    "RadiusofGyrationPlot",
     "Scholar2ResultLLM",
     "SerpGitTool",
     "SetUpAndRunTool",

mdagent/tools/base_tools/analysis_tools/__init__.py

@@ -1,13 +1,16 @@
 from .plot_tools import SimulationOutputFigures
 from .ppi_tools import PPIDistance
+from .rgy import RadiusofGyrationAverage, RadiusofGyrationPerFrame, RadiusofGyrationPlot
 from .rmsd_tools import RMSDCalculator
-from .vis_tools import CheckDirectoryFiles, VisFunctions, VisualizeProtein
+from .vis_tools import VisFunctions, VisualizeProtein
 
 __all__ = [
     "PPIDistance",
     "RMSDCalculator",
+    "RadiusofGyrationPerFrame",
+    "RadiusofGyrationPlot",
     "SimulationOutputFigures",
-    "CheckDirectoryFiles",
     "VisualizeProtein",
     "VisFunctions",
+    "RadiusofGyrationAverage",
 ]