spectrag 0.1.0__py3-none-any.whl

spectrag/__init__.py

@@ -0,0 +1,22 @@
+"""SpectRAG: spectral graph-augmented retrieval for RAG (arXiv:2610.36991).
+
+SpectRAGIngestion builds the index (offline, once per document set).
+SpectRAGRetrieval finds the context for a question and can answer it (online, per question).
+
+    from spectrag import SpectRAGIngestion, SpectRAGRetrieval
+    SpectRAGIngestion("my_index").ingest("docs/")
+    print(SpectRAGRetrieval("my_index").search("my question")["context"])
+"""
+
+from spectrag._version import INDEX_FORMAT, __version__
+from spectrag.config import DEFAULTS, make_settings
+from spectrag.errors import (DocumentError, EmbedderError, IndexError_, NoDocumentsError,
+                             SettingsError, SpectRAGError, UnknownSettingError)
+from spectrag.ingestion import SpectRAGIngestion
+from spectrag.log import enable_logging
+from spectrag.retrieval import SpectRAGRetrieval
+
+__all__ = ["SpectRAGIngestion", "SpectRAGRetrieval", "make_settings", "DEFAULTS",
+           "enable_logging", "SpectRAGError", "SettingsError", "UnknownSettingError",
+           "IndexError_", "EmbedderError", "DocumentError", "NoDocumentsError",
+           "__version__", "INDEX_FORMAT"]

spectrag/_version.py

@@ -0,0 +1,7 @@
+"""The package version (read by pyproject.toml) and the index format version."""
+
+__version__ = "0.1.0"
+
+# Bump INDEX_FORMAT when the files of an index change in a way older code cannot read.
+# 1 = JSON edge lists (before the Phase 1 audit), 2 = binary .npy/.npz edges and signals.
+INDEX_FORMAT = 2

spectrag/cli.py

@@ -0,0 +1,115 @@
+"""The `spectrag` command.
+
+    spectrag ingest docs/ --index my_index [--set chunk_tokens=256 --set embedder=ollama]
+    spectrag search "my question" --index my_index [--budget 4000] [--json]
+    spectrag answer "my question" --index my_index --model gemma3:4b
+    spectrag info --index my_index
+    spectrag graph --index my_index --out graph.html [--question "my question"]
+"""
+
+import argparse
+import json
+import sys
+
+from spectrag._version import __version__
+from spectrag.errors import SpectRAGError
+from spectrag.log import enable_logging
+
+
+def parse_sets(pairs):
+    """["alpha=0.5", "selector=mmr"] -> {"alpha": 0.5, "selector": "mmr"} (values read as JSON)."""
+    settings = {}
+    for pair in pairs or []:
+        if "=" not in pair:
+            raise SystemExit("--set needs key=value, got %r" % pair)
+        key, value = pair.split("=", 1)
+        try:
+            settings[key] = json.loads(value)
+        except ValueError:
+            settings[key] = value  # plain text, e.g. selector=mmr
+    return settings
+
+
+def command_ingest(args):
+    from spectrag import SpectRAGIngestion
+    info = SpectRAGIngestion(args.index, parse_sets(args.set)).ingest(args.path)
+    print(json.dumps(info, indent=1))
+
+
+def command_search(args):
+    from spectrag import SpectRAGRetrieval
+    result = SpectRAGRetrieval(args.index, parse_sets(args.set)).search(args.question, args.budget)
+    if args.json:
+        print(json.dumps({key: result[key] for key in ("question", "tokens", "chunks", "context")}, indent=1))
+        return
+    for chunk in result["chunks"]:
+        print("%2d. %-30s score %.3f  %d tokens  (%s)" % (chunk["rank"], chunk["doc_id"], chunk["score"],
+                                                         chunk["tokens"], chunk["reason"]))
+    print("\n" + result["context"])
+
+
+def command_answer(args):
+    from spectrag import SpectRAGRetrieval
+    from spectrag.utils.llm import ollama_llm
+    search = SpectRAGRetrieval(args.index, parse_sets(args.set))
+    result = search.answer(args.question, ollama_llm(args.model), args.budget)
+    print(result["answer"])
+    print("\nSources: " + ", ".join(sorted({c["doc_id"] for c in result["chunks"]})))
+
+
+def command_info(args):
+    from spectrag.utils.store import check_index_folder, load_index
+    from pathlib import Path
+    check_index_folder(Path(args.index))
+    print(json.dumps(load_index(args.index, only=("info",))["info"], indent=1))
+
+
+def command_graph(args):
+    from spectrag.utils.graph_view import save_graph_view
+    path = save_graph_view(args.index, args.out, args.question, parse_sets(args.set), args.max_nodes)
+    print("Graph view written to %s (open it in a browser)" % path)
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(prog="spectrag", description="SpectRAG: spectral graph RAG")
+    parser.add_argument("--version", action="version", version="spectrag " + __version__)
+    parser.add_argument("-v", "--verbose", action="store_true", help="print progress messages")
+    commands = parser.add_subparsers(dest="command", required=True)
+    ingest = commands.add_parser("ingest", help="build an index from a file or folder")
+    ingest.add_argument("path")
+    for name, helptext in [("search", "find the context for a question"), ("answer", "answer with an LLM")]:
+        sub = commands.add_parser(name, help=helptext)
+        sub.add_argument("question")
+        sub.add_argument("--budget", type=int, default=None, help="token budget (default 12000)")
+        if name == "search":
+            sub.add_argument("--json", action="store_true", help="print JSON")
+        else:
+            sub.add_argument("--model", default="gemma3:4b", help="Ollama model")
+    info = commands.add_parser("info", help="show what an index contains")
+    graph = commands.add_parser("graph", help="write an interactive graph view (HTML)")
+    graph.add_argument("--out", default="graph.html", help="HTML file to write")
+    graph.add_argument("--question", default=None, help="show this question's retrieval path")
+    graph.add_argument("--max-nodes", type=int, default=1500, dest="max_nodes")
+    for sub in (ingest, commands.choices["search"], commands.choices["answer"], info, graph):
+        sub.add_argument("--index", required=True, help="index folder")
+    for sub in (ingest, commands.choices["search"], commands.choices["answer"], graph):
+        sub.add_argument("--set", action="append", metavar="KEY=VALUE", help="change a setting")
+    return parser
+
+
+def main(argv=None):
+    args = build_parser().parse_args(argv)
+    if args.verbose:
+        enable_logging()
+    handlers = {"ingest": command_ingest, "search": command_search, "answer": command_answer,
+                "info": command_info, "graph": command_graph}
+    try:
+        handlers[args.command](args)
+    except SpectRAGError as error:
+        print("spectrag: error: %s" % error, file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

spectrag/config.py

@@ -0,0 +1,96 @@
+"""Default settings for SpectRAG.
+
+Values marked "paper" come from the SpectRAG paper (Tables 1-3).
+Values marked "assumed" are not given in the paper, so we picked a sensible value.
+"""
+
+import copy
+import difflib
+
+from spectrag.errors import UnknownSettingError
+from spectrag.validation import validate_settings
+
+# Edge-signal weights from the paper's "signal configuration 1" (Table 9).
+EDGE_WEIGHTS_CONFIG_1 = {
+    "semantic": 0.300,
+    "section": 0.250,
+    "entity": 0.200,
+    "adjacency": 0.084,
+    "utility": 0.083,
+    "citation": 0.083,
+}
+
+# Edge-signal weights from the paper's "signal configuration 2" (Table 10).
+EDGE_WEIGHTS_CONFIG_2 = {
+    "semantic": 0.150,
+    "section": 0.200,
+    "entity": 0.250,
+    "adjacency": 0.100,
+    "utility": 0.280,
+    "citation": 0.020,
+}
+
+DEFAULTS = {
+    # --- Chunking (paper) ---
+    "chunk_tokens": 1200,
+    "overlap_tokens": 100,
+    "min_chunk_tokens": 32,
+    "semantic_percentile": 90,
+    "semantic_embedder": None,   # embedder for semantic-chunking breaks; None = "embedder" (paper)
+    # --- Embedding ---
+    "embedder": "hash",          # "hash" (offline), "ollama" (nomic-embed-text), or your own function
+    "embed_dims": 768,           # paper: nomic-embed-text is 768-d (used by the hash embedder)
+    "ollama_model": "nomic-embed-text",
+    "ollama_url": "http://localhost:11434",
+    # --- Vector search (Section 6.2) ---
+    "ann_backend": "auto",       # "exact", "faiss", "hnsw", or "auto" (exact up to 20k chunks)
+    "hnsw_m": 32,                # HNSW links per node
+    "hnsw_ef_construction": 200,
+    "hnsw_ef_search": 128,
+    "dense_top_n": 1000,         # with an approximate backend: dense scores for the top n only
+    # --- Graph (weights: paper, the rest: assumed) ---
+    "edge_weights": EDGE_WEIGHTS_CONFIG_1,
+    "semantic_k": 10,            # assumed
+    "edge_threshold": 0.05,      # assumed
+    "graph_k": 10,               # assumed: each chunk keeps at most its 10 strongest edges
+    # --- Spectral (paper) ---
+    "spectral_dims": 32,
+    "n_clusters": None,          # assumed: None means about sqrt(n / 2) clusters
+    # --- Retrieval (paper) ---
+    "seed_count": 20,
+    "hops": 2,
+    "max_candidates": 200,
+    "token_budget": 12000,
+    # --- Retrieval (assumed) ---
+    "rrf_k": 60,                 # the usual RRF constant
+    "bm25_k1": 1.5,              # standard BM25 values (Robertson and Zaragoza)
+    "bm25_b": 0.75,
+    "bm25_weight": 0.5,          # hybrid relevance = 0.5 * dense + 0.5 * BM25
+    "alpha": 0.3,
+    "ppr_beta": 0.15,
+    "ppr_steps": 30,
+    "diffusion_t": 1.0,          # heat-kernel scale t for spectral closeness (Eq. 17)
+    "proximity_mode": "hybrid",  # "hybrid" (paper), "ppr", "spectral", "heat", "none"
+    "selector": "spectrag",      # "spectrag", "mmr", "protected_mmr", "anchor_expand_ce", "topk"
+    "mmr_lambda": 0.7,           # MMR: weight of relevance vs. novelty
+    "core_size": 3,
+    "precision_floor": 0.4,
+    "select_weights": {"a": 1.0, "b": 0.3, "c": 0.2, "d": 0.1, "lambda": 0.5},
+    "length_power": 0.5,
+}
+
+
+def make_settings(overrides=None):
+    """Return a fresh copy of DEFAULTS with the user's overrides applied and checked.
+
+    Raises UnknownSettingError (a KeyError) for a misspelled name, and SettingsError
+    (a ValueError) for a value of the wrong type or out of range.
+    """
+    settings = copy.deepcopy(DEFAULTS)
+    for key, value in (overrides or {}).items():
+        if key not in settings:
+            close = difflib.get_close_matches(key, DEFAULTS, n=1)
+            hint = " (did you mean %r?)" % close[0] if close else ""
+            raise UnknownSettingError("Unknown setting: %r%s" % (key, hint))
+        settings[key] = copy.deepcopy(value)
+    return validate_settings(settings)

spectrag/errors.py

@@ -0,0 +1,40 @@
+"""SpectRAG's exceptions. Each one also inherits the matching built-in exception, so code that
+catches ValueError, KeyError, FileNotFoundError or ConnectionError keeps working.
+
+    try:
+        search = SpectRAGRetrieval("my_index")
+    except spectrag.SpectRAGError as error:
+        print(error)
+"""
+
+
+class SpectRAGError(Exception):
+    """Base class for every error raised on purpose by SpectRAG."""
+
+
+class SettingsError(SpectRAGError, ValueError):
+    """A setting has the wrong type, is out of range, or is not one of the allowed choices."""
+
+
+class UnknownSettingError(SettingsError, KeyError):
+    """A setting name that SpectRAG does not know (often a typo)."""
+
+    def __str__(self):  # KeyError would otherwise print the message with quotes around it
+        return str(self.args[0]) if self.args else ""
+
+
+class IndexError_(SpectRAGError, FileNotFoundError):
+    """The index folder is missing, incomplete, or was made by an incompatible version.
+    (The trailing underscore avoids hiding Python's built-in IndexError.)"""
+
+
+class EmbedderError(SpectRAGError, ConnectionError):
+    """The embedder cannot be reached, or its vectors do not match the index."""
+
+
+class DocumentError(SpectRAGError, ValueError):
+    """Documents given to ingest_documents() are not in the expected form."""
+
+
+class NoDocumentsError(DocumentError):
+    """There was no document with text to ingest."""

spectrag/ingestion.py

@@ -0,0 +1,141 @@
+"""SpectRAGIngestion: the offline pipeline (paper Table 1). Run it once per document set.
+
+    ingest = SpectRAGIngestion("my_index")
+    info = ingest.ingest("examples/docs")
+"""
+
+import shutil
+import time
+from pathlib import Path
+
+import numpy as np
+
+from spectrag._version import INDEX_FORMAT, __version__
+from spectrag.config import make_settings
+from spectrag.errors import DocumentError, IndexError_, NoDocumentsError
+from spectrag.log import log_step, logger
+from spectrag.utils.chunker import chunk_documents
+from spectrag.utils.edges import build_edge_signals, count_edges
+from spectrag.utils.embedder import make_embedder
+from spectrag.utils.entities import add_entities
+from spectrag.utils.graph_build import build_graph, build_spectral
+from spectrag.utils.loader import load_docs
+from spectrag.utils.questions import add_utility_questions
+from spectrag.utils.references import add_references
+from spectrag.utils.store import edges_to_array, save_index
+from spectrag.utils.vector_index import build_vector_index, save_vector_index
+
+
+class SpectRAGIngestion:
+    """Load -> chunk -> embed -> edge signals -> weighted graph -> spectral -> save."""
+
+    def __init__(self, index_dir, settings=None, entity_fn=None, question_fn=None):
+        self.index_dir = index_dir
+        self.settings = make_settings(settings)
+        self.embed = make_embedder(self.settings)
+        # Semantic chunking can use a faster embedder for sentence distances (None = the same).
+        split_choice = self.settings["semantic_embedder"]
+        self.split_embed = self.embed if split_choice is None else \
+            make_embedder(dict(self.settings, embedder=split_choice))
+        self.entity_fn = entity_fn        # optional: text -> list of entities
+        self.question_fn = question_fn    # optional: text -> list of utility questions
+
+    def make_chunks(self, docs):
+        """Split documents into chunks with entities, references and utility questions."""
+        chunks = chunk_documents(docs, self.settings, embed_fn=self.split_embed)
+        add_entities(chunks, self.entity_fn)
+        add_references(chunks)
+        add_utility_questions(chunks, self.question_fn)
+        return chunks
+
+    def make_graph(self, signals):
+        """Fuse the signals (Eq. 10), drop weak edges, keep each chunk's strongest graph_k."""
+        return build_graph(signals, self.settings)
+
+    def make_spectral(self, edges, node_count):
+        """Laplacian (Eq. 11) -> spectral coordinates (Eq. 12-13) -> clusters."""
+        return build_spectral(edges, node_count, self.settings)
+
+    def ingest(self, path):
+        """Ingest a file or a folder (.txt, .md, .pdf), save the index, return a summary."""
+        docs = load_docs(path)
+        if not docs:
+            raise NoDocumentsError("No .txt, .md or .pdf file with text found in " + str(path))
+        return self.ingest_documents(docs)
+
+    def ingest_documents(self, docs):
+        """Ingest documents already in memory: [{"doc_id": ..., "text": ...}, ...]."""
+        check_documents(docs)
+        start = begin = time.time()
+        chunks = self.make_chunks(docs)
+        start = log_step("chunking", start, documents=len(docs), chunks=len(chunks))
+        embeddings = self.embed([chunk["text"] for chunk in chunks], kind="document")
+        vector_index = build_vector_index(embeddings, self.settings)  # exact, FAISS or HNSW
+        start = log_step("embedding", start, backend=vector_index["backend"])
+        signals, question_vectors, owners = build_edge_signals(
+            chunks, embeddings, self.embed, self.settings, vector_index)
+        start = log_step("edge signals", start, **count_edges(signals))
+        edges, weights = self.make_graph(signals)
+        coords, eigenvalues, clusters = self.make_spectral(edges, len(chunks))
+        start = log_step("graph and spectral", start, edges=len(edges), dims=coords.shape[1])
+        for chunk, cluster in zip(chunks, clusters):
+            chunk["cluster"] = int(cluster)
+        info = {"documents": len({c["doc_id"] for c in chunks}), "chunks": len(chunks),
+                "signal_edges": count_edges(signals), "graph_edges": len(edges),
+                "weights_used": {name: round(w, 4) for name, w in weights.items()},
+                "spectral_dims": coords.shape[1], "ann_backend": vector_index["backend"],
+                "clusters": len(set(clusters.tolist())),
+                "format_version": INDEX_FORMAT, "spectrag_version": __version__,
+                "seconds": round(time.time() - begin, 2)}
+        partial = prepare_folder(self.index_dir)
+        save_index(partial, {
+            "chunks": chunks, "embeddings": embeddings, "edges": edges_to_array(edges),
+            "signals": {name: edges_to_array(found) for name, found in signals.items()},
+            "spectral": coords, "eigenvalues": eigenvalues,
+            "question_vectors": question_vectors,
+            "question_owners": np.asarray(owners, dtype=np.int64),
+            "settings": self.settings, "info": info,
+        })
+        save_vector_index(vector_index, partial)
+        replace_folder(partial, self.index_dir)
+        log_step("saved index", start, folder=self.index_dir)
+        return info
+
+
+def check_documents(docs):
+    """Every document needs a text and a unique doc_id."""
+    if not docs:
+        raise NoDocumentsError("No documents with text to ingest")
+    seen = set()
+    for number, doc in enumerate(docs):
+        if not isinstance(doc, dict) or not isinstance(doc.get("text"), str) or "doc_id" not in doc:
+            raise DocumentError("Document %d must be a dict with 'doc_id' and 'text' (a string)" % number)
+        if doc["doc_id"] in seen:
+            raise DocumentError("Two documents have the same doc_id: %r" % (doc["doc_id"],))
+        seen.add(doc["doc_id"])
+    if not any(doc["text"].strip() for doc in docs):
+        raise NoDocumentsError("Every document is empty")
+
+
+def is_index(folder):
+    return (folder / "info.json").is_file() and (folder / "chunks.json").is_file()
+
+
+def prepare_folder(index_dir):
+    """A fresh '<index_dir>.partial' folder to write into, so a crash never leaves half an index."""
+    target = Path(index_dir)
+    if target.exists() and (not target.is_dir() or (any(target.iterdir()) and not is_index(target))):
+        raise IndexError_("%s already exists and is not a SpectRAG index; choose another folder" % target)
+    partial = target.with_name(target.name + ".partial")
+    if partial.exists():
+        shutil.rmtree(partial)
+    return partial
+
+
+def replace_folder(partial, index_dir):
+    """Move the finished index into place (replacing an older SpectRAG index)."""
+    target = Path(index_dir)
+    if target.exists():
+        logger.info("replacing the existing index in %s", target)
+        shutil.rmtree(target)
+    partial.rename(target)

spectrag/log.py

@@ -0,0 +1,30 @@
+"""Logging. SpectRAG is silent by default (like any library); switch messages on with:
+
+    import spectrag
+    spectrag.enable_logging()          # INFO: one line per pipeline step, with its time
+    spectrag.enable_logging("DEBUG")   # DEBUG: also one line per query
+"""
+
+import logging
+import time
+
+logger = logging.getLogger("spectrag")
+logger.addHandler(logging.NullHandler())
+
+
+def enable_logging(level="INFO"):
+    """Print SpectRAG's log messages to the terminal (stderr)."""
+    if not any(getattr(h, "spectrag", False) for h in logger.handlers):
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s"))
+        handler.spectrag = True
+        logger.addHandler(handler)
+    logger.setLevel(level)
+    return logger
+
+
+def log_step(name, start, **details):
+    """Log one finished step: 'chunking: 1.23 s (chunks=8)'. Returns the time now."""
+    extra = ", ".join("%s=%s" % item for item in details.items())
+    logger.info("%s: %.2f s%s", name, time.time() - start, " (%s)" % extra if extra else "")
+    return time.time()

spectrag/retrieval.py

@@ -0,0 +1,154 @@
+"""SpectRAGRetrieval: the online pipeline (paper Table 2, Algorithm 1). Run it per question.
+
+    search = SpectRAGRetrieval("my_index")
+    result = search.search("Which company did the co-founder of Aurora Labs join?")
+    print(result["context"])
+"""
+
+import time
+from pathlib import Path
+
+from spectrag._version import INDEX_FORMAT
+from spectrag.config import make_settings
+from spectrag.errors import EmbedderError, IndexError_
+from spectrag.log import logger
+from spectrag.utils.bm25 import build_bm25
+from spectrag.utils.diffusion import score_candidates
+from spectrag.utils.embedder import make_embedder
+from spectrag.utils.graph import degrees, neighbors
+from spectrag.utils.graph_build import build_graph, build_spectral
+from spectrag.utils.prompt import build_prompt, format_context
+from spectrag.utils.seeding import find_seeds
+from spectrag.utils.selectors import select
+from spectrag.utils.store import array_to_edges, check_index_folder, load_index
+from spectrag.utils.vector_index import load_vector_index
+
+# These change the graph itself, so they need rebuild_graph() instead of a per-call override.
+GRAPH_SETTINGS = ("edge_weights", "edge_threshold", "graph_k", "spectral_dims", "n_clusters")
+
+
+class SpectRAGRetrieval:
+    """Seeds -> graph expansion -> proximity re-scoring -> token-budget selection -> context."""
+
+    def __init__(self, index_dir, settings=None, embed_fn=None, rerank_fn=None):
+        """Load an index made by SpectRAGIngestion.
+
+        settings:  settings to change, e.g. {"token_budget": 4000, "selector": "mmr"}.
+                   Graph settings (edge_threshold, graph_k, ...) rebuild the graph from the
+                   saved signals, without embedding anything again.
+        embed_fn:  pass the same function again if the index was built with your own embedder.
+        rerank_fn: optional (question, texts) -> scores, e.g. a cross-encoder, used by the
+                   "anchor_expand_ce" selector.
+        """
+        self.index_dir = Path(index_dir)
+        check_index_folder(self.index_dir)
+        self.index = load_index(index_dir, skip=("signals",))  # raw signals are loaded on demand
+        made_with = self.index["info"].get("format_version", 1)
+        if made_with > INDEX_FORMAT:
+            raise IndexError_("This index has format %d; this SpectRAG version reads up to %d. "
+                              "Upgrade spectrag." % (made_with, INDEX_FORMAT))
+        self.index["edges"] = array_to_edges(self.index["edges"])
+        self.index["question_owners"] = [int(owner) for owner in self.index["question_owners"]]
+        built_with = make_settings(self.index["settings"])
+        self.settings = make_settings(dict(built_with, **(settings or {})))
+        if embed_fn is not None:
+            self.settings["embedder"] = embed_fn
+        elif str(self.settings["embedder"]).startswith("<function"):
+            raise ValueError("This index was built with your own embed function. "
+                             "Pass it again: SpectRAGRetrieval(index_dir, embed_fn=...)")
+        self.embed = make_embedder(self.settings)
+        self.rerank_fn = rerank_fn
+        self.bm25 = build_bm25([chunk["text"] for chunk in self.index["chunks"]])
+        self.vector_index = load_vector_index(self.index_dir, self.index["embeddings"], self.settings)
+        if any(self.settings[key] != built_with[key] for key in GRAPH_SETTINGS):
+            self.rebuild_graph()
+        self.refresh_graph()
+
+    def refresh_graph(self):
+        """Neighbour table (CSR) and degrees for the current edges."""
+        count = len(self.index["chunks"])
+        self.table = neighbors(self.index["edges"], count)
+        self.index["degree"] = degrees(self.index["edges"], count)
+
+    def rebuild_graph(self, overrides=None):
+        """Change graph settings without re-embedding: fuse the saved signals again (Eq. 10)
+        and recompute spectral coordinates and clusters (Eq. 11-13). In memory only."""
+        self.settings = make_settings(dict(self.settings, **(overrides or {})))
+        saved = load_index(self.index_dir, only=("signals",))["signals"]
+        signals = {name: array_to_edges(edges) for name, edges in saved.items()}
+        edges, _ = build_graph(signals, self.settings)
+        coords, eigenvalues, clusters = build_spectral(edges, len(self.index["chunks"]), self.settings)
+        self.index.update({"edges": edges, "spectral": coords, "eigenvalues": eigenvalues})
+        for chunk, cluster in zip(self.index["chunks"], clusters):
+            chunk["cluster"] = int(cluster)
+        self.refresh_graph()
+
+    def settings_for(self, overrides):
+        """The settings for one call: the current settings plus per-call overrides."""
+        if not overrides:
+            return self.settings
+        changed = [key for key in GRAPH_SETTINGS
+                   if key in overrides and overrides[key] != self.settings[key]]
+        if changed:
+            raise ValueError("Use rebuild_graph() to change graph settings: " + ", ".join(changed))
+        return make_settings(dict(self.settings, **overrides))
+
+    def retrieve(self, question, settings=None, query_vector=None):
+        """Stages 1-4. Returns (seed result, scored candidates, picks)."""
+        settings = self.settings_for(settings)
+        start = time.time()
+        if query_vector is None:
+            query_vector = self.embed([question], kind="query")[0]
+        if len(query_vector) != self.index["embeddings"].shape[1]:
+            raise EmbedderError("The query vector has %d dimensions but the index has %d: search "
+                                "with the same embedder the index was built with."
+                                % (len(query_vector), self.index["embeddings"].shape[1]))
+        seeds = find_seeds(question, query_vector, self.index, self.bm25, settings,
+                           self.vector_index)
+        scored = score_candidates(seeds["seeds"], seeds["relevance"], self.index,
+                                  self.table, settings)
+        rerank = None
+        if settings["selector"] == "anchor_expand_ce" and self.rerank_fn is not None:
+            texts = [self.index["chunks"][c]["text"] for c in scored["candidates"]]
+            rerank = self.rerank_fn(question, texts)
+        picks = select(scored, self.index, settings, self.table, rerank)
+        logger.debug("query %r: %d seeds, %d candidates, %d picked in %.1f ms", question[:60],
+                     len(seeds["seeds"]), len(scored["candidates"]), len(picks),
+                     1000 * (time.time() - start))
+        return seeds, scored, picks
+
+    def describe(self, k, reason, rank, scored):
+        """One selected chunk with all its scores (the 'ranked source metadata')."""
+        chunk = self.index["chunks"][scored["candidates"][k]]
+        info = {key: chunk[key] for key in ["id", "doc_id", "position", "section", "tokens", "text"]}
+        info.update({"rank": rank, "reason": reason, "hop": scored["hop"][k],
+                     "score": float(scored["scores"][k]),
+                     "relevance": float(scored["relevance"][k]),
+                     "proximity": float(scored["proximity"][k])})
+        return info
+
+    def search(self, question, token_budget=None, settings=None, query_vector=None):
+        """Find the context for a question. Returns a dict with 'context', 'chunks', ...
+
+        settings: per-call overrides, e.g. {"selector": "mmr", "alpha": 0.5}.
+        """
+        overrides = dict(settings or {})
+        if token_budget is not None:
+            overrides["token_budget"] = token_budget
+        seeds, scored, picks = self.retrieve(question, overrides, query_vector)
+        chunks = [self.describe(k, reason, rank, scored)
+                  for rank, (k, reason) in enumerate(picks, start=1)]
+        context = format_context(chunks)
+        return {"question": question, "context": context, "chunks": chunks,
+                "prompt": build_prompt(question, context),
+                "tokens": sum(chunk["tokens"] for chunk in chunks),
+                "seeds": seeds["seeds"], "candidates": scored["candidates"]}
+
+    def answer(self, question, llm_fn, token_budget=None, settings=None):
+        """search() + generation (Algorithm 1 line 13). llm_fn is any function: prompt -> text.
+
+        Returns the search result with an extra "answer" key; result["chunks"] are the sources.
+        """
+        result = self.search(question, token_budget, settings)
+        result["answer"] = llm_fn(result["prompt"]).strip()
+        return result

spectrag/utils/__init__.py

@@ -0,0 +1 @@
+"""Small helper functions used by SpectRAGIngestion and SpectRAGRetrieval."""

spectrag/utils/bm25.py

@@ -0,0 +1,53 @@
+"""BM25 keyword scoring (paper Table 2, "BM25 lexical retrieval"; Robertson and Zaragoza [11]).
+
+score(chunk, query) = sum over query words of
+    idf(word) * count * (k1 + 1) / (count + k1 * (1 - b + b * length / average_length))
+"""
+
+import math
+import re
+from collections import Counter
+
+import numpy as np
+
+
+def tokenize(text):
+    """Lowercase words: 'Helios Energy!' -> ['helios', 'energy']."""
+    return re.findall(r"\w+", text.lower())
+
+
+def build_bm25(texts):
+    """Build the keyword index once: word -> (chunk numbers, counts) as numpy arrays."""
+    lists = {}
+    lengths = []
+    for chunk_id, text in enumerate(texts):
+        counts = Counter(tokenize(text))
+        lengths.append(sum(counts.values()))
+        for word, count in counts.items():
+            lists.setdefault(word, ([], []))
+            lists[word][0].append(chunk_id)
+            lists[word][1].append(count)
+    postings = {word: (np.array(ids), np.array(counts, dtype=float))
+                for word, (ids, counts) in lists.items()}
+    average = sum(lengths) / len(lengths) if lengths else 0.0
+    return {"postings": postings, "lengths": np.array(lengths, dtype=float),
+            "average_length": average or 1.0}
+
+
+def idf(doc_freq, doc_count):
+    """Rare words count more. This form (as in Lucene) is never negative."""
+    return math.log(1 + (doc_count - doc_freq + 0.5) / (doc_freq + 0.5))
+
+
+def bm25_scores(bm25, query, k1=1.5, b=0.75):
+    """BM25 score of every chunk for the query (0 for chunks that share no word)."""
+    lengths = bm25["lengths"]
+    scores = np.zeros(len(lengths))
+    length_norm = k1 * (1 - b + b * lengths / bm25["average_length"])  # one value per chunk
+    for word in set(tokenize(query)):
+        if word not in bm25["postings"]:
+            continue
+        ids, counts = bm25["postings"][word]  # only chunks that contain the word
+        weight = idf(len(ids), len(lengths))
+        scores[ids] += weight * counts * (k1 + 1) / (counts + length_norm[ids])
+    return scores