mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-09 13:22:10 +02:00
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:
parent
e5206bddd0
commit
2bdc930b2a
16 changed files with 1013 additions and 13 deletions
|
|
@ -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
|
||||
|
|
|
|||
44
trustgraph-base/trustgraph/base/keyword_index_client.py
Normal file
44
trustgraph-base/trustgraph/base/keyword_index_client.py
Normal 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,
|
||||
)
|
||||
132
trustgraph-base/trustgraph/base/keyword_index_service.py
Normal file
132
trustgraph-base/trustgraph/base/keyword_index_service.py
Normal 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})'
|
||||
)
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue