IA Utils — Hybrid Retrieval (RAG)¶

Retrieval to ground an agent's own output in a document collection, not a quantum measurement result — see Vector Healing if you're looking for healing a numeric sequence instead.

An agent answering from a pile of documents has to find the right passage before it can quote it. The simplest approach, TF-IDF cosine similarity, only catches passages that share the query's actual words — a query phrased differently from the source text ("does information leak out?" vs. a paper that only ever says "traversable") can miss the right passage entirely. ia_utils.rag fixes this with two stages: a cheap first pass pools candidates from TF-IDF and a dense sentence embedding (which catches paraphrases the exact-word match misses), then a slower, more accurate cross-encoder model re-reads that small pool and picks the real top results.

Step 1. Index three documents and ask a question¶

from ia_utils.rag import build_index, search

docs = [
    ("The traversable wormhole construction couples two boundaries of an "
     "eternal BTZ black hole with a negative average null energy stress "
     "tensor, rendering the Einstein-Rosen bridge traversable for a probe.",
     "gao_wormhole.pdf"),
    ("Error mitigation techniques for short-depth quantum circuits "
     "extrapolate the noisy expectation value to the zero-noise limit "
     "without extra qubits.", "temme_mitigation.pdf"),
    ("A cross-encoder reads the query and passage together through one "
     "transformer, unlike a bi-encoder which embeds them separately and "
     "compares vectors afterward.", "reranking_note.pdf"),
]
index = build_index(docs)
results = search("does information leak through a wormhole", index, top=2)
[(r["rank"], r["source"], round(r["score"], 3)) for r in results]

[(1, 'gao_wormhole.pdf', -6.315), (2, 'reranking_note.pdf', -11.244)]

build_index chunks each (text, source) pair (1200 characters per chunk, 200-character overlap — a document doesn't have to fit in one chunk) and builds both a TF-IDF matrix and a dense embedding for every chunk. search returns the top results as plain dicts — score is the cross-encoder's own score by default (higher is better, but not bounded to [0, 1] the way a cosine similarity is), and cosine is always included alongside it so the two can be compared directly, like in Step 2.

Step 2. Compare against TF-IDF alone¶

from ia_utils.rag import search

query = "does information leak through a wormhole"
tfidf_only = search(query, index, top=2, rerank=False)
hybrid = search(query, index, top=2, rerank=True)
[(r["source"], r["score_type"]) for r in tfidf_only + hybrid]

[('gao_wormhole.pdf', 'cosine'), ('reranking_note.pdf', 'cosine'),
 ('gao_wormhole.pdf', 'rerank'), ('reranking_note.pdf', 'rerank')]

rerank=False skips the dense pool and the cross-encoder entirely — a plain TF-IDF ranking, score_type is "cosine" instead of "rerank". On a three-document toy collection like this one, both approaches often agree (as they do here) simply because there's nothing left for the dense pool to add — every document is already inside the candidate pool either way. The gap opens up on a real, larger collection: see Details below for a real, measured case where it mattered.

Step 3. Find a specific known phrase directly¶

from ia_utils.rag import search_exact

hits = search_exact("negative average null energy", index)
[(h["source"], h["match"]) for h in hits]

[('gao_wormhole.pdf', 'negative average null energy')]

Semantic search ranks by topical similarity, which can bury a short, specific, load-bearing phrase — an exact clause, a fixed parameter value, a named condition — under chunks that are merely more topically central. search_exact skips ranking entirely: no embedding, no reranker, no model download, just a literal (case-insensitive) substring match over every chunk's raw text, with context characters of surrounding text kept on each side. Pass regex=True to treat pattern as a real regular expression instead of a literal string.

Details¶

Where this came from¶

Promoted from quantumrag (a local literature-grounding tool used through Dense-Evolution-Discovery to check physics/chemistry claims against real, independently-verified arXiv papers before citing them), after validation there across roughly 30 topic collections and ~6300 chunks total. Only the retrieval mechanism moved here — quantumrag's own paper corpus and built indexes stay local, out of the package.

A real case where reranking changed the answer¶

On quantumrag's own quantum_info collection (30 real papers, hundreds of chunks), the query "how does error mitigation reduce noise in short depth circuits" returned, with TF-IDF alone, three chunks from the same paper (Temme et al. 2017) as its top 3 — technically on-topic, but redundant. With the hybrid pool + rerank, the third result changed to a chunk from a different, genuinely relevant paper (a differentiable-Kraus-tensor-networks review citing the same error-mitigation result) instead of repeating the first paper a third time. That is the effect this module is for: it shows up once a collection is large and varied enough for the extra recall and re-ordering to matter, not necessarily on a handful of documents.

The optional `rag` extra¶

pip install dense-evolution[rag] (scikit-learn + sentence-transformers). chunk_text is the one function in this module that needs neither — every other function, build_index/save_index/load_index/search included, needs at least scikit-learn (TF-IDF is built on sklearn.feature_extraction.text.TfidfVectorizer), even with compute_embeddings=False or rerank=False. build_index(..., compute_embeddings=False) skips only the sentence-transformers bi-encoder pass, not scikit-learn itself.

Model choice and where the two stages come from¶

Stage 1's dense bi-encoder (sentence-transformers/all-MiniLM-L6-v2, build_index's embedding_model) and stage 2's cross-encoder (cross-encoder/ms-marco-MiniLM-L-6-v2, search's cross_encoder_model) are both pretrained, general-purpose models — nothing is trained or fine-tuned by this module. The two-stage split itself follows Karpukhin et al. 2020, "Dense Passage Retrieval for Open-Domain Question Answering" (arXiv:2004.04906) for the dense-retrieval stage, and Nogueira & Cho 2019, "Passage Re-ranking with BERT" (arXiv:1901.04085) for the cross-encoder rerank stage — both real, established information-retrieval results, not new mechanisms invented for this module.

Persisting an index¶

save_index(index, index_dir) writes chunks.json, vectorizer.pkl, matrix.pkl, and (if compute_embeddings=True) embeddings.npy under index_dir. load_index(index_dir) reads them back into an identical RagIndex — rebuilding an index from scratch on every query is wasteful once a collection stops changing.

rag ¶

Hybrid two-stage retrieval (sparse TF-IDF UNION dense bi-encoder, pretrained cross-encoder rerank) for grounding an agent's output in a local document collection, instead of trusting an unverified citation.

Promoted from quantumrag (Desktop/Fullwork/quantumrag), a local RAG tool used through Dense-Evolution-Discovery to check physics/chemistry claims against real, independently-verified arXiv papers before citing them -- validated there across ~30 topic collections (~6300 chunks) before this promotion. Only the retrieval mechanism moves here, not quantumrag's own paper corpus or built indexes -- those stay local, out of the package.

Stage 1 pools candidates from TF-IDF cosine (exact-term overlap, cheap, whole collection) UNION a dense bi-encoder (catches paraphrases/synonyms TF-IDF misses on its own -- Karpukhin et al. 2020, "Dense Passage Retrieval for Open-Domain Question Answering", arXiv:2004.04906). Stage 2 reranks that pool with a pretrained cross-encoder (Nogueira & Cho 2019, "Passage Re-ranking with BERT", arXiv:1901.04085) -- a slower model that reads query+chunk together instead of comparing two separate vectors, applied only to the (small) pooled candidates, never the whole collection.

Needs the optional rag extra (scikit-learn + sentence-transformers) -- same pattern as native_hf.libcint_bridge (pyscf) and qmmm.region (rdkit): the whole module is off-limits without it, since even the TF-IDF-only path (build_index(compute_embeddings=False), search(rerank=False)) is built on scikit-learn's TfidfVectorizer/cosine_similarity. chunk_text (pure Python, no vectorizer/model involved) is the one function usable with no extra at all -- it doesn't import this module's sklearn/sentence-transformers names.

build_index ¶

build_index(
    documents,
    embedding_model: str = DEFAULT_EMBEDDING_MODEL,
    compute_embeddings: bool = True,
    chunk_size: int = CHUNK_SIZE_CHARS,
    overlap: int = CHUNK_OVERLAP_CHARS,
) -> RagIndex

documents: iterable of (text, source) pairs, one per already-extracted document (PDF/markdown/plain text parsing happens before this call -- this module only chunks and indexes text it's handed). compute_embeddings can be set False to skip the (slower) bi-encoder pass and build a TF-IDF-only index -- still needs the 'rag' extra for scikit-learn, just not sentence-transformers' model download; search() then falls back to plain cosine regardless of rerank.

Source code in tools/ia_utils/rag.py

def build_index(
    documents,
    embedding_model: str = DEFAULT_EMBEDDING_MODEL,
    compute_embeddings: bool = True,
    chunk_size: int = CHUNK_SIZE_CHARS,
    overlap: int = CHUNK_OVERLAP_CHARS,
) -> RagIndex:
    """
    documents: iterable of (text, source) pairs, one per already-extracted
    document (PDF/markdown/plain text parsing happens before this call --
    this module only chunks and indexes text it's handed). compute_embeddings
    can be set False to skip the (slower) bi-encoder pass and build a
    TF-IDF-only index -- still needs the 'rag' extra for scikit-learn, just
    not sentence-transformers' model download; search() then falls back to
    plain cosine regardless of `rerank`.
    """
    try:
        from sklearn.feature_extraction.text import TfidfVectorizer
    except ImportError as _import_error:  # pragma: no cover -- only reachable without scikit-learn installed, which CI here always has
        raise ImportError(_MISSING_RAG_EXTRA.format(_import_error)) from _import_error

    all_chunks = []
    for text, source in documents:
        all_chunks.extend(chunk_text(text, source, chunk_size, overlap))
    if not all_chunks:
        raise ValueError("no chunks produced -- 'documents' was empty or every text was blank")

    texts = [c["text"] for c in all_chunks]
    vectorizer = TfidfVectorizer(stop_words="english", max_features=20000, ngram_range=(1, 2))
    matrix = vectorizer.fit_transform(texts)

    embeddings = None
    if compute_embeddings:
        embeddings = _get_embedder(embedding_model).encode(
            texts, normalize_embeddings=True, show_progress_bar=False
        )

    return RagIndex(chunks=all_chunks, vectorizer=vectorizer, matrix=matrix, embeddings=embeddings)

search ¶

search(
    query: str,
    index: RagIndex,
    top: int = 3,
    rerank: bool = True,
    pool: int = DEFAULT_POOL,
    embedding_model: str = DEFAULT_EMBEDDING_MODEL,
    cross_encoder_model: str = DEFAULT_CROSS_ENCODER_MODEL,
) -> list

Returns the top top chunks as a list of dicts (rank, source, text, score, score_type, cosine), highest-scoring first.

rerank=True (default): stage 1 pools pool candidates from TF-IDF cosine UNION dense-embedding similarity (if the index has embeddings -- plain TF-IDF pool otherwise), stage 2 reorders that pool with the cross-encoder and returns its top top ("score_type": "rerank"). rerank=False: plain TF-IDF cosine ranking, no sentence-transformers model load needed (still needs the 'rag' extra's scikit-learn for cosine_similarity itself) -- useful as a baseline, or when only embeddings are unavailable for this index but a scored ranking is still wanted.

Source code in tools/ia_utils/rag.py

def search(
    query: str,
    index: RagIndex,
    top: int = 3,
    rerank: bool = True,
    pool: int = DEFAULT_POOL,
    embedding_model: str = DEFAULT_EMBEDDING_MODEL,
    cross_encoder_model: str = DEFAULT_CROSS_ENCODER_MODEL,
) -> list:
    """
    Returns the top `top` chunks as a list of dicts (rank, source, text,
    score, score_type, cosine), highest-scoring first.

    rerank=True (default): stage 1 pools `pool` candidates from TF-IDF
    cosine UNION dense-embedding similarity (if the index has embeddings --
    plain TF-IDF pool otherwise), stage 2 reorders that pool with the
    cross-encoder and returns its top `top` ("score_type": "rerank").
    rerank=False: plain TF-IDF cosine ranking, no sentence-transformers model
    load needed (still needs the 'rag' extra's scikit-learn for cosine_similarity
    itself) -- useful as a baseline, or when only embeddings are unavailable
    for this index but a scored ranking is still wanted.
    """
    try:
        from sklearn.metrics.pairwise import cosine_similarity
    except ImportError as _import_error:  # pragma: no cover -- only reachable without scikit-learn installed, which CI here always has
        raise ImportError(_MISSING_RAG_EXTRA.format(_import_error)) from _import_error

    query_vec = index.vectorizer.transform([query])
    cosine_scores = cosine_similarity(query_vec, index.matrix)[0]
    tfidf_order = cosine_scores.argsort()[::-1]

    if rerank and index.embeddings is not None:
        query_emb = _get_embedder(embedding_model).encode([query], normalize_embeddings=True)[0]
        dense_scores = index.embeddings @ query_emb
        dense_pool = dense_scores.argsort()[::-1][:pool]
        candidate_idx = sorted(set(tfidf_order[:pool].tolist()) | set(dense_pool.tolist()))
    elif rerank:
        candidate_idx = tfidf_order[:pool].tolist()
    else:
        candidate_idx = []

    if rerank and candidate_idx:
        pairs = [(query, index.chunks[i]["text"][:1024]) for i in candidate_idx]
        rerank_scores = _get_reranker(cross_encoder_model).predict(pairs)
        order = rerank_scores.argsort()[::-1][:top]
        top_idx = [candidate_idx[j] for j in order]
        scores = {candidate_idx[j]: float(rerank_scores[j]) for j in order}
        score_type = "rerank"
    else:
        top_idx = tfidf_order[:top].tolist()
        scores = {i: float(cosine_scores[i]) for i in top_idx}
        score_type = "cosine"

    return [
        {
            "rank": rank,
            "source": index.chunks[i]["source"],
            "text": index.chunks[i]["text"],
            "score": scores[i],
            "score_type": score_type,
            "cosine": float(cosine_scores[i]),
        }
        for rank, i in enumerate(top_idx, 1)
    ]

search_exact ¶

search_exact(
    pattern: str,
    index: RagIndex,
    regex: bool = False,
    max_hits: int = 10,
    context: int = 300,
) -> list

Substring/regex search over an index's raw chunk text -- no embedding, no reranker, no model download, so this needs neither the 'rag' extra nor a built index's embeddings (works even on an index built with compute_embeddings=False).

Semantic search (search(), above) ranks by topical similarity, which can bury a short, specific, load-bearing phrase (an exact clause, a fixed parameter value, a named condition) under chunks that are merely more topically central. Use this instead when you already know roughly what wording you're looking for and semantic search isn't surfacing it high enough.

regex=False (default) treats pattern as a literal, case-insensitive substring (re.escape'd internally); regex=True treats it as a real regular expression, case-sensitive, exactly as written. Stops after max_hits matches (a truncated field on the last result marks this, rather than silently returning a partial, unmarked list); context characters of surrounding text are kept on each side of the match.

Returns a list of {source, chunk_index, match, snippet, start, end} dicts, at most one per chunk -- the chunk's first match only, same as the original quantumrag CLI this was ported from, not an exhaustive enumeration of every occurrence within a chunk. Empty list (not an error) if nothing matches.

Source code in tools/ia_utils/rag.py

def search_exact(pattern: str, index: RagIndex, regex: bool = False, max_hits: int = 10, context: int = 300) -> list:
    """Substring/regex search over an index's raw chunk text -- no
    embedding, no reranker, no model download, so this needs neither the
    'rag' extra nor a built index's embeddings (works even on an index
    built with compute_embeddings=False).

    Semantic search (search(), above) ranks by topical similarity, which
    can bury a short, specific, load-bearing phrase (an exact clause, a
    fixed parameter value, a named condition) under chunks that are
    merely more topically central. Use this instead when you already know
    roughly what wording you're looking for and semantic search isn't
    surfacing it high enough.

    regex=False (default) treats `pattern` as a literal, case-insensitive
    substring (re.escape'd internally); regex=True treats it as a real
    regular expression, case-sensitive, exactly as written. Stops after
    `max_hits` matches (a `truncated` field on the last result marks this,
    rather than silently returning a partial, unmarked list); `context`
    characters of surrounding text are kept on each side of the match.

    Returns a list of {source, chunk_index, match, snippet, start, end}
    dicts, at most one per chunk -- the chunk's first match only, same as
    the original quantumrag CLI this was ported from, not an exhaustive
    enumeration of every occurrence within a chunk. Empty list (not an
    error) if nothing matches.
    """
    flags = 0 if regex else re.IGNORECASE
    compiled = re.compile(pattern if regex else re.escape(pattern), flags)
    hits = []
    for i, chunk in enumerate(index.chunks):
        text = chunk["text"]
        m = compiled.search(text)
        if not m:
            continue
        lo = max(0, m.start() - context)
        hi = min(len(text), m.end() + context)
        snippet = text[lo:hi]
        hits.append({
            "source": chunk["source"],
            "chunk_index": i,
            "match": m.group(0),
            "snippet": ("..." if lo > 0 else "") + snippet + ("..." if hi < len(text) else ""),
            "start": m.start(),
            "end": m.end(),
        })
        if len(hits) >= max_hits:
            hits[-1]["truncated"] = True
            break
    return hits

IA Utils — Hybrid Retrieval (RAG)¶

Step 1. Index three documents and ask a question¶

Step 2. Compare against TF-IDF alone¶

Step 3. Find a specific known phrase directly¶

See Also¶

Details¶

Where this came from¶

A real case where reranking changed the answer¶

The optional rag extra¶

Model choice and where the two stages come from¶

Persisting an index¶

rag ¶

build_index ¶

search ¶

search_exact ¶

The optional `rag` extra¶