Skip to content

D
DynaMate
Commit 2184d1bf by Cassandra Masschelein Committed by GitHub
Options

[update] increase chat flow between user and agent

parents 38f281e3 4d795428
Showing with 152 additions and 99 deletions
README.md
......@@ -34,13 +34,13 @@ OPENROUTER_API_KEY=your_key_here
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):
```
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! 🧬
......@@ -189,12 +189,14 @@ source setup.sh
```
Now you are ready to use DynaMate!
## 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
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:
```
--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>
--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.
"""
pdb_id: str
"PDB ID."
model: str
"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 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"
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:
# Inject the task prompt once before the loop
remaining_steps_string = "\n".join(remaining_steps)
self.messages.append({
"role": "user",
"content": (
"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."
),
})
for iteration in range(1, self.MAX_ITERATION + 1):
try:
remaining_steps_string = "\n".join(remaining_steps)
prompt = (
"The following essential steps are remaining in the pipeline:\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})
response = self._call_llm(self.messages)
tool_calls = response.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)
for tool_call in tool_calls:
exec_result = self._process_tool_call(tool_call)
self._process_tool_results(tool_call.function.name, exec_result, remaining_steps)
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:
self.logger.error(str(e))
raise
self.messages.append({"role": "user", "content": f"An error occurred: {e}"})
return remaining_steps
def _run_bfe(self, prompt):
iteration = 1
while iteration < self.MAX_ITERATION_BFE:
try:
self.messages.append({"role": "user", "content": prompt})
# Inject the task prompt once before the loop
self.messages.append({"role": "user", "content": prompt})
for iteration in range(1, self.MAX_ITERATION_BFE + 1):
try:
response = self._call_llm(self.messages)
tool_calls = response.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)
for tool_call in tool_calls:
exec_result = self._process_tool_call(tool_call)
self._process_tool_results_bfe(tool_call.function.name, exec_result)
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:
self.logger.error(str(e))
raise
......@@ -271,7 +277,7 @@ class MDAgent(BaseAgent):
def _format_expected_files(self, templates):
values = {
"pdb_id": self.pdb_id.upper(),
"ligand_name": self.ligand_name.upper(),
"ligand_name": self.ligand_name.upper() if self.ligand_name else "",
}
formatted = []
......
src/agents/prep_agent.py
......@@ -68,34 +68,88 @@ class PrepAgent(BaseAgent):
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):
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}")
tool_calls = response.tool_calls
if tool_calls:
self.logger.info(f"Length of tool calls: {len(tool_calls)}")
self.messages.append(response)
for tool_call in tool_calls:
self._process_tool_call(tool_call)
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"))
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):
lig_response = None
......@@ -225,11 +279,18 @@ class PrepAgent(BaseAgent):
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.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()
......
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.
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.
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
run molecular dynamics (MD) simulations using GROMACS.
......
src/scripts/equil_Gromacs.sh
......@@ -34,19 +34,15 @@ else
fi
#------- 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
echo "'em.gro' created"
echo "11 0" | $GMX energy -f em.edr -o potential.xvg >> $LOG_FILE 2>&1
else
echo "Error: Failed to create 'em.gro'" >> $LOG_FILE 2>&1
exit 1
fi
if [ -f em.gro ]; then
echo "'em.gro' created"
echo "11 0" | $GMX energy -f em.edr -o potential.xvg >> $LOG_FILE 2>&1
else
echo "'em.gro' already exists. Skipping energy minimisation." >> $LOG_FILE 2>&1
echo "Error: Failed to create 'em.gro'" >> $LOG_FILE 2>&1
exit 1
fi
#----------Create posres files-----------
......@@ -261,36 +257,26 @@ if grep -q "Cl-" index.ndx || grep -q "Na+" index.ndx; then
fi
fi
#--------------- 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
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
else
echo "Error: Failed to create 'nvt.gro'" >> $LOG_FILE 2>&1
exit 1
fi
if [ -f nvt.gro ]; then
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
else
echo "'nvt.gro' already exists. Skipping NVT." >> $LOG_FILE 2>&1
echo "Error: Failed to create 'nvt.gro'" >> $LOG_FILE 2>&1
exit 1
fi
#--------------- 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
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 "Density \n 0" | $GMX energy -f npt.edr -o density.xvg >> $LOG_FILE 2>&1
else
echo "Error: Failed to create 'npt.gro'" >> $LOG_FILE 2>&1
exit 1
fi
if [ -f npt.gro ]; then
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 "Density \n 0" | $GMX energy -f npt.edr -o density.xvg >> $LOG_FILE 2>&1
else
echo "'npt.gro' already exists. Skipping NPT." >> $LOG_FILE 2>&1
echo "Error: Failed to create 'npt.gro'" >> $LOG_FILE 2>&1
exit 1
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
#------- 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
else
echo "'md.gro' already exists. Skipping production MD."
fi
\ No newline at end of file
$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
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment