feat: hybrid retrieval (BM25 + vector RRF fusion) for document-RAG (#875) (#1030)

Adds a sparse keyword retrieval path beside the existing vector path in
document-RAG, fused by weighted Reciprocal Rank Fusion on chunk_id, behind
--retrieval-mode (vector | keyword | hybrid, default vector).

The keyword index is a new pluggable service (KeywordIndexService /
KeywordIndexClientSpec); the first backend is SQLite FTS5, consuming Chunk
messages off the ingestion stream and answering BM25 queries from one
process, since the index is a single local file. Query text is sanitized
into per-term quoted phrases (raw text is not valid FTS5 syntax), which
also makes dotted clause numbers and error codes exact-match without a
trigram index. Indexes are scoped per (workspace, collection) and dropped
on collection deletion.

The keyword-index client spec is only registered when the sparse path is
enabled, so existing flow definitions without keyword-index queues are
untouched; with retrieval_mode=vector the retrieval path is unchanged. In
hybrid mode a keyword-path failure degrades to vector-only.

trustgraph-flow/pyproject.toml

@@ -76,6 +76,7 @@ document-embeddings = "trustgraph.embeddings.document_embeddings:run"
 document-rag = "trustgraph.retrieval.document_rag:run"
 embeddings-fastembed = "trustgraph.embeddings.fastembed:run"
 embeddings-ollama = "trustgraph.embeddings.ollama:run"
+kw-index-fts5 = "trustgraph.storage.kw_index.fts5:run"
 graph-embeddings-query-milvus = "trustgraph.query.graph_embeddings.milvus:run"
 graph-embeddings-query-pinecone = "trustgraph.query.graph_embeddings.pinecone:run"
 graph-embeddings-query-qdrant = "trustgraph.query.graph_embeddings.qdrant:run"

trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py

@@ -31,8 +31,33 @@ logger = logging.getLogger(__name__)
 # This is only the fallback default: an explicit fetch_limit overrides it.
 OVERFETCH_FACTOR = 3
 
+# Reciprocal Rank Fusion constant. The standard value from Cormack et al.
+# (SIGIR 2009); higher values flatten the contribution of top ranks.
+RRF_K = 60
+
 LABEL="http://www.w3.org/2000/01/rdf-schema#label"
 
+def rrf_fuse(ranked_lists, weights, limit):
+    """Fuse ranked ChunkMatch lists by weighted Reciprocal Rank Fusion.
+
+    score(chunk) = sum over lists of weight / (RRF_K + rank), so fusion
+    needs only each list's ordering, never its native score scale — BM25
+    and cosine scores are incomparable. Returns the surviving matches
+    (first-seen object per chunk_id) in fused order, truncated to limit.
+    """
+    scores = {}
+    first_seen = {}
+    for matches, weight in zip(ranked_lists, weights):
+        for rank, match in enumerate(matches, start=1):
+            if not match.chunk_id:
+                continue
+            scores[match.chunk_id] = (
+                scores.get(match.chunk_id, 0.0) + weight / (RRF_K + rank)
+            )
+            first_seen.setdefault(match.chunk_id, match)
+    ordered = sorted(scores, key=lambda cid: -scores[cid])
+    return [first_seen[cid] for cid in ordered[:limit]]
+
 class Query:
 
     def __init__(
@@ -85,15 +110,8 @@ class Query:
 
         return qembeds
 
-    async def get_docs(self, concepts):
-        """
-        Get documents (chunks) matching the extracted concepts.
-
-        Returns:
-            tuple: (docs, chunk_ids) where:
-                - docs: list of document content strings
-                - chunk_ids: list of chunk IDs that were successfully fetched
-        """
+    async def get_vector_matches(self, concepts):
+        """Dense path: embed concepts, query the vector store, dedupe."""
         vectors = await self.get_vectors(concepts)
 
         if self.verbose:
@@ -123,6 +141,56 @@ class Query:
                     seen.add(match.chunk_id)
                     chunk_matches.append(match)
 
+        return chunk_matches
+
+    async def get_keyword_matches(self, query):
+        """Sparse path: BM25 search on the raw query text."""
+        if self.verbose:
+            logger.debug("Getting chunks from keyword index...")
+
+        return await self.rag.kw_index_client.query(
+            query=query, limit=self.fetch_limit,
+            collection=self.collection,
+        )
+
+    async def get_docs(self, concepts, query=""):
+        """
+        Get documents (chunks) matching the query, via the retrieval mode's
+        paths: dense (concept embeddings), sparse (BM25 over the raw query
+        text), or both fused by RRF. `query` is only consulted by the sparse
+        path; existing vector-mode callers may omit it.
+
+        Returns:
+            tuple: (docs, chunk_ids) where:
+                - docs: list of document content strings
+                - chunk_ids: list of chunk IDs that were successfully fetched
+        """
+        mode = self.rag.retrieval_mode
+
+        if mode == "keyword":
+            chunk_matches = await self.get_keyword_matches(query)
+        elif mode == "hybrid":
+            # The paths are independent; a keyword-index failure degrades
+            # to vector-only rather than failing the whole query.
+            async def keyword_or_empty():
+                try:
+                    return await self.get_keyword_matches(query)
+                except Exception as e:
+                    logger.warning(f"Keyword path failed, using vector only: {e}")
+                    return []
+
+            vector_matches, keyword_matches = await asyncio.gather(
+                self.get_vector_matches(concepts),
+                keyword_or_empty(),
+            )
+            chunk_matches = rrf_fuse(
+                [vector_matches, keyword_matches],
+                [self.rag.vector_weight, self.rag.keyword_weight],
+                self.fetch_limit,
+            )
+        else:
+            chunk_matches = await self.get_vector_matches(concepts)
+
         if self.verbose:
             logger.debug(f"Got {len(chunk_matches)} chunks, fetching content from Garage...")
 
@@ -154,6 +222,10 @@ class DocumentRag:
             verbose=False,
             rerank_diversity_mode="none",
             rerank_diversity_lambda=0.7,
+            kw_index_client=None,
+            retrieval_mode="vector",
+            vector_weight=1.0,
+            keyword_weight=1.0,
     ):
 
         self.verbose = verbose
@@ -169,6 +241,19 @@ class DocumentRag:
         self.rerank_diversity_mode = rerank_diversity_mode
         self.rerank_diversity_lambda = rerank_diversity_lambda
 
+        # Optional sparse (BM25) retrieval path. "vector" keeps the current
+        # dense-only behaviour; "keyword"/"hybrid" need a keyword index
+        # client wired.
+        if retrieval_mode != "vector" and kw_index_client is None:
+            raise ValueError(
+                f"retrieval_mode={retrieval_mode!r} requires a keyword "
+                f"index client"
+            )
+        self.kw_index_client = kw_index_client
+        self.retrieval_mode = retrieval_mode
+        self.vector_weight = vector_weight
+        self.keyword_weight = keyword_weight
+
         if self.verbose:
             logger.debug("DocumentRag initialized")
 
@@ -249,8 +334,13 @@ class DocumentRag:
             fetch_limit=fetch_count, track_usage=track_usage,
         )
 
-        # Extract concepts from query (grounding step)
-        concepts = await q.extract_concepts(query)
+        # Extract concepts from query (grounding step). Concepts only feed
+        # the dense path's embeddings; in keyword-only mode the LLM call
+        # would be paid and discarded, so ground on the raw query instead.
+        if self.retrieval_mode == "keyword":
+            concepts = [query]
+        else:
+            concepts = await q.extract_concepts(query)
 
         # Emit grounding explainability after concept extraction
         if explain_callback:
@@ -266,7 +356,7 @@ class DocumentRag:
             )
             await explain_callback(gnd_triples, gnd_uri)
 
-        docs, chunk_ids = await q.get_docs(concepts)
+        docs, chunk_ids = await q.get_docs(concepts, query)
 
         # Emit exploration explainability after chunks retrieved
         # (full candidate set, before any reranking)

trustgraph-flow/trustgraph/retrieval/document_rag/rag.py

@@ -14,6 +14,7 @@ from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
 from ... base import PromptClientSpec, EmbeddingsClientSpec
 from ... base import DocumentEmbeddingsClientSpec
 from ... base import RerankerClientSpec
+from ... base import KeywordIndexClientSpec
 from ... base import LibrarianSpec
 
 # Module logger
@@ -35,6 +36,9 @@ class Processor(FlowProcessor):
         fetch_limit = params.get("fetch_limit", 0)
         rerank_diversity_mode = params.get("rerank_diversity_mode", "none")
         rerank_diversity_lambda = params.get("rerank_diversity_lambda", 0.7)
+        retrieval_mode = params.get("retrieval_mode", "vector")
+        vector_weight = params.get("vector_weight", 1.0)
+        keyword_weight = params.get("keyword_weight", 1.0)
 
         super(Processor, self).__init__(
             **params | {
@@ -43,6 +47,9 @@ class Processor(FlowProcessor):
                 "fetch_limit": fetch_limit,
                 "rerank_diversity_mode": rerank_diversity_mode,
                 "rerank_diversity_lambda": rerank_diversity_lambda,
+                "retrieval_mode": retrieval_mode,
+                "vector_weight": vector_weight,
+                "keyword_weight": keyword_weight,
             }
         )
 
@@ -50,6 +57,9 @@ class Processor(FlowProcessor):
         self.fetch_limit = fetch_limit
         self.rerank_diversity_mode = rerank_diversity_mode
         self.rerank_diversity_lambda = rerank_diversity_lambda
+        self.retrieval_mode = retrieval_mode
+        self.vector_weight = vector_weight
+        self.keyword_weight = keyword_weight
 
         self.register_specification(
             ConsumerSpec(
@@ -87,6 +97,19 @@ class Processor(FlowProcessor):
             )
         )
 
+        # Only registered when the sparse path is enabled: the spec binds
+        # keyword-index topics from the flow definition, so registering it
+        # unconditionally would break flow classes that don't declare them.
+        # With the default retrieval_mode=vector, existing deployments are
+        # untouched.
+        if retrieval_mode != "vector":
+            self.register_specification(
+                KeywordIndexClientSpec(
+                    request_name = "keyword-index-request",
+                    response_name = "keyword-index-response",
+                )
+            )
+
         self.register_specification(
             ProducerSpec(
                 name = "response",
@@ -130,6 +153,13 @@ class Processor(FlowProcessor):
                 verbose=True,
                 rerank_diversity_mode=self.rerank_diversity_mode,
                 rerank_diversity_lambda=self.rerank_diversity_lambda,
+                # None when the spec wasn't registered (vector mode) or its
+                # topics were absent from this flow's definition (optional
+                # spec skipped) — DocumentRag validates per-request.
+                kw_index_client = flow("keyword-index-request"),
+                retrieval_mode=self.retrieval_mode,
+                vector_weight=self.vector_weight,
+                keyword_weight=self.keyword_weight,
             )
 
             if v.doc_limit:
@@ -299,6 +329,30 @@ class Processor(FlowProcessor):
             help='MMR relevance/diversity tradeoff, higher values prefer relevance'
         )
 
+        parser.add_argument(
+            '--retrieval-mode',
+            choices=['vector', 'keyword', 'hybrid'],
+            default='vector',
+            help='Chunk retrieval strategy: dense vector search (default), '
+                 'BM25 keyword search, or both fused by reciprocal rank '
+                 'fusion. keyword/hybrid need keyword-index queues in the '
+                 'flow definition'
+        )
+
+        parser.add_argument(
+            '--vector-weight',
+            type=float,
+            default=1.0,
+            help='Vector path weight in hybrid rank fusion (default: 1.0)'
+        )
+
+        parser.add_argument(
+            '--keyword-weight',
+            type=float,
+            default=1.0,
+            help='Keyword path weight in hybrid rank fusion (default: 1.0)'
+        )
+
 def run():
 
     Processor.launch(default_ident, __doc__)

trustgraph-flow/trustgraph/storage/kw_index/fts5/__init__.py

@@ -0,0 +1 @@
+from . service import *

trustgraph-flow/trustgraph/storage/kw_index/fts5/__main__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from . service import run
+
+if __name__ == '__main__':
+    run()

trustgraph-flow/trustgraph/storage/kw_index/fts5/service.py

@@ -0,0 +1,180 @@
+
+"""
+Keyword index over chunk text, backed by SQLite FTS5.  Consumes Chunk
+messages off the ingestion stream and answers BM25 keyword queries; both
+sides live in one service because the index is a single local file.  One
+FTS5 table per (workspace, collection) keeps BM25 corpus statistics and
+collection deletion scoped correctly.
+"""
+
+import asyncio
+import logging
+import re
+import sqlite3
+from pathlib import Path
+
+from .... base import KeywordIndexService, CollectionConfigHandler
+from .... schema import ChunkMatch
+
+# Module logger
+logger = logging.getLogger(__name__)
+
+default_ident = "kw-index"
+default_index_path = "/data/kw-index.db"
+
+# FTS5 table names embed workspace/collection; quoting handles the rest, but
+# strip anything outside the character set other stores allow in names so a
+# hostile name can't smuggle quote characters.
+_NAME_SAFE = re.compile(r"[^A-Za-z0-9_-]")
+
+def _table(workspace, collection):
+    ws = _NAME_SAFE.sub("_", workspace)
+    coll = _NAME_SAFE.sub("_", collection)
+    return f"kw_{ws}_{coll}"
+
+def to_match_query(text):
+    """User text -> FTS5 MATCH expression.
+
+    Raw text is not valid FTS5 syntax ("7.3.2" is a syntax error, the "-" in
+    "AURA-7" is column-filter syntax), so each whitespace token is quoted as
+    a phrase and the phrases are OR-ed: BM25 scores accumulate over matching
+    terms, and a quoted phrase of sub-tokens ("7.3.2" -> [7 3 2]) still
+    matches the exact dotted term without also matching "7.3.1".
+    """
+    tokens = [t for t in text.split() if t.strip('"')]
+    if not tokens:
+        return None
+    return " OR ".join('"' + t.replace('"', '""') + '"' for t in tokens)
+
+class Processor(CollectionConfigHandler, KeywordIndexService):
+
+    def __init__(self, **params):
+
+        index_path = params.get("index_path", default_index_path)
+
+        super(Processor, self).__init__(
+            **params | {
+                "index_path": index_path,
+            }
+        )
+
+        Path(index_path).parent.mkdir(parents=True, exist_ok=True)
+
+        # Writes are serialized on one connection by the lock; reads get
+        # their own connection so a query never queues behind the chunk
+        # ingestion backlog. WAL lets the reader proceed while a write
+        # commits, and NORMAL sync is safe with WAL (an index is
+        # re-derivable from the chunk store anyway). All sqlite work runs
+        # in a thread so the event loop is never blocked.
+        self.db = sqlite3.connect(index_path, check_same_thread=False)
+        self.db.execute("PRAGMA journal_mode=WAL")
+        self.db.execute("PRAGMA synchronous=NORMAL")
+        self.read_db = sqlite3.connect(index_path, check_same_thread=False)
+        self._lock = asyncio.Lock()
+
+        # Register for config push notifications
+        self.register_config_handler(self.on_collection_config, types=["collection"])
+
+        logger.info(f"Keyword index at {index_path}")
+
+    def _index(self, table, chunk_id, body):
+        self.db.execute(
+            f'CREATE VIRTUAL TABLE IF NOT EXISTS "{table}" '
+            f'USING fts5(chunk_id UNINDEXED, body)'
+        )
+        # Re-ingesting a chunk replaces its previous row rather than
+        # accumulating duplicates.
+        self.db.execute(
+            f'DELETE FROM "{table}" WHERE chunk_id = ?', (chunk_id,)
+        )
+        self.db.execute(
+            f'INSERT INTO "{table}" (chunk_id, body) VALUES (?, ?)',
+            (chunk_id, body),
+        )
+        self.db.commit()
+
+    def _query(self, table, match, limit):
+        try:
+            rows = self.read_db.execute(
+                f'SELECT chunk_id, bm25("{table}") FROM "{table}" '
+                f'WHERE "{table}" MATCH ? ORDER BY bm25("{table}") LIMIT ?',
+                (match, limit),
+            ).fetchall()
+        except sqlite3.OperationalError as e:
+            if "no such table" in str(e):
+                # Nothing indexed for this collection yet
+                return []
+            raise
+        # bm25() is lower-is-better (negative); negate so ChunkMatch.score
+        # is higher-is-better like the vector path.
+        return [ChunkMatch(chunk_id=r[0], score=-r[1]) for r in rows]
+
+    async def index_chunk(self, workspace, message):
+
+        if not self.collection_exists(workspace, message.metadata.collection):
+            logger.warning(
+                f"Collection {message.metadata.collection} for workspace {workspace} "
+                f"does not exist in config (likely deleted while data was in-flight). "
+                f"Dropping message."
+            )
+            return
+
+        chunk_id = message.document_id
+        if not chunk_id:
+            return
+
+        body = message.chunk.decode("utf-8", errors="replace")
+        if not body.strip():
+            return
+
+        table = _table(workspace, message.metadata.collection)
+
+        async with self._lock:
+            await asyncio.to_thread(self._index, table, chunk_id, body)
+
+    async def query_keyword_index(self, workspace, request):
+
+        match = to_match_query(request.query)
+        if match is None:
+            return []
+
+        limit = request.limit if request.limit > 0 else 20
+        table = _table(workspace, request.collection)
+
+        # No lock: reads run on their own connection and WAL keeps them
+        # consistent alongside the writer.
+        return await asyncio.to_thread(self._query, table, match, limit)
+
+    async def create_collection(self, workspace: str, collection: str, metadata: dict):
+        """FTS5 tables are created lazily on first indexed chunk."""
+        logger.info(
+            f"Collection create request for {workspace}/{collection} - "
+            f"table created lazily on first write"
+        )
+
+    async def delete_collection(self, workspace: str, collection: str):
+        """Drop the FTS5 table for this collection via config push."""
+        table = _table(workspace, collection)
+
+        def drop():
+            self.db.execute(f'DROP TABLE IF EXISTS "{table}"')
+            self.db.commit()
+
+        async with self._lock:
+            await asyncio.to_thread(drop)
+        logger.info(f"Deleted keyword index table: {table}")
+
+    @staticmethod
+    def add_args(parser):
+
+        KeywordIndexService.add_args(parser)
+
+        parser.add_argument(
+            '--index-path',
+            default=default_index_path,
+            help=f'SQLite FTS5 index file (default: {default_index_path})'
+        )
+
+def run():
+
+    Processor.launch(default_ident, __doc__)