Prompting improvements

charge_backend/moleculedb/molecule_naming.py

@@ -19,13 +19,16 @@
 import json
 import os
 import requests
-import pandas as pd
+import functools
 
+from charge_backend.moleculedb.purchasable import moldb_connect
 from typing import Literal, TypeAlias
 
 MolNameFormat: TypeAlias = Literal["brand", "iupac", "formula", "smiles"]
 
+_MOLECULE_DATABASE_PATH = os.getenv("FLASK_MOLECULE_DB", "/data/db/molecules.db")
 _DATABASE_PATH = os.getenv("FLASK_INCHI_DB", "/data/inchi_mapping.json")
+MOLDB_CONNECTION = None
 if os.path.exists(_DATABASE_PATH):
     with open(_DATABASE_PATH, "rb") as fp:
         DATABASE = json.load(fp)
@@ -54,6 +57,41 @@ def inchi_lookup(inchi: str, prefer_iupac: bool = False) -> str | None:
     return None
 
 
+@functools.cache
+def ai_preferred_name_lookup(inchi):
+    global MOLDB_CONNECTION
+    if MOLDB_CONNECTION is None:
+        if os.path.exists(_MOLECULE_DATABASE_PATH):
+            MOLDB_CONNECTION = moldb_connect(_MOLECULE_DATABASE_PATH)
+    if MOLDB_CONNECTION is None:
+        return None
+    cursor = MOLDB_CONNECTION.cursor()
+
+    row = cursor.execute(
+        """
+        SELECT n.ai_preferred_name
+        FROM names n
+        JOIN molecules m ON n.molecule_id = m.molecule_id
+        WHERE m.key = ?
+        """,
+        (inchi,),
+    ).fetchone()
+
+    return row[0].strip() if row and row[0] and row[0].strip() else None
+
+
+def smiles_preferred_name(smiles):
+    if Chem is None:
+        return None
+
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:  # Invalid SMILES
+        return None
+
+    inchi = str(Chem.MolToInchi(mol))
+    return ai_preferred_name_lookup(inchi)
+
+
 def smiles_to_html(
     smiles: str,
     molecule_name_format: MolNameFormat = "brand",
@@ -67,9 +105,16 @@ def smiles_to_html(
     if mol is None:  # Invalid SMILES
         return smiles
 
-    # First, try to find a canonical or IUPAC name
+    # First, try to find a canonical, IUPAC name
     if molecule_name_format in ("brand", "iupac"):
         inchi = str(Chem.MolToInchi(mol))
+
+        # default to AI preferred name
+        if molecule_name_format == "brand":
+          name = ai_preferred_name_lookup(inchi)
+          if name:
+              return name
+
         name = inchi_lookup(inchi, molecule_name_format == "iupac")
         if name:
             return name

charge_backend/retrosynthesis/ai.py

@@ -17,6 +17,7 @@
 from charge_backend.flask_experiment import FlaskExperiment, GraphContext
 from charge_backend.moleculedb.molecule_naming import (
     smiles_to_html,
+    smiles_preferred_name,
 )
 from charge_backend.retrosynthesis.template import (
     generate_nodes_for_molecular_graph,
@@ -25,6 +26,9 @@
 from charge_backend.moleculedb.purchasable import is_purchasable
 from charge_backend.retrosynthesis.mapping import build_mapped_reaction_dict_or_none
 from charge_backend.retrosynthesis.database import find_exact_reactions
+from charge_backend.retrosynthesis.functional_groups import (
+    functional_groups_from_smiles,
+)
 
 from charge_backend.retrosynthesis.retrosynthesis_task import (
     TemplateFreeRetrosynthesisTask as RetrosynthesisTask,
@@ -33,29 +37,38 @@
 
 from lc_conductor import ToolRuntime
 
+RETROSYNTH_PROMPT_TEMPLATE = """You are a chemistry retrosynthesis assistant. Perform single-step retrosynthesis only.
 
-RETROSYNTH_UNCONSTRAINED_USER_PROMPT_TEMPLATE = (
-    "Provide a retrosynthetic pathway for the target molecule `{target_molecule}`. "
-    + "If there are `*`, the `*` indicate the boundaries of the polymer repeat unit."
-    + "The pathway should be provided as a tuple of reactants as SMILES and the product as SMILES. "
-    + "Perform only single step retrosynthesis. Make sure the SMILES strings are valid. "
-    + "Use tools to verify the SMILES strings and diagnose any issues that arise."
-    + "Do the evaluation step-by-step. Propose a retrosynthetic step, then evaluate it. "
-    + "If the evaluation fails, propose a new retrosynthetic step and evaluate it again. "
-    + "Find the best possible retrosynthetic step, and use tools to see if the "
-    + "proposed reactants are synthesizable. "
-)
+Target:
+- Preferred Name: {preferred_name}
+- SMILES: `{smiles}`
+- Functional Groups: {fgs}
+- Polymer rule: if `*` appears, it marks polymer repeat-unit boundaries
 
-RETROSYNTH_CONSTRAINED_USER_PROMPT_TEMPLATE = (
-    "Provide a retrosynthetic pathway for the target molecule `{target_molecule}`. "
-    + "If there are `*`, the `*` indicate the boundaries of the polymer repeat unit."
-    + "The pathway should be provided as a tuple of reactants as SMILES and the product as SMILES. "
-    + "Perform only single step retrosynthesis. Make sure the SMILES strings are valid. "
-    + "Use tools to verify the SMILES strings and diagnose any issues that arise. "
-    + "The following reactant cannot be used in the retrosynthetic step: {constrained_reactant}. "
-    + "Do the evaluation step-by-step. Propose a retrosynthetic step, then evaluate it. "
-    + "If the evaluation fails, propose a new retrosynthetic step and evaluate it again. "
-)
+Task:
+Find the best one-step retrosynthetic path to the target. Use available tools to verify each candidate; if a required tool is unavailable, perform the
+same check by chemical reasoning.
+
+Requirements:
+1. Identify the key bond formation, functional group transformation, or disconnection that most directly explains the target.
+2. Propose candidate reactants for a single retrosynthetic step.
+3. Verify that each proposed reactant SMILES is syntactically valid.
+4. Check whether the proposed reactants are chemically plausible and reasonably synthesizable.
+5. Evaluate the implied forward reaction. The reactants should regenerate the target in one step without adding, deleting, or rearranging unrelated atoms.
+6. If `predict_reaction_products` is available, use it to predict products from the proposed reactants, then canonicalize and compare the predicted
+product with the target. If there is any inconsistency log it and try some other set of reactants.
+7. If prediction tools are unavailable, perform the same forward-product equivalence check by chemical reasoning.
+8. If a candidate fails validation or there is any inconsistency, diagnose the issue, log it, and try another candidate.
+9. Choose the best validated step.
+10. Return the selected reactants and regenerated product as SMILES.
+
+Ranking criteria:
+- exact or near-exact forward-product equivalence to target
+- chemical plausibility
+- reactants are buyable, or can be reduced to buyable precursors in few plausible steps
+- one-step feasibility
+
+"""
 
 
 async def ai_based_retrosynthesis(
@@ -106,17 +119,21 @@ async def ai_based_retrosynthesis(
         callback_handler.agent_key = agent_key
         callback_handler.on_agent_update = history_callback
 
+    preferred_name = smiles_preferred_name(current_node.smiles)
+    functional_groups = functional_groups_from_smiles(current_node.smiles)
+
+    user_prompt = RETROSYNTH_PROMPT_TEMPLATE.format(
+        preferred_name=preferred_name,
+        smiles=current_node.smiles,
+        fgs=", ".join(functional_groups) if functional_groups else "None",
+    )
+
     if constraint:
-        user_prompt = RETROSYNTH_CONSTRAINED_USER_PROMPT_TEMPLATE.format(
-            target_molecule=current_node.smiles,
-            constrained_reactant=constraint,
-        )
-    else:
-        user_prompt = RETROSYNTH_UNCONSTRAINED_USER_PROMPT_TEMPLATE.format(
-            target_molecule=current_node.smiles
+        user_prompt += (
+            f"\n\nConstraint - The following reactants cannot be used in the retrosynthetic step: "
+            f"{constraint}."
         )
 
-    user_prompt += "\nDouble check the reactants with the `predict_reaction_products` tool to see if the products are equivalent to the given product. If there is any inconsistency (canonicalize both sides of the equation first), log it and try some other set of reactants."
     if query is not None:
         user_prompt += (
             f"\n\nAdditionally, adhere to the following requirements:\n{query}\n\n"

charge_backend/retrosynthesis/functional_groups.py

@@ -0,0 +1,65 @@
+try:
+    from rdkit import Chem
+except ImportError:
+    Chem = None
+
+FUNCTIONAL_GROUP_SMARTS = {
+    # Carbonyl compounds
+    "aldehyde": "[CX3H1](=O)[#6]",
+    "ketone": "[#6][CX3](=O)[#6]",
+    "carboxylic acid": "C(=O)[OX2H1]",
+    "ester": "C(=O)O[#6]",
+    "amide": "C(=O)N",
+    "acyl halide": "C(=O)[F,Cl,Br,I]",
+    "anhydride": "C(=O)OC(=O)",
+    # Alcohols / ethers
+    "alcohol": "[OX2H][CX4]",
+    "phenol": "[OX2H][c]",
+    "ether": "[#6]-O-[#6]",
+    "epoxide": "[OX2r3]1CC1",
+    # Nitrogen
+    "amine": "[NX3;H2,H1,H0;!$(NC=O)]",
+    "imine": "[CX3]=[NX2]",
+    "nitrile": "C#N",
+    "nitro": "[N+](=O)[O-]",
+    "azo": "[N]=[N]",
+    "isocyanate": "N=C=O",
+    "isothiocyanate": "N=C=S",
+    "urea": "N-C(=O)-N",
+    "carbamate": "O-C(=O)-N",
+    # Sulfur
+    "thiol": "[SX2H]",
+    "thioether": "[#6]-S-[#6]",
+    "sulfoxide": "[SX3](=O)",
+    "sulfone": "[SX4](=O)(=O)",
+    "sulfonamide": "S(=O)(=O)N",
+    "sulfonic acid": "S(=O)(=O)[OX2H]",
+    # Phosphorus
+    "phosphate": "P(=O)(O)(O)",
+    # Unsaturation
+    "alkene": "C=C",
+    "alkyne": "C#C",
+    # Rings
+    "aromatic ring": "a1aaaaa1",
+    # Halogens
+    "organohalide": "[#6][F,Cl,Br,I]",
+    # Hydrocarbon
+    "alkyl substituent": "[CX4][a]",
+}
+
+
+def functional_groups_from_smiles(smiles):
+    """Return simple SMARTS-matched functional group names."""
+    if Chem is None:
+        return []
+
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return []
+
+    groups = []
+    for name, smarts in FUNCTIONAL_GROUP_SMARTS.items():
+        pattern = Chem.MolFromSmarts(smarts)
+        if pattern is not None and mol.HasSubstructMatch(pattern):
+            groups.append(name)
+    return groups