trustgraph-flow 2.6.2__tar.gz → 2.6.4__tar.gz

{trustgraph_flow-2.6.2 → trustgraph_flow-2.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.6.2
+Version: 2.6.4
 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
@@ -17,6 +17,7 @@ Requires-Dist: cryptography
 Requires-Dist: faiss-cpu
 Requires-Dist: falkordb
 Requires-Dist: fastembed
+Requires-Dist: flashrank
 Requires-Dist: ibis
 Requires-Dist: jsonschema
 Requires-Dist: langchain

{trustgraph_flow-2.6.2 → trustgraph_flow-2.6.4}/pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "faiss-cpu",
     "falkordb",
     "fastembed",
+    "flashrank",
     "ibis",
     "jsonschema",
     "langchain",
@@ -83,6 +84,7 @@ graph-embeddings-write-pinecone = "trustgraph.storage.graph_embeddings.pinecone:
 graph-embeddings-write-qdrant = "trustgraph.storage.graph_embeddings.qdrant:run"
 graph-embeddings = "trustgraph.embeddings.graph_embeddings:run"
 graph-rag = "trustgraph.retrieval.graph_rag:run"
+reranker-flashrank = "trustgraph.reranker.flashrank:run"
 kg-extract-agent = "trustgraph.extract.kg.agent:run"
 kg-extract-definitions = "trustgraph.extract.kg.definitions:run"
 kg-extract-rows = "trustgraph.extract.kg.rows:run"

{trustgraph_flow-2.6.2 → trustgraph_flow-2.6.4}/trustgraph/bootstrap/initialisers/default_flow_start.py

@@ -18,6 +18,10 @@ description : str (default "Default")
     Human-readable description passed to flow-svc.
 parameters : dict (optional)
     Optional parameter overrides passed to start-flow.
+list_timeout : int (default 10)
+    Timeout in seconds for the list-flows request.
+start_timeout : int (default 30)
+    Timeout in seconds for the start-flow request.
 """
 
 from trustgraph.schema import FlowRequest
@@ -34,6 +38,8 @@ class DefaultFlowStart(Initialiser):
             blueprint=None,
             description="Default",
             parameters=None,
+            list_timeout=10,
+            start_timeout=30,
             **kwargs,
     ):
         super().__init__(**kwargs)
@@ -46,6 +52,8 @@ class DefaultFlowStart(Initialiser):
         self.blueprint = blueprint
         self.description = description
         self.parameters = dict(parameters) if parameters else {}
+        self.list_timeout = list_timeout
+        self.start_timeout = start_timeout
 
     async def run(self, ctx, old_flag, new_flag):
 
@@ -70,7 +78,7 @@ class DefaultFlowStart(Initialiser):
                 FlowRequest(
                     operation="list-flows",
                 ),
-                timeout=10,
+                timeout=self.list_timeout,
             )
             if list_resp.error:
                 raise RuntimeError(
@@ -99,7 +107,7 @@ class DefaultFlowStart(Initialiser):
                     description=self.description,
                     parameters=self.parameters,
                 ),
-                timeout=30,
+                timeout=self.start_timeout,
             )
             if resp.error:
                 raise RuntimeError(

{trustgraph_flow-2.6.2 → trustgraph_flow-2.6.4}/trustgraph/bootstrap/initialisers/workspace_init.py

@@ -14,7 +14,9 @@ seed_file : str (required when source=="seed-file")
     Path to a JSON seed file with the same shape TemplateSeed consumes.
 overwrite : bool (default False)
     On re-run (flag change), if True overwrite all keys; if False,
-    upsert-missing-only (preserves in-workspace customisations).
+    upsert-missing-only (preserves in-workspace customisations)
+iam_timeout : int (default 10)
+    Timeout in seconds for the IAM create-workspace request.
 
 Raises (in ``run``)
 -------------------
@@ -41,7 +43,9 @@ class WorkspaceInit(Initialiser):
             source="template",
             seed_file=None,
             overwrite=False,
+            iam_timeout=10,
             **kwargs,
+            
     ):
         super().__init__(**kwargs)
 
@@ -59,6 +63,7 @@ class WorkspaceInit(Initialiser):
         self.source = source
         self.seed_file = seed_file
         self.overwrite = overwrite
+        self.iam_timeout = iam_timeout
 
     async def run(self, ctx, old_flag, new_flag):
         await self._create_workspace(ctx)
@@ -120,10 +125,10 @@ class WorkspaceInit(Initialiser):
                     workspace_record=WorkspaceInput(
                         id=self.workspace,
                         name=self.workspace.title(),
-                        enabled=True,
+                        enabled=True,  
                     ),
                 ),
-                timeout=10,
+                timeout=self.iam_timeout,
             )
             if resp.error:
                 if resp.error.type == "duplicate":

trustgraph_flow-2.6.4/trustgraph/flow_version.py

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

{trustgraph_flow-2.6.2 → trustgraph_flow-2.6.4}/trustgraph/gateway/dispatch/manager.py

@@ -37,6 +37,7 @@ from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
 from . document_embeddings_query import DocumentEmbeddingsQueryRequestor
 from . row_embeddings_query import RowEmbeddingsQueryRequestor
 from . mcp_tool import McpToolRequestor
+from . reranker import RerankerRequestor
 from . text_load import TextLoad
 from . document_load import DocumentLoad
 
@@ -74,6 +75,7 @@ request_response_dispatchers = {
     "structured-diag": StructuredDiagRequestor,
     "row-embeddings": RowEmbeddingsQueryRequestor,
     "sparql": SparqlQueryRequestor,
+    "reranker": RerankerRequestor,
 }
 
 system_dispatchers = {

trustgraph_flow-2.6.4/trustgraph/gateway/dispatch/reranker.py

@@ -0,0 +1,31 @@
+
+from ... schema import RerankerRequest, RerankerResponse
+from ... messaging import TranslatorRegistry
+
+from . requestor import ServiceRequestor
+
+class RerankerRequestor(ServiceRequestor):
+    def __init__(
+            self, backend, request_queue, response_queue, timeout,
+            consumer, subscriber,
+    ):
+
+        super(RerankerRequestor, self).__init__(
+            backend=backend,
+            request_queue=request_queue,
+            response_queue=response_queue,
+            request_schema=RerankerRequest,
+            response_schema=RerankerResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
+            timeout=timeout,
+        )
+
+        self.request_translator = TranslatorRegistry.get_request_translator("reranker")
+        self.response_translator = TranslatorRegistry.get_response_translator("reranker")
+
+    def to_request(self, body):
+        return self.request_translator.decode(body)
+
+    def from_response(self, message):
+        return self.response_translator.encode_with_completion(message)

{trustgraph_flow-2.6.2 → trustgraph_flow-2.6.4}/trustgraph/gateway/registry.py

@@ -311,7 +311,7 @@ register(Operation(
 ))
 register(Operation(
     name="list-my-workspaces",
-    capability="workspaces:list-own",
+    capability=AUTHENTICATED,
     resource_level=ResourceLevel.SYSTEM,
     extract_resource=_empty_resource,
     extract_parameters=_no_parameters,
@@ -518,6 +518,7 @@ _FLOW_SERVICES = {
     "structured-diag": "structured-query:read",
     "row-embeddings": "row-embeddings:read",
     "sparql": "sparql:read",
+    "reranker": "reranker",
 }
 for _kind, _cap in _FLOW_SERVICES.items():
     _register_flow_kind("flow-service", _kind, _cap)

{trustgraph_flow-2.6.2 → trustgraph_flow-2.6.4}/trustgraph/iam/service/iam.py

@@ -72,13 +72,13 @@ _READER_CAPS = {
     "row-embeddings:read",
     "llm",
     "embeddings",
+    "reranker",
     "mcp",
     "config:read",
     "flows:read",
     "collections:read",
     "knowledge:read",
     "keys:self",
-    "workspaces:list-own",
 }
 
 _WRITER_CAPS = _READER_CAPS | {

trustgraph_flow-2.6.4/trustgraph/reranker/flashrank/__init__.py

@@ -0,0 +1,2 @@
+
+from . processor import *

trustgraph_flow-2.6.4/trustgraph/reranker/flashrank/__main__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from . processor import run
+
+if __name__ == '__main__':
+    run()

trustgraph_flow-2.6.4/trustgraph/reranker/flashrank/processor.py

@@ -0,0 +1,109 @@
+
+"""
+Reranker service using flashrank.
+Scores query-document pairs and returns the top results ranked by
+relevance.
+"""
+
+import asyncio
+import logging
+
+from ... base import RerankerService
+from ... schema import RerankerResult
+
+from flashrank import Ranker, RerankRequest
+
+logger = logging.getLogger(__name__)
+
+default_ident = "reranker"
+
+default_model = "ms-marco-MiniLM-L-12-v2"
+
+class Processor(RerankerService):
+
+    def __init__(self, **params):
+
+        model = params.get("model", default_model)
+
+        super(Processor, self).__init__(
+            **params | { "model": model }
+        )
+
+        self.default_model = model
+
+        self.cached_model_name = None
+        self.ranker = None
+
+        self._load_model(model)
+
+    def _load_model(self, model_name):
+        if self.cached_model_name != model_name:
+            logger.info(f"Loading flashrank model: {model_name}")
+            self.ranker = Ranker(model_name=model_name)
+            self.cached_model_name = model_name
+            logger.info(f"flashrank model {model_name} loaded successfully")
+        else:
+            logger.debug(f"Using cached model: {model_name}")
+
+    def _run_rerank(self, query, passages):
+        request = RerankRequest(query=query, passages=passages)
+        return self.ranker.rerank(request)
+
+    async def on_rerank(self, queries, documents, limit, model=None):
+
+        if not queries or not documents:
+            return []
+
+        use_model = model or self.default_model
+
+        if self.cached_model_name != use_model:
+            await asyncio.to_thread(self._load_model, use_model)
+
+        passages = [
+            {"id": d.document_id, "text": d.document_text}
+            for d in documents
+        ]
+
+        best_scores = {}
+
+        for q in queries:
+            ranked = await asyncio.to_thread(
+                self._run_rerank, q.query_text, passages,
+            )
+
+            for r in ranked:
+                doc_id = r["id"]
+                score = r["score"]
+                score = float(score)
+                if doc_id not in best_scores or score > best_scores[doc_id][1]:
+                    best_scores[doc_id] = (q.query_id, score)
+
+        results = sorted(
+            best_scores.items(),
+            key=lambda x: x[1][1],
+            reverse=True,
+        )[:limit]
+
+        return [
+            RerankerResult(
+                document_id=doc_id,
+                query_id=query_id,
+                score=score,
+            )
+            for doc_id, (query_id, score) in results
+        ]
+
+    @staticmethod
+    def add_args(parser):
+
+        RerankerService.add_args(parser)
+
+        parser.add_argument(
+            '-m', '--model',
+            default=default_model,
+            help=f'Reranker model (default: {default_model})'
+        )
+
+def run():
+
+    Processor.launch(default_ident, __doc__)