diff --git a/.gitignore b/.gitignore index 8f50ecd..71ee47a 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,4 @@ software/ Launch/ Plotting/ extra_logs/ +Launch_scripts/ diff --git a/docker/Dockerfile b/docker/Dockerfile index 1df482e..38568b8 100644 --- a/docker/Dockerfile +++ b/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 diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index 37e0dd5..2b96523 100644 --- a/docker/entrypoint.sh +++ b/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 "$@" \ No newline at end of file diff --git a/gmx_MMPBSA/env.yml b/gmx_MMPBSA/env.yml index e394468..1f9a9ef 100644 --- a/gmx_MMPBSA/env.yml +++ b/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 diff --git a/main.py b/main.py index 3548357..ebe1801 100644 --- a/main.py +++ b/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 diff --git a/src/agents/agent.py b/src/agents/agent.py index 45ec8e5..a9eeeb5 100644 --- a/src/agents/agent.py +++ b/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} diff --git a/src/agents/md_agent.py b/src/agents/md_agent.py index f388590..82e8d6b 100644 --- a/src/agents/md_agent.py +++ b/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." ), }) diff --git a/src/agents/prep_agent.py b/src/agents/prep_agent.py index 3df7aa0..d9b57ff 100644 --- a/src/agents/prep_agent.py +++ b/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."}, diff --git a/src/constants.py b/src/constants.py index 759f994..c5421e7 100644 --- a/src/constants.py +++ b/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" diff --git a/src/scripts/run_tleap.sh b/src/scripts/run_tleap.sh index 8416710..b090866 100755 --- a/src/scripts/run_tleap.sh +++ b/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 \ No newline at end of file +fi diff --git a/src/tools/MMPBSA.py b/src/tools/MMPBSA.py index af0cd81..9287f37 100644 --- a/src/tools/MMPBSA.py +++ b/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, diff --git a/src/tools/RAG_tools.py b/src/tools/RAG_tools.py index c106490..b6a1ea6 100644 --- a/src/tools/RAG_tools.py +++ b/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() diff --git a/src/tools/pdb_tools.py b/src/tools/pdb_tools.py index 38c4ba0..22b80ca 100644 --- a/src/tools/pdb_tools.py +++ b/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"))]