[update] increase chat flow between user and agent

README.md

@@ -34,13 +34,13 @@ OPENROUTER_API_KEY=your_key_here
 
 
 3. Run the agent:
 3. Run the agent:
 ```
 ```
-docker run --env-file .env dynamate --pdb-id <pdb-id> --ligand <ligand-name (optional)> --model <model_name> --temp <simulation temperature (K), default: chosen by the agent> --duration <simulation duration (ns), default: chosen by the agent>
+docker run --env-file .env dynamate --model <model_name> --pdb-id <pdb-id> --ligand <ligand-name (optional)> --temp <simulation temperature (K), default: chosen by the agent> --duration <simulation duration (ns), default: chosen by the agent>
 ```
 ```
 
 
 4. Interactive mode (for debugging or exploration):
 4. Interactive mode (for debugging or exploration):
 ```
 ```
 docker run -it --rm --env-file .env dynamate /bin/bash
 docker run -it --rm --env-file .env dynamate /bin/bash
-python main.py --pdb-id <pdb-id> --model <model_name>
+python main.py --model <model_name> --pdb-id <pdb-id>
 ```
 ```
 
 
 Happy molecular dynamics simulations! 🧬
 Happy molecular dynamics simulations! 🧬
@@ -189,12 +189,14 @@ source setup.sh
 ```
 ```
 Now you are ready to use DynaMate! 
 Now you are ready to use DynaMate! 
 ## Usage
 ## Usage
-To launch the script specify the PDB (or upload it), possible ligand name, and model name in the command line arguments. For example, to launch the MD run with the protein 5UEZ, ligand 89G, and model GPT-5 mini:
+To launch the script specify the model name in the command line arguments. For example, to launch the agent with GPT-5 mini:
 ```bash
 ```bash
-python main.py --pdb_id 5UEZ --ligand 89G --model openrouter/openai/gpt-5-mini
+python main.py --model openrouter/openai/gpt-5-mini
 ```
 ```
 You can optionally specify:
 You can optionally specify:
 ```
 ```
+--pdb-id <protein you would like to run MD for, default: prompted at runtime>
+--ligand <ligand you would like to run MD with, default: prompted at runtime>
 --temp <simulation temperature (K), default: chosen by the agent>
 --temp <simulation temperature (K), default: chosen by the agent>
 --duration <simulation duration (ns), default: chosen by the agent>
 --duration <simulation duration (ns), default: chosen by the agent>
 ```
 ```

main.py

@@ -34,12 +34,12 @@ class CommandLineArgs:
     Tyro automatically generates a command line interface from this class.
     Tyro automatically generates a command line interface from this class.
     """
     """
 
 
-    pdb_id: str
-    "PDB ID."
-
     model: str
     model: str
     "Model name to use for the MD pipeline."
     "Model name to use for the MD pipeline."
 
 
+    pdb_id: str | None = None
+    "PDB ID (optional; agent will ask interactively if not provided)."
+
     ligand: str | None = None
     ligand: str | None = None
     "Ligand ID (optional; defaults to no ligand)."
     "Ligand ID (optional; defaults to no ligand)."
 
 

src/agents/md_agent.py

@@ -199,64 +199,70 @@ class MDAgent(BaseAgent):
             self.completed_summary += f"{name} failed;\n"
             self.completed_summary += f"{name} failed;\n"
 
 
     def _run_agent(self, remaining_steps: list[str]):
     def _run_agent(self, remaining_steps: list[str]):
-        iteration = 1
-
-        # let the LLM continuously propose next tool steps, up to MAX_ITERATIONS
-        while remaining_steps and iteration < self.MAX_ITERATION:
-            try:
+        # Inject the task prompt once before the loop
         remaining_steps_string = "\n".join(remaining_steps)
         remaining_steps_string = "\n".join(remaining_steps)
-                prompt = (
-                    "The following essential steps are remaining in the pipeline:\n"
+        self.messages.append({
+            "role": "user",
+            "content": (
+                "Execute the molecular dynamics pipeline. The required steps are:\n"
                 f"{remaining_steps_string}\n\n"
                 f"{remaining_steps_string}\n\n"
-                    f"Completed: {self.completed_steps}\n"
-                    "Choose the next best tool to execute. Do not ask the user anything."
-                )
-                self.messages.append({"role": "user", "content": prompt})
+                "Work through each step in order using the available tools. "
+                "Ask the user if you need clarification on any step."
+            ),
+        })
 
 
+        for iteration in range(1, self.MAX_ITERATION + 1):
+            try:
                 response = self._call_llm(self.messages)
                 response = self._call_llm(self.messages)
                 tool_calls = response.tool_calls
                 tool_calls = response.tool_calls
 
 
                 if tool_calls:
                 if tool_calls:
-                    self.logger.info(f"Length of tool calls: {len(tool_calls)}")
+                    self.logger.info(f"Iteration {iteration}: {len(tool_calls)} tool call(s)")
                     self.messages.append(response)
                     self.messages.append(response)
                     for tool_call in tool_calls:
                     for tool_call in tool_calls:
                         exec_result = self._process_tool_call(tool_call)
                         exec_result = self._process_tool_call(tool_call)
                         self._process_tool_results(tool_call.function.name, exec_result, remaining_steps)
                         self._process_tool_results(tool_call.function.name, exec_result, remaining_steps)
-
                 else:
                 else:
-                    assistant_message = {"role": "assistant", "content": response.content}
-                    self.messages.append(assistant_message)
+                    self.messages.append({"role": "assistant", "content": response.content})
+                    if response.content:
+                        print(f"\nAgent: {response.content}")
+                    user_input = input("You (press Enter to finish): ").strip()
+                    if not user_input:
+                        self.logger.info(f"Agent finished after {iteration} iteration(s).")
+                        break
+                    self.messages.append({"role": "user", "content": user_input})
 
 
-                iteration += 1
-                self.logger.info(f"Logging agent iteration {iteration}")
             except Exception as e:
             except Exception as e:
                 self.logger.error(str(e))
                 self.logger.error(str(e))
-                raise
+                self.messages.append({"role": "user", "content": f"An error occurred: {e}"})
 
 
         return remaining_steps
         return remaining_steps
 
 
     def _run_bfe(self, prompt):
     def _run_bfe(self, prompt):
-        iteration = 1
-        while iteration < self.MAX_ITERATION_BFE:
-            try:
+        # Inject the task prompt once before the loop
         self.messages.append({"role": "user", "content": prompt})
         self.messages.append({"role": "user", "content": prompt})
 
 
+        for iteration in range(1, self.MAX_ITERATION_BFE + 1):
+            try:
                 response = self._call_llm(self.messages)
                 response = self._call_llm(self.messages)
                 tool_calls = response.tool_calls
                 tool_calls = response.tool_calls
 
 
                 if tool_calls:
                 if tool_calls:
-                    self.logger.info(f"Length of tool calls: {len(tool_calls)}")
+                    self.logger.info(f"BFE iteration {iteration}: {len(tool_calls)} tool call(s)")
                     self.messages.append(response)
                     self.messages.append(response)
                     for tool_call in tool_calls:
                     for tool_call in tool_calls:
                         exec_result = self._process_tool_call(tool_call)
                         exec_result = self._process_tool_call(tool_call)
                         self._process_tool_results_bfe(tool_call.function.name, exec_result)
                         self._process_tool_results_bfe(tool_call.function.name, exec_result)
-
                 else:
                 else:
-                    assistant_message = {"role": "assistant", "content": response.content}
-                    self.messages.append(assistant_message)
+                    self.messages.append({"role": "assistant", "content": response.content})
+                    if response.content:
+                        print(f"\nAgent: {response.content}")
+                    user_input = input("You (press Enter to finish): ").strip()
+                    if not user_input:
+                        self.logger.info(f"BFE agent finished after {iteration} iteration(s).")
+                        break
+                    self.messages.append({"role": "user", "content": user_input})
 
 
-                iteration += 1
-                self.logger.info(f"Logging agent iteration {iteration}")
             except Exception as e:
             except Exception as e:
                 self.logger.error(str(e))
                 self.logger.error(str(e))
                 raise
                 raise
@@ -271,7 +277,7 @@ class MDAgent(BaseAgent):
     def _format_expected_files(self, templates):
     def _format_expected_files(self, templates):
         values = {
         values = {
             "pdb_id": self.pdb_id.upper(),
             "pdb_id": self.pdb_id.upper(),
-            "ligand_name": self.ligand_name.upper(),
+            "ligand_name": self.ligand_name.upper() if self.ligand_name else "",
         }
         }
 
 
         formatted = []
         formatted = []

src/agents/prep_agent.py

@@ -68,34 +68,88 @@ class PrepAgent(BaseAgent):
 
 
         self.messages.append(system_prompt)
         self.messages.append(system_prompt)
 
 
+    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).",
+        })
+        response = self._call_llm(self.messages)
+        self.messages.append({"role": "assistant", "content": response.content})
+        print(f"\nAgent: {response.content}")
+
+        while True:
+            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",
+                    "content": (
+                        f"Extract the PDB ID (4-character alphanumeric code) and ligand ID "
+                        f"(3-character alphanumeric code, if present) from this user response:\n\n"
+                        f"\"{user_input}\"\n\n"
+                        f"Reply with only valid JSON in this exact format: "
+                        f'{{\"pdb_id\": \"XXXX\", \"ligand\": \"XXX\"}} or '
+                        f'{{\"pdb_id\": \"XXXX\", \"ligand\": null}} if no ligand was mentioned.'
+                    ),
+                }
+            ]
+            parse_response = self._call_llm(parse_messages)
+
+            try:
+                raw = parse_response.content.strip()
+                raw = re.sub(r'^```(?:json)?\s*|\s*```$', '', raw, flags=re.MULTILINE).strip()
+                extracted = json.loads(raw)
+                pdb_id = (extracted.get("pdb_id") or "").strip().upper() or None
+                ligand = (extracted.get("ligand") or "").strip().upper() or None
+            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)")
+
+            confirmation = input("You: ").strip().lower()
+            if confirmation in ("yes", "y"):
+                break
+
+            print("Agent: No problem, please provide the PDB ID and ligand again.")
+
+        self.pdb_id = pdb_id
+        if self.ligand_name is None:
+            self.ligand_name = ligand
+
     def _get_pdb_file_path(self, prompt):
     def _get_pdb_file_path(self, prompt):
-        pdb_file_path = None
-        num_calls = 0
+        # Inject the task prompt once before the loop
+        self.messages.append({"role": "user", "content": prompt})
 
 
-        while pdb_file_path is None and num_calls < 5:
-            response = self._prompt_llm(prompt)
-            num_calls += 1
+        for _ in range(5):
+            response = self._call_llm(self.messages)
             self.logger.info(f"Response: {response.content}")
             self.logger.info(f"Response: {response.content}")
 
 
             tool_calls = response.tool_calls
             tool_calls = response.tool_calls
             if tool_calls:
             if tool_calls:
                 self.logger.info(f"Length of tool calls: {len(tool_calls)}")
                 self.logger.info(f"Length of tool calls: {len(tool_calls)}")
                 self.messages.append(response)
                 self.messages.append(response)
-
                 for tool_call in tool_calls:
                 for tool_call in tool_calls:
                     self._process_tool_call(tool_call)
                     self._process_tool_call(tool_call)
-
             else:
             else:
-                assistant_message = {"role": "assistant", "content": response.content}
-                self.messages.append(assistant_message)
+                self.messages.append({"role": "assistant", "content": response.content})
+                if response.content:
+                    print(f"\nAgent: {response.content}")
+                user_input = input("You: ").strip()
+                if user_input:
+                    self.messages.append({"role": "user", "content": user_input})
 
 
-            # Check if a PDB file exists in sandbox
             pdb_files = list(self.sandbox_dir.glob("*.pdb"))
             pdb_files = list(self.sandbox_dir.glob("*.pdb"))
-
             if pdb_files:
             if pdb_files:
-                pdb_file_path = str(pdb_files[0])
+                return str(pdb_files[0])
 
 
-        return pdb_file_path
+        return None
 
 
     def _find_ligand(self):
     def _find_ligand(self):
         lig_response = None
         lig_response = None
@@ -225,11 +279,18 @@ class PrepAgent(BaseAgent):
 
 
         self._setup_system_prompt()
         self._setup_system_prompt()
 
 
-        user_input = self.pdb_id
-        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}")
+        interactive = self.pdb_id is None
+        if interactive:
+            self._ask_for_system()
+
+        if interactive:
+            prompt = "Please proceed to fetch and prepare the PDB file for the system we just discussed."
+        else:
+            prompt = f"Fetch and prepare the PDB file for {self.pdb_id}."
+
+        self.logger.info(f"Starting prep for PDB: {self.pdb_id}, Ligand: {self.ligand_name or 'none'}")
         self.pdb_file_path = self._get_pdb_file_path(prompt)
         self.pdb_file_path = self._get_pdb_file_path(prompt)
-        self.logger.info(f"I now have access to the structure information for protein {self.pdb_file_path}")
+        self.logger.info(f"Structure ready: {self.pdb_file_path}")
 
 
         self._find_ligand()
         self._find_ligand()
 
 

src/prompts/default.py

@@ -3,7 +3,8 @@ PREP_SYSTEM_PROMPT = """"You are a helpful science assistant        designed to
                 Classify the user request and prepare the input files for an appropriate molecular dynamics pipeline. 
                 Classify the user request and prepare the input files for an appropriate molecular dynamics pipeline. 
                 The user will either specify a PDB ID or upload the file into {sandbox_dir}. 
                 The user will either specify a PDB ID or upload the file into {sandbox_dir}. 
                 Depending on the user inputs you should define what a sucessful MD pipeline would involve.
                 Depending on the user inputs you should define what a sucessful MD pipeline would involve.
-                Call the relevant tools when needed to prepare the system for molecular dynamics. Do not ask the user anything."""
+                Call the relevant tools when needed to prepare the system for molecular dynamics.
+                You may ask the user clarifying questions when necessary."""
 
 
 MD_SYSTEM_PROMPT = """You are an MD execution assistant. You have access to tools that prepare and
 MD_SYSTEM_PROMPT = """You are an MD execution assistant. You have access to tools that prepare and
             run molecular dynamics (MD) simulations using GROMACS.
             run molecular dynamics (MD) simulations using GROMACS.

src/scripts/equil_Gromacs.sh

@@ -34,19 +34,15 @@ else
 fi
 fi
 
 
 #------- ENERGY MINIMISATION ------------
 #------- ENERGY MINIMISATION ------------
-if ! ls em.gro 1> /dev/null 2>&1; then
-	$GMX grompp -f em.mdp -c $INPUT_GRO -p topol.top -o em.tpr >> $LOG_FILE 2>&1
-	$GMX mdrun -v -deffnm em >> $LOG_FILE 2>&1
+$GMX grompp -f em.mdp -c $INPUT_GRO -p topol.top -o em.tpr >> $LOG_FILE 2>&1
+$GMX mdrun -v -deffnm em >> $LOG_FILE 2>&1
 
 
-	if [ -f em.gro ]; then
+if [ -f em.gro ]; then
 	echo "'em.gro' created"
 	echo "'em.gro' created"
 	echo "11 0" | $GMX energy -f em.edr -o potential.xvg >> $LOG_FILE 2>&1
 	echo "11 0" | $GMX energy -f em.edr -o potential.xvg >> $LOG_FILE 2>&1
-	else
+else
 	echo "Error: Failed to create 'em.gro'" >> $LOG_FILE 2>&1
 	echo "Error: Failed to create 'em.gro'" >> $LOG_FILE 2>&1
 	exit 1
 	exit 1
-	fi
-else
-	echo "'em.gro' already exists. Skipping energy minimisation." >> $LOG_FILE 2>&1
 fi
 fi
 
 
 #----------Create posres files-----------
 #----------Create posres files-----------
@@ -261,36 +257,26 @@ if grep -q "Cl-" index.ndx || grep -q "Na+" index.ndx; then
 	fi
 	fi
 fi
 fi
 #--------------- NVT --------------------
 #--------------- NVT --------------------
+$GMX grompp -f nvt.mdp -c em.gro -r em.gro -p topol.top -o nvt.tpr -n index.ndx -maxwarn 2 >> $LOG_FILE 2>&1
+$GMX mdrun -v -deffnm nvt >> $LOG_FILE 2>&1
 
 
-if ! ls nvt.gro 1> /dev/null 2>&1; then
-
-	$GMX grompp -f nvt.mdp -c em.gro -r em.gro -p topol.top -o nvt.tpr -n index.ndx -maxwarn 2 >> $LOG_FILE 2>&1
-	$GMX mdrun -v -deffnm nvt >> $LOG_FILE 2>&1
-        
-	if [ -f nvt.gro ]; then
+if [ -f nvt.gro ]; then
 	echo "'nvt.gro' created" >> $LOG_FILE 2>&1
 	echo "'nvt.gro' created" >> $LOG_FILE 2>&1
 	echo -e "Temperature \n 0" | $GMX energy -f nvt.edr -o temperature.xvg >> $LOG_FILE 2>&1
 	echo -e "Temperature \n 0" | $GMX energy -f nvt.edr -o temperature.xvg >> $LOG_FILE 2>&1
-	else
+else
 	echo "Error: Failed to create 'nvt.gro'" >> $LOG_FILE 2>&1
 	echo "Error: Failed to create 'nvt.gro'" >> $LOG_FILE 2>&1
 	exit 1
 	exit 1
-	fi
-else
-	echo "'nvt.gro' already exists. Skipping NVT." >> $LOG_FILE 2>&1
 fi
 fi
 
 
 #--------------- NPT --------------------
 #--------------- NPT --------------------
-if ! ls npt.gro 1> /dev/null 2>&1; then
-	$GMX grompp -f npt.mdp -c nvt.gro -t nvt.cpt -r nvt.gro -p topol.top -o npt.tpr -n index.ndx -maxwarn 2 >> $LOG_FILE 2>&1
-	$GMX mdrun -v -deffnm npt >> $LOG_FILE 2>&1
+$GMX grompp -f npt.mdp -c nvt.gro -t nvt.cpt -r nvt.gro -p topol.top -o npt.tpr -n index.ndx -maxwarn 2 >> $LOG_FILE 2>&1
+$GMX mdrun -v -deffnm npt >> $LOG_FILE 2>&1
 
 
-	if [ -f npt.gro ]; then
+if [ -f npt.gro ]; then
 	echo "'npt.gro' created" >> $LOG_FILE 2>&1
 	echo "'npt.gro' created" >> $LOG_FILE 2>&1
 	echo -e "Pressure \n 0" | $GMX energy -f npt.edr -o pressure.xvg >> $LOG_FILE 2>&1
 	echo -e "Pressure \n 0" | $GMX energy -f npt.edr -o pressure.xvg >> $LOG_FILE 2>&1
 	echo -e "Density \n 0" | $GMX energy -f npt.edr -o density.xvg >> $LOG_FILE 2>&1
 	echo -e "Density \n 0" | $GMX energy -f npt.edr -o density.xvg >> $LOG_FILE 2>&1
-	else
+else
 	echo "Error: Failed to create 'npt.gro'" >> $LOG_FILE 2>&1
 	echo "Error: Failed to create 'npt.gro'" >> $LOG_FILE 2>&1
 	exit 1
 	exit 1
-	fi
-else
-	echo "'npt.gro' already exists. Skipping NPT." >> $LOG_FILE 2>&1
 fi
 fi
\ No newline at end of file

src/scripts/prod_Gromacs.sh

@@ -71,9 +71,5 @@ if grep -q "Cl-" index.ndx || grep -q "Na+" index.ndx; then
 fi
 fi
 
 
 #------- PRODUCTION MD ------------
 #------- PRODUCTION MD ------------
-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
-    echo "y" | $GMX mdrun -v -deffnm md >> $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
+echo "y" | $GMX mdrun -v -deffnm md >> $LOG_FILE 2>&1
\ No newline at end of file
-else
-    echo "'md.gro' already exists. Skipping production MD."
-fi
\ No newline at end of file