Graph rxn mapper

.gitignore

@@ -1,3 +1,10 @@
+flask_tools/pipette/_java_build/
+.podman-run/
+.podman-root/
+.podman-runroot/
+.podman-tmp/
+.rdt-podman.env
+
 # Logs
 logs
 *.log

flask_tools/pipette/README.md

@@ -8,6 +8,9 @@ The current pipeline includes:
 - reaction SMILES parsing
 - basic SMILES validation
 - exact-match checker interfaces for reaction databases
+- graph based balancing - Attempt to balance reaction by adding copies of reactants. Good for dimerization reactions
+  - If this is enabled, the reaction fixing LLM is called here instead of later
+- atom mapping
 - reaction fixing LLM call
   - Run if no exact match is found
   - If new reaction is returned, goes back to start. Only allowed to run once in a pipeline
@@ -150,20 +153,14 @@ The tools can also be disabled by setting `tool_list: null` in the the config.
           "skipped_reason": null
         },
         {
-          "name": "llm_reaction_fix",
+          "name": "RDTAtomMapper",
           "status": "pass",
           "data": {
-            "original_reaction_smiles": "Cn1cnc2c1c(=O)[nH]c(=O)n2C.CI>>CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
-            "fixed_reaction_smiles": "CI.Cn1cnc2c1c(=O)[nH]c(=O)n2C>>Cn1c(=O)c2c(ncn2C)n(C)c1=O.[H+].[I-]",
-            "removed_agents": [],
-            "added_reactants": [],
-            "added_products": [
-              "[H+]",
-              "[I-]"
-            ]
+            "input_reaction_smiles": "CI.Cn1cnc2c1c(=O)[nH]c(=O)n2C>>Cn1c(=O)c2c(ncn2C)n(C)c1=O.I",
+            "mapped_reaction_smiles": "[O:10]=[C:9]1[NH:11][C:12](=[O:13])[N:14]([C:7]=2[N:6]=[CH:5][N:4]([C:8]12)[CH3:3])[CH3:15].[I:2][CH3:1]>>[N:14]1([C:7]=2[N:6]=[CH:5][N:4]([CH3:3])[C:8]2[C:9]([N:11]([CH3:1])[C:12]1=[O:13])=[O:10])[CH3:15].[IH:2]",
+            "product_to_reactant": []
           },
-          "comment": "N-methylation of the xanthine NH with methyl iodide requires HI as the byproduct, represented as [H+] and [I-]. No agents were present to remove.",
-          "skipped_reason": null
+          "comment": "RDT atom mapping completed."
         },
         {
           "name": "basic_smiles_validation",
@@ -236,3 +233,46 @@ config = PipetteConfig.from_yaml("my-config.yaml")
 `pytest`
 Or
 `pytest -m llm_query` to run the tests that use LLM
+
+# ReactionDecoder / RDT
+
+`pipette` includes a Python wrapper around the Java-based [ReactionDecoder Tool](https://github.com/asad/ReactionDecoder) (RDT). There is accomplished with a short Java wrapper script that calls ReactionDecoder, and gets called by the Python wrapper script.
+
+The wrapper uses `RDT`'s built-in defaults for mapping options.
+
+To compile the Java wrapper script, you must first build a fat jar of RDT.
+
+You must have Java 25 installed. Check your existing java with `java --version`. The `install_rdt.sh` script will install Java and Maven, and then compile the RDT java wrapper.
+
+```bash
+./scripts/install_rdt.sh
+export PIPETTE_RDT_JAR=/absolute/path/to/ReactionDecoder/target/rdt-4.0.0-jar-with-dependencies.jar
+export PIPETTE_RDT_HELPER_BUILD_DIR=/absolute/path/to/flask-tools/flask_tools/pipette/_java_build
+```
+
+Use it from Python:
+
+```python
+from flask_tools.pipette.verifiers.rdt import (
+  map_reaction_smiles_with_rdt,
+  map_reaction_smiles_list_with_rdt,
+)
+
+mapped = map_reaction_smiles_with_rdt("CC(=O)O.OCC>>CC(=O)OCC.O")
+mapped_many = map_reaction_smiles_list_with_rdt(
+  [
+    "CC(=O)O.OCC>>CC(=O)OCC.O",
+    "CCO>>CC=O",
+  ]
+)
+```
+
+Or from the CLI:
+
+```bash
+pipette-rdt --rxn-smi 'CC(=O)O.OCC>>CC(=O)OCC.O'
+# Outputs [O:3]=[C:2]([OH:4])[CH3:1].[OH:5][CH2:6][CH3:7]>>[O:5]([C:2]([CH3:1])=[O:4])[CH2:6][CH3:7].[OH2:3]
+pipette-rdt --file reactions.txt --json
+```
+
+If the input uses `reactants>agents>products`, the wrapper strips agents for RDT, maps the core reaction, and then reinserts the original agents into the returned reaction SMILES.

flask_tools/pipette/__init__.py

@@ -5,7 +5,6 @@
 ## SPDX-License-Identifier: Apache-2.0
 ###############################################################################
 
-from .grade_rxn import grade_reaction
 from .constants import FinalGrade, ReactionGrade, ToolResult, ToolStatus
 
 __all__ = [
@@ -15,3 +14,11 @@
     "ToolStatus",
     "grade_reaction",
 ]
+
+
+def __getattr__(name: str):
+    if name == "grade_reaction":
+        from .grade_rxn import grade_reaction
+
+        return grade_reaction
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

flask_tools/pipette/assets/ai_judge_no_dft.yaml

@@ -1,4 +1,5 @@
 llm_judge:
+  enable_atom_mapping_dict_in_prompt: false
   allow_fail:
     - exact_match
 settings:

flask_tools/pipette/assets/ai_judge_no_dft_atom_mapper.yaml

@@ -0,0 +1,14 @@
+llm_judge:
+  allow_fail:
+    - exact_match
+tool_list:
+  - "basic_smiles_validation"
+  - "exact_match"
+  - "charge_conservation"
+  - "mass_conservation"
+  - "reaction_energy"
+settings:
+  stop_on_hard_fail: true
+  mass_tolerance_atoms: 0
+  reaction_energy_max_ev_mol: 0.2
+  use_dft: false

flask_tools/pipette/assets/ai_judge_with_dft.yaml

@@ -1,4 +1,9 @@
-tool_list: all
+tool_list:
+  - "basic_smiles_validation"
+  - "exact_match"
+  - "charge_conservation"
+  - "mass_conservation"
+  - "reaction_energy"
 tools_settings:
   reaction_energy:
     database: null

flask_tools/pipette/config.py

@@ -27,6 +27,11 @@ def package_config_path(filename: str) -> Path:
     return Path(__file__).with_name("assets") / filename
 
 
+def graph_rxn_mapper_prompt_path(filename: str) -> Path:
+    # Relative to top level of pipette module / graph_rxn_mapper / prmopts
+    return Path(__file__).with_name("graph_rxn_mapper") / "prompts" / filename
+
+
 def _validate_mapping_format(data: object, *, name: str) -> dict[str, Any]:
     if data is None:
         return {}
@@ -47,6 +52,23 @@ def _resolve_optional_path(path_value: object, *, base_dir: Path) -> Path | None
     return candidate
 
 
+def _resolve_defaultable_cwd_path(
+    path_value: object,
+    *,
+    default_path: Path,
+    name: str,
+) -> Path:
+    if path_value is None or path_value == "default":
+        return default_path
+    if not isinstance(path_value, str):
+        raise ValueError(f"{name} must be a string, 'default', or null.")
+
+    candidate = Path(path_value).expanduser()
+    if not candidate.is_absolute():
+        candidate = (Path.cwd() / candidate).resolve()
+    return candidate
+
+
 @dataclass
 class PipelineConfig:
     stop_on_hard_fail: bool = True
@@ -136,6 +158,7 @@ def _llm_kwargs_from_mapping(
 @dataclass
 class LLMJudgeConfig(LLMConfig):
     allow_fail: Literal["all"] | list[str] = field(default_factory=list)
+    enable_atom_mapping_dict_in_prompt: bool = False
     prompt_path: Path = field(
         default_factory=lambda: package_config_path("judge-prompt.txt")
     )
@@ -156,8 +179,17 @@ def from_mapping(
                 raise ValueError(
                     "llm_judge.allow_fail must be 'all' or a list of tool names."
                 )
+        enable_atom_mapping_dict_in_prompt = mapping.get(
+            "enable_atom_mapping_dict_in_prompt",
+            cls.enable_atom_mapping_dict_in_prompt,
+        )
+        if not isinstance(enable_atom_mapping_dict_in_prompt, bool):
+            raise ValueError(
+                "llm_judge.enable_atom_mapping_dict_in_prompt must be a boolean."
+            )
         return cls(
             allow_fail=allow_fail if allow_fail == "all" else list(allow_fail),
+            enable_atom_mapping_dict_in_prompt=enable_atom_mapping_dict_in_prompt,
             **cls._llm_kwargs_from_mapping(
                 mapping,
                 name="llm_judge",
@@ -214,9 +246,81 @@ def from_mapping(
         )
 
 
+@dataclass
+class LLMAtomMappingConfig:
+    url: str = DEFAULT_LLM_BASE_URL
+    model: str = "gpt-5.4"
+    reasoning_effort: ReasoningEffort = "medium"
+    api_key: str | None = None
+    system_prompt_path: Path = field(
+        default_factory=lambda: graph_rxn_mapper_prompt_path("atom_mapping_system.md")
+    )
+    user_prompt_path: Path = field(
+        default_factory=lambda: graph_rxn_mapper_prompt_path("atom_mapping_user.md")
+    )
+    skill_prompt_path: Path = field(
+        default_factory=lambda: graph_rxn_mapper_prompt_path("atom_mapping_skill.md")
+    )
+
+    @classmethod
+    def from_mapping(
+        cls,
+        data: object,
+        *,
+        base_dir: Path,
+    ) -> LLMAtomMappingConfig:
+        mapping = _validate_mapping_format(data, name="tools_settings.llm_atom_mapping")
+        del base_dir
+
+        url = mapping.get("url")
+        if url is not None and not isinstance(url, str):
+            raise ValueError(
+                "tools_settings.llm_atom_mapping.url must be a string when provided."
+            )
+
+        model = mapping.get("model", cls.model)
+        if not isinstance(model, str):
+            raise ValueError("tools_settings.llm_atom_mapping.model must be a string.")
+
+        reasoning_effort = mapping.get("reasoning_effort", cls.reasoning_effort)
+        if reasoning_effort not in {"low", "medium", "high"}:
+            raise ValueError(
+                "tools_settings.llm_atom_mapping.reasoning_effort must be 'low', 'medium', or 'high'."
+            )
+
+        api_key = mapping.get("api_key")
+        if api_key is not None and not isinstance(api_key, str):
+            raise ValueError(
+                "tools_settings.llm_atom_mapping.api_key must be a string when provided."
+            )
+
+        return cls(
+            url=resolve_llm_base_url(url),
+            model=model,
+            reasoning_effort=reasoning_effort,
+            api_key=api_key,
+            system_prompt_path=_resolve_defaultable_cwd_path(
+                mapping.get("system_prompt_path"),
+                default_path=graph_rxn_mapper_prompt_path("atom_mapping_system.md"),
+                name="tools_settings.llm_atom_mapping.system_prompt_path",
+            ),
+            user_prompt_path=_resolve_defaultable_cwd_path(
+                mapping.get("user_prompt_path"),
+                default_path=graph_rxn_mapper_prompt_path("atom_mapping_user.md"),
+                name="tools_settings.llm_atom_mapping.user_prompt_path",
+            ),
+            skill_prompt_path=_resolve_defaultable_cwd_path(
+                mapping.get("skill_prompt_path"),
+                default_path=graph_rxn_mapper_prompt_path("atom_mapping_skill.md"),
+                name="tools_settings.llm_atom_mapping.skill_prompt_path",
+            ),
+        )
+
+
 @dataclass
 class ToolsConfig:
     reaction_energy: ReactionEnergyConfig = field(default_factory=ReactionEnergyConfig)
+    llm_atom_mapping: LLMAtomMappingConfig = field(default_factory=LLMAtomMappingConfig)
 
     @classmethod
     def from_mapping(
@@ -226,11 +330,16 @@ def from_mapping(
         base_dir: Path,
     ) -> ToolsConfig:
         mapping = _validate_mapping_format(data, name="tools_settings")
+
         return cls(
             reaction_energy=ReactionEnergyConfig.from_mapping(
                 mapping.get("reaction_energy"),
                 base_dir=base_dir,
-            )
+            ),
+            llm_atom_mapping=LLMAtomMappingConfig.from_mapping(
+                mapping.get("llm_atom_mapping"),
+                base_dir=base_dir,
+            ),
         )
 
 

flask_tools/pipette/constants.py

@@ -52,6 +52,22 @@ def resolve_llm_base_url(explicit_base_url: str | None = None) -> str:
 DEFAULT_LLM_BASE_URL = resolve_llm_base_url()
 
 
+class SmilesContainer(str):
+    """A class so most tools can assume rxn_smiles passed to tool.run() is a plain string, but some can play around
+    with the presence of the reagents
+    """
+
+    def __new__(cls, value, original_smiles=None):
+        instance = super().__new__(cls, value)
+        return instance
+
+    def __init__(self, value, reagents_smi: str | None = None):
+        self.reagents_smi = reagents_smi
+
+    def __repr__(self):
+        return f"SmilesContainer({str.__repr__(self)}, reagents_smiles={self.reagents_smi!r})"
+
+
 class ToolStatus(str, Enum):
     PASS = "pass"  # Reaction passed this tool
     FAIL = "fail"  # Reaction failed to pass this tool
@@ -72,7 +88,7 @@ class ToolResult(BaseModel):
     status: ToolStatus
     data: (
         SerializeAsAny[ToolResultDetails] | None
-    )  # None if tool had an error or wasn't run. SerializeAsAny or else model_dump only outputs the parent class ToolResultDetails' fields which are nothing.
+    )  # None if tool had an error or wasn't run. Must use SerializeAsAny or else model_dump only outputs the parent class ToolResultDetails' fields which are no fields.
     comment: str = ""
     skipped_reason: str | None = (
         None  # If a priority checker skipped this tool, like in an exact rule pipeline, or a traceback if there was an error

flask_tools/pipette/grade_rxn.py

@@ -32,12 +32,13 @@
 from typing import TYPE_CHECKING
 
 from flask_tools.pipette.config import PipetteConfig, load_config, ConfigType
-from flask_tools.pipette.constants import ReactionGrade, ToolResult
+from flask_tools.pipette.constants import ReactionGrade
 from flask_tools.pipette.pipeline import build_default_pipeline
 from flask_tools.pipette.reaction_fixer import ReactionFixResultDetails
 
 if TYPE_CHECKING:
     from .judge import AsyncLLMJudge
+    from flask_tools.pipette.constants import ToolResult
 
 REACTION_SMILES_COLUMNS = (
     "rxn_smiles",
@@ -52,8 +53,9 @@ def _get_possible_fixed_rxn_smi(reaction_grade: ReactionGrade) -> str | None:
     tool_res: ToolResult
     for tool_res in reaction_grade.results:
         if tool_res.name == "llm_reaction_fix":
-            d: ReactionFixResultDetails = tool_res.data  # noqa
-            return d.fixed_reaction_smiles
+            d: ReactionFixResultDetails | None = tool_res.data
+            if d:
+                return d.fixed_reaction_smiles
     return None
 
 
@@ -78,7 +80,7 @@ def _build_output_records(
         {
             "rxn_smiles": rxn_smiles,
             "cleaned_rxn_smiles": _get_possible_fixed_rxn_smi(result) or rxn_smiles,
-            "grade": result.model_dump(mode="json"),
+            "grade": result.model_dump(mode="json", exclude_none=True),
         }
         for rxn_smiles, result in zip(rxn_smiles_list, results, strict=True)
     ]
@@ -187,6 +189,7 @@ def main() -> list[dict]:
         f"or '{ConfigType.LLM_JUDGE_NO_DFT}.",
     )
     parser.add_argument(
+        "-v",
         "--verbose",
         action="store_true",
         help="Prints out json object",