dee/core/blast.py

"""BLAST, in a place both the REST layer and the agent can reach.

The engine has been able to BLAST since the identify feature: `server.py` runs
blastp against `nr` on a background thread with a hash-keyed cache. The agent
could not call any of it — `agent_tools.py` importing from `server.py` would be
a circular import (server imports the agent), so the capability sat behind the
Flask layer. This module is the same extraction `scoring.py` already did for
the ESM-2 scorer, for the same reason.

Two things here are not incidental.

THE SEQUENCE LEAVES THE SPACE
    Every other tool in this engine computes locally or sends at most
    (gene_symbol, organism). BLAST is the exception: it posts the user's
    actual residues to NCBI. `server.py` already enforces that server-side
    (`blast_consent`, added by an audit — "must be enforced server-side, not
    just by the consent modal"). The agent path uses the confirm gate for the
    same purpose, so the model cannot route around it either. Nothing in this
    module asks for consent — it is the caller's job, and both callers do it.

IT IS SLOW, AND SAYING SO BEATS HANGING
    NCBI BLAST is 30 s to several minutes, and the queue is shared. So the
    agent path waits a bounded time and then reports "still running" with the
    job id rather than holding a run open. A tool that returns "not finished
    yet" is honest; one that blocks for four minutes looks broken.
"""
from __future__ import annotations

import hashlib
import logging
import re
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

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

# How long the AGENT path waits before handing back a job id. The REST path
# polls and does not use this.
AGENT_WAIT_SECONDS = 75.0

# blastp/nr for protein, blastn/nt for nucleotide. Both are the public web
# service; there is no local database on this box.
_PROGRAMS = {
    "protein":    {"program": "blastp", "database": "nr"},
    "nucleotide": {"program": "blastn", "database": "nt"},
}

_AA_ONLY = set("EFILPQZXJBOU")     # residues that cannot be DNA/RNA


@dataclass
class BlastJob:
    job_id: str
    seq_hash: str
    kind: str = "protein"
    status: str = "pending"
    started_at: float = field(default_factory=time.time)
    finished_at: Optional[float] = None
    hits: List[Dict[str, Any]] = field(default_factory=list)
    error: Optional[str] = None

    def elapsed(self) -> float:
        end = self.finished_at if self.finished_at else time.time()
        return round(end - self.started_at, 1)

    def done(self) -> bool:
        return self.status in ("done", "error")

    def public(self) -> Dict[str, Any]:
        return {"job_id": self.job_id, "status": self.status, "kind": self.kind,
                "elapsed_seconds": self.elapsed(), "hits": self.hits,
                "error": self.error}


_CACHE: Dict[str, BlastJob] = {}     # (kind, seq_hash) -> job
_BY_ID: Dict[str, BlastJob] = {}
_LOCK = threading.Lock()


def hash_sequence(seq: str) -> str:
    return hashlib.sha256(seq.encode("utf-8")).hexdigest()[:16]


def extract_organism(hit_def: str) -> str:
    """'GFP [Aequorea victoria]' -> 'Aequorea victoria'."""
    m = re.search(r"\[([^\]]+)\]", hit_def or "")
    return m.group(1) if m else ""


def guess_kind(seq: str) -> str:
    """protein or nucleotide, from the alphabet.

    Deliberately conservative: only letters that CANNOT be nucleotides make it
    protein. A short ACGT-only peptide is genuinely ambiguous, and calling it
    nucleotide is the safer error — blastn on a peptide returns nothing, while
    blastp on DNA returns confident nonsense.
    """
    s = re.sub(r"[^A-Za-z]", "", seq or "").upper()
    if not s:
        return "nucleotide"
    return "protein" if set(s) & _AA_ONLY else "nucleotide"


def _run(job: BlastJob, seq: str, hitlist: int, expect: float) -> None:
    try:
        from Bio.Blast import NCBIWWW, NCBIXML

        cfg = _PROGRAMS[job.kind]
        job.status = "submitting"
        handle = NCBIWWW.qblast(program=cfg["program"], database=cfg["database"],
                                sequence=seq, hitlist_size=hitlist, expect=expect)
        job.status = "parsing"
        hits: List[Dict[str, Any]] = []
        for record in NCBIXML.parse(handle):
            for alignment in record.alignments[:hitlist]:
                if not alignment.hsps:
                    continue
                hsp = alignment.hsps[0]          # best HSP for this subject
                hits.append({
                    "title": (alignment.hit_def or "")[:180],
                    "accession": getattr(alignment, "accession", "") or "",
                    "organism": extract_organism(alignment.hit_def or ""),
                    "identity_pct": round(100.0 * hsp.identities /
                                          max(1, hsp.align_length), 1),
                    "coverage_pct": round(100.0 * hsp.align_length /
                                          max(1, len(seq)), 1),
                    "evalue": float(hsp.expect),
                    "bit_score": round(float(hsp.bits), 1),
                    "align_length": int(hsp.align_length),
                })
            break                                 # one query in, one record out
        job.hits = hits
        job.status = "done"
    except Exception as exc:                      # noqa: BLE001
        logger.exception("BLAST failed")
        # The message is shown to a scientist, so it names the service rather
        # than leaking a stack shape.
        job.error = f"NCBI BLAST did not return a result ({type(exc).__name__})."
        job.status = "error"
    finally:
        job.finished_at = time.time()


def submit(seq: str, *, kind: Optional[str] = None, hitlist: int = 8,
           expect: float = 1e-5) -> BlastJob:
    """Start (or reuse) a BLAST for this sequence. Returns immediately."""
    seq = re.sub(r"[^A-Za-z]", "", seq or "").upper()
    kind = kind if kind in _PROGRAMS else guess_kind(seq)
    key = f"{kind}:{hash_sequence(seq)}"
    with _LOCK:
        cached = _CACHE.get(key)
        if cached is not None:
            return cached
        job = BlastJob(job_id=hashlib.sha256(key.encode()).hexdigest()[:12],
                       seq_hash=hash_sequence(seq), kind=kind)
        _CACHE[key] = job
        _BY_ID[job.job_id] = job
    threading.Thread(target=_run, args=(job, seq, hitlist, expect),
                     daemon=True).start()
    return job


def get(job_id: str) -> Optional[BlastJob]:
    with _LOCK:
        return _BY_ID.get(job_id)


def wait(job: BlastJob, timeout: Optional[float] = None,
         poll: float = 0.25) -> BlastJob:
    """Block until the job settles or `timeout` passes. Never raises.

    `timeout=None` reads AGENT_WAIT_SECONDS *at call time*, not as a default
    bound at import. Written the obvious way — `timeout=AGENT_WAIT_SECONDS` in
    the signature — the value is frozen when the module loads, so the constant
    cannot be tuned per deployment and a test that lowers it still waits the
    full 75 s. (It did: this file's suite took 77 s until this was fixed.)
    """
    if timeout is None:
        timeout = AGENT_WAIT_SECONDS
    deadline = time.time() + max(0.0, timeout)
    while time.time() < deadline and not job.done():
        time.sleep(poll)
    return job


def reset_cache() -> None:
    """Tests only."""
    with _LOCK:
        _CACHE.clear()
        _BY_ID.clear()