mondalsou/lead_optimization_agent

1"""2Lead Optimization Agent — Utility Functions3============================================4Thin wrappers around RDKit, exposing ADMET analysis as tool-callable5functions for the Claude agent. Fully local — no external API required.6"""7import math8from typing import Optional9 10try:11    from rdkit import Chem12    from rdkit.Chem import Descriptors, rdMolDescriptors, QED, Crippen13    from rdkit.Chem.FilterCatalog import FilterCatalog, FilterCatalogParams14    RDKIT_AVAILABLE = True15except ImportError:16    RDKIT_AVAILABLE = False17    print("Warning: RDKit not available. Install with: conda install -c conda-forge rdkit")18 19 20# ─────────────────────────────────────────────────────────────────────────────21# Internal: Local ADMET Computation (pure RDKit, no network)22# ─────────────────────────────────────────────────────────────────────────────23 24def _des(val: float, lo: float, hi: float, lo_d: float = 1.0, hi_d: float = 0.0) -> float:25    """Linear desirability: clamps to [lo_d, hi_d] outside [lo, hi] range."""26    if lo_d >= hi_d:  # decreasing: higher val → lower desirability27        if val <= lo:28            return lo_d29        if val >= hi:30            return hi_d31    else:              # increasing: higher val → higher desirability32        if val <= lo:33            return lo_d34        if val >= hi:35            return hi_d36    t = (val - lo) / (hi - lo)37    return lo_d + t * (hi_d - lo_d)38 39 40def _has_smarts(mol, smarts: str) -> bool:41    try:42        pat = Chem.MolFromSmarts(smarts)43        return pat is not None and mol.HasSubstructMatch(pat)44    except Exception:45        return False46 47 48 49def _cyp_flags(mol) -> dict:50    """Heuristic CYP450 substrate/inhibitor flags using SMARTS."""51    # CYP3A4 substrate: large + aromatic + heteroatom (most drugs)52    mw = Descriptors.ExactMolWt(mol)53    has_ar = _has_smarts(mol, "a")54    has_n_or_o = _has_smarts(mol, "[N,O]")55    cyp3a4 = mw > 300 and has_ar and has_n_or_o56 57    # CYP2D6 substrate: basic N within 5 bonds of aromatic ring58    cyp2d6 = _has_smarts(mol, "[NX3;H0,H1,H2;!$(NC=O)]~*~*~a") or \59              _has_smarts(mol, "[NX3;H0,H1,H2;!$(NC=O)]~*~a")60 61    # Inhibitor risk: imidazole, pyridine N, or extended pi system62    cyp_inhib = (63        _has_smarts(mol, "n1ccnc1")   # imidazole-like64        or _has_smarts(mol, "n1ccccc1")  # pyridine65        or _has_smarts(mol, "c1ccc2ccccc2c1")  # naphthalene / fused rings66    )67 68    return {69        "cyp3a4_substrate": cyp3a4,70        "cyp2d6_substrate": cyp2d6,71        "cyp2c9_substrate": has_ar and _has_smarts(mol, "C(=O)[OH]"),  # carboxylic acid + ring72        "cyp_inhibitor_risk": cyp_inhib,73    }74 75 76def _esol_log_s(mol, mw: float, clogp: float, rb: int) -> float:77    """78    ESOL (Delaney 2004) solubility estimate:79    logS = 0.16 − 0.63·cLogP − 0.0062·MW + 0.066·RB − 0.74·AP80    AP = fraction aromatic carbons (proxy for aromaticity penalty)81    """82    ap = sum(1 for a in mol.GetAtoms()83             if a.GetIsAromatic() and a.GetAtomicNum() == 6) / max(mol.GetNumHeavyAtoms(), 1)84    return 0.16 - 0.63 * clogp - 0.0062 * mw + 0.066 * rb - 0.74 * ap85 86 87def _solubility_class(log_s: float) -> str:88    if log_s >= -1:   return "Highly Soluble"89    if log_s >= -2:   return "Soluble"90    if log_s >= -4:   return "Moderately Soluble"91    if log_s >= -6:   return "Poorly Soluble"92    return "Insoluble"93 94 95def _gi_absorption(tpsa: float, rb: int, mw: float, clogp: float, hbd: int) -> dict:96    """97    GI absorption: High if Veber rules pass AND Lipinski-like.98    Veber (2002): TPSA ≤ 140 AND rotatable bonds ≤ 1099    Also penalise if multiple Lipinski violations.100    """101    veber = tpsa <= 140 and rb <= 10102    lipinski_ok = mw <= 500 and clogp <= 5 and hbd <= 5103    if veber and lipinski_ok:104        absorption = "High"105        score = 0.85 + 0.15 * max(0, (140 - tpsa) / 140)106    elif veber or lipinski_ok:107        absorption = "Moderate"108        score = 0.55109    else:110        absorption = "Low"111        score = 0.20112    return {"absorption": absorption, "bioavailability_score": round(score, 2)}113 114 115def _bbb_probability(clogp: float, mw: float, tpsa: float, hbd: int) -> dict:116    """117    BBB penetration probability from rule-based weighted desirability.118    Reference ranges: Pajouhesh & Lenz (2005), Wager CNS MPO.119    """120    # Each parameter contributes 0-1 desirability121    d_logp = _des(clogp, 0.0, 5.0, 0.1, 1.0)   # cLogP 0-5 (peak ~2-3)122    if clogp > 3:123        d_logp = _des(clogp, 3.0, 5.0, 1.0, 0.3)124 125    d_mw   = _des(mw,   200.0, 450.0, 0.7, 0.0)  # MW < 400 strongly preferred126    d_tpsa = _des(tpsa, 40.0,  90.0,  1.0, 0.0)  # TPSA < 60 ideal127    d_hbd  = _des(float(hbd), 0.0, 3.0, 1.0, 0.0)128 129    # Weighted average (TPSA and HBD are most important for BBB)130    prob = 0.20 * d_logp + 0.20 * d_mw + 0.35 * d_tpsa + 0.25 * d_hbd131    prob = round(max(0.05, min(0.98, prob)), 3)132    penetrates = prob >= 0.5133 134    confidence = "High" if prob > 0.75 or prob < 0.25 else "Moderate"135    return {"penetrates": penetrates, "probability": prob, "confidence": confidence}136 137 138def _cns_mpo(clogp: float, mw: float, tpsa: float, hbd: int) -> dict:139    """140    CNS MPO score — Wager et al. 2010 (Pfizer), 5-parameter version.141    Each parameter contributes 0-1 desirability, summed to max 5.142    (pKa excluded — cannot be reliably computed from SMILES alone.)143    """144    d1 = _des(clogp,       3.0, 5.0, 1.0, 0.0)   # cLogP ≤ 3 → 1, ≥ 5 → 0145    d2 = _des(clogp,       1.0, 3.0, 0.0, 1.0)   # cLogD ≈ cLogP (neutral approx)146    if clogp > 3.0:147        d2 = _des(clogp,   3.0, 5.0, 1.0, 0.0)148    d3 = _des(mw,        360.0, 500.0, 1.0, 0.0)  # MW ≤ 360 → 1, ≥ 500 → 0149    d4 = _des(tpsa,       40.0,  90.0, 1.0, 0.0)  # TPSA ≤ 40 → 1, ≥ 90 → 0150    d5 = _des(float(hbd),  0.5,   3.5, 1.0, 0.0)  # HBD ≤ 0.5 → 1, ≥ 3.5 → 0151 152    score = round(d1 + d2 + d3 + d4 + d5, 2)153 154    if score >= 4.0:155        cns_class = "CNS+"156    elif score >= 3.0:157        cns_class = "CNS Borderline"158    else:159        cns_class = "CNS-"160 161    return {"score": score, "cns_class": cns_class}162 163 164def _lipinski_check(mw: float, clogp: float, hbd: int, hba: int) -> tuple:165    """Returns (rules list, summary string)."""166    rules = [167        {"name": "MW ≤ 500",        "value": mw,    "pass": mw    <= 500},168        {"name": "cLogP ≤ 5",       "value": clogp, "pass": clogp <= 5},169        {"name": "HBD ≤ 5",         "value": hbd,   "pass": hbd   <= 5},170        {"name": "HBA ≤ 10",        "value": hba,   "pass": hba   <= 10},171    ]172    n_fail = sum(1 for r in rules if not r["pass"])173    summary = "Pass" if n_fail == 0 else ("Borderline" if n_fail == 1 else "Fail")174    return rules, summary175 176 177def _pains_alerts(mol) -> list:178    """Return PAINS alert names using RDKit FilterCatalog."""179    try:180        params = FilterCatalogParams()181        params.AddCatalog(FilterCatalogParams.FilterCatalogs.PAINS)182        catalog = FilterCatalog(params)183        entries = catalog.GetMatches(mol)184        return [{"name": e.GetDescription(), "description": "PAINS structural alert"} for e in entries]185    except Exception:186        return []187 188 189def _sa_score(mol, mw: float) -> tuple:190    """191    Synthetic accessibility heuristic.192    Rough inverse of complexity: penalise rings, stereocenters, MW.193    Returns (score 1-10, class string). Lower = easier to synthesise.194    """195    ring_count = rdMolDescriptors.CalcNumRings(mol)196    spiro = rdMolDescriptors.CalcNumSpiroAtoms(mol)197    stereo = len(Chem.FindMolChiralCenters(mol, includeUnassigned=True))198    heavy = mol.GetNumHeavyAtoms()199 200    complexity = (0.02 * mw + 0.5 * ring_count + 1.0 * spiro + 0.7 * stereo)201    score = min(10.0, max(1.0, complexity))202    score = round(score, 1)203 204    if score <= 3:205        sa_class = "Easy"206    elif score <= 6:207        sa_class = "Moderate"208    else:209        sa_class = "Difficult"210 211    return score, sa_class212 213 214def _fail_fast_score(lipinski_summary: str, num_alerts: int,215                     mw: float, tpsa: float, clogp: float) -> float:216    """217    Composite risk score 0-10 (higher = more concerning).218    """219    score = 0.0220    score += {"Pass": 0, "Borderline": 2, "Fail": 5}.get(lipinski_summary, 0)221    score += min(num_alerts * 2, 4)222    if mw > 600:   score += 2223    if tpsa > 140: score += 1224    if clogp > 5:  score += 1225    if clogp < 0:  score += 1226    return round(min(score, 10.0), 1)227 228 229def _decision(fail_fast: float, bbb_prob: float, qed_score: float,230              cns_mpo: float, num_alerts: int) -> tuple:231    """232    Returns (decision: str, rationale: str).233    Decision: 'Progress' | 'Optimize' | 'Kill'234    """235    if fail_fast >= 7 or num_alerts >= 2:236        return (237            "Kill",238            f"High risk score ({fail_fast}/10) or PAINS alerts ({num_alerts}) "239            "indicate major liabilities. Not worth further optimization."240        )241    if qed_score >= 0.6 and cns_mpo >= 4.0 and bbb_prob >= 0.5:242        return (243            "Progress",244            f"Solid drug-likeness (QED {qed_score:.2f}), "245            f"CNS MPO {cns_mpo:.1f}, BBB probability {bbb_prob:.2f}. "246            "Molecule shows promise — continue optimization."247        )248    return (249        "Optimize",250        f"Some properties need improvement: QED {qed_score:.2f}, "251        f"CNS MPO {cns_mpo:.1f}, BBB prob {bbb_prob:.2f}. "252        "Targeted structural changes can address the gaps."253    )254 255 256def _optimization_suggestions(clogp: float, mw: float, tpsa: float,257                               hbd: int, bbb_prob: float, cns_mpo: float,258                               qed_score: float, rb: int) -> list:259    """Generate context-specific optimization suggestions."""260    tips = []261    if bbb_prob < 0.5:262        if tpsa > 90:263            tips.append({"text": "Reduce polar surface area (TPSA > 90 Ų) — consider replacing amide/OH with ester or nitrile bioisostere"})264        if hbd > 1:265            tips.append({"text": f"Reduce H-bond donors ({hbd}) — replace NH/OH groups with methylated analogues or bioisosteres"})266        if mw > 400:267            tips.append({"text": f"Reduce molecular weight ({mw:.0f} Da) — remove non-essential substituents"})268    if clogp < 1.0:269        tips.append({"text": "Increase lipophilicity (cLogP too low for good CNS penetration) — add small alkyl or fluoro groups"})270    elif clogp > 4.5:271        tips.append({"text": "Reduce lipophilicity (cLogP too high risks off-target effects) — add polar groups or replace alkyl with heteroatom"})272    if qed_score < 0.6:273        tips.append({"text": f"Drug-likeness (QED {qed_score:.2f}) is suboptimal — simplify structure and reduce MW toward 300-400 Da range"})274    if cns_mpo < 4.0:275        tips.append({"text": f"CNS MPO score ({cns_mpo:.1f}) below threshold — target TPSA < 60, HBD ≤ 1, cLogP 1-3, MW < 400"})276    if rb > 8:277        tips.append({"text": f"Too many rotatable bonds ({rb}) — constrain flexible chains by ring formation or rigidification"})278    if not tips:279        tips.append({"text": "Properties are well-balanced — fine-tune specific targets while preserving overall profile"})280    return tips[:4]  # cap at 4 suggestions281 282 283def analyze_local(smiles: str) -> dict:284    """285    Compute a full ADMET profile for a SMILES string using local RDKit.286    Returns a dict in the same shape as the Drug Discovery Triage API response,287    so extract_key_scores() works without any changes.288    """289    if not RDKIT_AVAILABLE:290        return {"error": "RDKit not installed. Run: conda install -c conda-forge rdkit"}291 292    mol = Chem.MolFromSmiles(smiles.strip())293    if mol is None:294        return {"error": f"Invalid SMILES: RDKit could not parse '{smiles[:60]}'"}295 296    # ── Physicochemical ───────────────────────────────────────────────────────297    canonical = Chem.MolToSmiles(mol)298    mw        = round(Descriptors.ExactMolWt(mol), 2)299    clogp     = round(Crippen.MolLogP(mol), 2)300    tpsa      = round(rdMolDescriptors.CalcTPSA(mol), 1)301    hbd       = rdMolDescriptors.CalcNumHBD(mol)302    hba       = rdMolDescriptors.CalcNumHBA(mol)303    rb        = rdMolDescriptors.CalcNumRotatableBonds(mol)304    fsp3      = round(rdMolDescriptors.CalcFractionCSP3(mol), 3)305    ring_ct   = rdMolDescriptors.CalcNumRings(mol)306    ar_ring   = rdMolDescriptors.CalcNumAromaticRings(mol)307 308    # ── QED ───────────────────────────────────────────────────────────────────309    qed_val   = round(QED.qed(mol), 3)310    if qed_val >= 0.67:311        qed_class = "Drug-like"312    elif qed_val >= 0.34:313        qed_class = "Moderate"314    else:315        qed_class = "Poor"316 317    # ── Lipinski ──────────────────────────────────────────────────────────────318    lip_rules, lip_summary = _lipinski_check(mw, clogp, hbd, hba)319 320    # ── PAINS alerts ──────────────────────────────────────────────────────────321    alerts = _pains_alerts(mol)322 323    # ── ADMET ─────────────────────────────────────────────────────────────────324    log_s        = round(_esol_log_s(mol, mw, clogp, rb), 2)325    sol_class    = _solubility_class(log_s)326    gi           = _gi_absorption(tpsa, rb, mw, clogp, hbd)327    bbb          = _bbb_probability(clogp, mw, tpsa, hbd)328    cns          = _cns_mpo(clogp, mw, tpsa, hbd)329    cyp          = _cyp_flags(mol)330 331    # ── Synthetic accessibility & scoring ─────────────────────────────────────332    sa_score, sa_class = _sa_score(mol, mw)333    fail_fast    = _fail_fast_score(lip_summary, len(alerts), mw, tpsa, clogp)334    decision, rationale = _decision(fail_fast, bbb["probability"], qed_val,335                                    cns["score"], len(alerts))336    suggestions  = _optimization_suggestions(clogp, mw, tpsa, hbd,337                                             bbb["probability"], cns["score"],338                                             qed_val, rb)339 340    return {341        # ── Identity ──────────────────────────────────────────────────────────342        "canonical_smiles": canonical,343        # ── Physicochemical ───────────────────────────────────────────────────344        "molecular_weight": mw,345        "clogp":            clogp,346        "tpsa":             tpsa,347        "hbd":              hbd,348        "hba":              hba,349        "rotatable_bonds":  rb,350        "fraction_sp3":     fsp3,351        "ring_count":       ring_ct,352        "aromatic_ring_count": ar_ring,353        # ── Drug-likeness ─────────────────────────────────────────────────────354        "qed": {"qed_score": qed_val, "qed_class": qed_class},355        "lipinski_rules":   lip_rules,356        "lipinski_summary": lip_summary,357        "alerts":           alerts,358        "synthetic_accessibility_score": sa_score,359        "synthetic_accessibility_class": sa_class,360        "fail_fast_score":  fail_fast,361        "decision":         decision,362        "decision_rationale": rationale,363        "optimization_suggestions": suggestions,364        # ── ADMET ─────────────────────────────────────────────────────────────365        "admet": {366            "solubility": {367                "log_s":            log_s,368                "solubility_class": sol_class,369            },370            "gi_absorption": gi,371            "bbb_penetration": bbb,372            "cns_mpo": cns,373            "cyp_metabolism": cyp,374            "toxicity": {"endpoints": {}},375        },376    }377 378 379# ─────────────────────────────────────────────────────────────────────────────380# Tool 1: Analyze Molecule  (local RDKit — no network required)381# ─────────────────────────────────────────────────────────────────────────────382def call_admet_api(smiles: str, timeout: int = 90) -> dict:383    """384    Analyze a SMILES string and return full ADMET profile.385    Computed locally via RDKit — instant, no API key or network needed.386    (timeout arg kept for API compatibility but unused)387    """388    return analyze_local(smiles.strip())389 390 391# ─────────────────────────────────────────────────────────────────────────────392# Tool 2: Validate SMILES  (local RDKit, no API cost)393# ─────────────────────────────────────────────────────────────────────────────394def is_valid_smiles(smiles: str) -> dict:395    """396    Validate a SMILES string using RDKit (fast, offline).397    The agent should always call this before analyze_molecule.398    """399    if not RDKIT_AVAILABLE:400        return {"valid": True, "note": "RDKit unavailable — skipping local validation"}401    try:402        mol = Chem.MolFromSmiles(smiles.strip())403        if mol is None:404            return {"valid": False, "error": "RDKit could not parse this SMILES string"}405        if mol.GetNumHeavyAtoms() == 0:406            return {"valid": False, "error": "Empty SMILES — no atoms found"}407        return {408            "valid": True,409            "canonical_smiles": Chem.MolToSmiles(mol),410            "num_heavy_atoms": mol.GetNumHeavyAtoms(),411        }412    except Exception as e:413        return {"valid": False, "error": str(e)}414 415 416# ─────────────────────────────────────────────────────────────────────────────417# Score Extraction  (flatten full API response → concise dict)418# ─────────────────────────────────────────────────────────────────────────────419def extract_key_scores(response: dict) -> dict:420    """421    Flatten the full API response into the metrics the agent needs to reason about.422    Includes physicochemical properties, ADMET, toxicity, and the API's own decision.423    """424    if "error" in response:425        return response426 427    admet = response.get("admet", {})428    cyp   = admet.get("cyp_metabolism", {})429    alerts = response.get("alerts", [])430    alert_names = [a.get("name", "") for a in alerts] if isinstance(alerts, list) else []431    suggestions = [s.get("text", "") for s in response.get("optimization_suggestions", [])]432 433    qed_block = response.get("qed", {})434    qed_score = (qed_block.get("qed_score") or 0) if isinstance(qed_block, dict) else 0435 436    return {437        # ── Identity ──────────────────────────────────────────────────────────438        "canonical_smiles":   response.get("canonical_smiles", ""),439        # ── Physicochemical ───────────────────────────────────────────────────440        "molecular_weight":   round(response.get("molecular_weight", 0), 1),441        "clogp":              round(response.get("clogp", 0), 2),442        "tpsa":               round(response.get("tpsa", 0), 1),443        "hbd":                response.get("hbd", 0),444        "hba":                response.get("hba", 0),445        "rotatable_bonds":    response.get("rotatable_bonds", 0),446        # ── Drug-likeness ─────────────────────────────────────────────────────447        "qed_score":          round(float(qed_score), 3),448        "qed_class":          qed_block.get("qed_class", "") if isinstance(qed_block, dict) else "",449        "lipinski_summary":   response.get("lipinski_summary", "Unknown"),450        "fail_fast_score":    round(response.get("fail_fast_score", 0), 1),451        "decision":           response.get("decision", "Unknown"),452        "decision_rationale": response.get("decision_rationale", ""),453        # ── ADMET ─────────────────────────────────────────────────────────────454        "solubility_class":   admet.get("solubility", {}).get("solubility_class", "Unknown"),455        "log_s":              admet.get("solubility", {}).get("log_s", None),456        "gi_absorption":      admet.get("gi_absorption", {}).get("absorption", "Unknown"),457        "bbb_penetrates":     admet.get("bbb_penetration", {}).get("penetrates", False),458        "bbb_probability":    round(admet.get("bbb_penetration", {}).get("probability", 0), 3),459        "cns_mpo_score":      round(admet.get("cns_mpo", {}).get("score", 0), 2),460        "cns_class":          admet.get("cns_mpo", {}).get("cns_class", "Unknown"),461        # ── Toxicity ──────────────────────────────────────────────────────────462        "alerts":             alert_names,463        "num_alerts":         len(alert_names),464        "cyp3a4_substrate":   cyp.get("cyp3a4_substrate", False),465        "cyp_inhibitor_risk": cyp.get("cyp_inhibitor_risk", False),466        # ── API's own suggestions (bonus context for the agent) ───────────────467        "api_suggestions":    suggestions,468    }469 470 471# ─────────────────────────────────────────────────────────────────────────────472# Tool 3: Compare Candidates473# ─────────────────────────────────────────────────────────────────────────────474def compare_candidates(smiles_list: list, labels: list = None) -> dict:475    """476    Analyze multiple SMILES and return side-by-side key scores.477    Used by the agent to rank candidates across optimization rounds.478    """479    if not labels:480        labels = [f"Candidate {i}" for i in range(len(smiles_list))]481 482    results = []483    for smiles, label in zip(smiles_list, labels):484        resp = call_admet_api(smiles)485        entry = extract_key_scores(resp) if "error" not in resp else resp486        entry["label"] = label487        results.append(entry)488 489    return {"candidates": results, "count": len(results)}490 491 492# ─────────────────────────────────────────────────────────────────────────────493# Tool Dispatcher  (called by the agent loop)494# ─────────────────────────────────────────────────────────────────────────────495def tool_executor(tool_name: str, tool_input: dict) -> dict:496    """Route a Claude tool_use block to the correct function."""497    if tool_name == "analyze_molecule":498        resp = call_admet_api(tool_input["smiles"])499        return extract_key_scores(resp) if "error" not in resp else resp500 501    elif tool_name == "validate_smiles":502        return is_valid_smiles(tool_input["smiles"])503 504    elif tool_name == "compare_candidates":505        return compare_candidates(506            tool_input["smiles_list"],507            tool_input.get("labels"),508        )509 510    return {"error": f"Unknown tool: '{tool_name}'"}511