Feature/docking tool - #152
Conversation
Adds run_docking @tool + construct_docking_graph (mirrors the gRASPA pattern): dock a candidate (SMILES/name/PubChem CID) into a receptor and return binding affinity + poses. Bundled vancomycin target; custom receptor via .pdbqt + box. - schemas/docking_schema.py; tools/docking_core.py (+mock_docking, lazy deps); tools/docking_tools.py; graphs/docking_agent.py; tests/test_docking_tools.py - 'docking' optional extra (meeko); Vina via conda-forge, lazy-imported + test-skipped - ruff clean; full suite 289 passed; live Vina dock verified with the extra
Wire construct_docking_graph into workflow_map (llm_agent.py) and the CLI ALL_WORKFLOW_TYPES so the tool is selectable via '-w docking'; document it in configuration_with_toml.md.
There was a problem hiding this comment.
Can you check if we need a separate graph for this? Looks to me that we can use the single_agent graph, and bind the new tools to it. Feel free to add new prompts under prompts/ as you see needed, as the default prompt for single_agent may not work for your application.
Then, under example/, please create a new example to run ChemGraph with these added tools/prompts. Best if you can provide actual prompt examples to showcase what problems these added tools can solve.
There was a problem hiding this comment.
Good call, you're right, a separate graph wasn't needed. I've removed graphs/docking_agent.py and reverted the workflow wiring (llm_agent.py workflow_map, cli/commands.py, and the docs). The tool now binds to the existing single_agent graph:
ChemGraph(workflow_type="single_agent", system_prompt=docking_prompt, tools=[run_docking])
I added a docking-specific prompt at prompt/docking_prompt.py (since the default single_agent prompt doesn't cover docking), and a runnable example at examples/docking/run_chemgraph.py + README.md that docks a candidate into a receptor and reports the best affinity.
There was a problem hiding this comment.
Is this a force field/parameter file?
There was a problem hiding this comment.
No, it's a prepared receptor structure, not a force-field/parameter file. It's PDBQT format: atomic coordinates plus Gasteiger partial charges and AutoDock atom types, which is what Vina uses as the rigid receptor. This one is vancomycin (chain A of RCSB PDB 1FVM). Since it's illustrative rather than a core default, I've moved it out of the package into examples/docking/ and documented its provenance in the example README. Users point the tool at their own receptor (a .pdbqt, or a SMILES/name/CID).
| # (validated by redocking the native D-Ala-D-Ala ligand from PDB 1FVM). | ||
| _FILES = Path(__file__).parent / "files" / "docking" | ||
| _VANCOMYCIN_RECEPTOR = _FILES / "vancomycin_receptor.pdbqt" | ||
| _VANCOMYCIN_CENTER = [-3.436, 5.510, 22.100] |
There was a problem hiding this comment.
Are these parameters (_CENTER, _SIZE) fixed for different workflows? Is it worth considering them as parameters inside the core tool?
There was a problem hiding this comment.
Agreed, they shouldn't be hardcoded. I removed the fixed _CENTER/_SIZE constants. The search box is now auto-detected (site_detection="auto" tries reference-ligand → fpocket → blind), and center, box_size, and reference_ligand are optional parameters in the input schema so the user can override detection when they know the pocket. The only remaining constant is _BOX_PADDING (Å added around detected atoms), which is easy to promote to a parameter too if you'd prefer.
| ] | ||
| docking = [ | ||
| "meeko", | ||
| # AutoDock Vina is required at runtime but is NOT pip-installable in practice |
There was a problem hiding this comment.
Can you also add this as an error message to where AutoDock Vina is imported? E.g If user run the tools but the environment using missing AutoDock Vina, the error message will be returned like "AuroDock Vina is not installed. Please install it via ...".
There was a problem hiding this comment.
Done. AutoDock Vina is now imported lazily inside run_docking_core, wrapped so a missing install raises a clear, actionable message:
"AutoDock Vina is not installed. Install it with: conda install -c conda-forge vina
(it is not pip-installable; the docking extra installs Meeko but not Vina.)"
|
Thanks for the review. Pushed a commit addressing all four points: (1) No separate graph : removed docking_agent.py and reverted the workflow wiring; the tool now binds to single_agent via ChemGraph(workflow_type="single_agent", system_prompt=docking_prompt, tools=[run_docking]). Added prompt/docking_prompt.py and a runnable examples/docking/ (script + README). |
2 participants
Summary
Adds a molecular docking tool (AutoDock Vina) so an agent can estimate how strongly
a small molecule binds a target and return its best pose. Introduces
run_docking(
@tool) plus a dedicatedconstruct_docking_graph, mirroring the existing gRASPAtool/graph pattern, and registers a
dockingworkflow (-w docking).The candidate may be a SMILES, a molecule name, or a PubChem CID (resolved
automatically, reusing
molecule_name_to_smiles_core+pubchempy+rdkit). Thereceptor defaults to a bundled vancomycin target, or accepts a prepared rigid
receptor
.pdbqtplus a search box. Motivated by the Intro-to-HPC bootcamp(Project 2: vancomycin / D-Ala-D-Ala resistance).
Dependencies: the
dockingextra installsmeeko(pip). AutoDock Vina isrequired at runtime but is installed from conda-forge (
conda install -c conda-forge vina), since the PyPIvinabuilds from source. Vina is lazy-imported with a clearerror, and the Vina test skips when it's absent — so core install and CI are
unaffected (same spirit as the
tblitehandling).Related issues
N/A
Type of change
How was this tested?
ruff check .— clean.python -m pytest tests/test_docking_tools.py— hermetic tests (SMILESresolution,
mock_docking, graph build) pass; a live AutoDock Vina dock into thebundled vancomycin receptor runs with the
dockingextra + conda Vina(4 passed), and skips otherwise (3 passed / 1 skipped, matching CI).
pytest tests/ -k "not tblite"— 290 passed, 27 skipped (the workflow-parametrized tests now also cover
docking).confirmed
engine="vina"and per-pose scores in the written poses PDBQT.Checklist
mainand targetsmainruff check .passespytest tests/ -k "not tblite"passes (plus extras tests if Academy/backends touched)