spectrafuse 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. spectrafuse/__init__.py +27 -0
  2. spectrafuse/_version.py +7 -0
  3. spectrafuse/cli.py +115 -0
  4. spectrafuse/config.py +96 -0
  5. spectrafuse/errors.py +52 -0
  6. spectrafuse/ingestion.py +118 -0
  7. spectrafuse/log.py +30 -0
  8. spectrafuse/retrieval.py +158 -0
  9. spectrafuse/utils/__init__.py +1 -0
  10. spectrafuse/utils/bm25.py +53 -0
  11. spectrafuse/utils/chunker.py +112 -0
  12. spectrafuse/utils/clusters.py +63 -0
  13. spectrafuse/utils/diffusion.py +105 -0
  14. spectrafuse/utils/edge_list.py +18 -0
  15. spectrafuse/utils/edges.py +38 -0
  16. spectrafuse/utils/edges_content.py +110 -0
  17. spectrafuse/utils/edges_structure.py +86 -0
  18. spectrafuse/utils/embedder.py +89 -0
  19. spectrafuse/utils/embedder_id.py +78 -0
  20. spectrafuse/utils/entities.py +72 -0
  21. spectrafuse/utils/evaluation.py +112 -0
  22. spectrafuse/utils/expansion.py +68 -0
  23. spectrafuse/utils/graph.py +90 -0
  24. spectrafuse/utils/graph_build.py +27 -0
  25. spectrafuse/utils/graph_template.py +105 -0
  26. spectrafuse/utils/graph_view.py +126 -0
  27. spectrafuse/utils/index_folder.py +78 -0
  28. spectrafuse/utils/llm.py +62 -0
  29. spectrafuse/utils/loader.py +61 -0
  30. spectrafuse/utils/metrics.py +61 -0
  31. spectrafuse/utils/prompt.py +62 -0
  32. spectrafuse/utils/questions.py +40 -0
  33. spectrafuse/utils/references.py +43 -0
  34. spectrafuse/utils/sections.py +98 -0
  35. spectrafuse/utils/seeding.py +106 -0
  36. spectrafuse/utils/selection.py +93 -0
  37. spectrafuse/utils/selectors.py +104 -0
  38. spectrafuse/utils/semantic_split.py +54 -0
  39. spectrafuse/utils/spectral.py +72 -0
  40. spectrafuse/utils/store.py +94 -0
  41. spectrafuse/utils/tokens.py +64 -0
  42. spectrafuse/utils/tuning.py +78 -0
  43. spectrafuse/utils/units.py +106 -0
  44. spectrafuse/utils/vector_index.py +124 -0
  45. spectrafuse/validation.py +84 -0
  46. spectrafuse-0.2.0.dist-info/METADATA +606 -0
  47. spectrafuse-0.2.0.dist-info/RECORD +52 -0
  48. spectrafuse-0.2.0.dist-info/WHEEL +5 -0
  49. spectrafuse-0.2.0.dist-info/entry_points.txt +2 -0
  50. spectrafuse-0.2.0.dist-info/licenses/LICENSE +21 -0
  51. spectrafuse-0.2.0.dist-info/licenses/NOTICE +23 -0
  52. spectrafuse-0.2.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,27 @@
1
+ """SpectraFuse: an independent implementation inspired by SpectRAG (Sahinoglu, Mohammed,
2
+ Kazangirler and Kazangirler, arXiv:2610.36991), spectral graph-augmented retrieval for RAG.
3
+ Not the official SpectRAG and not affiliated with its authors; the official implementation is
4
+ https://github.com/SightAndCite/SpectRAG. The two classes keep the paper's method name.
5
+
6
+ SpectRAGIngestion builds the index (offline, once per document set).
7
+ SpectRAGRetrieval finds the context for a question and can answer it (online, per question).
8
+
9
+ from spectrafuse import SpectRAGIngestion, SpectRAGRetrieval
10
+ SpectRAGIngestion("my_index").ingest("docs/")
11
+ print(SpectRAGRetrieval("my_index").search("my question")["context"])
12
+ """
13
+
14
+ from spectrafuse._version import INDEX_FORMAT, __version__
15
+ from spectrafuse.config import DEFAULTS, make_settings
16
+ from spectrafuse.errors import (DocumentError, DocumentNotFoundError, EmbedderError, IndexError_,
17
+ LLMError, MissingPackageError, NoDocumentsError, SettingsError,
18
+ SpectRAGError, UnknownSettingError)
19
+ from spectrafuse.ingestion import SpectRAGIngestion
20
+ from spectrafuse.log import enable_logging
21
+ from spectrafuse.retrieval import SpectRAGRetrieval
22
+
23
+ __all__ = ["SpectRAGIngestion", "SpectRAGRetrieval", "make_settings", "DEFAULTS",
24
+ "enable_logging", "SpectRAGError", "SettingsError", "UnknownSettingError",
25
+ "IndexError_", "EmbedderError", "DocumentError", "NoDocumentsError",
26
+ "DocumentNotFoundError", "MissingPackageError", "LLMError",
27
+ "__version__", "INDEX_FORMAT"]
@@ -0,0 +1,7 @@
1
+ """The package version (read by pyproject.toml) and the index format version."""
2
+
3
+ __version__ = "0.2.0"
4
+
5
+ # Bump INDEX_FORMAT when the files of an index change in a way older code cannot read.
6
+ # 1 = JSON edge lists (before the Phase 1 audit), 2 = binary .npy/.npz edges and signals.
7
+ INDEX_FORMAT = 2
spectrafuse/cli.py ADDED
@@ -0,0 +1,115 @@
1
+ """The `spectrafuse` command.
2
+
3
+ spectrafuse ingest docs/ --index my_index [--set chunk_tokens=256 --set embedder=ollama]
4
+ spectrafuse search "my question" --index my_index [--budget 4000] [--json]
5
+ spectrafuse answer "my question" --index my_index --model gemma3:4b
6
+ spectrafuse info --index my_index
7
+ spectrafuse graph --index my_index --out graph.html [--question "my question"]
8
+ """
9
+
10
+ import argparse
11
+ import json
12
+ import sys
13
+
14
+ from spectrafuse._version import __version__
15
+ from spectrafuse.errors import SpectRAGError
16
+ from spectrafuse.log import enable_logging
17
+
18
+
19
+ def parse_sets(pairs):
20
+ """["alpha=0.5", "selector=mmr"] -> {"alpha": 0.5, "selector": "mmr"} (values read as JSON)."""
21
+ settings = {}
22
+ for pair in pairs or []:
23
+ if "=" not in pair:
24
+ raise SystemExit("--set needs key=value, got %r" % pair)
25
+ key, value = pair.split("=", 1)
26
+ try:
27
+ settings[key] = json.loads(value)
28
+ except ValueError:
29
+ settings[key] = value # plain text, e.g. selector=mmr
30
+ return settings
31
+
32
+
33
+ def command_ingest(args):
34
+ from spectrafuse import SpectRAGIngestion
35
+ info = SpectRAGIngestion(args.index, parse_sets(args.set)).ingest(args.path)
36
+ print(json.dumps(info, indent=1))
37
+
38
+
39
+ def command_search(args):
40
+ from spectrafuse import SpectRAGRetrieval
41
+ result = SpectRAGRetrieval(args.index, parse_sets(args.set)).search(args.question, args.budget)
42
+ if args.json:
43
+ print(json.dumps({key: result[key] for key in ("question", "tokens", "chunks", "context")}, indent=1))
44
+ return
45
+ for chunk in result["chunks"]:
46
+ print("%2d. %-30s score %.3f %d tokens (%s)" % (chunk["rank"], chunk["doc_id"], chunk["score"],
47
+ chunk["tokens"], chunk["reason"]))
48
+ print("\n" + result["context"])
49
+
50
+
51
+ def command_answer(args):
52
+ from spectrafuse import SpectRAGRetrieval
53
+ from spectrafuse.utils.llm import ollama_llm
54
+ search = SpectRAGRetrieval(args.index, parse_sets(args.set))
55
+ result = search.answer(args.question, ollama_llm(args.model), args.budget)
56
+ print(result["answer"])
57
+ print("\nSources: " + ", ".join(sorted({c["doc_id"] for c in result["chunks"]})))
58
+
59
+
60
+ def command_info(args):
61
+ from spectrafuse.utils.store import check_index_folder, load_index
62
+ from pathlib import Path
63
+ check_index_folder(Path(args.index))
64
+ print(json.dumps(load_index(args.index, only=("info",))["info"], indent=1))
65
+
66
+
67
+ def command_graph(args):
68
+ from spectrafuse.utils.graph_view import save_graph_view
69
+ path = save_graph_view(args.index, args.out, args.question, parse_sets(args.set), args.max_nodes)
70
+ print("Graph view written to %s (open it in a browser)" % path)
71
+
72
+
73
+ def build_parser():
74
+ parser = argparse.ArgumentParser(prog="spectrafuse", description="SpectraFuse: spectral graph RAG, an independent implementation inspired by SpectRAG (arXiv:2610.36991). Official SpectRAG: github.com/SightAndCite/SpectRAG")
75
+ parser.add_argument("--version", action="version", version="spectrafuse " + __version__)
76
+ parser.add_argument("-v", "--verbose", action="store_true", help="print progress messages")
77
+ commands = parser.add_subparsers(dest="command", required=True)
78
+ ingest = commands.add_parser("ingest", help="build an index from a file or folder")
79
+ ingest.add_argument("path")
80
+ for name, helptext in [("search", "find the context for a question"), ("answer", "answer with an LLM")]:
81
+ sub = commands.add_parser(name, help=helptext)
82
+ sub.add_argument("question")
83
+ sub.add_argument("--budget", type=int, default=None, help="token budget (default 12000)")
84
+ if name == "search":
85
+ sub.add_argument("--json", action="store_true", help="print JSON")
86
+ else:
87
+ sub.add_argument("--model", default="gemma3:4b", help="Ollama model")
88
+ info = commands.add_parser("info", help="show what an index contains")
89
+ graph = commands.add_parser("graph", help="write an interactive graph view (HTML)")
90
+ graph.add_argument("--out", default="graph.html", help="HTML file to write")
91
+ graph.add_argument("--question", default=None, help="show this question's retrieval path")
92
+ graph.add_argument("--max-nodes", type=int, default=1500, dest="max_nodes")
93
+ for sub in (ingest, commands.choices["search"], commands.choices["answer"], info, graph):
94
+ sub.add_argument("--index", required=True, help="index folder")
95
+ for sub in (ingest, commands.choices["search"], commands.choices["answer"], graph):
96
+ sub.add_argument("--set", action="append", metavar="KEY=VALUE", help="change a setting")
97
+ return parser
98
+
99
+
100
+ def main(argv=None):
101
+ args = build_parser().parse_args(argv)
102
+ if args.verbose:
103
+ enable_logging()
104
+ handlers = {"ingest": command_ingest, "search": command_search, "answer": command_answer,
105
+ "info": command_info, "graph": command_graph}
106
+ try:
107
+ handlers[args.command](args)
108
+ except SpectRAGError as error:
109
+ print("spectrafuse: error: %s" % error, file=sys.stderr)
110
+ return 1
111
+ return 0
112
+
113
+
114
+ if __name__ == "__main__":
115
+ sys.exit(main())
spectrafuse/config.py ADDED
@@ -0,0 +1,96 @@
1
+ """Default settings for SpectraFuse.
2
+
3
+ Values marked "paper" come from the SpectRAG paper (Tables 1-3).
4
+ Values marked "assumed" are not given in the paper, so we picked a sensible value.
5
+ """
6
+
7
+ import copy
8
+ import difflib
9
+
10
+ from spectrafuse.errors import UnknownSettingError
11
+ from spectrafuse.validation import validate_settings
12
+
13
+ # Edge-signal weights from the paper's "signal configuration 1" (Table 9).
14
+ EDGE_WEIGHTS_CONFIG_1 = {
15
+ "semantic": 0.300,
16
+ "section": 0.250,
17
+ "entity": 0.200,
18
+ "adjacency": 0.084,
19
+ "utility": 0.083,
20
+ "citation": 0.083,
21
+ }
22
+
23
+ # Edge-signal weights from the paper's "signal configuration 2" (Table 10).
24
+ EDGE_WEIGHTS_CONFIG_2 = {
25
+ "semantic": 0.150,
26
+ "section": 0.200,
27
+ "entity": 0.250,
28
+ "adjacency": 0.100,
29
+ "utility": 0.280,
30
+ "citation": 0.020,
31
+ }
32
+
33
+ DEFAULTS = {
34
+ # --- Chunking (paper) ---
35
+ "chunk_tokens": 1200,
36
+ "overlap_tokens": 100,
37
+ "min_chunk_tokens": 32,
38
+ "semantic_percentile": 90,
39
+ "semantic_embedder": None, # embedder for semantic-chunking breaks; None = "embedder" (paper)
40
+ # --- Embedding ---
41
+ "embedder": "hash", # "hash" (offline), "ollama" (nomic-embed-text), or your own function
42
+ "embed_dims": 768, # paper: nomic-embed-text is 768-d (used by the hash embedder)
43
+ "ollama_model": "nomic-embed-text",
44
+ "ollama_url": "http://localhost:11434",
45
+ # --- Vector search (Section 6.2) ---
46
+ "ann_backend": "auto", # "exact", "faiss", "hnsw", or "auto" (exact up to 20k chunks)
47
+ "hnsw_m": 32, # HNSW links per node
48
+ "hnsw_ef_construction": 200,
49
+ "hnsw_ef_search": 128,
50
+ "dense_top_n": 1000, # with an approximate backend: dense scores for the top n only
51
+ # --- Graph (weights: paper, the rest: assumed) ---
52
+ "edge_weights": EDGE_WEIGHTS_CONFIG_1,
53
+ "semantic_k": 10, # assumed
54
+ "edge_threshold": 0.05, # assumed
55
+ "graph_k": 10, # assumed: each chunk keeps at most its 10 strongest edges
56
+ # --- Spectral (paper) ---
57
+ "spectral_dims": 32,
58
+ "n_clusters": None, # assumed: None means about sqrt(n / 2) clusters
59
+ # --- Retrieval (paper) ---
60
+ "seed_count": 20,
61
+ "hops": 2,
62
+ "max_candidates": 200,
63
+ "token_budget": 12000,
64
+ # --- Retrieval (assumed) ---
65
+ "rrf_k": 60, # the usual RRF constant
66
+ "bm25_k1": 1.5, # standard BM25 values (Robertson and Zaragoza)
67
+ "bm25_b": 0.75,
68
+ "bm25_weight": 0.5, # hybrid relevance = 0.5 * dense + 0.5 * BM25
69
+ "alpha": 0.3,
70
+ "ppr_beta": 0.15,
71
+ "ppr_steps": 30,
72
+ "diffusion_t": 1.0, # heat-kernel scale t for spectral closeness (Eq. 17)
73
+ "proximity_mode": "hybrid", # "hybrid" (paper), "ppr", "spectral", "heat", "none"
74
+ "selector": "spectrag", # "spectrag", "mmr", "protected_mmr", "anchor_expand_ce", "topk"
75
+ "mmr_lambda": 0.7, # MMR: weight of relevance vs. novelty
76
+ "core_size": 3,
77
+ "precision_floor": 0.4,
78
+ "select_weights": {"a": 1.0, "b": 0.3, "c": 0.2, "d": 0.1, "lambda": 0.5},
79
+ "length_power": 0.5,
80
+ }
81
+
82
+
83
+ def make_settings(overrides=None):
84
+ """Return a fresh copy of DEFAULTS with the user's overrides applied and checked.
85
+
86
+ Raises UnknownSettingError (a KeyError) for a misspelled name, and SettingsError
87
+ (a ValueError) for a value of the wrong type or out of range.
88
+ """
89
+ settings = copy.deepcopy(DEFAULTS)
90
+ for key, value in (overrides or {}).items():
91
+ if key not in settings:
92
+ close = difflib.get_close_matches(key, DEFAULTS, n=1)
93
+ hint = " (did you mean %r?)" % close[0] if close else ""
94
+ raise UnknownSettingError("Unknown setting: %r%s" % (key, hint))
95
+ settings[key] = copy.deepcopy(value)
96
+ return validate_settings(settings)
spectrafuse/errors.py ADDED
@@ -0,0 +1,52 @@
1
+ """SpectraFuse's exceptions. Each one also inherits the matching built-in exception, so code that
2
+ catches ValueError, KeyError, FileNotFoundError or ConnectionError keeps working.
3
+
4
+ try:
5
+ search = SpectRAGRetrieval("my_index")
6
+ except spectrafuse.SpectRAGError as error:
7
+ print(error)
8
+ """
9
+
10
+
11
+ class SpectRAGError(Exception):
12
+ """Base class for every error raised on purpose by SpectraFuse."""
13
+
14
+
15
+ class SettingsError(SpectRAGError, ValueError):
16
+ """A setting has the wrong type, is out of range, or is not one of the allowed choices."""
17
+
18
+
19
+ class UnknownSettingError(SettingsError, KeyError):
20
+ """A setting name that SpectraFuse does not know (often a typo)."""
21
+
22
+ def __str__(self): # KeyError would otherwise print the message with quotes around it
23
+ return str(self.args[0]) if self.args else ""
24
+
25
+
26
+ class IndexError_(SpectRAGError, FileNotFoundError):
27
+ """The index folder is missing, incomplete, or was made by an incompatible version.
28
+ (The trailing underscore avoids hiding Python's built-in IndexError.)"""
29
+
30
+
31
+ class EmbedderError(SpectRAGError, ConnectionError):
32
+ """The embedder cannot be reached, or its vectors do not match the index."""
33
+
34
+
35
+ class DocumentError(SpectRAGError, ValueError):
36
+ """Documents given to ingest_documents() are not in the expected form."""
37
+
38
+
39
+ class NoDocumentsError(DocumentError):
40
+ """There was no document with text to ingest."""
41
+
42
+
43
+ class DocumentNotFoundError(DocumentError, FileNotFoundError):
44
+ """The file or folder given to ingest() does not exist."""
45
+
46
+
47
+ class MissingPackageError(SpectRAGError, ImportError):
48
+ """An optional package is needed (e.g. faiss-cpu or pypdf); the message says what to install."""
49
+
50
+
51
+ class LLMError(SpectRAGError, ConnectionError):
52
+ """The LLM (e.g. a local Ollama model) cannot be reached or did not reply in time."""
@@ -0,0 +1,118 @@
1
+ """SpectRAGIngestion: the offline pipeline (paper Table 1). Run it once per document set.
2
+
3
+ ingest = SpectRAGIngestion("my_index")
4
+ info = ingest.ingest("examples/docs")
5
+ """
6
+
7
+ import time
8
+
9
+ import numpy as np
10
+
11
+ from spectrafuse._version import INDEX_FORMAT, __version__
12
+ from spectrafuse.config import make_settings
13
+ from spectrafuse.errors import DocumentError, NoDocumentsError
14
+ from spectrafuse.log import log_step
15
+ from spectrafuse.utils.chunker import chunk_documents
16
+ from spectrafuse.utils.edges import build_edge_signals, count_edges
17
+ from spectrafuse.utils.embedder import make_embedder
18
+ from spectrafuse.utils.embedder_id import embed_probe, saved_settings
19
+ from spectrafuse.utils.entities import add_entities
20
+ from spectrafuse.utils.graph_build import build_graph, build_spectral
21
+ from spectrafuse.utils.index_folder import prepare_folder, replace_folder
22
+ from spectrafuse.utils.loader import load_docs
23
+ from spectrafuse.utils.questions import add_utility_questions
24
+ from spectrafuse.utils.references import add_references
25
+ from spectrafuse.utils.store import edges_to_array, save_index
26
+ from spectrafuse.utils.vector_index import build_vector_index, save_vector_index
27
+
28
+
29
+ class SpectRAGIngestion:
30
+ """Load -> chunk -> embed -> edge signals -> weighted graph -> spectral -> save."""
31
+
32
+ def __init__(self, index_dir, settings=None, entity_fn=None, question_fn=None):
33
+ self.index_dir = index_dir
34
+ self.settings = make_settings(settings)
35
+ self.embed = make_embedder(self.settings)
36
+ # Semantic chunking can use a faster embedder for sentence distances (None = the same).
37
+ split_choice = self.settings["semantic_embedder"]
38
+ self.split_embed = self.embed if split_choice is None else \
39
+ make_embedder(dict(self.settings, embedder=split_choice))
40
+ self.entity_fn = entity_fn # optional: text -> list of entities
41
+ self.question_fn = question_fn # optional: text -> list of utility questions
42
+
43
+ def make_chunks(self, docs):
44
+ """Split documents into chunks with entities, references and utility questions."""
45
+ chunks = chunk_documents(docs, self.settings, embed_fn=self.split_embed)
46
+ add_entities(chunks, self.entity_fn)
47
+ add_references(chunks)
48
+ add_utility_questions(chunks, self.question_fn)
49
+ return chunks
50
+
51
+ def make_graph(self, signals):
52
+ """Fuse the signals (Eq. 10), drop weak edges, keep each chunk's strongest graph_k."""
53
+ return build_graph(signals, self.settings)
54
+
55
+ def make_spectral(self, edges, node_count):
56
+ """Laplacian (Eq. 11) -> spectral coordinates (Eq. 12-13) -> clusters."""
57
+ return build_spectral(edges, node_count, self.settings)
58
+
59
+ def ingest(self, path):
60
+ """Ingest a file or a folder (.txt, .md, .pdf), save the index, return a summary."""
61
+ docs = load_docs(path)
62
+ if not docs:
63
+ raise NoDocumentsError("No .txt, .md or .pdf file with text found in " + str(path))
64
+ return self.ingest_documents(docs)
65
+
66
+ def ingest_documents(self, docs):
67
+ """Ingest documents already in memory: [{"doc_id": ..., "text": ...}, ...]."""
68
+ check_documents(docs)
69
+ start = begin = time.time()
70
+ chunks = self.make_chunks(docs)
71
+ start = log_step("chunking", start, documents=len(docs), chunks=len(chunks))
72
+ embeddings = self.embed([chunk["text"] for chunk in chunks], kind="document")
73
+ vector_index = build_vector_index(embeddings, self.settings) # exact, FAISS or HNSW
74
+ start = log_step("embedding", start, backend=vector_index["backend"])
75
+ signals, question_vectors, owners = build_edge_signals(
76
+ chunks, embeddings, self.embed, self.settings, vector_index)
77
+ start = log_step("edge signals", start, **count_edges(signals))
78
+ edges, weights = self.make_graph(signals)
79
+ coords, eigenvalues, clusters = self.make_spectral(edges, len(chunks))
80
+ start = log_step("graph and spectral", start, edges=len(edges), dims=coords.shape[1])
81
+ for chunk, cluster in zip(chunks, clusters):
82
+ chunk["cluster"] = int(cluster)
83
+ info = {"documents": len({c["doc_id"] for c in chunks}), "chunks": len(chunks),
84
+ "signal_edges": count_edges(signals), "graph_edges": len(edges),
85
+ "weights_used": {name: round(w, 4) for name, w in weights.items()},
86
+ "spectral_dims": coords.shape[1], "ann_backend": vector_index["backend"],
87
+ "clusters": len(set(clusters.tolist())),
88
+ "format_version": INDEX_FORMAT, "spectrafuse_version": __version__,
89
+ "seconds": round(time.time() - begin, 2)}
90
+ partial = prepare_folder(self.index_dir)
91
+ save_index(partial, {
92
+ "chunks": chunks, "embeddings": embeddings, "edges": edges_to_array(edges),
93
+ "signals": {name: edges_to_array(found) for name, found in signals.items()},
94
+ "spectral": coords, "eigenvalues": eigenvalues,
95
+ "question_vectors": question_vectors,
96
+ "question_owners": np.asarray(owners, dtype=np.int64),
97
+ "settings": saved_settings(self.settings), "info": info, # functions saved by name
98
+ "embedder_probe": embed_probe(self.embed), # fingerprint, checked when opened
99
+ })
100
+ save_vector_index(vector_index, partial)
101
+ replace_folder(partial, self.index_dir)
102
+ log_step("saved index", start, folder=self.index_dir)
103
+ return info
104
+
105
+
106
+ def check_documents(docs):
107
+ """Every document needs a text and a unique doc_id."""
108
+ if not docs:
109
+ raise NoDocumentsError("No documents with text to ingest")
110
+ seen = set()
111
+ for number, doc in enumerate(docs):
112
+ if not isinstance(doc, dict) or not isinstance(doc.get("text"), str) or "doc_id" not in doc:
113
+ raise DocumentError("Document %d must be a dict with 'doc_id' and 'text' (a string)" % number)
114
+ if doc["doc_id"] in seen:
115
+ raise DocumentError("Two documents have the same doc_id: %r" % (doc["doc_id"],))
116
+ seen.add(doc["doc_id"])
117
+ if not any(doc["text"].strip() for doc in docs):
118
+ raise NoDocumentsError("Every document is empty")
spectrafuse/log.py ADDED
@@ -0,0 +1,30 @@
1
+ """Logging. SpectraFuse is silent by default (like any library); switch messages on with:
2
+
3
+ import spectrafuse
4
+ spectrafuse.enable_logging() # INFO: one line per pipeline step, with its time
5
+ spectrafuse.enable_logging("DEBUG") # DEBUG: also one line per query
6
+ """
7
+
8
+ import logging
9
+ import time
10
+
11
+ logger = logging.getLogger("spectrafuse")
12
+ logger.addHandler(logging.NullHandler())
13
+
14
+
15
+ def enable_logging(level="INFO"):
16
+ """Print SpectraFuse's log messages to the terminal (stderr)."""
17
+ if not any(getattr(h, "spectrafuse", False) for h in logger.handlers):
18
+ handler = logging.StreamHandler()
19
+ handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s"))
20
+ handler.spectrafuse = True
21
+ logger.addHandler(handler)
22
+ logger.setLevel(level)
23
+ return logger
24
+
25
+
26
+ def log_step(name, start, **details):
27
+ """Log one finished step: 'chunking: 1.23 s (chunks=8)'. Returns the time now."""
28
+ extra = ", ".join("%s=%s" % item for item in details.items())
29
+ logger.info("%s: %.2f s%s", name, time.time() - start, " (%s)" % extra if extra else "")
30
+ return time.time()
@@ -0,0 +1,158 @@
1
+ """SpectRAGRetrieval: the online pipeline (paper Table 2, Algorithm 1). Run it per question.
2
+
3
+ search = SpectRAGRetrieval("my_index")
4
+ result = search.search("Which company did the co-founder of Aurora Labs join?")
5
+ print(result["context"])
6
+ """
7
+
8
+ import time
9
+
10
+ from spectrafuse._version import INDEX_FORMAT
11
+ from spectrafuse.config import make_settings
12
+ from spectrafuse.errors import EmbedderError, IndexError_, SettingsError
13
+ from spectrafuse.log import logger
14
+ from spectrafuse.utils.bm25 import build_bm25
15
+ from spectrafuse.utils.diffusion import score_candidates
16
+ from spectrafuse.utils.embedder import make_embedder
17
+ from spectrafuse.utils.embedder_id import (CUSTOM, check_embedder, is_saved_function,
18
+ read_saved_settings)
19
+ from spectrafuse.utils.graph import degrees, neighbors
20
+ from spectrafuse.utils.graph_build import build_graph, build_spectral
21
+ from spectrafuse.utils.index_folder import readable_folder
22
+ from spectrafuse.utils.prompt import build_prompt, format_context
23
+ from spectrafuse.utils.seeding import find_seeds
24
+ from spectrafuse.utils.selectors import select
25
+ from spectrafuse.utils.store import array_to_edges, check_index_folder, load_index
26
+ from spectrafuse.utils.vector_index import load_vector_index
27
+
28
+ # These change the graph itself, so they need rebuild_graph() instead of a per-call override.
29
+ GRAPH_SETTINGS = ("edge_weights", "edge_threshold", "graph_k", "spectral_dims", "n_clusters")
30
+
31
+
32
+ class SpectRAGRetrieval:
33
+ """Seeds -> graph expansion -> proximity re-scoring -> token-budget selection -> context."""
34
+
35
+ def __init__(self, index_dir, settings=None, embed_fn=None, rerank_fn=None):
36
+ """Load an index made by SpectRAGIngestion.
37
+
38
+ settings: settings to change, e.g. {"token_budget": 4000, "selector": "mmr"}.
39
+ Graph settings (edge_threshold, graph_k, ...) rebuild the graph from the
40
+ saved signals, without embedding anything again.
41
+ embed_fn: pass the same function again if the index was built with your own embedder.
42
+ rerank_fn: optional (question, texts) -> scores, e.g. a cross-encoder, used by the
43
+ "anchor_expand_ce" selector.
44
+ """
45
+ self.index_dir = readable_folder(index_dir) # '<index>.old' only after an interrupted ingest
46
+ check_index_folder(self.index_dir)
47
+ self.index = load_index(self.index_dir, skip=("signals",)) # raw signals: loaded on demand
48
+ made_with = self.index["info"].get("format_version", 1)
49
+ if made_with > INDEX_FORMAT:
50
+ raise IndexError_("This index has format %d; this SpectraFuse version reads up to %d. "
51
+ "Upgrade spectrafuse." % (made_with, INDEX_FORMAT))
52
+ self.index["edges"] = array_to_edges(self.index["edges"])
53
+ self.index["question_owners"] = [int(owner) for owner in self.index["question_owners"]]
54
+ built_with = make_settings(read_saved_settings(self.index["settings"]))
55
+ self.settings = make_settings(dict(built_with, **(settings or {})))
56
+ if embed_fn is not None:
57
+ self.settings["embedder"] = embed_fn
58
+ elif is_saved_function(self.settings["embedder"]):
59
+ raise SettingsError("This index was built with your own embed function (%s). "
60
+ "Pass it again: SpectRAGRetrieval(index_dir, embed_fn=...)"
61
+ % self.settings["embedder"][len(CUSTOM):])
62
+ self.embed = make_embedder(self.settings)
63
+ check_embedder(self.embed, self.index.get("embedder_probe"), built_with["embedder"])
64
+ self.rerank_fn = rerank_fn
65
+ self.bm25 = build_bm25([chunk["text"] for chunk in self.index["chunks"]])
66
+ self.vector_index = load_vector_index(self.index_dir, self.index["embeddings"], self.settings)
67
+ if any(self.settings[key] != built_with[key] for key in GRAPH_SETTINGS):
68
+ self.rebuild_graph()
69
+ self.refresh_graph()
70
+
71
+ def refresh_graph(self):
72
+ """Neighbour table (CSR) and degrees for the current edges."""
73
+ count = len(self.index["chunks"])
74
+ self.table = neighbors(self.index["edges"], count)
75
+ self.index["degree"] = degrees(self.index["edges"], count)
76
+
77
+ def rebuild_graph(self, overrides=None):
78
+ """Change graph settings without re-embedding: fuse the saved signals again (Eq. 10)
79
+ and recompute spectral coordinates and clusters (Eq. 11-13). In memory only."""
80
+ self.settings = make_settings(dict(self.settings, **(overrides or {})))
81
+ saved = load_index(self.index_dir, only=("signals",))["signals"]
82
+ signals = {name: array_to_edges(edges) for name, edges in saved.items()}
83
+ edges, _ = build_graph(signals, self.settings)
84
+ coords, eigenvalues, clusters = build_spectral(edges, len(self.index["chunks"]), self.settings)
85
+ self.index.update({"edges": edges, "spectral": coords, "eigenvalues": eigenvalues})
86
+ for chunk, cluster in zip(self.index["chunks"], clusters):
87
+ chunk["cluster"] = int(cluster)
88
+ self.refresh_graph()
89
+
90
+ def settings_for(self, overrides):
91
+ """The settings for one call: the current settings plus per-call overrides."""
92
+ if not overrides:
93
+ return self.settings
94
+ changed = [key for key in GRAPH_SETTINGS
95
+ if key in overrides and overrides[key] != self.settings[key]]
96
+ if changed:
97
+ raise SettingsError("Use rebuild_graph() to change graph settings: " + ", ".join(changed))
98
+ return make_settings(dict(self.settings, **overrides))
99
+
100
+ def retrieve(self, question, settings=None, query_vector=None):
101
+ """Stages 1-4. Returns (seed result, scored candidates, picks)."""
102
+ settings = self.settings_for(settings)
103
+ start = time.time()
104
+ if query_vector is None:
105
+ query_vector = self.embed([question], kind="query")[0]
106
+ if len(query_vector) != self.index["embeddings"].shape[1]:
107
+ raise EmbedderError("The query vector has %d dimensions but the index has %d: search "
108
+ "with the same embedder the index was built with."
109
+ % (len(query_vector), self.index["embeddings"].shape[1]))
110
+ seeds = find_seeds(question, query_vector, self.index, self.bm25, settings,
111
+ self.vector_index)
112
+ scored = score_candidates(seeds["seeds"], seeds["relevance"], self.index,
113
+ self.table, settings)
114
+ rerank = None
115
+ if settings["selector"] == "anchor_expand_ce" and self.rerank_fn is not None:
116
+ texts = [self.index["chunks"][c]["text"] for c in scored["candidates"]]
117
+ rerank = self.rerank_fn(question, texts)
118
+ picks = select(scored, self.index, settings, self.table, rerank)
119
+ logger.debug("query %r: %d seeds, %d candidates, %d picked in %.1f ms", question[:60],
120
+ len(seeds["seeds"]), len(scored["candidates"]), len(picks),
121
+ 1000 * (time.time() - start))
122
+ return seeds, scored, picks
123
+
124
+ def describe(self, k, reason, rank, scored):
125
+ """One selected chunk with all its scores (the 'ranked source metadata')."""
126
+ chunk = self.index["chunks"][scored["candidates"][k]]
127
+ info = {key: chunk[key] for key in ["id", "doc_id", "position", "section", "tokens", "text"]}
128
+ info.update({"rank": rank, "reason": reason, "hop": scored["hop"][k],
129
+ "score": float(scored["scores"][k]),
130
+ "relevance": float(scored["relevance"][k]),
131
+ "proximity": float(scored["proximity"][k])})
132
+ return info
133
+
134
+ def search(self, question, token_budget=None, settings=None, query_vector=None):
135
+ """Find the context for a question. Returns a dict with 'context', 'chunks', ...
136
+
137
+ settings: per-call overrides, e.g. {"selector": "mmr", "alpha": 0.5}.
138
+ """
139
+ overrides = dict(settings or {})
140
+ if token_budget is not None:
141
+ overrides["token_budget"] = token_budget
142
+ seeds, scored, picks = self.retrieve(question, overrides, query_vector)
143
+ chunks = [self.describe(k, reason, rank, scored)
144
+ for rank, (k, reason) in enumerate(picks, start=1)]
145
+ context = format_context(chunks)
146
+ return {"question": question, "context": context, "chunks": chunks,
147
+ "prompt": build_prompt(question, context),
148
+ "tokens": sum(chunk["tokens"] for chunk in chunks),
149
+ "seeds": seeds["seeds"], "candidates": scored["candidates"]}
150
+
151
+ def answer(self, question, llm_fn, token_budget=None, settings=None):
152
+ """search() + generation (Algorithm 1 line 13). llm_fn is any function: prompt -> text.
153
+
154
+ Returns the search result with an extra "answer" key; result["chunks"] are the sources.
155
+ """
156
+ result = self.search(question, token_budget, settings)
157
+ result["answer"] = llm_fn(result["prompt"]).strip()
158
+ return result
@@ -0,0 +1 @@
1
+ """Small helper functions used by SpectRAGIngestion and SpectRAGRetrieval."""