dee/core/agent_tools.py

"""Tools the Turing agent (dee/core/agent.py) can call via OpenRouter
function-calling — thin wrappers around the SAME dee/core/*.py functions
the REST API already calls, with the SAME input bounds those routes
already enforce (an LLM-supplied argument is no more trustworthy than a
raw POST body — every bound here mirrors dee/server.py's own route).

Scope, deliberately: only clean, self-contained actions that can complete
within one HTTP request. The full async /api/run pipeline (job-id +
polling, unbounded protein length, up to 50k search steps) is NOT wired
up — design_variant_library below is a bounded, synchronous sibling of it
instead, capped the same way POST /api/de/round2 already proves safe in
production (500 aa, same settings shape) rather than a new, unproven
bound. Round-2 proposal itself (needs prior bench measurements a fresh
chat won't have) and CRISPR's genome off-target / Ensembl exon lookups
(slow external NCBI/Ensembl fetches with a "still building" state the
REST UI can poll for but a single chat turn can't) stay out for now.

dee/core/ stays auth-agnostic by convention (see dee/auth.py's own
docstring) — execute_tool takes a plain `auth_anonymous: bool`, not the
AuthContext object, so this module never imports dee.auth.
"""

from __future__ import annotations

import logging
import re
from typing import Any, Callable, Dict, List

logger = logging.getLogger("dee.agent_tools")


def _tool_fetch_sequence(args: Dict[str, Any]) -> Dict[str, Any]:
    """Resolve a gene symbol / accession / raw paste to an actual DNA
    sequence — the same dee.core.resolve.resolve_target() that backs the
    "paste anything" box in the CRISPR UI (POST /api/crispr/resolve), now
    reachable from chat too. Previously a chat request naming a gene instead
    of pasting its sequence was a dead end (found 2026-07-11: "it's not
    capable of searching www") even though this exact lookup already existed
    and is fast/synchronous (2 Ensembl REST calls, cached) — unlike the
    genome-scale off-target/exon work that's deliberately kept out of chat
    for being slow enough to need polling (see module docstring)."""
    from dee.core import resolve as _resolve

    text = str(args.get("text") or "").strip()
    if not text:
        return {"ok": False, "error": "missing 'text'"}
    # Same 1 Mbp cap as POST /api/crispr/resolve.
    if len(text) > 1_000_000:
        return {"ok": False, "error": "Input too long. Cap is 1 Mbp."}

    organism = str(args.get("organism", "")).lower().strip()
    if organism not in ("", "ecoli", "human", "mouse"):
        organism = ""

    try:
        result = _resolve.resolve_target(text, organism=organism)
    except Exception:  # noqa: BLE001
        logger.exception("fetch_sequence resolve failed")
        return {"ok": False, "error": "Lookup failed — check the name/accession and try again."}

    if not result.get("ok"):
        return {"ok": False, "error": result.get("error") or "Couldn't resolve that input."}

    sequence = result.get("sequence", "")
    return {
        "ok": True,
        "kind": result.get("kind"),
        "sequence": sequence,
        "length": len(sequence),
        "gene_symbol": result.get("gene_symbol", ""),
        "label": result.get("label", ""),
        "source": result.get("source", ""),
    }


def _tool_design_crispr_guides(args: Dict[str, Any]) -> Dict[str, Any]:
    from dee.core.crispr import find_guides

    seq = str(args.get("sequence") or "").strip()
    if not seq:
        return {"ok": False, "error": "missing 'sequence'"}
    # Same 1 Mbp cap as POST /api/crispr/design.
    if len(seq) > 1_000_000:
        return {"ok": False, "error": (
            "Sequence too long. Cap is 1 Mbp — paste just the gene/region "
            "you're editing, not a whole chromosome."
        )}

    enzyme = str(args.get("enzyme", "cas9")).lower()
    if enzyme not in ("cas9", "cas12a"):
        enzyme = "cas9"
    mode = str(args.get("mode", "knockout")).lower()
    if mode not in ("knockout", "base_edit"):
        mode = "knockout"
    try:
        max_results = int(args.get("max_results", 20))
    except (TypeError, ValueError):
        max_results = 20
    # Smaller default/cap than the REST UI's 500 — this reply is read
    # inline in chat, not scrolled through a results table.
    max_results = max(1, min(50, max_results))

    try:
        guides = find_guides(seq, enzyme=enzyme, max_results=max_results, mode=mode)
    except ValueError as exc:
        return {"ok": False, "error": str(exc)}

    guide_dicts = [
        {
            "rank": g.rank,
            "strand": g.strand,
            "position": g.position,
            "spacer": g.spacer,
            "pam": g.pam,
            "composite_score": round(g.composite_score, 3),
            "on_target_score": round(g.on_target_score, 3),
            "gc_pct": round(g.gc_pct, 1),
            "notes": g.notes,
        }
        for g in guides
    ]

    # Same calibration as POST /api/crispr/design (server.py) — fold the
    # cross-user aggregate into each guide's score, then re-rank. The chat
    # tool skipped this entirely before (found 2026-07-11); no-op until a
    # rebuild has populated public.outcome_priors for "crispr".
    field_prior_n = 0
    try:
        from dee.core import outcomes as _O
        prior = _O.load_cached_tool_prior("crispr")
        if prior and prior.effects:
            field_prior_n = len(prior.effects)
            guide_dicts = _O.calibrate_crispr_guides(guide_dicts, prior)
    except Exception:  # noqa: BLE001
        logger.exception("field-prior calibration skipped in design_crispr_guides")

    return {
        "ok": True,
        "n_guides": len(guide_dicts),
        "enzyme": enzyme,
        "mode": mode,
        "field_prior_keys": field_prior_n,
        "guides": guide_dicts,
    }


def _tool_design_primers(args: Dict[str, Any]) -> Dict[str, Any]:
    from dee.core import primers as _P

    template = re.sub(r"\s+", "", str(args.get("template") or "")).upper()
    if not template:
        return {"ok": False, "error": "Paste a template to design primers against."}
    if not re.fullmatch(r"[ACGTN]+", template):
        return {"ok": False, "error": "Template must be DNA (A/C/G/T/N) only."}
    # Same cap as POST /api/primers/design.
    if len(template) > _P.MAX_TEMPLATE:
        return {"ok": False, "error": (
            f"Template longer than {_P.MAX_TEMPLATE:,} bp — trim to the region of interest."
        )}

    def _oi(v):
        try:
            return int(v)
        except (TypeError, ValueError):
            return None

    try:
        result = _P.design_primers(template, _oi(args.get("target_start")), _oi(args.get("target_end")))
    except Exception as exc:  # noqa: BLE001 — mirrors the REST route's own catch-all
        logger.exception("design_primers tool call failed")
        return {"ok": False, "error": "Primer design failed — check the inputs."}
    if not result.get("ok"):
        return result

    pairs = result.get("pairs") or []
    return {
        "ok": True,
        "n_forward": result.get("n_forward"),
        "n_reverse": result.get("n_reverse"),
        "warnings": result.get("warnings") or [],
        # Trimmed to what's useful read inline in chat — drop internal
        # candidate metadata (hairpin/self-dimer scores, etc.) the REST
        # results table shows but a chat reply doesn't need.
        "pairs": [
            {
                "forward_seq": p.get("forward", {}).get("seq"),
                "forward_tm": p.get("forward", {}).get("tm"),
                "reverse_seq": p.get("reverse", {}).get("seq"),
                "reverse_tm": p.get("reverse", {}).get("tm"),
                "product_size": p.get("product_size"),
            }
            for p in pairs
        ],
    }


# Cap mirrors POST /api/de/round2's _DE_ROUND2_MAX_AA (dee/server.py) — that
# route already proves 500 aa synchronous ESM-2 scoring is safe in
# production on this deploy's CPU tier; reusing the same number rather than
# inventing an unproven bound for this tool.
_MAX_PROTEIN_AA = 500


def _tool_design_variant_library(args: Dict[str, Any]) -> Dict[str, Any]:
    raw = re.sub(r"\s+", "", str(args.get("sequence") or "")).upper()
    if not raw:
        return {"ok": False, "error": "missing 'sequence'"}

    # Auto-detect + translate DNA input. fetch_sequence (and the CRISPR/
    # primer tools) all deal in DNA — a model chaining "fetch this gene,
    # then DE it" naturally ends up holding DNA, not protein, and this tool
    # used to require protein only with no bridge, no translate tool to
    # call either. The DNA just failed plain alphabet validation with no
    # path forward, which from the user's side looked exactly like the
    # model forgetting the sequence it had just fetched (found 2026-07-12).
    # Best-effort and silent-fallback rather than a hard error: a protein
    # made entirely of A/C/G/T-coded residues (Ala/Cys/Gly/Thr) is legal
    # and technically indistinguishable from DNA by alphabet alone, so if
    # this doesn't validate as a real CDS, fall through to treating the
    # input as protein instead of failing outright.
    protein = raw
    if len(raw) >= 30 and len(raw) % 3 == 0 and re.match(r"^[ACGTN]+$", raw):
        try:
            from dee.core.sequence import translate_dna
            protein = translate_dna(raw, identifier="chat-input").protein
        except Exception:  # noqa: BLE001
            protein = raw

    if not re.match(r"^[ACDEFGHIKLMNPQRSTVWY*]+$", protein):
        return {"ok": False, "error": "Sequence must be a protein (standard amino-acid letters), not DNA."}
    if len(protein) > _MAX_PROTEIN_AA:
        return {"ok": False, "error": (
            f"Sequence too long for a chat reply — cap is {_MAX_PROTEIN_AA} aa. "
            "Use the Directed Evolution tool in the sidebar for longer "
            "proteins; it runs in the background instead of one request."
        )}

    host = str(args.get("host", "e_coli")).lower()
    if host not in ("e_coli", "yeast", "human"):
        host = "e_coli"
    try:
        percentile = float(args.get("percentile", 85.0))
    except (TypeError, ValueError):
        percentile = 85.0
    percentile = max(1.0, min(99.0, percentile))
    try:
        k = int(args.get("k", 10))
    except (TypeError, ValueError):
        k = 10
    # Smaller cap than the REST route's 200 — a chat reply lists variants
    # inline, it doesn't render a scrollable results table.
    k = max(1, min(20, k))
    try:
        max_mutations = int(args.get("max_mutations", 3))
    except (TypeError, ValueError):
        max_mutations = 3
    max_mutations = max(1, min(6, max_mutations))
    min_mutations = min(2, max_mutations)

    from dee.core import scoring
    from dee.models.scorer import top_percentile_pool
    from dee.optimizer.search import SearchConfig, evolve

    try:
        scorer = scoring.get_scorer("small")
        # Bounded wait, not scoring.score_guarded's default indefinite
        # block — this is one HTTP request's worth of budget, not the REST
        # pipeline's async job with its own polling UI. Fail fast and
        # clearly instead of holding the chat turn open.
        scores_df = scoring.score_guarded(scorer, protein, wait_timeout=20.0)
    except scoring.ScoringBusyError:
        return {"ok": False, "kind": "busy", "error": (
            "The design engine is busy scoring another request right now — try again in a few seconds."
        )}
    except Exception:  # noqa: BLE001
        logger.exception("design_variant_library scoring failed")
        return {"ok": False, "error": "Scoring failed — check the sequence and try again."}

    pool = top_percentile_pool(scores_df, percentile=percentile)

    # Same soft blend as the REST /api/run job (server.py): fold the
    # cross-user aggregate into round-1 scores by amino-acid substitution
    # type, before search. Previously ONLY the REST job did this — the chat
    # tool went straight from scoring to search, so a chat-driven design
    # never reflected the aggregate no matter how much data existed (found
    # 2026-07-11 while checking why it couldn't be tested from chat). No-op
    # until a rebuild has actually populated public.mutation_priors.
    field_prior_n = 0
    try:
        from dee.core import aggregate as _agg
        _gp = _agg.load_cached_global_prior()
        if _gp and _gp.effects:
            field_prior_n = len(_gp.effects)
            pool = pool.copy()
            pool["delta_ll"] = [
                float(r.delta_ll) + 0.3 * _gp.effects.get(
                    (str(r.wt_aa).upper(), str(r.mut_aa).upper()), 0.0)
                for r in pool.itertuples(index=False)
            ]
    except Exception:  # noqa: BLE001
        logger.exception("global-prior blend skipped in design_variant_library")

    try:
        # Reduced search budget vs. the REST defaults (restarts=8,
        # steps=1200) — the pool is already scored, this part is pure
        # numpy/pandas simulated annealing (no more ESM-2 calls), but kept
        # small anyway to leave headroom in the request's total time budget
        # for the OpenRouter round trips before and after this tool call.
        variants = evolve(pool, SearchConfig(
            k=k, max_mutations=max_mutations, min_mutations=min_mutations,
            n_restarts=4, steps_per_restart=600,
        ))
    except ValueError as exc:
        return {"ok": False, "error": str(exc)}

    # Reverse-translate each variant to actual, orderable DNA — same step
    # the REST /api/run job uses (variants_to_dataframe: host-aware codon
    # optimization + forbidden-restriction-site scrubbing). Previously this
    # tool stopped at mutation labels ("W44K") with no sequence at all, so a
    # chat-driven design had nothing to copy/paste or order (found
    # 2026-07-12). Pure computation, no network — safe for one chat turn;
    # best-effort so a rare encoding failure degrades to labels-only instead
    # of losing the whole result.
    dna_rows: List[Dict[str, Any]] = []
    try:
        from dee.core.codon import variants_to_dataframe, DEFAULT_FORBIDDEN_SITES
        dna_rows = variants_to_dataframe(
            protein, variants, host=host, forbidden_sites=DEFAULT_FORBIDDEN_SITES,
        ).to_dict(orient="records")
    except Exception:  # noqa: BLE001
        logger.exception("DNA encoding failed in design_variant_library")

    return {
        "ok": True,
        "model": "small",
        "host": host,
        "n_scored_positions": len(scores_df),
        "n_pool": len(pool),
        "field_prior_substitutions": field_prior_n,
        "variants": [
            dict(
                {"rank": v.rank, "mutations": list(v.mutation_labels), "fitness": round(v.fitness, 3)},
                **({
                    "dna": dna_rows[i]["Optimized_DNA_Seq"],
                    "length_bp": dna_rows[i]["Length_bp"],
                    "protein": dna_rows[i]["Mutant_AA_Seq"],
                    # 0-indexed AMINO ACID positions that changed — NOT a
                    # DNA-level diff against the WT sequence. WT and each
                    # variant are reverse-translated + restriction-site-
                    # scrubbed independently, so their DNA can legitimately
                    # differ at unmutated codons too (a different silent
                    # codon choice to clear a forbidden site) — diffing the
                    # two DNA strings would mark those as "the mutation,"
                    # which is simply wrong. The client maps each position p
                    # to nt range [p*3, p*3+3) — reverse_translate() is a
                    # strict 1:1 codon-per-residue mapping with no leading
                    # offset, so this is exact, not an approximation.
                    "mutated_positions": [m.position for m in v.mutations],
                } if i < len(dna_rows) else {}),
            )
            for i, v in enumerate(variants)
        ],
    }


# ─── Registry + OpenRouter tool specs ───────────────────────────────────
# requires_signin mirrors each function's REST route exactly: /api/crispr/
# design, /api/primers/design, and /api/run (the REST sibling of
# design_variant_library) all require a full account, not just the
# anonymous-trial quota /api/agent/step itself is gated by.
_TOOLS: Dict[str, Dict[str, Any]] = {
    "fetch_sequence": {"fn": _tool_fetch_sequence, "requires_signin": True},
    "design_crispr_guides": {"fn": _tool_design_crispr_guides, "requires_signin": True},
    "design_primers": {"fn": _tool_design_primers, "requires_signin": True},
    "design_variant_library": {"fn": _tool_design_variant_library, "requires_signin": True},
}

TOOL_SPECS: List[Dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "fetch_sequence",
            "description": (
                "Resolve a gene symbol (e.g. GFP, TP53), an accession "
                "(Ensembl ENST.../ENSG..., or RefSeq NM_.../NR_...), or a "
                "raw pasted sequence into an actual DNA sequence, with a "
                "human-readable label. Call this FIRST whenever the user "
                "names a gene/protein/accession instead of pasting a "
                "sequence themselves, then pass the returned 'sequence' "
                "into design_crispr_guides / design_primers / "
                "design_variant_library as needed — don't ask the user to "
                "go paste it manually if this can fetch it. Gene-symbol "
                "lookups need 'organism' set to human or mouse; accessions "
                "and raw sequences don't. Requires the user to be signed in."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Gene symbol, accession, or raw DNA/FASTA. Max 1,000,000 bp.",
                    },
                    "organism": {
                        "type": "string",
                        "enum": ["human", "mouse", "ecoli"],
                        "description": "Required for a gene-symbol lookup; ignored for accessions and raw sequences.",
                    },
                },
                "required": ["text"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "design_crispr_guides",
            "description": (
                "Design and rank CRISPR guide RNAs against a DNA sequence for "
                "gene knockout or base editing. Requires the user to be signed in."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "sequence": {
                        "type": "string",
                        "description": "DNA sequence to target (FASTA headers/whitespace tolerated). Max 1,000,000 bp.",
                    },
                    "enzyme": {
                        "type": "string",
                        "enum": ["cas9", "cas12a"],
                        "description": "Cas enzyme. Defaults to cas9.",
                    },
                    "mode": {
                        "type": "string",
                        "enum": ["knockout", "base_edit"],
                        "description": "Editing mode. Defaults to knockout.",
                    },
                    "max_results": {
                        "type": "integer",
                        "description": "Max guides to return, 1-50. Defaults to 20.",
                    },
                },
                "required": ["sequence"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "design_primers",
            "description": (
                "Design forward/reverse PCR primer pairs to amplify a region of "
                "a DNA template. Requires the user to be signed in."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "template": {
                        "type": "string",
                        "description": f"DNA template (A/C/G/T/N only). Max {60_000:,} bp.",
                    },
                    "target_start": {
                        "type": "integer",
                        "description": "1-based start of the region the product must span. Omit to amplify the whole template.",
                    },
                    "target_end": {
                        "type": "integer",
                        "description": "1-based end of the region the product must span. Omit to amplify the whole template.",
                    },
                },
                "required": ["template"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "design_variant_library",
            "description": (
                "Score every possible single-amino-acid substitution in a protein "
                "with ESM-2 and propose a ranked library of multi-mutant variants. "
                f"Chat-sized version of the Directed Evolution tool — capped at "
                f"{_MAX_PROTEIN_AA} aa and a smaller search budget so it completes "
                "within one reply; for longer proteins or a deeper search, use the "
                "Directed Evolution tool in the sidebar instead. Requires the user "
                "to be signed in."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "sequence": {
                        "type": "string",
                        "description": (
                            f"Protein sequence (standard amino acids), OR a DNA coding "
                            f"sequence (CDS) — auto-translated, so the raw output of "
                            f"fetch_sequence can be passed straight through with no "
                            f"separate translation step. Max {_MAX_PROTEIN_AA} aa."
                        ),
                    },
                    "host": {
                        "type": "string",
                        "enum": ["e_coli", "yeast", "human"],
                        "description": "Expression host, used for codon-usage-aware downstream steps. Defaults to e_coli.",
                    },
                    "k": {
                        "type": "integer",
                        "description": "Number of ranked variants to return, 1-20. Defaults to 10.",
                    },
                    "max_mutations": {
                        "type": "integer",
                        "description": "Max simultaneous substitutions per variant, 1-6. Defaults to 3.",
                    },
                    "percentile": {
                        "type": "number",
                        "description": "Keep single-site mutations above this ΔLL percentile before searching multi-mutants, 1-99. Defaults to 85.",
                    },
                },
                "required": ["sequence"],
            },
        },
    },
]


def execute_tool(name: str, arguments: Dict[str, Any], auth_anonymous: bool) -> Dict[str, Any]:
    """Run one tool call. Never raises — always returns a JSON-serializable
    dict, {"ok": False, "error": ...} on any failure the model can read
    and recover from."""
    spec = _TOOLS.get(name)
    if spec is None:
        return {"ok": False, "error": f"unknown tool {name!r}"}
    if spec["requires_signin"] and auth_anonymous:
        return {
            "ok": False,
            "kind": "signin_required",
            "error": "This requires a free account. Sign in to continue.",
        }
    try:
        return spec["fn"](arguments)
    except Exception:  # noqa: BLE001 — a tool must never crash the agent loop
        logger.exception("tool %r failed", name)
        return {"ok": False, "error": f"{name} failed — check the inputs and try again."}