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.

tests/unit/test_base/test_optional_request_response_spec.py

@@ -0,0 +1,81 @@
+"""
+Tests for RequestResponseSpec's optional flag: an optional client spec
+binds only when the flow definition declares its topics, so a definition
+predating the topics skips the binding (flow(name) then returns None)
+instead of raising KeyError during Flow construction — which would wedge
+the processor's start-flow retry loop.
+"""
+
+import pytest
+from unittest.mock import MagicMock
+
+from trustgraph.base.request_response_spec import RequestResponseSpec
+
+
+class StubImpl:
+    """Captures constructor kwargs; stands in for RequestResponse."""
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+
+def make_spec(optional):
+    return RequestResponseSpec(
+        request_name="keyword-index-request",
+        request_schema=object,
+        response_name="keyword-index-response",
+        response_schema=object,
+        impl=StubImpl,
+        optional=optional,
+    )
+
+
+def make_flow():
+    flow = MagicMock()
+    flow.id = "f-id"
+    flow.name = "f-name"
+    flow.workspace = "ws"
+    flow.consumer = {}
+    return flow
+
+
+FULL_TOPICS = {
+    "topics": {
+        "keyword-index-request": "request:tg:keyword-index:ws:f",
+        "keyword-index-response": "response:tg:keyword-index:ws:f",
+    }
+}
+
+
+class TestOptionalRequestResponseSpec:
+
+    def test_optional_spec_skips_binding_when_topics_absent(self):
+        flow = make_flow()
+        make_spec(optional=True).add(flow, MagicMock(), {"topics": {}})
+        assert flow.consumer == {}
+
+    def test_optional_spec_skips_when_only_one_topic_present(self):
+        flow = make_flow()
+        definition = {
+            "topics": {
+                "keyword-index-request": "request:tg:keyword-index:ws:f",
+            }
+        }
+        make_spec(optional=True).add(flow, MagicMock(), definition)
+        assert flow.consumer == {}
+
+    def test_optional_spec_binds_when_topics_present(self):
+        flow = make_flow()
+        make_spec(optional=True).add(flow, MagicMock(), FULL_TOPICS)
+        client = flow.consumer["keyword-index-request"]
+        assert isinstance(client, StubImpl)
+        assert client.kwargs["request_topic"] == \
+            "request:tg:keyword-index:ws:f"
+
+    def test_default_spec_still_requires_topics(self):
+        # Non-optional specs keep the existing contract: a missing topic
+        # is a definition error, surfaced immediately.
+        with pytest.raises(KeyError):
+            make_spec(optional=False).add(
+                make_flow(), MagicMock(), {"topics": {}},
+            )

tests/unit/test_retrieval/test_document_rag_hybrid.py

@@ -0,0 +1,211 @@
+"""
+Tests for the retrieval-mode dispatch in DocumentRag (issue: hybrid
+BM25 + vector retrieval).
+
+Covered behaviours:
+
+  1. Default: retrieval_mode="vector" never touches the keyword client and
+     produces the same chunks as before — the sparse path is strictly opt-in.
+  2. keyword: only the keyword index is queried (no vector-store query, no
+     embedding of concepts); chunk order follows the BM25 ranking.
+  3. hybrid: both paths run and are fused by weighted RRF on chunk_id; a
+     keyword-path failure degrades to vector-only instead of failing the
+     query.
+  4. Constructing with keyword/hybrid but no keyword client is an error.
+
+Pure orchestration tests: all subsidiary clients are stubs.
+"""
+
+import pytest
+from unittest.mock import AsyncMock
+
+from trustgraph.retrieval.document_rag.document_rag import (
+    DocumentRag, rrf_fuse, RRF_K,
+)
+from trustgraph.base import PromptResult
+from trustgraph.schema import ChunkMatch
+
+
+CONTENT = {
+    "v1": "vector chunk one",
+    "v2": "vector chunk two",
+    "k1": "keyword chunk one",
+    "both": "chunk found by both paths",
+}
+
+
+def build_clients(vector_ids, keyword_ids):
+    prompt_client = AsyncMock()
+    embeddings_client = AsyncMock()
+    doc_embeddings_client = AsyncMock()
+    kw_index_client = AsyncMock()
+    fetch_chunk = AsyncMock()
+
+    async def mock_prompt(template_id, variables=None, **kwargs):
+        if template_id == "extract-concepts":
+            return PromptResult(response_type="text", text="concept")
+        return PromptResult(response_type="text", text="")
+
+    prompt_client.prompt.side_effect = mock_prompt
+    prompt_client.document_prompt.return_value = PromptResult(
+        response_type="text", text="answer",
+    )
+
+    embeddings_client.embed.return_value = [[0.1, 0.2]]
+
+    doc_embeddings_client.query.return_value = [
+        ChunkMatch(chunk_id=c) for c in vector_ids
+    ]
+    kw_index_client.query.return_value = [
+        ChunkMatch(chunk_id=c, score=1.0) for c in keyword_ids
+    ]
+
+    fetch_chunk.side_effect = lambda chunk_id: CONTENT[chunk_id]
+
+    return (
+        prompt_client, embeddings_client, doc_embeddings_client,
+        kw_index_client, fetch_chunk,
+    )
+
+
+def build_rag(vector_ids, keyword_ids, **kwargs):
+    prompt, embeddings, doc_embeddings, kw, fetch = build_clients(
+        vector_ids, keyword_ids,
+    )
+    rag = DocumentRag(
+        prompt_client=prompt,
+        embeddings_client=embeddings,
+        doc_embeddings_client=doc_embeddings,
+        fetch_chunk=fetch,
+        kw_index_client=kw,
+        **kwargs,
+    )
+    return rag, doc_embeddings, kw, embeddings, prompt
+
+
+# ---------------------------------------------------------------------------
+# rrf_fuse
+# ---------------------------------------------------------------------------
+
+class TestRrfFuse:
+
+    def test_chunk_in_both_lists_outranks_single_list_leaders(self):
+        a = ChunkMatch("a")
+        b = ChunkMatch("b")
+        both = ChunkMatch("both")
+        fused = rrf_fuse([[a, both], [both, b]], [1.0, 1.0], 10)
+        assert [m.chunk_id for m in fused][0] == "both"
+        assert {m.chunk_id for m in fused} == {"a", "b", "both"}
+
+    def test_weights_bias_the_fusion(self):
+        a, b = ChunkMatch("a"), ChunkMatch("b")
+        fused = rrf_fuse([[a], [b]], [1.0, 10.0], 10)
+        assert [m.chunk_id for m in fused] == ["b", "a"]
+
+    def test_limit_truncates(self):
+        matches = [ChunkMatch(f"c{i}") for i in range(5)]
+        assert len(rrf_fuse([matches], [1.0], 2)) == 2
+
+    def test_cross_list_accumulation_beats_single_top_rank(self):
+        # b sums 1/(K+2) + 1/(K+3) across two lists, beating the single
+        # 1/(K+1) that a gets — the accumulation property that
+        # distinguishes RRF from a best-rank merge.
+        a, b, x, y = (ChunkMatch(c) for c in "abxy")
+        fused = rrf_fuse([[a, b], [x, y, b]], [1.0, 1.0], 10)
+        assert fused[0].chunk_id == "b"
+        assert 1 / (RRF_K + 2) + 1 / (RRF_K + 3) > 1 / (RRF_K + 1)
+
+    def test_empty_chunk_ids_are_skipped(self):
+        fused = rrf_fuse([[ChunkMatch(""), ChunkMatch("a")]], [1.0], 10)
+        assert [m.chunk_id for m in fused] == ["a"]
+
+
+# ---------------------------------------------------------------------------
+# Mode dispatch through DocumentRag.query()
+# ---------------------------------------------------------------------------
+
+@pytest.mark.asyncio
+async def test_vector_mode_never_touches_keyword_client():
+    rag, doc_embeddings, kw, _, prompt = build_rag(
+        ["v1", "v2"], ["k1"], retrieval_mode="vector",
+    )
+    await rag.query("question")
+
+    kw.query.assert_not_called()
+    doc_embeddings.query.assert_called()
+    docs = prompt.document_prompt.call_args.kwargs["documents"]
+    assert docs == [CONTENT["v1"], CONTENT["v2"]]
+
+
+@pytest.mark.asyncio
+async def test_default_mode_is_vector_with_no_keyword_client():
+    prompt, embeddings, doc_embeddings, _, fetch = build_clients(
+        ["v1"], [],
+    )
+    rag = DocumentRag(
+        prompt_client=prompt,
+        embeddings_client=embeddings,
+        doc_embeddings_client=doc_embeddings,
+        fetch_chunk=fetch,
+    )
+    await rag.query("question")
+    docs = prompt.document_prompt.call_args.kwargs["documents"]
+    assert docs == [CONTENT["v1"]]
+
+
+@pytest.mark.asyncio
+async def test_keyword_mode_skips_vector_store_and_embeddings():
+    rag, doc_embeddings, kw, embeddings, prompt = build_rag(
+        ["v1", "v2"], ["k1", "both"], retrieval_mode="keyword",
+    )
+    await rag.query("what does clause 7.3.2 say")
+
+    doc_embeddings.query.assert_not_called()
+    embeddings.embed.assert_not_called()
+    # No dense path -> no concept-extraction LLM call either
+    prompt.prompt.assert_not_called()
+    # The sparse path searches the raw query text, not extracted concepts
+    assert kw.query.call_args.kwargs["query"] == "what does clause 7.3.2 say"
+    docs = prompt.document_prompt.call_args.kwargs["documents"]
+    assert docs == [CONTENT["k1"], CONTENT["both"]]
+
+
+@pytest.mark.asyncio
+async def test_hybrid_mode_fuses_both_paths():
+    # both appears in both rankings, so RRF must put it first
+    rag, doc_embeddings, kw, _, prompt = build_rag(
+        ["v1", "both"], ["both", "k1"], retrieval_mode="hybrid",
+    )
+    await rag.query("question")
+
+    doc_embeddings.query.assert_called()
+    kw.query.assert_called()
+    docs = prompt.document_prompt.call_args.kwargs["documents"]
+    assert docs[0] == CONTENT["both"]
+    assert set(docs) == {CONTENT["both"], CONTENT["v1"], CONTENT["k1"]}
+
+
+@pytest.mark.asyncio
+async def test_hybrid_degrades_to_vector_when_keyword_path_fails():
+    rag, doc_embeddings, kw, _, prompt = build_rag(
+        ["v1", "v2"], [], retrieval_mode="hybrid",
+    )
+    kw.query.side_effect = RuntimeError("keyword index down")
+
+    await rag.query("question")
+
+    docs = prompt.document_prompt.call_args.kwargs["documents"]
+    assert docs == [CONTENT["v1"], CONTENT["v2"]]
+
+
+def test_non_vector_mode_without_client_is_an_error():
+    prompt, embeddings, doc_embeddings, _, fetch = build_clients([], [])
+    for mode in ("keyword", "hybrid"):
+        with pytest.raises(ValueError):
+            DocumentRag(
+                prompt_client=prompt,
+                embeddings_client=embeddings,
+                doc_embeddings_client=doc_embeddings,
+                fetch_chunk=fetch,
+                retrieval_mode=mode,
+            )

tests/unit/test_storage/test_kw_index_fts5_storage.py

@@ -0,0 +1,157 @@
+"""
+Unit tests for trustgraph.storage.kw_index.fts5.service — the SQLite FTS5
+keyword index. Covers the MATCH-expression sanitizer (raw user text is not
+valid FTS5 syntax), exact-term retrieval for the motivating cases (dotted
+clause numbers, error codes, hyphenated identifiers), chunk re-ingestion
+replacing rather than duplicating, (workspace, collection) scoping, and
+collection deletion.
+"""
+
+import tempfile
+from pathlib import Path
+
+import pytest
+from unittest.mock import AsyncMock
+from unittest import IsolatedAsyncioTestCase
+
+from trustgraph.schema import Chunk, Metadata, KeywordIndexRequest
+from trustgraph.storage.kw_index.fts5.service import (
+    Processor, to_match_query, _table,
+)
+
+
+class TestMatchQuerySanitizer:
+
+    def test_plain_words_are_quoted_and_or_joined(self):
+        assert to_match_query("return policy") == '"return" OR "policy"'
+
+    def test_dotted_and_hyphenated_terms_survive(self):
+        # Raw "7.3.2" is an FTS5 syntax error; "AURA-7" parses "-" as a
+        # column filter. Quoting neutralizes both.
+        assert to_match_query("clause 7.3.2 AURA-7") == (
+            '"clause" OR "7.3.2" OR "AURA-7"'
+        )
+
+    def test_embedded_quotes_are_escaped(self):
+        assert to_match_query('say "hello"') == '"say" OR """hello"""'
+
+    def test_empty_and_quote_only_queries_yield_none(self):
+        assert to_match_query("") is None
+        assert to_match_query("   ") is None
+        assert to_match_query('"') is None
+
+
+def make_processor(index_path):
+    # A real file, not :memory: — the service holds separate write and read
+    # connections, which only share a database through the filesystem.
+    processor = Processor(
+        taskgroup=AsyncMock(),
+        id="test-kw-index",
+        index_path=index_path,
+    )
+    # Config-pushed collection state isn't wired in unit tests
+    processor.collection_exists = lambda workspace, collection: True
+    return processor
+
+
+def chunk(chunk_id, text, collection="default"):
+    return Chunk(
+        metadata=Metadata(id="doc1", collection=collection),
+        chunk=text.encode("utf-8"),
+        document_id=chunk_id,
+    )
+
+
+CHUNKS = [
+    ("c1", "Clause 7.3.2 states that indemnification obligations survive."),
+    ("c2", "Clause 7.3.1 covers limitation of liability."),
+    ("c3", "Error E4032 occurs when the connection pool is exhausted."),
+]
+
+
+class TestFts5KeywordIndex(IsolatedAsyncioTestCase):
+
+    async def asyncSetUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.processor = make_processor(str(Path(self._tmp.name) / "kw.db"))
+        for chunk_id, text in CHUNKS:
+            await self.processor.index_chunk("ws", chunk("ws-" + chunk_id, text))
+
+    async def asyncTearDown(self):
+        self.processor.db.close()
+        self.processor.read_db.close()
+        self._tmp.cleanup()
+
+    async def query(self, text, collection="default", limit=0):
+        return await self.processor.query_keyword_index(
+            "ws", KeywordIndexRequest(
+                query=text, limit=limit, collection=collection,
+            ),
+        )
+
+    async def test_exact_dotted_term_matches_only_its_clause(self):
+        matches = await self.query("7.3.2")
+        assert [m.chunk_id for m in matches] == ["ws-c1"]
+
+    async def test_error_code_matches(self):
+        matches = await self.query("E4032")
+        assert [m.chunk_id for m in matches] == ["ws-c3"]
+
+    async def test_scores_are_higher_is_better(self):
+        matches = await self.query("clause indemnification")
+        assert matches[0].chunk_id == "ws-c1"
+        assert all(m.score > 0 for m in matches)
+        # c1 matches both terms so it must outrank c2
+        by_id = {m.chunk_id: m.score for m in matches}
+        assert by_id["ws-c1"] > by_id["ws-c2"]
+
+    async def test_reingesting_a_chunk_replaces_it(self):
+        await self.processor.index_chunk(
+            "ws", chunk("ws-c1", "Completely different content now.")
+        )
+        assert await self.query("indemnification 7.3.2") == []
+        matches = await self.query("completely different")
+        assert [m.chunk_id for m in matches] == ["ws-c1"]
+
+    async def test_collections_are_isolated(self):
+        await self.processor.index_chunk(
+            "ws", chunk("other-c1", "indemnification text", collection="other")
+        )
+        default_ids = [m.chunk_id for m in await self.query("indemnification")]
+        other_ids = [
+            m.chunk_id
+            for m in await self.query("indemnification", collection="other")
+        ]
+        assert "other-c1" not in default_ids
+        assert other_ids == ["other-c1"]
+
+    async def test_workspaces_are_isolated(self):
+        matches = await self.processor.query_keyword_index(
+            "someone-else", KeywordIndexRequest(
+                query="indemnification", collection="default",
+            ),
+        )
+        assert matches == []
+
+    async def test_unindexed_collection_returns_empty_not_error(self):
+        assert await self.query("anything", collection="never-written") == []
+
+    async def test_hostile_query_text_is_inert(self):
+        # FTS5 operators and SQL fragments arrive as quoted phrases
+        assert await self.query('body: DROP TABLE OR NOT NEAR(') == []
+
+    async def test_limit_is_applied(self):
+        matches = await self.query("clause", limit=1)
+        assert len(matches) == 1
+
+    async def test_delete_collection_drops_the_index(self):
+        await self.processor.delete_collection("ws", "default")
+        assert await self.query("clause") == []
+
+    async def test_dropped_message_when_collection_missing(self):
+        self.processor.collection_exists = lambda w, c: False
+        await self.processor.index_chunk(
+            "ws", chunk("ws-c9", "should be dropped")
+        )
+        self.processor.collection_exists = lambda w, c: True
+        assert await self.query("dropped") == []