[feat] improved parameter selection

main.py

@@ -28,75 +28,6 @@ def _ensure_api_key(env_var: str, prompt_name: str) -> str | None:
     return key
     return key
 
 
 
 
-def update_mdp_file(src_path, dest_path, md_temp, md_duration):
-    def extract_comment(part):
-        comment = ""
-
-        if ";" in part:
-            _, comment = part.split(";", 1)
-            comment = ";" + comment
-
-        return comment
-
-    with open(src_path, "r") as f:
-        lines = f.readlines()
-
-    dt = None
-
-    for line in lines:
-        if line.strip().startswith("dt"):
-            dt = float(line.split("=")[1].split(";")[0].strip())
-            break
-    if dt is None:
-        raise ValueError(f"dt not found in {src_path}")
-
-    # Calculate nsteps
-    # md_duration is in ns, dt is in ps, so convert: 1 ns = 1000 ps
-    nsteps = int((md_duration * 1000) / dt)
-
-    # Update the lines
-    new_lines = []
-
-    for line in lines:
-        if line.strip().startswith("ref_t"):
-            # replace the value with md_temp
-            parts = line.split("=")
-            comment = extract_comment(parts[1])
-            new_line = f"{parts[0]}= {md_temp} {md_temp} {comment}\n"
-
-        elif line.strip().startswith("nsteps"):
-            parts = line.split("=")
-            comment = extract_comment(parts[1])
-
-            new_line = f"{parts[0]}= {nsteps} {comment}\n"
-        else:
-            new_line = line
-
-        new_lines.append(new_line)
-
-    with open(dest_path, "w") as f:
-        f.writelines(new_lines)
-
-
-def process_mdp_files(mdp_dir, sandbox_dir, md_temp, md_duration):
-    if not mdp_dir.exists():
-        raise FileNotFoundError(f"MDP files directory not found: {mdp_dir}")
-
-    mdp_files = os.listdir(mdp_dir)
-
-    for file in mdp_files:
-        full_file_name = os.path.join(mdp_dir, file)
-        dest_file_name = os.path.join(sandbox_dir, file)
-
-        if not os.path.isfile(full_file_name):
-            continue
-
-        if file in ("md.mdp", "npt.mdp", "nvt.mdp"):
-            update_mdp_file(full_file_name, dest_file_name, md_temp, md_duration)
-        else:
-            shutil.copy(full_file_name, sandbox_dir)
-
-
 @dataclass
 @dataclass
 class CommandLineArgs:
 class CommandLineArgs:
     """
     """
@@ -112,6 +43,12 @@ class CommandLineArgs:
     ligand: str | None = None
     ligand: str | None = None
     "Ligand ID (optional; defaults to no ligand)."
     "Ligand ID (optional; defaults to no ligand)."
 
 
+    temp: float | None = None
+    "Simulation temperature in Kelvin."
+
+    duration: float | None = None
+    "Simulation length in nanoseconds."
+
     model_supports_system_messages: bool = True
     model_supports_system_messages: bool = True
 
 
 
 
@@ -123,11 +60,7 @@ def main(config: CommandLineArgs):
     sandbox_dir = constants.DATA_DIR / run_name
     sandbox_dir = constants.DATA_DIR / run_name
     sandbox_dir.mkdir(parents=True, exist_ok=True)
     sandbox_dir.mkdir(parents=True, exist_ok=True)
 
 
-    root_logger.info("autoMD - your assistant for running molecular dynamics")
-    root_logger.info("================")
-    root_logger.info(
-        "I can read, fetch, and prepare PDB files to run MD simulations. Is there a particular system I can help you with today?"
-    )
+    root_logger.info("DynaMate - your assistant for running molecular dynamics")
 
 
     try:
     try:
         _ensure_api_key("OPENROUTER_API_KEY", "OPENROUTER_API_KEY")
         _ensure_api_key("OPENROUTER_API_KEY", "OPENROUTER_API_KEY")
@@ -135,31 +68,27 @@ def main(config: CommandLineArgs):
         root_logger.error(str(e))
         root_logger.error(str(e))
         return
         return
 
 
-    root_logger.info("\n=== 1. Starting PrepAgent (Planning & Parameter Determination) ===")
+    root_logger.info("\n=== Starting PrepAgent (Planning & Parameter Determination) ===")
     prep_agent = PrepAgent(
     prep_agent = PrepAgent(
         model_name=config.model,
         model_name=config.model,
         temperature=constants.TEMPERATURE,
         temperature=constants.TEMPERATURE,
         sandbox_dir=sandbox_dir,
         sandbox_dir=sandbox_dir,
         pdb_id=config.pdb_id,
         pdb_id=config.pdb_id,
         ligand_name=config.ligand,
         ligand_name=config.ligand,
+        md_temp=config.temp,
+        md_duration=config.duration,
         model_supports_system_messages=config.model_supports_system_messages,
         model_supports_system_messages=config.model_supports_system_messages,
     )
     )
-    # Start the prep agent to understand request from user and generate a plan
     prep_agent.setup_tools()
     prep_agent.setup_tools()
     pdb_file_path, ligand_name, plan, llm_cost = prep_agent.run()
     pdb_file_path, ligand_name, plan, llm_cost = prep_agent.run()
     root_logger.info("PrepAgent completed. Plan generated.")
     root_logger.info("PrepAgent completed. Plan generated.")
 
 
-    # Copy required mdp_files into sandbox for the agent and update ref_t and nsteps
     md_temp, md_duration = plan["parameters"]["temperature_k"], plan["parameters"]["duration_ns"]
     md_temp, md_duration = plan["parameters"]["temperature_k"], plan["parameters"]["duration_ns"]
 
 
-    md_duration=0.01
-
     root_logger.info(f"Applying parameters: Temp={md_temp}K, Duration={md_duration}ns")
     root_logger.info(f"Applying parameters: Temp={md_temp}K, Duration={md_duration}ns")
-    process_mdp_files(constants.MDP_FILES, sandbox_dir, md_temp, md_duration)
 
 
-    root_logger.info("\n=== 2. Starting MDAgent (Execution & Tool Loop) ===")
+    root_logger.info("\n=== Starting MDAgent (Execution & Tool Loop) ===")
 
 
-    # Copy PDB into run directory
     md_agent = MDAgent(
     md_agent = MDAgent(
         model_name=config.model,
         model_name=config.model,
         temperature=constants.TEMPERATURE,
         temperature=constants.TEMPERATURE,
@@ -167,6 +96,8 @@ def main(config: CommandLineArgs):
         structure_path=pdb_file_path,
         structure_path=pdb_file_path,
         pdb_id=Path(pdb_file_path).stem,
         pdb_id=Path(pdb_file_path).stem,
         ligand_name=ligand_name,
         ligand_name=ligand_name,
+        md_temp=md_temp,
+        md_duration=md_duration,
         model_supports_system_messages=config.model_supports_system_messages,
         model_supports_system_messages=config.model_supports_system_messages,
         plan=plan,
         plan=plan,
     )
     )
@@ -175,14 +106,12 @@ def main(config: CommandLineArgs):
     # Run the MD pipeline (handles user input and retries internally)
     # Run the MD pipeline (handles user input and retries internally)
     result = md_agent.run()
     result = md_agent.run()
 
 
-    # Print result and summary
     if not result:
     if not result:
         root_logger.error("=== MD Pipeline failed or incomplete ===")
         root_logger.error("=== MD Pipeline failed or incomplete ===")
     else:
     else:
         root_logger.info("=== MD Pipeline completed successfully ===")
         root_logger.info("=== MD Pipeline completed successfully ===")
 
 
 
 
-
 if __name__ == "__main__":
 if __name__ == "__main__":
     config = tyro.cli(CommandLineArgs)
     config = tyro.cli(CommandLineArgs)
     main(config)
     main(config)

setup.sh

 #!/bin/bash
 #!/bin/bash
 
 
 # load your environment
 # load your environment
-source ~/miniforge3/bin/activate dynagent
+source ~/miniforge3/bin/activate dynamate
 
 
 # Activate GROMACS environment
 # Activate GROMACS environment
 if [ -f /usr/local/gromacs/bin/GMXRC ]; then
 if [ -f /usr/local/gromacs/bin/GMXRC ]; then

src/agents/agent.py

@@ -28,6 +28,8 @@ class BaseAgent(ABC):
         sandbox_dir: str,
         sandbox_dir: str,
         pdb_id: str | None = None,
         pdb_id: str | None = None,
         ligand_name: str | None = None,
         ligand_name: str | None = None,
+        md_temp: float | None = None,
+        md_duration: float | None = None,
         model_supports_system_messages: bool = True,
         model_supports_system_messages: bool = True,
     ):
     ):
         self.model_name = model_name
         self.model_name = model_name
@@ -35,6 +37,8 @@ class BaseAgent(ABC):
         self.sandbox_dir = Path(sandbox_dir)
         self.sandbox_dir = Path(sandbox_dir)
         self.pdb_id = pdb_id
         self.pdb_id = pdb_id
         self.ligand_name = ligand_name
         self.ligand_name = ligand_name
+        self.md_temp = md_temp
+        self.md_duration = md_duration
 
 
         self.model_supports_system_messages = model_supports_system_messages
         self.model_supports_system_messages = model_supports_system_messages
 
 
@@ -66,7 +70,6 @@ class BaseAgent(ABC):
         if not messages:
         if not messages:
             return []
             return []
 
 
-        # Start from the end
         block = []
         block = []
         i = len(messages) - 1
         i = len(messages) - 1
 
 
@@ -77,11 +80,7 @@ class BaseAgent(ABC):
             role = m.get("role")
             role = m.get("role")
 
 
             # Stop conditions for different last-message types:
             # Stop conditions for different last-message types:
-
-            # 1. If the last message is a 'tool' message, we must also include the
-            #    assistant tool call that triggered it.
             if role == "tool":
             if role == "tool":
-                # include until we find the matching assistant tool call
                 i -= 1
                 i -= 1
                 while i >= 0 and messages[i].get("role") != "assistant":
                 while i >= 0 and messages[i].get("role") != "assistant":
                     block.insert(0, messages[i])
                     block.insert(0, messages[i])
@@ -89,22 +88,13 @@ class BaseAgent(ABC):
                 if i >= 0:
                 if i >= 0:
                     block.insert(0, messages[i])
                     block.insert(0, messages[i])
                 return block
                 return block
-
-            # 2. If last message is assistant tool call (function call),
-            #    we keep only that one so the model knows what it's continuing.
             if role == "assistant" and m.get("tool_calls"):
             if role == "assistant" and m.get("tool_calls"):
                 return block
                 return block
-
-            # 3. If the last message is a normal assistant message,
-            #    keep it and the preceding user message.
             if role == "assistant":
             if role == "assistant":
-                # include the previous user message
                 if i - 1 >= 0 and messages[i-1].get("role") == "user":
                 if i - 1 >= 0 and messages[i-1].get("role") == "user":
                     block.insert(0, messages[i-1])
                     block.insert(0, messages[i-1])
                 return block
                 return block
 
 
-            # 4. If the last message is a user message,
-            #    include previous assistant message.
             if role == "user":
             if role == "user":
                 if i - 1 >= 0 and messages[i-1].get("role") == "assistant":
                 if i - 1 >= 0 and messages[i-1].get("role") == "assistant":
                     block.insert(0, messages[i-1])
                     block.insert(0, messages[i-1])
@@ -168,7 +158,6 @@ class BaseAgent(ABC):
 
 
         tool_output = None
         tool_output = None
 
 
-        # try:
         self._validate_tool_path(tool_input)
         self._validate_tool_path(tool_input)
         func = TOOL_MAP.get(tool_name)
         func = TOOL_MAP.get(tool_name)
 
 
@@ -179,15 +168,6 @@ class BaseAgent(ABC):
         passed = self._additional_check_for_errors_tool_output(tool_name, tool_output)
         passed = self._additional_check_for_errors_tool_output(tool_name, tool_output)
 
 
         return {"ok": passed, "output": tool_output}
         return {"ok": passed, "output": tool_output}
-        # except Exception as e:
-        #     tb = traceback.format_exc()
-        #     error_msg = (
-        #         f"Tool '{tool_name}' failed with exception:\n{e}\n\n"
-        #         f"Traceback:\n{tb}\n"
-        #         "Decide whether to retry, fix the issue, or skip this step."
-        #     )
-        #     self.logger.error(error_msg)
-        #     return {"ok": False, "output": error_msg}
 
 
 
 
     def _format_tool_usage_ouput(self, id_, tool_name, arguments, output):
     def _format_tool_usage_ouput(self, id_, tool_name, arguments, output):
@@ -203,7 +183,7 @@ class BaseAgent(ABC):
         def to_dict_safe(msg):
         def to_dict_safe(msg):
             if isinstance(msg, dict):
             if isinstance(msg, dict):
                 return msg
                 return msg
-            elif hasattr(msg, "model_dump"):  # Pydantic/LiteLLM object
+            elif hasattr(msg, "model_dump"):
                 return msg.model_dump()
                 return msg.model_dump()
             else:
             else:
                 return str(msg)
                 return str(msg)

src/agents/md_agent.py

@@ -7,7 +7,7 @@ from src.agents.agent import BaseAgent, ToolOutputError
 from src.prompts import MD_SYSTEM_PROMPT
 from src.prompts import MD_SYSTEM_PROMPT
 from src.tools import tool_schema
 from src.tools import tool_schema
 
 
-litellm.drop_params = True  # avoid problems with setting temp on GPT-5
+litellm.drop_params = True 
 
 
 
 
 class MDAgent(BaseAgent):
 class MDAgent(BaseAgent):
@@ -31,11 +31,20 @@ class MDAgent(BaseAgent):
         structure_path,
         structure_path,
         pdb_id=None,
         pdb_id=None,
         ligand_name=None,
         ligand_name=None,
+        md_temp=None,
+        md_duration=None,
         model_supports_system_messages=True,
         model_supports_system_messages=True,
         plan: Dict[str, Any] = None,
         plan: Dict[str, Any] = None,
     ):
     ):
         super().__init__(
         super().__init__(
-            model_name, temperature, sandbox_dir, pdb_id, ligand_name, model_supports_system_messages
+            model_name,
+            temperature,
+            sandbox_dir,
+            pdb_id,
+            ligand_name,
+            md_temp,
+            md_duration,
+            model_supports_system_messages
         )
         )
 
 
         self.structure_path = Path(structure_path)
         self.structure_path = Path(structure_path)
@@ -249,7 +258,7 @@ class MDAgent(BaseAgent):
         success = self._pipeline_successful()
         success = self._pipeline_successful()
 
 
         if self.ligand_name and success:
         if self.ligand_name and success:
-            user_prompt = "Would you like me to calculate the free energy of binding for your protein-ligand system using the MMPBSA tool? (yes/no) \n"
+            user_prompt = "\n==========\nWould you like me to calculate the free energy of binding for your protein-ligand system using the MMPBSA tool? (yes/no) \n\n"
 
 
             user_answer = input(user_prompt).strip().lower()
             user_answer = input(user_prompt).strip().lower()
 
 

src/agents/prep_agent.py

@@ -9,7 +9,7 @@ from src.agents.agent import BaseAgent
 from src.prompts import PREP_SYSTEM_PROMPT
 from src.prompts import PREP_SYSTEM_PROMPT
 
 
 
 
-litellm.drop_params = True  # avoid problems with setting temp on GPT-5
+litellm.drop_params = True
 
 
 
 
 class Tool(BaseModel):
 class Tool(BaseModel):
@@ -26,10 +26,19 @@ class PrepAgent(BaseAgent):
         sandbox_dir,
         sandbox_dir,
         pdb_id=None,
         pdb_id=None,
         ligand_name=None,
         ligand_name=None,
+        md_temp=None,
+        md_duration=None,
         model_supports_system_messages=True,
         model_supports_system_messages=True,
     ):
     ):
         super().__init__(
         super().__init__(
-            model_name, temperature, sandbox_dir, pdb_id, ligand_name, model_supports_system_messages
+            model_name,
+            temperature,
+            sandbox_dir,
+            pdb_id,
+            ligand_name,
+            md_temp,
+            md_duration,
+            model_supports_system_messages
         )
         )
 
 
         self.messages: List[Dict[str, Any]] = []
         self.messages: List[Dict[str, Any]] = []
@@ -122,7 +131,7 @@ class PrepAgent(BaseAgent):
                     sys.exit(1)
                     sys.exit(1)
 
 
     def _find_simulation_temperature(self):
     def _find_simulation_temperature(self):
-        temperature = None
+        temperature = self.md_temp
 
 
         while temperature is None:
         while temperature is None:
             user_input = "Please pick a suitable temperature (in Kelvins) for running the molecular dynamics simulation. Also provide a rational for why you selected this temperature."
             user_input = "Please pick a suitable temperature (in Kelvins) for running the molecular dynamics simulation. Also provide a rational for why you selected this temperature."
@@ -146,7 +155,7 @@ class PrepAgent(BaseAgent):
         return temperature
         return temperature
 
 
     def _calculate_duration(self):
     def _calculate_duration(self):
-        duration = None
+        duration = self.md_duration
 
 
         while duration is None:
         while duration is None:
             user_input = "Please pick a suitable simulation duration (in nanoseconds) for running a short molecular dynamics of this system. Also provide a rational for why you selected this duration. Note that dynamics simulations of 10 ns take 1 hour to complete, so keep the experiment brief (less than 1 ns duration)."
             user_input = "Please pick a suitable simulation duration (in nanoseconds) for running a short molecular dynamics of this system. Also provide a rational for why you selected this duration. Note that dynamics simulations of 10 ns take 1 hour to complete, so keep the experiment brief (less than 1 ns duration)."
@@ -218,18 +227,20 @@ class PrepAgent(BaseAgent):
         prompt = f"I would like to run molecular dynamics for the system {user_input}. If a PDB has not been uploaded, use the tools available to fetch and prepare the PDB for {user_input}."
         prompt = f"I would like to run molecular dynamics for the system {user_input}. If a PDB has not been uploaded, use the tools available to fetch and prepare the PDB for {user_input}."
         self.logger.info(f"User input: {prompt}")
         self.logger.info(f"User input: {prompt}")
         self.pdb_file_path = self._get_pdb_file_path(prompt)
         self.pdb_file_path = self._get_pdb_file_path(prompt)
-        self.logger.info(f"Thank you, I now have access to the structure information for protein {self.pdb_file_path}")
+        self.logger.info(f"I now have access to the structure information for protein {self.pdb_file_path}")
 
 
         self._find_ligand()
         self._find_ligand()
 
 
-        temperature = self._find_simulation_temperature()
-        self.logger.info(f"Using simulation temperature: {temperature} K")
+        if self.md_temp is None:
+            self.md_temp = self._find_simulation_temperature()
+        self.logger.info(f"Using simulation temperature: {self.md_temp} K")
 
 
-        duration = self._calculate_duration()
-        self.logger.info(f"Using simulation duration: {duration} ns")
+        if self.md_duration is None:
+            self.md_duration = self._calculate_duration()
+        self.logger.info(f"Using simulation duration: {self.md_duration} ns")
 
 
         # Build plan steps depending on ligand
         # Build plan steps depending on ligand
-        plan = self._generate_plan(temperature, duration)
+        plan = self._generate_plan(self.md_temp, self.md_duration)
         self.agent_plan = json.dumps(plan, indent=2)
         self.agent_plan = json.dumps(plan, indent=2)
         self.logger.info(f"Generated plan: {self.agent_plan}")
         self.logger.info(f"Generated plan: {self.agent_plan}")
 
 

src/constants.py

@@ -14,4 +14,6 @@ ENV_FILE = Path(__file__).resolve().parent.parent / ".env"
 DATA_DIR = Path(__file__).resolve().parent.parent / "sandbox"
 DATA_DIR = Path(__file__).resolve().parent.parent / "sandbox"
 
 
 AGENT_LOGS = Path(__file__).resolve().parent.parent / "agent_logs"
 AGENT_LOGS = Path(__file__).resolve().parent.parent / "agent_logs"
-JSON_LOG_FILE = AGENT_LOGS / "agent_runs.jsonl"
+JSON_LOG_FILE = AGENT_LOGS / "agent_runs.jsonl"
\ No newline at end of file
+
+MMPBSA_ENV_DIR = Path("/path/to/your/envs/mmpbsa")
\ No newline at end of file

src/scripts/mdp_files/md.mdp

 title                   = Protein-ligand complex MD simulation 
 title                   = Protein-ligand complex MD simulation 
 ; Run parameters
 ; Run parameters
 integrator              = md        ; leap-frog integrator
 integrator              = md        ; leap-frog integrator
-nsteps                  = 5000   ; 2 * 50,000 = 100 ps (0.01 ns)
+nsteps                  = 50000   ; 2 * 500,000 = 1000 ps (0.1 ns)
 dt                      = 0.002     ; 2 fs
 dt                      = 0.002     ; 2 fs
 ; Output control
 ; Output control
 nstenergy               = 5000     ; save energies every 10.0 ps
 nstenergy               = 5000     ; save energies every 10.0 ps
@@ -43,4 +43,4 @@ pbc                     = xyz       ; 3-D PBC
 ; Dispersion correction is not used for proteins with the C36 additive FF
 ; Dispersion correction is not used for proteins with the C36 additive FF
 DispCorr                = no 
 DispCorr                = no 
 ; Velocity generation
 ; Velocity generation
-gen_vel                 = no        ; continuing from NPT equilibration 
+gen_vel                 = no        ; continuing from NPT equilibration 
\ No newline at end of file

src/scripts/prod_Gromacs.sh

 #!/bin/bash
 #!/bin/bash
-if [ "$#" -lt 2 ]; then
-    echo "Usage: $0 input_gro npt_cpt_file"
+if [ "$#" -lt 3 ]; then
+    echo "Usage: $0 input_gro npt_cpt_file log_file [ligand_name]"
     exit 1
     exit 1
 fi
 fi
 
 
@@ -11,6 +11,65 @@ LOG_FILE="$3"
 
 
 > $LOG_FILE 
 > $LOG_FILE 
 
 
+# Optional fourth argument
+if [ "$#" -ge 4 ]; then
+    LIGNAME="$4"
+else
+    LIGNAME=""
+fi
+
+#------- EDIT MD.MDP FILE ------------
+
+# if no ions are present, update md.mdp such that Water_and_ions group is changed to Water only
+# if ligand, Protein Water becomes Protein_ligand Water
+if ! grep -q "Cl-" index.ndx && ! grep -q "Na+" index.ndx; then
+	md_file="md.mdp"
+    original="Protein Water_and_ions"
+	
+	if [ -n "$LIGNAME" ]; then
+		replacement="Protein_$LIGNAME Water"
+		if grep "Protein_$LIGNAME" index.ndx; then
+			echo "Protein_$LIGNAME already in index.ndx"
+		else
+			echo -e "1 | 13\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
+		fi
+	else
+		replacement="Protein Water"
+	fi
+
+	if grep "$original" "$md_file"; then
+		sed -i "s|$original|$replacement|" "$md_file"
+		echo "$replacement added successfully to tc-grps group in $md_file."
+	else
+		echo "tc-grps line was not found in $md_file."
+	fi
+else
+	echo "Ions present. Keeping Water_and_ions group."
+fi
+
+# if ions are present, but ligand exists, update md.mdp accordingly
+if grep -q "Cl-" index.ndx || grep -q "Na+" index.ndx; then
+    if [ -n "$LIGNAME" ]; then	
+        md_file="md.mdp" 
+
+        original="Protein Water_and_ions"
+        replacement="Protein_$LIGNAME Water_and_ions"
+        
+        if grep "Protein_$LIGNAME" index.ndx; then
+            echo "Protein_$LIGNAME already in index.ndx"
+        else
+            echo -e "1 | 13\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
+        fi
+
+        if grep "$original" "$md_file"; then
+            sed -i "s|$original|$replacement|" "$md_file"
+            echo "Protein_$LIGNAME added successfully to tc-grps group in $md_file."
+        else
+            echo "tc-grps line "Protein Water_and_ions" was not found in $md_file."
+        fi
+    fi
+fi
+
 #------- PRODUCTION MD ------------
 #------- PRODUCTION MD ------------
 if ! ls md.gro 1> /dev/null 2>&1; then
 if ! ls md.gro 1> /dev/null 2>&1; then
     $GMX grompp -f md.mdp -c $INPUT_GRO -t $NPT_CPT_FILE -p topol.top -n index.ndx -o md.tpr >> $LOG_FILE 2>&1
     $GMX grompp -f md.mdp -c $INPUT_GRO -t $NPT_CPT_FILE -p topol.top -n index.ndx -o md.tpr >> $LOG_FILE 2>&1

src/tools/MMPBSA.py

@@ -6,18 +6,19 @@ import parmed as pmd # type: ignore
 import sys
 import sys
 import re
 import re
 import os
 import os
+from src import constants
 
 
-def run_gmxMMPBSA(sandbox_dir: str, pdb_id: str, nsteps:str, nstxout_compressed:str, temp=str) -> str:
+def run_gmxMMPBSA(sandbox_dir: str, pdb_id: str, nsteps:str, nstxout_compressed:str, md_temp:str) -> str:
     nframes=int(nsteps)/int(nstxout_compressed)
     nframes=int(nsteps)/int(nstxout_compressed)
     os.makedirs(f"{sandbox_dir}/gmx_MMPBSA", exist_ok=True)
     os.makedirs(f"{sandbox_dir}/gmx_MMPBSA", exist_ok=True)
-    os.chdir(f"{sandbox_dir}/gmx_MMPBSA")
-    mmpbsa_infile = open('mmpbsa.in', 'w' )
+    MMPBSA_dir=f"{sandbox_dir}/gmx_MMPBSA"
+    mmpbsa_infile = open(f"{MMPBSA_dir}/mmpbsa.in", 'w' )
     mmpbsa_infile.write(f'''&general
     mmpbsa_infile.write(f'''&general
 sys_name={pdb_id}
 sys_name={pdb_id}
 startframe=1
 startframe=1
 endframe={int(float(nframes))}
 endframe={int(float(nframes))}
 interval=5
 interval=5
-temperature={int(float(temp))}
+temperature={int(float(md_temp))}
 verbose=2
 verbose=2
 /
 /
 &pb
 &pb
@@ -68,17 +69,13 @@ verbose=2
 ''')
 ''')
     mmpbsa_infile.close()
     mmpbsa_infile.close()
 
 
-    #run_gmxMMPBSA("6JJ3","10000000","5000","300")
-
     tpr_file=f"{sandbox_dir}/md.tpr"
     tpr_file=f"{sandbox_dir}/md.tpr"
     xtc_file=f"{sandbox_dir}/md_noPBC.xtc"
     xtc_file=f"{sandbox_dir}/md_noPBC.xtc"
     index_file=f"{sandbox_dir}/index.ndx"
     index_file=f"{sandbox_dir}/index.ndx"
     topol_file=f"{sandbox_dir}/topol.top"
     topol_file=f"{sandbox_dir}/topol.top"
 
 
-    GMXMMPBSA_PATH = "/home/hackathon/miniforge3/envs/gmxMMPBSA/bin/gmx_MMPBSA"
-
     cmd = [
     cmd = [
-        GMXMMPBSA_PATH,
+        constants.MMPBSA_ENV_DIR,
         "-O",
         "-O",
         "-i", "mmpbsa.in",
         "-i", "mmpbsa.in",
         "-cs", tpr_file,
         "-cs", tpr_file,
@@ -91,5 +88,25 @@ verbose=2
         "-nogui"
         "-nogui"
     ]
     ]
 
 
-    subprocess.run(cmd, check=True)
-    return f"MMPBSA complete! Files created: {sandbox_dir}/gmx_MMPBSA/FINAL_RESULTS_MMPBSA.dat and {sandbox_dir}/gmx_MMPBSA/FINAL_RESULTS_MMPBSA.csv"
+    result = subprocess.run(cmd, cwd=MMPBSA_dir, stdout=sys.stdout, stderr=sys.stderr, text=True)
+
+    MMPBSA_output = ""
+
+    log_file_path = Path(MMPBSA_dir) / "gmx_MMPBSA.log"
+
+    if log_file_path.exists():
+        try:
+            MMPBSA_output = log_file_path.read_text(encoding="utf-8")
+        except Exception as e:
+            MMPBSA_output = f"Could not read gmx_MMPBSA log file: {e}"
+
+    if result.returncode != 0:
+        return (f"Equilibration script failed with return code {result.returncode}.\n"
+                f"--- Full gmx_MMPBSA Log ---\n"
+                f"{MMPBSA_output}\n"
+                f"--- Shell Script Stderr ---\n"
+                f"{result.stderr or 'None captured directly'}")
+    else:
+        return (f"MMPBSA complete! Files created: {MMPBSA_dir}/FINAL_RESULTS_MMPBSA.dat and {MMPBSA_dir}/FINAL_RESULTS_MMPBSA.csv."
+                f"Full gmx_MMPBSA output:\n"
+                f"{MMPBSA_output}")
\ No newline at end of file

src/tools/map.py

@@ -38,15 +38,15 @@ TOOL_MAP = {
     ),
     ),
     # GROMACS-related
     # GROMACS-related
     "gromacs_equil": lambda s, i: gromacs_equil(
     "gromacs_equil": lambda s, i: gromacs_equil(
-        s.sandbox_dir, i["input_gro"], 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_file=i.get("ligand_file")
     ),
     ),
     "gromacs_production": lambda s, i: gromacs_production(
     "gromacs_production": lambda s, i: gromacs_production(
-        s.sandbox_dir, i["input_gro"], i["npt_cpt_file"], ligand_name=i.get("ligand_name")
+        s.sandbox_dir, i["input_gro"], i["npt_cpt_file"], i["md_temp"], i["md_duration"], ligand_name=i.get("ligand_name")
     ),
     ),
     "gromacs_analysis": lambda s, i: gromacs_analysis(s.sandbox_dir, i["input_xtc"], ligand_name=i.get("ligand_name")),
     "gromacs_analysis": lambda s, i: gromacs_analysis(s.sandbox_dir, i["input_xtc"], ligand_name=i.get("ligand_name")),
     # MMPBSA-related
     # MMPBSA-related
     "run_gmxMMPBSA": lambda s, i: run_gmxMMPBSA(
     "run_gmxMMPBSA": lambda s, i: run_gmxMMPBSA(
-        s.sandbox_dir, i["pdb_id"], i["nsteps"], i["nstxout_compressed"], i["temp"],
+        s.sandbox_dir, i["pdb_id"], i["nsteps"], i["nstxout_compressed"], i["md_temp"],
     ),
     ),
     # # RAG tools
     # # RAG tools
     "search_papers": lambda _, i: search_papers(i["query"]),
     "search_papers": lambda _, i: search_papers(i["query"]),

src/tools/pdb_tools.py

@@ -342,7 +342,7 @@ def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = No
 
 
         logger.info("Atom renaming of ligand completed.")
         logger.info("Atom renaming of ligand completed.")
 
 
-        return f"Successfully Prepared PDB structure with a ligand and saved the extracted ligand PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb and the protonated ligand PDB file to {sandbox_dir}/{ligand_name}_h.pdb. Ligand was protonated at pH=7 and atom names were cleaned (renumbered)."
+        return f"Successfully Prepared PDB structure with a ligand and saved the extracted ligand PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb and the protonated ligand PDB file to {sandbox_dir}/{ligand_name}_h.pdb. Ligand was protonated at pH=7 and atom names were cleaned (renumbered)"
 
 
     return f"Successfully Prepared PDB structure without a ligand and saved the extracted PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb"
     return f"Successfully Prepared PDB structure without a ligand and saved the extracted PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb"
 
 

src/tools/tool_schema.py

@@ -251,12 +251,6 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                             "This file must exist in the provided sandbox_dir ({sandbox_dir})."
                             "This file must exist in the provided sandbox_dir ({sandbox_dir})."
                         ),
                         ),
                     },
                     },
-                    # "charge_ligand": {
-                    #     "type": ["integer", "null"],
-                    #     "description": (
-                    #         "The net charge of the ligand. If not provided, the charge will be calculated from the structure."
-                    #     ),
-                    # },
                     "ligand_name": {
                     "ligand_name": {
                         "type": "string",
                         "type": "string",
                         "description": (
                         "description": (
@@ -363,6 +357,12 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                             "This file must be located within sandbox_dir."
                             "This file must be located within sandbox_dir."
                         ),
                         ),
                     },
                     },
+                    "md_temp": {
+                        "type": "string",
+                        "description": (
+                            "Temperature in Kelvin for the MD simulation."
+                        ),
+                    },
                     "ligand_name": {
                     "ligand_name": {
                         "type": "string",
                         "type": "string",
                         "description": (
                         "description": (
@@ -378,7 +378,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                         ),
                         ),
                     },
                     },
                 },
                 },
-                "required": ["sandbox_dir", "input_gro"],
+                "required": ["sandbox_dir", "input_gro", "md_temp"],
             },
             },
         ),
         ),
         Tool(
         Tool(
@@ -416,6 +416,18 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                             "Must be located within sandbox_dir."
                             "Must be located within sandbox_dir."
                         ),
                         ),
                     },
                     },
+                    "md_temp": {
+                        "type": "string",
+                        "description": (
+                            "Temperature in Kelvin for the MD simulation."
+                        ),
+                    },
+                    "md_duration": {
+                        "type": "string",
+                        "description": (
+                            "Duration of the MD simulation in nanoseconds."
+                        ),
+                    },
                     "ligand_name": {
                     "ligand_name": {
                         "type": "string",
                         "type": "string",
                         "description": (
                         "description": (
@@ -423,7 +435,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                         ),
                         ),
                     }
                     }
                 },
                 },
-                "required": ["sandbox_dir", "input_gro", "npt_cpt_file"],
+                "required": ["sandbox_dir", "input_gro", "npt_cpt_file", "md_temp", "md_duration"],
             },
             },
         ),
         ),
         Tool(
         Tool(
@@ -502,12 +514,12 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
                         "type": ["string"],
                         "type": ["string"],
                         "description": ("Number of MD steps that elapse between writing position coordinates using lossy compression, found in the md.mdp file located in sandbox_dir."),
                         "description": ("Number of MD steps that elapse between writing position coordinates using lossy compression, found in the md.mdp file located in sandbox_dir."),
                     },
                     },
-                    "temp": {
+                    "md_temp": {
                         "type": ["string"],
                         "type": ["string"],
                         "description": ("Temperature used during the MD simulation, found in the md.mdp file located in sandbox_dir. This value is an integer with base 10."),
                         "description": ("Temperature used during the MD simulation, found in the md.mdp file located in sandbox_dir. This value is an integer with base 10."),
                     },
                     },
                 },
                 },
-                "required": ["sandbox_dir", "pdb_id", "nsteps", "nstxout_compressed", "temp"],
+                "required": ["sandbox_dir", "pdb_id", "nsteps", "nstxout_compressed", "md_temp"],
             },
             },
         ),
         ),
         Tool(
         Tool(