mcp_server/engine.py

"""Deterministic engine: the math behind the MCP tools. No LLM.

Loads the shipped embeddings + descriptions once (warm). For a pathway not in the shipped
`.npz`, BioLORD is lazy-loaded and the description (or prettified name) is embedded on the fly
— so the core needs no vector DB (exact-key lookup).
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import numpy as np

import config
from core.schema import ComparisonResult, CutoffParams, GseaResult, PairDrillDown
from core.signatures import Signature, build_signatures
from core.similarity import (
    best_match_matrix,
    cluster,
    cosine_matrix,
    nes_correlation,
)


def _prettify(name: str) -> str:
    return name.replace("_", " ").lower()


def _jaccard(a: set, b: set) -> Optional[float]:
    if not a and not b:
        return None
    union = a | b
    return len(a & b) / len(union) if union else None


class Engine:
    def __init__(
        self,
        embeddings_npz: Path = config.EMBEDDINGS_NPZ,
        descriptions_json: Path = config.DESCRIPTIONS_JSON,
    ):
        data = np.load(embeddings_npz, allow_pickle=False)
        self._emb: dict[str, np.ndarray] = {
            n: v.astype(float) for n, v in zip(data["names"].tolist(), data["vectors"])
        }
        self._desc: dict[str, str] = (
            json.loads(Path(descriptions_json).read_text()) if Path(descriptions_json).exists() else {}
        )
        self._model = None  # lazy BioLORD, only if an unknown pathway appears
        self._embedder_id = config.EMBEDDER_DEFAULT

    # --- embedding lookup ---------------------------------------------------
    def _embed(self, name: str) -> np.ndarray:
        vec = self._emb.get(name)
        if vec is not None:
            return vec
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self._embedder_id, device=config.EMBED_DEVICE)
        text = self._desc.get(name) or _prettify(name)
        vec = np.asarray(self._model.encode([text], normalize_embeddings=True)[0], dtype=float)
        self._emb[name] = vec  # cache for the rest of this session
        return vec

    def embed_query(self, text: str) -> np.ndarray:
        """Embed a free-text query / concept with the warm BioLORD (no caching; transient)."""
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self._embedder_id, device=config.EMBED_DEVICE)
        return np.asarray(self._model.encode([text], normalize_embeddings=True)[0], dtype=float)

    # --- MCP tool implementations ------------------------------------------
    def compare(self, results: list[dict], params: Optional[dict] = None) -> dict:
        cparams = CutoffParams.model_validate(params or {})
        warnings: list[str] = []

        sigs: list[Signature] = []
        for rdict in results:
            sigs.extend(build_signatures(GseaResult.model_validate(rdict), cparams, self._embed, warnings))

        labels = [s.label for s in sigs]
        if not sigs:
            return ComparisonResult(
                signatures=[], similarity=[], order=[], linkage=[], clusters={},
                nes_correlation=[], pairs=[], warnings=warnings or ["no signatures after filtering"],
            ).model_dump()

        sim = (best_match_matrix(sigs) if cparams.aggregation == "best_match"
               else cosine_matrix(np.stack([s.vector for s in sigs])))
        order, linkage_rows, flat = cluster(sim, config.CLUSTER_CUT_FRACTION)
        clusters = {labels[i]: int(flat[i]) for i in range(len(labels))}
        nes_corr = nes_correlation(sigs)
        pairs = self._drill_down(sigs, sim, nes_corr)

        return ComparisonResult(
            signatures=labels,
            similarity=sim.tolist(),
            order=order,
            linkage=linkage_rows,
            clusters=clusters,
            nes_correlation=nes_corr,
            pairs=pairs,
            warnings=warnings,
        ).model_dump()

    def _drill_down(
        self, sigs: list[Signature], sim: np.ndarray, nes_corr: list, top: int = config.TOP_PAIRS,
    ) -> list[PairDrillDown]:
        m = len(sigs)
        ranked = sorted(
            ((sim[i][j], i, j) for i in range(m) for j in range(i + 1, m)),
            reverse=True,
        )[:top]
        out: list[PairDrillDown] = []
        for _, i, j in ranked:
            a, b = sigs[i], sigs[j]
            sa, sb = set(a.pathways), set(b.pathways)
            le_a = {g for genes in a.leading_edge for g in genes}
            le_b = {g for genes in b.leading_edge for g in genes}
            out.append(PairDrillDown(
                a=a.label, b=b.label,
                semantic=float(sim[i][j]),
                nes_corr=nes_corr[i][j],
                shared_pathways=sorted(sa & sb),
                unique_a=sorted(sa - sb),
                unique_b=sorted(sb - sa),
                top_shared_themes=self._top_themes(a, b),
                leading_edge_jaccard=_jaccard(le_a, le_b),
            ))
        return out

    @staticmethod
    def _top_themes(a: Signature, b: Signature, k: int = 3) -> list[dict]:
        """Top semantically-matched pathway pairs across the two signatures."""
        M = np.clip(np.asarray(a.embeddings) @ np.asarray(b.embeddings).T, -1.0, 1.0)
        flat = sorted(
            ((float(M[x, y]), x, y) for x in range(M.shape[0]) for y in range(M.shape[1])),
            reverse=True,
        )
        themes, seen_a, seen_b = [], set(), set()
        for s, x, y in flat:
            if x in seen_a or y in seen_b:
                continue
            themes.append({"a": a.pathways[x], "b": b.pathways[y], "sim": round(s, 4)})
            seen_a.add(x); seen_b.add(y)
            if len(themes) >= k:
                break
        return themes

    def describe(self, names: list[str]) -> dict:
        return {n: self._desc.get(n) or _prettify(n) for n in names}

    def health(self) -> dict:
        return {
            "status": "ok",
            "n_embeddings": len(self._emb),
            "embed_dim": config.EMBED_DIM,
            "embedder": self._embedder_id,
            "model_loaded": self._model is not None,
        }