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.
This commit is contained in:
Sunny Yang 2026-07-07 05:54:02 -06:00 committed by GitHub
parent e5206bddd0
commit 2bdc930b2a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 1013 additions and 13 deletions

View file

@ -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

View file

@ -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,
)

View file

@ -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})'
)

View file

@ -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
)

View file

@ -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