core/discovery.py

"""Aurelius core — relationship scoring and hidden-connection discovery.

Two query types beyond pathfinding, both source-agnostic (they speak only
the GraphSource protocol, so they work identically on live Wikipedia and
on an ingested protein or finance graph):

relate(a, b)   — how strongly are two nodes connected, and through what?
                 Evidence: direct edges, directed 2-hop paths (a→x→b),
                 co-targets (a→x←b), co-sources (x→a, x→b when backlinks
                 exist), and embedding similarity. Returns a 0-100 strength
                 plus the actual intermediaries, not just a number.

discover(a)    — Swanson ABC literature-based discovery: rank nodes C that
                 share many intermediaries B with A (A→B→C) but have NO
                 direct A→C edge. High bridge-count + high embedding
                 similarity + no direct link = a candidate hidden
                 connection. On stored sources the similarity half uses
                 the fused [text ; α·struct] vector, which is what lets a
                 structurally-close-but-textually-far node surface.
"""

from __future__ import annotations

import asyncio

import numpy as np

from .embedding import EmbeddingCache, cosine_similarity
from .representation import fuse
from .source import GraphSource
from .types import NodeRef

# Caps that keep live-mode (API-backed) discovery bounded.
_RELATE_NEIGHBOR_CAP   = 400
_DISCOVER_BRIDGE_CAP   = 20     # B nodes expanded (concurrently) per discover
_DISCOVER_CANDIDATES   = 800    # C pool cap before ranking
_DISCOVER_EMBED_CAP    = 200    # candidates text-embedded in live mode


def _store_of(source: GraphSource):
    """The GraphStore behind an ingested adapter, or None for live ones."""
    return getattr(source, "store", None) if hasattr(source, "ingested") else None


async def _fused_embedding(source: GraphSource, ref: NodeRef,
                           emb_cache: EmbeddingCache) -> np.ndarray | None:
    """Fused vector from the store when available, else live text embed."""
    store = _store_of(source)
    if store is not None:
        t = store.get_embedding(source.name, ref.id, "text")
        s = store.get_embedding(source.name, ref.id, "struct")
        if t is not None or s is not None:
            return fuse(t, s, source.name)
    info = await source.node_info(ref, rich=True)
    vecs = await emb_cache.embed([ref.key()], [info.text or ref.title])
    return vecs[0] if vecs and vecs[0].size else None


# ══════════════════════════════════════════════════════════════
# relate(a, b)
# ══════════════════════════════════════════════════════════════

async def relate(source: GraphSource, a_query: str, b_query: str) -> dict:
    a = await source.resolve(a_query)
    b = await source.resolve(b_query)
    if not a or not b:
        missing = a_query if not a else b_query
        return {"error": f"Cannot find: '{missing}'"}

    nb_a, nb_b = await asyncio.gather(
        source.neighbors(a), source.neighbors(b))
    out_a = {e.dst.id: e.dst for e in nb_a[:_RELATE_NEIGHBOR_CAP]}
    out_b = {e.dst.id: e.dst for e in nb_b[:_RELATE_NEIGHBOR_CAP]}

    in_a: dict[str, NodeRef] = {}
    in_b: dict[str, NodeRef] = {}
    if source.supports_backlinks:
        bk_a, bk_b = await asyncio.gather(
            source.back_neighbors(a), source.back_neighbors(b))
        in_a = {e.src.id: e.src for e in bk_a}
        in_b = {e.src.id: e.src for e in bk_b}

    direct_ab = b.id in out_a
    direct_ba = a.id in out_b

    # Directed 2-hop a→x→b: x is an out-neighbor of a AND an in-neighbor
    # of b (or, without backlinks, unverifiable — skipped).
    paths_ab = [out_a[x] for x in (set(out_a) & set(in_b))] if in_b else []
    paths_ba = [out_b[x] for x in (set(out_b) & set(in_a))] if in_a else []
    co_targets = [out_a[x] for x in (set(out_a) & set(out_b))]   # a→x←b
    co_sources = [in_a[x] for x in (set(in_a) & set(in_b))]      # x→a, x→b

    emb_cache = EmbeddingCache()
    ea, eb = await asyncio.gather(
        _fused_embedding(source, a, emb_cache),
        _fused_embedding(source, b, emb_cache))
    sim = cosine_similarity(ea, eb)

    # Composite strength: direct edges dominate, then 2-hop evidence
    # (saturating), then shared-neighbourhood evidence, then similarity.
    def _sat(count: int, scale: float) -> float:
        return 1.0 - float(np.exp(-count / scale))

    strength = (
        (0.35 if (direct_ab or direct_ba) else 0.0)
        + 0.30 * _sat(len(paths_ab) + len(paths_ba), 5.0)
        + 0.20 * _sat(len(co_targets) + len(co_sources), 20.0)
        + 0.15 * max(0.0, sim)
    )

    def _refs(refs: list[NodeRef], cap: int = 12) -> list[dict]:
        return [{"id": r.id, "title": r.title} for r in refs[:cap]]

    return {
        "source": source.name,
        "a": {"id": a.id, "title": a.title},
        "b": {"id": b.id, "title": b.title},
        "direct": {"a_to_b": direct_ab, "b_to_a": direct_ba},
        "paths_a_to_b": _refs(paths_ab),
        "paths_b_to_a": _refs(paths_ba),
        "n_paths": len(paths_ab) + len(paths_ba),
        "co_targets": _refs(co_targets),
        "n_co_targets": len(co_targets),
        "co_sources": _refs(co_sources),
        "n_co_sources": len(co_sources),
        "similarity": round(sim, 3),
        "strength": round(100 * min(1.0, strength), 1),
    }


# ══════════════════════════════════════════════════════════════
# discover(a)
# ══════════════════════════════════════════════════════════════

async def discover(source: GraphSource, a_query: str, k: int = 12) -> dict:
    a = await source.resolve(a_query)
    if not a:
        return {"error": f"Cannot find: '{a_query}'"}

    emb_cache = EmbeddingCache()
    ea = await _fused_embedding(source, a, emb_cache)

    nb_a = await source.neighbors(a)
    direct: dict[str, NodeRef] = {e.dst.id: e.dst for e in nb_a}
    if not direct:
        return {"error": f"'{a.title}' has no outbound links to walk."}

    # Choose the B set: rank a's neighbours by text similarity to a so the
    # bridges we expand are the *relevant* ones, then expand concurrently.
    b_refs = list(direct.values())
    if len(b_refs) > _DISCOVER_BRIDGE_CAP and ea is not None:
        infos = await source.node_infos(b_refs)
        keys  = [r.key() for r in b_refs]
        embs  = await emb_cache.embed(
            keys, [i.text or r.title for i, r in zip(infos, b_refs)])
        # Text-vs-fused dims can differ (stored fused vectors are longer);
        # compare on the shared text prefix length.
        d = min(ea.shape[0], embs[0].shape[0]) if embs[0].size else 0
        sims = [cosine_similarity(e[:d], ea[:d]) if e.size else 0.0
                for e in embs]
        order = np.argsort(sims)[::-1]
        b_refs = [b_refs[i] for i in order[:_DISCOVER_BRIDGE_CAP]]
    else:
        b_refs = b_refs[:_DISCOVER_BRIDGE_CAP]

    results = await asyncio.gather(
        *(source.neighbors(b) for b in b_refs), return_exceptions=True)

    # Aggregate candidates C with their bridges B (A→B→C, no A→C).
    bridges_of: dict[str, list[NodeRef]] = {}
    cand_ref: dict[str, NodeRef] = {}
    for b_ref, edges in zip(b_refs, results):
        if isinstance(edges, BaseException):
            continue
        for e in edges:
            c = e.dst
            if c.id == a.id or c.id in direct:
                continue
            bridges_of.setdefault(c.id, []).append(b_ref)
            cand_ref[c.id] = c
            if len(cand_ref) >= _DISCOVER_CANDIDATES * 4:
                break

    if not cand_ref:
        return {"a": {"id": a.id, "title": a.title}, "source": source.name,
                "candidates": []}

    # Rank: bridge support first, then embedding similarity on the top pool.
    pool = sorted(cand_ref, key=lambda c: -len(bridges_of[c]))[:_DISCOVER_CANDIDATES]

    store = _store_of(source)
    sims: dict[str, float] = {}
    if store is not None and ea is not None:
        for c in pool:
            t = store.get_embedding(source.name, c, "text")
            s = store.get_embedding(source.name, c, "struct")
            ec = fuse(t, s, source.name)
            d = min(ea.shape[0], ec.shape[0]) if ec is not None else 0
            sims[c] = cosine_similarity(ec[:d], ea[:d]) if d else 0.0
    elif ea is not None:
        head = pool[:_DISCOVER_EMBED_CAP]
        refs = [cand_ref[c] for c in head]
        infos = await source.node_infos(refs)
        embs = await emb_cache.embed(
            [r.key() for r in refs],
            [i.text or r.title for i, r in zip(infos, refs)])
        d0 = ea.shape[0]
        for c, e in zip(head, embs):
            d = min(d0, e.shape[0]) if e.size else 0
            sims[c] = cosine_similarity(e[:d], ea[:d]) if d else 0.0

    max_bridges = max(len(bridges_of[c]) for c in pool)
    scored = []
    for c in pool:
        support = len(bridges_of[c]) / max_bridges
        scored.append((0.55 * support + 0.45 * max(0.0, sims.get(c, 0.0)), c))
    scored.sort(reverse=True)

    candidates = [{
        "id": c,
        "title": cand_ref[c].title,
        "score": round(100 * sc, 1),
        "n_bridges": len(bridges_of[c]),
        "similarity": round(sims.get(c, 0.0), 3),
        "bridges": [{"id": b.id, "title": b.title}
                    for b in bridges_of[c][:6]],
    } for sc, c in scored[:k]]

    return {"a": {"id": a.id, "title": a.title}, "source": source.name,
            "candidates": candidates}