From 2bdc930b2af871bc91e2c7986b859d608e9d884d Mon Sep 17 00:00:00 2001 From: Sunny Yang Date: Tue, 7 Jul 2026 05:54:02 -0600 Subject: [PATCH] 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. --- .../proc-group/groups/embeddings-store.yaml | 8 + .../test_optional_request_response_spec.py | 81 +++++++ .../test_document_rag_hybrid.py | 211 ++++++++++++++++++ .../test_kw_index_fts5_storage.py | 157 +++++++++++++ trustgraph-base/trustgraph/base/__init__.py | 2 + .../trustgraph/base/keyword_index_client.py | 44 ++++ .../trustgraph/base/keyword_index_service.py | 132 +++++++++++ .../trustgraph/base/request_response_spec.py | 14 +- .../trustgraph/schema/services/query.py | 21 ++ trustgraph-flow/pyproject.toml | 1 + .../retrieval/document_rag/document_rag.py | 114 +++++++++- .../trustgraph/retrieval/document_rag/rag.py | 54 +++++ .../trustgraph/storage/kw_index/__init__.py | 0 .../storage/kw_index/fts5/__init__.py | 1 + .../storage/kw_index/fts5/__main__.py | 6 + .../storage/kw_index/fts5/service.py | 180 +++++++++++++++ 16 files changed, 1013 insertions(+), 13 deletions(-) create mode 100644 tests/unit/test_base/test_optional_request_response_spec.py create mode 100644 tests/unit/test_retrieval/test_document_rag_hybrid.py create mode 100644 tests/unit/test_storage/test_kw_index_fts5_storage.py create mode 100644 trustgraph-base/trustgraph/base/keyword_index_client.py create mode 100644 trustgraph-base/trustgraph/base/keyword_index_service.py create mode 100644 trustgraph-flow/trustgraph/storage/kw_index/__init__.py create mode 100644 trustgraph-flow/trustgraph/storage/kw_index/fts5/__init__.py create mode 100644 trustgraph-flow/trustgraph/storage/kw_index/fts5/__main__.py create mode 100644 trustgraph-flow/trustgraph/storage/kw_index/fts5/service.py diff --git a/dev-tools/proc-group/groups/embeddings-store.yaml b/dev-tools/proc-group/groups/embeddings-store.yaml index b5d4a6c8..e76ac0a3 100644 --- a/dev-tools/proc-group/groups/embeddings-store.yaml +++ b/dev-tools/proc-group/groups/embeddings-store.yaml @@ -32,6 +32,14 @@ processors: id: graph-embeddings-write store_uri: http://localhost:6333 + # Keyword (BM25) index: ingest-write and query in one processor, since + # the FTS5 index is a single local file. + - class: trustgraph.storage.kw_index.fts5.Processor + params: + <<: *defaults + id: kw-index + index_path: /tmp/tg-kw-index.db + - class: trustgraph.query.row_embeddings.qdrant.Processor params: <<: *defaults diff --git a/tests/unit/test_base/test_optional_request_response_spec.py b/tests/unit/test_base/test_optional_request_response_spec.py new file mode 100644 index 00000000..f3eca975 --- /dev/null +++ b/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": {}}, + ) diff --git a/tests/unit/test_retrieval/test_document_rag_hybrid.py b/tests/unit/test_retrieval/test_document_rag_hybrid.py new file mode 100644 index 00000000..fc24f411 --- /dev/null +++ b/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, + ) diff --git a/tests/unit/test_storage/test_kw_index_fts5_storage.py b/tests/unit/test_storage/test_kw_index_fts5_storage.py new file mode 100644 index 00000000..2867f8d6 --- /dev/null +++ b/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") == [] diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 0e77f8ac..e5dd28de 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -44,6 +44,8 @@ from . agent_client import AgentClientSpec from . structured_query_client import StructuredQueryClientSpec from . reranker_client import RerankerClientSpec from . reranker_service import RerankerService +from . keyword_index_service import KeywordIndexService +from . keyword_index_client import KeywordIndexClientSpec, KeywordIndexClient from . row_embeddings_query_client import RowEmbeddingsQueryClientSpec from . collection_config_handler import CollectionConfigHandler from . audit_publisher import AuditPublisher diff --git a/trustgraph-base/trustgraph/base/keyword_index_client.py b/trustgraph-base/trustgraph/base/keyword_index_client.py new file mode 100644 index 00000000..f2cbd362 --- /dev/null +++ b/trustgraph-base/trustgraph/base/keyword_index_client.py @@ -0,0 +1,44 @@ + +import logging + +from . request_response_spec import RequestResponse, RequestResponseSpec +from .. schema import KeywordIndexRequest, KeywordIndexResponse + +# Module logger +logger = logging.getLogger(__name__) + +class KeywordIndexClient(RequestResponse): + async def query(self, query, limit=20, collection="default", timeout=30): + + resp = await self.request( + KeywordIndexRequest( + query = query, + limit = limit, + collection = collection + ), + timeout=timeout + ) + + logger.debug("Keyword index response: %s", resp) + + if resp.error: + raise RuntimeError(resp.error.message) + + # Return ChunkMatch objects with chunk_id and score + return resp.chunks + +class KeywordIndexClientSpec(RequestResponseSpec): + def __init__( + self, request_name, response_name, + ): + super(KeywordIndexClientSpec, self).__init__( + request_name = request_name, + request_schema = KeywordIndexRequest, + response_name = response_name, + response_schema = KeywordIndexResponse, + impl = KeywordIndexClient, + # Flow definitions predating the keyword index don't declare + # these topics; bind only where they exist so one stale + # definition can't wedge the processor. + optional = True, + ) diff --git a/trustgraph-base/trustgraph/base/keyword_index_service.py b/trustgraph-base/trustgraph/base/keyword_index_service.py new file mode 100644 index 00000000..99d15940 --- /dev/null +++ b/trustgraph-base/trustgraph/base/keyword_index_service.py @@ -0,0 +1,132 @@ +""" +Keyword index service base class. A single service owns both sides of the +lexical index: it consumes Chunk messages off the ingestion stream (the last +message in the pipeline that still carries chunk text) and answers keyword +search requests over what it has indexed. Unlike the vector stores, ingest +and query are not split into two processors: the first backend (SQLite FTS5) +is a single-file index that cannot be shared between containers, so one +process must own it. Backends with a server (Elasticsearch/OpenSearch) can +still be split later behind the same schema. +""" + +from __future__ import annotations + +from argparse import ArgumentParser + +import logging + +from .. schema import Chunk +from .. schema import KeywordIndexRequest, KeywordIndexResponse +from .. schema import Error +from .. exceptions import TooManyRequests + +from . flow_processor import FlowProcessor +from . consumer_spec import ConsumerSpec +from . producer_spec import ProducerSpec + +# Module logger +logger = logging.getLogger(__name__) + +default_ident = "kw-index" +default_concurrency = 10 + +class KeywordIndexService(FlowProcessor): + + def __init__(self, **params): + + id = params.get("id") + concurrency = params.get("concurrency", default_concurrency) + + super(KeywordIndexService, self).__init__( + **params | { "id": id } + ) + + self.register_specification( + ConsumerSpec( + name = "input", + schema = Chunk, + handler = self.on_chunk, + ) + ) + + self.register_specification( + ConsumerSpec( + name = "request", + schema = KeywordIndexRequest, + handler = self.on_request, + concurrency = concurrency, + ) + ) + + self.register_specification( + ProducerSpec( + name = "response", + schema = KeywordIndexResponse, + ) + ) + + async def on_chunk(self, msg, consumer, flow): + + try: + + request = msg.value() + + # Workspace comes from the flow the message arrived on. + await self.index_chunk(flow.workspace, request) + + except TooManyRequests as e: + raise e + + except Exception as e: + + logger.error(f"Exception in keyword index store: {e}", exc_info=True) + raise e + + async def on_request(self, msg, consumer, flow): + + try: + + request = msg.value() + + # Sender-produced ID + id = msg.properties()["id"] + + logger.debug(f"Handling keyword index query request {id}...") + + chunks = await self.query_keyword_index( + flow.workspace, request, + ) + + logger.debug("Sending keyword index query response...") + r = KeywordIndexResponse(chunks=chunks, error=None) + await flow("response").send(r, properties={"id": id}) + + logger.debug("Keyword index query request completed") + + except Exception as e: + + logger.error(f"Exception in keyword index query service: {e}", exc_info=True) + + logger.info("Sending error response...") + + r = KeywordIndexResponse( + error=Error( + type = "keyword-index-query-error", + message = str(e), + ), + chunks=[], + ) + + await flow("response").send(r, properties={"id": id}) + + @staticmethod + def add_args(parser: ArgumentParser) -> None: + + FlowProcessor.add_args(parser) + + parser.add_argument( + '-c', '--concurrency', + type=int, + default=default_concurrency, + help=f'Number of concurrent requests (default: {default_concurrency})' + ) diff --git a/trustgraph-base/trustgraph/base/request_response_spec.py b/trustgraph-base/trustgraph/base/request_response_spec.py index aa934a7f..7c8ad3bd 100644 --- a/trustgraph-base/trustgraph/base/request_response_spec.py +++ b/trustgraph-base/trustgraph/base/request_response_spec.py @@ -109,16 +109,28 @@ class RequestResponse(Subscriber): class RequestResponseSpec(Spec): def __init__( self, request_name, request_schema, response_name, - response_schema, impl=RequestResponse + response_schema, impl=RequestResponse, optional=False ): self.request_name = request_name self.request_schema = request_schema self.response_name = response_name self.response_schema = response_schema self.impl = impl + self.optional = optional def add(self, flow: Any, processor: Any, definition: dict[str, Any]) -> None: + # An optional client binds only when the flow definition declares + # its topics. Older definitions predating the topics would otherwise + # KeyError here during Flow construction, which wedges the whole + # processor in a start-flow retry loop; skipping instead leaves + # flow(name) returning None for the caller to handle per-request. + topics = definition.get("topics", {}) + if self.optional and ( + self.request_name not in topics + or self.response_name not in topics): + return + request_metrics = ProducerMetrics( processor = flow.id, flow = flow.name, name = self.request_name ) diff --git a/trustgraph-base/trustgraph/schema/services/query.py b/trustgraph-base/trustgraph/schema/services/query.py index 9c11a157..f390db8c 100644 --- a/trustgraph-base/trustgraph/schema/services/query.py +++ b/trustgraph-base/trustgraph/schema/services/query.py @@ -71,6 +71,27 @@ document_embeddings_response_queue = queue('document-embeddings', cls='response' ############################################################################ +# Keyword index query - lexical (BM25) search over chunk text, the sparse +# counterpart to the doc embeddings query above. Matches share the ChunkMatch +# shape so both retrieval paths key on chunk_id; score is "higher is better" +# in both (BM25 rank scores are negated by the service to match). + +@dataclass +class KeywordIndexRequest: + query: str = "" + limit: int = 0 + collection: str = "" + +@dataclass +class KeywordIndexResponse: + error: Error | None = None + chunks: list[ChunkMatch] = field(default_factory=list) + +keyword_index_request_queue = queue('keyword-index', cls='request') +keyword_index_response_queue = queue('keyword-index', cls='response') + +############################################################################ + # Row embeddings query - for semantic/fuzzy matching on row index values @dataclass diff --git a/trustgraph-flow/pyproject.toml b/trustgraph-flow/pyproject.toml index 648df8bd..e7c65732 100644 --- a/trustgraph-flow/pyproject.toml +++ b/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" diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py index f2087912..b64317ad 100644 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py +++ b/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) diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py index 80dfb6b1..f2c128db 100755 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py +++ b/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__) diff --git a/trustgraph-flow/trustgraph/storage/kw_index/__init__.py b/trustgraph-flow/trustgraph/storage/kw_index/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trustgraph-flow/trustgraph/storage/kw_index/fts5/__init__.py b/trustgraph-flow/trustgraph/storage/kw_index/fts5/__init__.py new file mode 100644 index 00000000..98f4d9da --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/kw_index/fts5/__init__.py @@ -0,0 +1 @@ +from . service import * diff --git a/trustgraph-flow/trustgraph/storage/kw_index/fts5/__main__.py b/trustgraph-flow/trustgraph/storage/kw_index/fts5/__main__.py new file mode 100644 index 00000000..da5a9021 --- /dev/null +++ b/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() diff --git a/trustgraph-flow/trustgraph/storage/kw_index/fts5/service.py b/trustgraph-flow/trustgraph/storage/kw_index/fts5/service.py new file mode 100644 index 00000000..03b147eb --- /dev/null +++ b/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__)