Cluster/Docker hardening, error handling, and MMPBSA env fixes

.gitignore

@@ -30,3 +30,4 @@ software/
 Launch/
 Plotting/
 extra_logs/
+Launch_scripts/

docker/Dockerfile

@@ -1,4 +1,4 @@
-FROM md_software:latest
+FROM bohdannaida/md_software:latest
 
 RUN apt-get update && apt-get install -y \
     openbabel \
@@ -15,6 +15,7 @@ RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -
     bash /tmp/miniconda.sh -b -p $CONDA_DIR && \
     rm /tmp/miniconda.sh
 ENV PATH="$CONDA_DIR/bin:$PATH"
+ENV CONDA_FETCH_THREADS=1
 
 # Accept Conda's Terms and Conditions
 
@@ -24,12 +25,18 @@ RUN conda tos accept --override-channels --channel https://repo.anaconda.com/pkg
 COPY environment.yml /tmp/environment.yml
 RUN conda env create -f /tmp/environment.yml && \
     conda clean -afy
+# Separate env for MM/PB(GB)SA binding affinity
+COPY gmx_MMPBSA/env.yml /tmp/gmxmmpbsa.yml
+RUN conda env create -f /tmp/gmxmmpbsa.yml && \
+    conda clean -afy
+ENV MMPBSA_ENV_DIR="/opt/conda/envs/gmxMMPBSA/bin/gmx_MMPBSA"
 
-ARG CONDA_ENV_NAME=dynagent
+ARG CONDA_ENV_NAME=dynamate
 ENV CONDA_DEFAULT_ENV=$CONDA_ENV_NAME
 ENV PATH="$CONDA_DIR/envs/$CONDA_ENV_NAME/bin:$PATH"
 
 
+
 # Create a new user for the agent
 ARG AGENT_USER=beautifulagent
 RUN useradd -m $AGENT_USER

docker/entrypoint.sh

@@ -1,5 +1,10 @@
 #!/bin/bash
 source /opt/conda/etc/profile.d/conda.sh
-conda activate dynagent
+conda activate dynamate
 source /opt/software/ambertools25/amber.sh
+
+# First arg is a flag (e.g. --model) -> run the agent; else exec it (e.g. /bin/bash)
+if [ "${1:0:1}" = '-' ]; then
+    set -- python /app/main.py "$@"
+fi
 exec "$@"

gmx_MMPBSA/env.yml

@@ -6,6 +6,7 @@ dependencies:
   - python=3.9
   - pip
   - ambertools<=23.3
+  - mpich
   - mpi4py<=3.1.5
   - gromacs<=2023.4
   - git

main.py

@@ -20,7 +20,7 @@ def _ensure_api_key(env_var: str, prompt_name: str) -> str | None:
     key = os.environ.get(env_var)
 
     if not key:
-        print(f"--- Missing API Key: {prompt_name} ---")
+        print(f"Missing API Key: {prompt_name}")
         key = input(f"Please enter your {prompt_name} API key: ").strip()
         if key:
             os.environ[env_var] = key

src/agents/agent.py

@@ -177,15 +177,19 @@ def _safe_execute_tool(self, tool_name: str, tool_input: Dict[str, Any]) -> dict
                 self.logger.error(error_msg)
                 return {"ok": False, "output": error_msg}
 
-        tool_output = None
-
         self._validate_tool_path(tool_input)
         func = TOOL_MAP.get(tool_name)
 
         if not func:
             raise ValueError(f"Unknown tool: {tool_name}")
 
-        tool_output = func(self, tool_input)
+        try:
+            tool_output = func(self, tool_input)
+        except Exception as e:
+            error_msg = f"Tool '{tool_name}' raised an exception: {type(e).__name__}: {e}\n{traceback.format_exc()}"
+            self.logger.error(error_msg)
+            return {"ok": False, "output": error_msg}
+
         passed = self._additional_check_for_errors_tool_output(tool_name, tool_output)
 
         return {"ok": passed, "output": tool_output}

src/agents/md_agent.py

@@ -212,7 +212,7 @@ def _run_agent(self, remaining_steps: list[str]):
                 "Execute the molecular dynamics pipeline. The required steps are:\n"
                 f"{remaining_steps_string}\n\n"
                 "Work through each step in order using the available tools. "
-                "Ask the user if you need clarification on any step."
+                "Do not ask the user anything."
             ),
         })
 

src/agents/prep_agent.py

@@ -70,7 +70,6 @@ def _setup_system_prompt(self) -> None:
 
     def _ask_for_system(self):
         """Have the LLM ask the user for the PDB ID and optional ligand, parse via LLM, then confirm."""
-        # Step 1: LLM asks the user
         self.messages.append({
             "role": "user",
             "content": "Ask the user what molecular system they would like to simulate (PDB ID or file upload) and whether they have a ligand to include (3-letter code).",
@@ -83,7 +82,6 @@ def _ask_for_system(self):
             user_input = input("You: ").strip()
             self.messages.append({"role": "user", "content": user_input})
 
-            # Step 2: LLM extracts PDB ID and ligand as JSON
             parse_messages = [
                 {
                     "role": "user",
@@ -108,7 +106,6 @@ def _ask_for_system(self):
             except (json.JSONDecodeError, AttributeError):
                 pdb_id, ligand = None, None
 
-            # Step 3: Confirm with user
             confirm_parts = [f"PDB ID: {pdb_id or 'not found'}"]
             confirm_parts.append(f"Ligand: {ligand}" if ligand else "Ligand: none")
             print(f"\nAgent: I understood the following — {', '.join(confirm_parts)}. Is that correct? (yes/no)")
@@ -162,10 +159,10 @@ def _find_ligand(self):
             lig_name = re.search(r"^[A-Z0-9]{3}$", user_input)
             if not lig_name:
                 self.logger.error(
-                    f"'{user_input}' is not a valid 3-character ligand code."
+                    f"'{user_input}' is not a valid 3-character ligand code. Continuing without a ligand."
                 )
-                self.ligand_name = input("Please enter the three character identifier for the ligand (or press Enter to skip): ").strip().upper() or None
-                continue
+                self.ligand_name = None
+                return
 
             self.ligand_name = lig_name.group()
             self.logger.info(f"User requested ligand: {self.ligand_name}")
@@ -183,7 +180,7 @@ def _find_ligand(self):
             self.logger.error(
                 f"Ligand '{self.ligand_name}' not found in {self.pdb_file_path}."
             )
-            self.ligand_name = input("Please enter the correct three character identifier for the ligand (or press Enter to skip): ").strip().upper() or None
+            self.ligand_name = None
 
     def _find_simulation_temperature(self):
         temperature = self.md_temp
@@ -234,10 +231,7 @@ def _calculate_duration(self):
     def _generate_plan(self, temperature, duration):
         if self.ligand_name:
             steps = [
-                {
-                    "step": "prepare_pdb_file_ligand",
-                    "description": "Clean and preprocess PDB file for protein-ligand system.",
-                },
+                {"step": "prepare_pdb_file_ligand", "description": "Clean and preprocess PDB file for protein-ligand system"},
                 {"step": "add_caps", "description": "Add N- and C-terminal capping groups."},
                 {"step": "rename_histidines", "description": "Rename HIS to HIE, HIP or HID."},
                 {"step": "param_ligand", "description": "Generate ligand parameters using antechamber or acpype."},
@@ -248,10 +242,7 @@ def _generate_plan(self, temperature, duration):
             ]
         else:
             steps = [
-                {
-                    "step": "prepare_pdb_file_ligand",
-                    "description": "Clean and preprocess PDB file for protein-only system.",
-                },
+                {"step": "prepare_pdb_file_ligand","description": "Clean and preprocess PDB file for protein-only system."},
                 {"step": "add_caps", "description": "Add N- and C-terminal capping groups."},
                 {"step": "rename_histidines", "description": "Rename HIS to HIE, HIP or HID."},
                 {"step": "run_tleap", "description": "Build system topology and solvate complex using tleap."},

src/constants.py

@@ -5,7 +5,7 @@
 SUMMARY_OUTPUT_TOKENS = 6000
 MAX_CONTEXT_TOKENS = 32000
 PAPER_DIR = Path(__file__).resolve().parent.parent / "my_papers"
-MODEL_NAME = "openrouter/openai/gpt-4.1-2025-04-14"
+MODEL_NAME = "openrouter/openai/gpt-5.5"
 TEMPERATURE = 0.1
 
 SCRIPTS_DIR = Path(__file__).resolve().parent / "scripts"

src/scripts/run_tleap.sh

@@ -13,6 +13,8 @@ PDB_ID=$3
 # Create tleap input file
 cat > leap.in << EOF
 source leaprc.protein.ff14SB
+# Load DNA parameters to support DNA residues (DA/DC/DG/DT)
+source leaprc.DNA.bsc1
 source leaprc.water.tip3p
 
 # Map PDB atom names to template atom names
@@ -38,4 +40,4 @@ if [ $? -ne 0 ]; then
     echo "tleap failed to generate output files."
 else
     echo "${PDBFILE} processed. Generated ${PDB_ID}.prmtop, ${PDB_ID}.inpcrd, and ${PDB_ID}_tleap.pdb."
-fi
+fi

src/tools/MMPBSA.py

@@ -73,9 +73,10 @@ def run_gmxMMPBSA(sandbox_dir: str, pdb_id: str, nsteps:str, nstxout_compressed:
     xtc_file=f"{sandbox_dir}/md_noPBC.xtc"
     index_file=f"{sandbox_dir}/index.ndx"
     topol_file=f"{sandbox_dir}/topol.top"
+    MMPBSA_ENV_DIR=os.environ.get("MMPBSA_ENV_DIR", str(constants.MMPBSA_ENV_DIR))
 
     cmd = [
-        constants.MMPBSA_ENV_DIR,
+        MMPBSA_ENV_DIR,
         "-O",
         "-i", "mmpbsa.in",
         "-cs", tpr_file,

src/tools/RAG_tools.py

@@ -23,15 +23,15 @@ def _ensure_openai_key() -> None:
 
 
 def _load_documents() -> Docs:
-    _ensure_openai_key()
-    docs = Docs()
-
     pdf_files = list(constants.PAPER_DIR.rglob("*.pdf")) if constants.PAPER_DIR.exists() else []
     total_files = len(pdf_files)
 
     if total_files == 0:
         return None
 
+    _ensure_openai_key()  # only needed once we actually have papers to embed
+    docs = Docs()
+
     pickled_docs = "my_docs.pkl"
 
     if not os.path.exists(pickled_docs):
@@ -56,7 +56,6 @@ def _load_documents() -> Docs:
 
 def search_papers(query: dict):
     global documents
-    _ensure_openai_key()
 
     if not documents:
         documents = _load_documents()

src/tools/pdb_tools.py

@@ -205,10 +205,14 @@ def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = No
             else:
                 index = ligand_pdb_file.split("_")[-1].split(".")[0]  # get index from filename
                 protonated_file = f"{sandbox_dir}/{ligand_name}_{index}_h.pdb"
-                list_protonated_files.append(f"{ligand_name}_{index}_h.pdb")
+
+            list_protonated_files.append(protonated_file)
 
             cmd = shlex.split(f"obabel {ligand_pdb_file} -O {protonated_file} -p7")
-            subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            if result.returncode != 0 or not os.path.exists(protonated_file):
+                error_detail = result.stderr.decode(errors="replace").strip() or result.stdout.decode(errors="replace").strip()
+                return f"obabel failed to protonate {ligand_pdb_file} (return code {result.returncode}): {error_detail}"
             with open(protonated_file, "r") as infile:
                 lines = infile.readlines()
                 filtered_lines = [line for line in lines if not (line.startswith("CONECT") or line.startswith("MASTER"))]