[feat] add ability to accept systems with multiple ligands

src/constants.py

-from pathlib import Path
-
-MAX_CHARACTERS_TO_LOG = 5000
-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"
-TEMPERATURE = 0.1
-
-SCRIPTS_DIR = Path(__file__).resolve().parent / "scripts"
-MDP_FILES = SCRIPTS_DIR / "mdp_files"
-
-ENV_FILE = Path(__file__).resolve().parent.parent / ".env"
-DATA_DIR = Path(__file__).resolve().parent.parent / "sandbox"
-
-AGENT_LOGS = Path(__file__).resolve().parent.parent / "agent_logs"
-JSON_LOG_FILE = AGENT_LOGS / "agent_runs.jsonl"
-
-MMPBSA_ENV_DIR = Path("/path/to/your/envs/mmpbsa")
\ No newline at end of file

src/scripts/run_tleap_ligand.sh

@@ -2,13 +2,13 @@
 # Usage: ./run_tleap.sh input.pdb
 
 if [ $# -ne 4 ]; then
-    echo "Usage: $0 sandbox_dir complex_pdb ligand_name prepi_file"
+    echo "Usage: $0 sandbox_dir complex_pdb frcmod_file prepi_file"
     exit 1
 fi
 
 SANDBOX_DIR=$1 # PDBFILE already has sandbox path
 PDBFILE=$2
-LIGNAME=$3
+FRCMOD_FILE=$3
 PREPI_FILE=$4
 
 # Create tleap input file
@@ -22,7 +22,7 @@ addPdbAtomMap { { "CH3"  "C" } { "HH31" "H1" } { "HH32" "H2" } { "HH33" "H3" } {
 
 # Load ligand parameters
 loadamberprep ${SANDBOX_DIR}/${PREPI_FILE}
-loadamberparams ${SANDBOX_DIR}/${LIGNAME}.frcmod
+loadamberparams ${SANDBOX_DIR}/${FRCMOD_FILE}
 
 # PDBFILE already has sandbox path
 mol = loadpdb ${PDBFILE}

src/tools/amber_tools.py

@@ -34,23 +34,29 @@ def run_tleap(sandbox_dir: str, input_pdb: str, pdb_id: str) -> str:
         return f"tleap ran successfully with output: {result.stdout}. \n New files added: {sandbox_dir}/topol.top, {sandbox_dir}/{pdb_id}.gro"
 
 
-def run_tleap_ligand(sandbox_dir: str, input_pdb: str, pdb_id: str, ligand_file: str, ligand_name: str) -> str:
+def run_tleap_ligand(sandbox_dir: str, input_pdb: str, pdb_id: str, ligand_files: str | list[str], ligand_name: str) -> str:
     """
     Run tleap preparation using run_tleap.sh, for a protein-ligand complex.
     """
+    # make sure it's a list
+    if isinstance(ligand_files, str):
+        ligand_files = [ligand_files]
+    else:
+        ligand_files = ligand_files
     # complex.pdb
     with (
         open(f"{sandbox_dir}/{input_pdb}", "r") as pdb_infile,
-        open(f"{sandbox_dir}/{ligand_file}", "r") as ligand_infile,
         open(f"{sandbox_dir}/complex.pdb", "w") as outfile,
     ):
         for line in pdb_infile:
             if not line.startswith("END"):
                 outfile.write(line)
 
+        with open(f"{sandbox_dir}/{ligand_files[0]}", "r") as ligand_infile:
             for line in ligand_infile:
                 if line.startswith("HETATM"):
                     outfile.write(line)
+            logger.info(f"Only the first ligand {ligand_files[0]} was added to complex.pdb as we will only simulate one ligand.")
 
         outfile.write("TER\n")
         outfile.write("END\n")
@@ -59,13 +65,14 @@ def run_tleap_ligand(sandbox_dir: str, input_pdb: str, pdb_id: str, ligand_file:
     # tleap with ligand
     script = constants.SCRIPTS_DIR / "run_tleap_ligand.sh"
 
-    if Path(f"{sandbox_dir}/{ligand_name}_fixed.prepi").exists():
-        prepi_file = f"{ligand_name}_fixed.prepi"
+    ligand_stem=Path(f"{ligand_files[0]}").stem
+    if Path(f"{sandbox_dir}/{ligand_stem}_fixed.prepi").exists():
+        prepi_file = f"{ligand_stem}_fixed.prepi"
     else:
-        prepi_file = f"{ligand_name}.prepi"
+        prepi_file = f"{ligand_stem}.prepi"
 
     result = subprocess.run(
-        [str(script), sandbox_dir, complex_pdb, ligand_name, prepi_file],
+        [str(script), sandbox_dir, complex_pdb, f"{ligand_stem}.frcmod", prepi_file],
         cwd=sandbox_dir,
         capture_output=True,
         text=True,

src/tools/gromacs_tools.py

@@ -9,12 +9,14 @@ import time
 
 logger = get_class_logger(__name__)
 
-def gromacs_equil(sandbox_dir: str, input_gro: str, md_temp: str, ligand_name=None, ligand_file=None) -> str:
+def gromacs_equil(sandbox_dir: str, input_gro: str, md_temp: str, ligand_name=None, ligand_files=None) -> str:
     # sometimes llm passes ligands as empty strings
     if not ligand_name:
         ligand_name = None
-    if not ligand_file:
+    if not ligand_files:
         ligand_file = None
+    else:
+        ligand_file=ligand_files[0]
         
     # ------------ Modify topol.top to include position restraints ------------
 

src/tools/ligand_tools.py

@@ -6,12 +6,82 @@ from src.utils import get_class_logger
 
 logger = get_class_logger(__name__)
 
+def fix_charges(input_file, output_file=None):
+            """
+            Adjust charges so the total is an integer.
+            Only modifies one atom's charge (last atom in coordinate/charge section).
+            """
+            with open(input_file) as f:
+                lines = f.readlines()
+
+            # Identify section that looks like atom definitions (before LOOP/IMPROPER)
+            atom_section_end = len(lines)
+            for i, line in enumerate(lines):
+                if re.match(r"^\s*(LOOP|IMPROPER|DONE|STOP)\b", line):
+                    atom_section_end = i
+                    break
+
+            atom_lines = []
+            charges = []
+            float_pattern = re.compile(r"([-+]?\d*\.\d+|\d+)(?!.*\S)")
+
+            for i, line in enumerate(lines[:atom_section_end]):
+                tokens = line.split()
+                if len(tokens) >= 8:  # typical atom line length
+                    last = tokens[-1]
+                    try:
+                        charge = float(last)
+                        atom_lines.append(i)
+                        charges.append(charge)
+                    except ValueError:
+                        pass
+
+            if not charges:
+                logger.warning("No atom charges found.")
+                return
+
+            total = sum(charges)
+            target = round(total)
+            delta = target - total
+
+            if abs(delta) < 1e-6:
+                logger.warning(f"Already integer total ({total:.6f}). No change.")
+                return
+
+            # Adjust last atom charge in section
+            last_atom_idx = atom_lines[-1]
+            old_charge = charges[-1]
+            new_charge = old_charge + delta
+
+            lines[last_atom_idx] = float_pattern.sub(f"{new_charge:.6f}", lines[last_atom_idx])
+
+            if output_file is None:
+                output_file = Path(input_file).with_name(Path(input_file).stem + "_fixed.res")
+
+            with open(output_file, "w") as f:
+                f.writelines(lines)
+
+            logger.info(f"Adjusted total charge: {total:.6f} → {target}")
+            logger.info(f"Atom line {last_atom_idx + 1}: {old_charge:.6f} → {new_charge:.6f}")
+            logger.info(f"Saved to: {output_file}")
+
+def param_ligand(sandbox_dir: str, ligand_files: str | list[str], ligand_name: str, charge_ligand: str | None = None) -> str:
+    if isinstance(ligand_files, str):
+        ligand_files = [ligand_files]
+    else:
+        ligand_files = ligand_files
+
+    ligand_file=ligand_files[0]
+    if len(ligand_files) > 1:
+        logger.info(f"For now, we will only parameterize the first ligand since we will only simulate one ligand: {ligand_file}")
+
+    logger.info(f"Parameterizing ligand file: {ligand_file}")
+    ligand_stem=Path(f"{ligand_file}").stem
 
-def param_ligand(sandbox_dir: str, ligand_file: str, ligand_name: str, charge_ligand: str | None = None) -> str:
     # Find charge of ligand if not provided
     charges = []
 
-    tmp_mol2file = f"{sandbox_dir}/ligand_tmp.mol2"
+    tmp_mol2file = f"{sandbox_dir}/{ligand_stem}.mol2"
     # Create temporary mol2 file using obabel
     cmd = shlex.split(f"obabel {sandbox_dir}/{ligand_file} -O {tmp_mol2file}")
     run_1 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
@@ -42,105 +112,52 @@ def param_ligand(sandbox_dir: str, ligand_file: str, ligand_name: str, charge_li
     # Clean up temporary mol2 file
     Path(tmp_mol2file).unlink(missing_ok=True)
 
-    logger.info(f"Charge of ligand {ligand_file} determined to be {charge_ligand}")
+    logger.info(f"Charge of ligand {ligand_stem} determined to be {charge_ligand}")
 
     # Create mol2 file using antechamber
     cmd = shlex.split(
-        f"antechamber -i {sandbox_dir}/{ligand_file} -fi pdb -o {sandbox_dir}/{ligand_name}.mol2 -fo mol2 -c bcc -nc {charge_ligand} -s 2"
+        f"antechamber -i {sandbox_dir}/{ligand_file} -fi pdb -o {sandbox_dir}/{ligand_stem}.mol2 -fo mol2 -c bcc -nc {charge_ligand} -s 2"
     )
     run_2 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
     if run_2.returncode != 0:
         error_text = "\n".join(filter(None, [run_2.stderr, run_2.stdout]))
         return f"Ligand parameterization failed with error: {error_text}"
 
-    logger.info(f"Mol2 file for ligand {ligand_name} created")
+    logger.info(f"Mol2 file for ligand {ligand_stem} created")
 
-    cmd = shlex.split(f"sed -i 's/UNL/{ligand_name}/g' {sandbox_dir}/{ligand_name}.mol2")
+    cmd = shlex.split(f"sed -i 's/UNL/{ligand_name}/g' {sandbox_dir}/{ligand_stem}.mol2")
     run_3 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
     if run_3.returncode != 0:
         error_text = "\n".join(filter(None, [run_3.stderr, run_3.stdout]))
         return f"Ligand parameterization failed with error: {error_text}"
 
-    # Parmed to make total charge an integer
-    # pmd.load_file(f"{sandbox_dir}/{ligand_name}.mol2").fix_charges(precision=4).save(f"{sandbox_dir}/{ligand_name}_fixed.mol2")
-
     # Create prepi file using antechamber
     cmd = shlex.split(
-        f"antechamber -i {sandbox_dir}/{ligand_name}.mol2 -fi mol2 -o {sandbox_dir}/{ligand_name}.prepi -fo prepi -c bcc"
+        f"antechamber -i {sandbox_dir}/{ligand_stem}.mol2 -fi mol2 -o {sandbox_dir}/{ligand_stem}.prepi -fo prepi -c bcc"
     )
     run_4 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
     if run_4.returncode != 0:
         error_text = "\n".join(filter(None, [run_4.stderr, run_4.stdout]))
         return f"Ligand parameterization failed with error: {error_text}"
+    logger.info(f"Prepi file for ligand {ligand_file} created: {ligand_stem}.prepi")
 
     # Create frcmod file using parmchk2
-    cmd = shlex.split(f"parmchk2 -i {sandbox_dir}/{ligand_name}.mol2 -f mol2 -o {sandbox_dir}/{ligand_name}.frcmod")
+    cmd = shlex.split(f"parmchk2 -i {sandbox_dir}/{ligand_stem}.mol2 -f mol2 -o {sandbox_dir}/{ligand_stem}.frcmod")
     run_5 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True, check=True)
     if run_5.returncode != 0:
         error_text = "\n".join(filter(None, [run_5.stderr, run_5.stdout]))
         return f"Ligand parameterization failed with error: {error_text}"
+    logger.info(f"Frcmod file for ligand {ligand_file} created: {ligand_stem}.frcmod")
 
     # Update charge prepi file
+    fix_charges(f"{sandbox_dir}/{ligand_stem}.prepi", f"{sandbox_dir}/{ligand_stem}_fixed.prepi")
 
-    def fix_charges(input_file, output_file=None):
-        """
-        Adjust charges so the total is an integer.
-        Only modifies one atom's charge (last atom in coordinate/charge section).
-        """
-        with open(input_file) as f:
-            lines = f.readlines()
-
-        # Identify section that looks like atom definitions (before LOOP/IMPROPER)
-        atom_section_end = len(lines)
-        for i, line in enumerate(lines):
-            if re.match(r"^\s*(LOOP|IMPROPER|DONE|STOP)\b", line):
-                atom_section_end = i
-                break
-
-        atom_lines = []
-        charges = []
-        float_pattern = re.compile(r"([-+]?\d*\.\d+|\d+)(?!.*\S)")
-
-        for i, line in enumerate(lines[:atom_section_end]):
-            tokens = line.split()
-            if len(tokens) >= 8:  # typical atom line length
-                last = tokens[-1]
-                try:
-                    charge = float(last)
-                    atom_lines.append(i)
-                    charges.append(charge)
-                except ValueError:
-                    pass
-
-        if not charges:
-            logger.warning("No atom charges found.")
-            return
-
-        total = sum(charges)
-        target = round(total)
-        delta = target - total
-
-        if abs(delta) < 1e-6:
-            logger.warning(f"Already integer total ({total:.6f}). No change.")
-            return
-
-        # Adjust last atom charge in section
-        last_atom_idx = atom_lines[-1]
-        old_charge = charges[-1]
-        new_charge = old_charge + delta
-
-        lines[last_atom_idx] = float_pattern.sub(f"{new_charge:.6f}", lines[last_atom_idx])
-
-        if output_file is None:
-            output_file = Path(input_file).with_name(Path(input_file).stem + "_fixed.res")
-
-        with open(output_file, "w") as f:
-            f.writelines(lines)
-
-        logger.info(f"Adjusted total charge: {total:.6f} → {target}")
-        logger.info(f"Atom line {last_atom_idx + 1}: {old_charge:.6f} → {new_charge:.6f}")
-        logger.info(f"Saved to: {output_file}")
-
-    fix_charges(f"{sandbox_dir}/{ligand_name}.prepi", f"{sandbox_dir}/{ligand_name}_fixed.prepi")
+    if Path(f"{sandbox_dir}/{ligand_stem}_fixed.prepi").exists():
+        prepi_file = f"{ligand_stem}_fixed.prepi"
+    else:
+        prepi_file = f"{ligand_stem}.prepi"
 
-    return "Ligand parameterisation complete. File saved to {sandbox_dir}/{ligand_name}_fixed.prepi"
+    if len(ligand_files) == 1:
+        return f"Ligand parameterisation complete. Parameters saved to {sandbox_dir}/{prepi_file} and {sandbox_dir}/{ligand_stem}.frcmod"
+    if len(ligand_files) > 1:
+        return f"Ligand parameterisation complete for the the first ligand: {ligand_files[0]}. Parameters saved to {sandbox_dir}/{ligand_stem}.frcmod and {sandbox_dir}/{prepi_file}"

src/tools/map.py

@@ -30,15 +30,15 @@ TOOL_MAP = {
     "add_caps": lambda s, i: add_caps(s.sandbox_dir, i["input_pdb"], i["pdb_id"]),
     "rename_histidines": lambda s, i: rename_histidines(s.sandbox_dir, i["input_pdb"], i["pdb_id"]),
     # Ligand handling
-    "param_ligand": lambda s, i: param_ligand(s.sandbox_dir, i["ligand_file"], i["ligand_name"]),
+    "param_ligand": lambda s, i: param_ligand(s.sandbox_dir, i["ligand_files"] if isinstance(i["ligand_files"], list) else [i["ligand_files"]], i["ligand_name"]),
     # AMBER-related
     "run_tleap": lambda s, i: run_tleap(s.sandbox_dir, i["input_pdb"], i["pdb_id"]),
     "run_tleap_ligand": lambda s, i: run_tleap_ligand(
-        s.sandbox_dir, i["input_pdb"], i["pdb_id"], i["ligand_file"], i["ligand_name"]
+        s.sandbox_dir, i["input_pdb"], i["pdb_id"], i["ligand_files"] if isinstance(i["ligand_files"], list) else [i["ligand_files"]], i["ligand_name"]
     ),
     # GROMACS-related
     "gromacs_equil": lambda s, i: gromacs_equil(
-        s.sandbox_dir, i["input_gro"], i["md_temp"], ligand_name=i.get("ligand_name"), ligand_file=i.get("ligand_file")
+        s.sandbox_dir, i["input_gro"], i["md_temp"], ligand_name=i.get("ligand_name"), ligand_files=i.get("ligand_files")
     ),
     "gromacs_production": lambda s, i: gromacs_production(
         s.sandbox_dir, i["input_gro"], i["npt_cpt_file"], i["md_temp"], i["md_duration"], ligand_name=i.get("ligand_name")

src/tools/tool_schema.py

@@ -243,12 +243,17 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                         "type": "string",
                         "description": f"The directory path where input files are located and output files are produced (this is the sandbox directory): {sandbox_dir}",
                     },
-                    "ligand_file": {
-                        "type": "string",
+                    "ligand_files": {
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
                         "description": (
-                            f"Input PDB file of the protonated ligand to be parameterized (without full path). "
-                            f"The ligand file is named {ligand_name}_h.pdb if previously protonated with the prepare_pdb_file_ligand tool."
-                            "This file must exist in the provided sandbox_dir ({sandbox_dir})."
+                            "List of input PDB file(s) of the protonated ligand(s) to be parameterized "
+                            "(without full path). If there is only one ligand, provide a list with one element, "
+                            f"e.g. ['{ligand_name}_h.pdb']. If there are multiple ligands, provide a list such as "
+                            f"['{ligand_name}_h_1.pdb', '{ligand_name}_h_2.pdb'] where the ligands are indexed starting at 1. "
+                            f"All files must exist in the sandbox directory ({sandbox_dir})."
                         ),
                     },
                     "ligand_name": {
@@ -258,7 +263,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                         ),
                     },
                 },
-                "required": ["sandbox_dir", "ligand_file", "ligand_name"],
+                "required": ["sandbox_dir", "ligand_files", "ligand_name"],
             },
         ),
         Tool(
@@ -298,12 +303,17 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                         "type": "string",
                         "description": (f"The PDB ID of the structure to be prepared: {pdb_id}. "),
                     },
-                    "ligand_file": {
-                        "type": "string",
+                    "ligand_files": {
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
                         "description": (
-                            f"Input PDB file of the ligand to be merged with the protein (without full path). "
-                            f"The ligand file is named {ligand_name}_h.pdb if previously protonated with the prepare_pdb_file_ligand tool."
-                            "This file must exist in the provided sandbox_dir."
+                            "List of input PDB file(s) of the protonated ligand(s) to be parameterized "
+                            "(without full path). If there is only one ligand, provide a list with one element, "
+                            f"e.g. ['{ligand_name}_h.pdb']. If there are multiple ligands, provide a list such as "
+                            f"['{ligand_name}_h_1.pdb', '{ligand_name}_h_2.pdb'] where the ligands are indexed starting at 1. "
+                            f"All files must exist in the sandbox directory ({sandbox_dir})."
                         ),
                     },
                     "ligand_name": {
@@ -313,13 +323,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                         ),
                     },
                 },
-                "required": [
-                    "sandbox_dir",
-                    "input_pdb",
-                    "pdb_id",
-                    "ligand_file",
-                    "ligand_name",
-                ],
+                "required": ["sandbox_dir", "input_pdb", "pdb_id", "ligand_files", "ligand_name"],
             },
         ),
         Tool(
@@ -369,12 +373,17 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                             "Three-character residue name of the ligand (capital letters or numbers). If no ligand was provided by user, do not input this argument."
                         ),
                     },
-                    "ligand_file": {
-                        "type": "string",
+                    "ligand_files": {
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
                         "description": (
-                            "Optional argument for the input PDB file of the ligand to be merged with the protein (without full path). "
-                            "The ligand file is typically named '<ligand_name>_h.pdb' if previously protonated. "
-                            "If no ligand was provided by user, do not input this argument."
+                             "List of input PDB file(s) of the protonated ligand(s) to be parameterized "
+                            "(without full path). If there is only one ligand, provide a list with one element, "
+                            f"e.g. ['{ligand_name}_h.pdb']. If there are multiple ligands, provide a list such as "
+                            f"['{ligand_name}_h_1.pdb', '{ligand_name}_h_2.pdb']. "
+                            f"All files must exist in the sandbox directory ({sandbox_dir})."
                         ),
                     },
                 },