core/source.py

"""Aurelius core — the GraphSource protocol and adapter registry.

An adapter makes one interconnected dataset navigable. Implement this
protocol (plus, for large static datasets, an ingest script that fills the
store) and the whole engine — navigator, relate, discover, the UI — works
on your graph with zero engine changes. That property is the product.

Two adapter modes:
  live      — neighbors/backlinks fetched from an upstream API per call
              (Wikipedia's links/linkshere, OpenAlex's refs/cited-by).
  ingested  — the graph was written into core.store by an ingest run;
              the adapter answers from the store (biomed, news, finance).
              StoreBackedSource below is the shared implementation for
              this mode.
"""

from __future__ import annotations

import abc
from typing import Optional

from .types import Edge, NodeInfo, NodeRef


class GraphSource(abc.ABC):
    """Contract every data-source adapter implements."""

    #: unique registry name, e.g. "wikipedia"
    name: str = ""
    #: human description for /api/sources and the UI picker
    description: str = ""
    #: edge types this source can emit (documentation + UI legend)
    edge_types: tuple[str, ...] = ("link",)
    #: False → back_neighbors() is unsupported or expensive; the navigator
    #: then skips goal-zone construction and runs forward-only.
    supports_backlinks: bool = True

    # ── identity ─────────────────────────────────────────────────────────
    @abc.abstractmethod
    async def resolve(self, query: str) -> Optional[NodeRef]:
        """Free-text query → a node, or None. Adapters own their notion of
        fuzzy matching / disambiguation (the v1 lesson: a resolved title is
        not necessarily a *search-worthy* target — handle stubs here)."""

    # ── structure ────────────────────────────────────────────────────────
    @abc.abstractmethod
    async def neighbors(self, n: NodeRef, *,
                        hunt_id: str | None = None,
                        priority_ids: set[str] | None = None) -> list[Edge]:
        """Outbound edges of n. `hunt_id`/`priority_ids` are optional
        fetch hints (keep paginating until hunt_id is found; order
        priority_ids first) — adapters that don't paginate ignore them."""

    async def back_neighbors(self, n: NodeRef, limit: int = 500) -> list[Edge]:
        """Inbound edges (edges whose dst is n). Only called when
        supports_backlinks is True."""
        raise NotImplementedError

    # ── content ──────────────────────────────────────────────────────────
    async def node_info(self, n: NodeRef, rich: bool = False) -> NodeInfo:
        """Embedding text + display summary for n. rich=True may spend an
        extra fetch for a fuller text (used for the two search endpoints,
        whose embeddings anchor every score in the run)."""
        return NodeInfo(text=n.title)

    async def node_infos(self, ns: list[NodeRef]) -> list[NodeInfo]:
        """Batch form of node_info (rich=False). Adapters that already hold
        the info from a neighbors() fetch should override to answer from
        cache without I/O — the navigator calls this once per expansion
        with every candidate."""
        return [await self.node_info(n) for n in ns]

    # ── recommendations ──────────────────────────────────────────────────
    async def suggest(self, query: str, limit: int = 8) -> list[dict]:
        """Type-ahead suggestions for the search box, drawn from THIS
        source's own vocabulary — the fix for the old Wikipedia-only
        autocomplete. Each item is
            {id, title, kind?, subtitle?}
        where `kind` groups results (company/paper/disease/…) and
        `subtitle` is a short human hint. Default: no suggestions (a live
        source with nothing cheap to offer simply returns []); adapters
        override with a domain-appropriate lookup."""
        return []

    # ── niceties (optional) ──────────────────────────────────────────────
    async def edge_display(self, src: NodeRef, dst: NodeRef) -> Optional[str]:
        """Human-facing rendering of an edge (Wikipedia: the piped link
        text). None = nothing special to show."""
        return None

    async def sample_pair(self) -> Optional[tuple[str, str]]:
        """Two queries that make a good demo pair, or None."""
        return None


# ══════════════════════════════════════════════════════════════
# Registry
# ══════════════════════════════════════════════════════════════

_REGISTRY: dict[str, GraphSource] = {}


def register(source: GraphSource) -> GraphSource:
    if not source.name:
        raise ValueError("GraphSource.name must be set")
    _REGISTRY[source.name] = source
    return source


def get_source(name: str) -> GraphSource:
    try:
        return _REGISTRY[name]
    except KeyError:
        raise KeyError(f"Unknown source '{name}'. Registered: {sorted(_REGISTRY)}")


def list_sources() -> list[GraphSource]:
    return list(_REGISTRY.values())