trustgraph-flow 2.7.0__tar.gz → 2.7.2__tar.gz

{trustgraph_flow-2.7.0 → trustgraph_flow-2.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.7.0
+Version: 2.7.2
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph

{trustgraph_flow-2.7.0 → trustgraph_flow-2.7.2}/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-2.7.2/trustgraph/flow_version.py

@@ -0,0 +1 @@
+__version__ = "2.7.2"

{trustgraph_flow-2.7.0 → trustgraph_flow-2.7.2}/trustgraph/gateway/auth.py

@@ -189,36 +189,44 @@ class IamAuth:
     # Lifecycle
     # ------------------------------------------------------------------
 
+    async def _fetch_signing_key(self):
+        """Single attempt to fetch the IAM signing public key.
+        Returns True on success."""
+        try:
+            async def _fetch(client):
+                return await client.get_signing_key_public()
+
+            pem = await self._with_client(_fetch)
+            if pem:
+                self._signing_public_pem = pem
+                logger.info(
+                    "IamAuth: fetched IAM signing public key "
+                    f"({len(pem)} bytes)"
+                )
+                return True
+        except Exception as e:
+            logger.debug(
+                f"IamAuth: signing key fetch failed: "
+                f"{type(e).__name__}: {e}"
+            )
+        return False
+
     async def start(self, max_retries=30, retry_delay=2.0):
         """Fetch the signing public key from IAM.  Retries on
         failure — the gateway may be starting before IAM is ready."""
 
-        async def _fetch(client):
-            return await client.get_signing_key_public()
-
         for attempt in range(max_retries):
-            try:
-                pem = await self._with_client(_fetch)
-                if pem:
-                    self._signing_public_pem = pem
-                    logger.info(
-                        "IamAuth: fetched IAM signing public key "
-                        f"({len(pem)} bytes)"
-                    )
-                    return
-            except Exception as e:
-                logger.info(
-                    f"IamAuth: waiting for IAM signing key "
-                    f"({type(e).__name__}: {e}); "
-                    f"retry {attempt + 1}/{max_retries}"
-                )
+            if await self._fetch_signing_key():
+                return
+            logger.info(
+                f"IamAuth: waiting for IAM signing key; "
+                f"retry {attempt + 1}/{max_retries}"
+            )
             await asyncio.sleep(retry_delay)
 
-        # Don't prevent startup forever.  A later authenticate() call
-        # will try again via the JWT path.
         logger.warning(
             "IamAuth: could not fetch IAM signing key at startup; "
-            "JWT validation will fail until it's available"
+            "JWT validation will retry on first request"
         )
 
     # ------------------------------------------------------------------
@@ -259,7 +267,7 @@ class IamAuth:
             self._annotate_request(request, identity)
             return identity
         if token.count(".") == 2:
-            identity = self._verify_jwt(token)
+            identity = await self._verify_jwt(token)
             self._annotate_request(request, identity)
             return identity
         raise _auth_failure()
@@ -281,7 +289,9 @@ class IamAuth:
         except Exception:
             pass
 
-    def _verify_jwt(self, token):
+    async def _verify_jwt(self, token):
+        if not self._signing_public_pem:
+            await self._fetch_signing_key()
         if not self._signing_public_pem:
             raise _auth_failure()
         try:

{trustgraph_flow-2.7.0 → trustgraph_flow-2.7.2}/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-2.7.0 → trustgraph_flow-2.7.2}/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,10 @@ 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)
+        fetch_chunk_timeout = params.get("fetch_chunk_timeout", 120)
+        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 +48,10 @@ class Processor(FlowProcessor):
                 "fetch_limit": fetch_limit,
                 "rerank_diversity_mode": rerank_diversity_mode,
                 "rerank_diversity_lambda": rerank_diversity_lambda,
+                "fetch_chunk_timeout": fetch_chunk_timeout,
+                "retrieval_mode": retrieval_mode,
+                "vector_weight": vector_weight,
+                "keyword_weight": keyword_weight,
             }
         )
 
@@ -50,6 +59,10 @@ class Processor(FlowProcessor):
         self.fetch_limit = fetch_limit
         self.rerank_diversity_mode = rerank_diversity_mode
         self.rerank_diversity_lambda = rerank_diversity_lambda
+        self.fetch_chunk_timeout = fetch_chunk_timeout
+        self.retrieval_mode = retrieval_mode
+        self.vector_weight = vector_weight
+        self.keyword_weight = keyword_weight
 
         self.register_specification(
             ConsumerSpec(
@@ -87,6 +100,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",
@@ -116,7 +142,7 @@ class Processor(FlowProcessor):
 
             logger.info(f"Handling input {id}...")
 
-            async def fetch_chunk(chunk_id, timeout=120):
+            async def fetch_chunk(chunk_id, timeout=self.fetch_chunk_timeout):
                 return await flow.librarian.fetch_document_text(
                     document_id=chunk_id, timeout=timeout,
                 )
@@ -130,6 +156,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 +332,38 @@ class Processor(FlowProcessor):
             help='MMR relevance/diversity tradeoff, higher values prefer relevance'
         )
 
+        parser.add_argument(
+            '--fetch-chunk-timeout',
+            type=int,
+            default=120,
+            help='Timeout in seconds for fetching a document chunk from the '
+                 'librarian (default: 120)'
+        )
+        
+        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-2.7.2/trustgraph/storage/kw_index/fts5/__init__.py

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

trustgraph_flow-2.7.2/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-2.7.2/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__)

{trustgraph_flow-2.7.0 → trustgraph_flow-2.7.2}/trustgraph_flow.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.7.0
+Version: 2.7.2
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Author-email: "trustgraph.ai" <security@trustgraph.ai>
 Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph

{trustgraph_flow-2.7.0 → trustgraph_flow-2.7.2}/trustgraph_flow.egg-info/SOURCES.txt

@@ -358,6 +358,10 @@ trustgraph/storage/graph_embeddings/qdrant/write.py
 trustgraph/storage/knowledge/__init__.py
 trustgraph/storage/knowledge/__main__.py
 trustgraph/storage/knowledge/store.py
+trustgraph/storage/kw_index/__init__.py
+trustgraph/storage/kw_index/fts5/__init__.py
+trustgraph/storage/kw_index/fts5/__main__.py
+trustgraph/storage/kw_index/fts5/service.py
 trustgraph/storage/row_embeddings/__init__.py
 trustgraph/storage/row_embeddings/qdrant/__init__.py
 trustgraph/storage/row_embeddings/qdrant/__main__.py

{trustgraph_flow-2.7.0 → trustgraph_flow-2.7.2}/trustgraph_flow.egg-info/entry_points.txt

@@ -35,6 +35,7 @@ kg-extract-rows = trustgraph.extract.kg.rows:run
 kg-extract-topics = trustgraph.extract.kg.topics:run
 kg-manager = trustgraph.cores:run
 kg-store = trustgraph.storage.knowledge:run
+kw-index-fts5 = trustgraph.storage.kw_index.fts5:run
 librarian = trustgraph.librarian:run
 mcp-tool = trustgraph.agent.mcp_tool:run
 metering = trustgraph.metering:run

trustgraph_flow-2.7.0/trustgraph/flow_version.py

@@ -1 +0,0 @@
-__version__ = "2.7.0"