trustgraph-flow 2.8.0__tar.gz → 2.8.2__tar.gz

{trustgraph_flow-2.8.0 → trustgraph_flow-2.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 2.8.0
+Version: 2.8.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.8.0 → trustgraph_flow-2.8.2}/pyproject.toml

@@ -91,7 +91,6 @@ kg-extract-agent = "trustgraph.extract.kg.agent:run"
 kg-extract-definitions = "trustgraph.extract.kg.definitions:run"
 kg-extract-rows = "trustgraph.extract.kg.rows:run"
 kg-extract-relationships = "trustgraph.extract.kg.relationships:run"
-kg-extract-topics = "trustgraph.extract.kg.topics:run"
 kg-extract-ontology = "trustgraph.extract.kg.ontology:run"
 kg-manager = "trustgraph.cores:run"
 kg-store = "trustgraph.storage.knowledge:run"

{trustgraph_flow-2.8.0 → trustgraph_flow-2.8.2}/trustgraph/bootstrap/bootstrapper/service.py

@@ -169,6 +169,7 @@ class Processor(AsyncProcessor):
             request_topic=config_request_queue,
             response_topic=config_response_queue,
             subscription=f"{self.id}--config--{uuid.uuid4()}",
+            timeout=self.config_timeout,
         )
 
     def _make_flow_client(self, workspace):

{trustgraph_flow-2.8.0 → trustgraph_flow-2.8.2}/trustgraph/cores/knowledge.py

@@ -5,7 +5,6 @@ from .. schema import LibrarianRequest, DocumentMetadata
 from .. knowledge import hash
 from .. exceptions import RequestError
 from .. tables.knowledge import KnowledgeTableStore
-from .. base import Publisher
 
 import base64
 import asyncio
@@ -491,32 +490,25 @@ class KnowledgeManager:
             )
         )
 
-        t_pub = None
-        ge_pub = None
+        t_prod = None
+        ge_prod = None
 
         try:
 
             logger.debug(f"Triples queue: {t_q}")
             logger.debug(f"Graph embeddings queue: {ge_q}")
 
-            t_pub = Publisher(
-                self.flow_config.pubsub, t_q,
-                schema=Triples,
+            t_prod = await self.flow_config.pubsub.create_producer(
+                topic=t_q, schema=Triples,
             )
-            ge_pub = Publisher(
-                self.flow_config.pubsub, ge_q,
-                schema=GraphEmbeddings
+            ge_prod = await self.flow_config.pubsub.create_producer(
+                topic=ge_q, schema=GraphEmbeddings,
             )
 
-            logger.debug("Starting publishers...")
-
-            await t_pub.start()
-            await ge_pub.start()
-
             async def publish_triples(t):
                 if hasattr(t, 'metadata') and hasattr(t.metadata, 'collection'):
                     t.metadata.collection = request.collection or "default"
-                await t_pub.send(None, t)
+                await t_prod.send(t)
 
             logger.debug("Publishing triples...")
 
@@ -529,7 +521,7 @@ class KnowledgeManager:
             async def publish_ge(g):
                 if hasattr(g, 'metadata') and hasattr(g.metadata, 'collection'):
                     g.metadata.collection = request.collection or "default"
-                await ge_pub.send(None, g)
+                await ge_prod.send(g)
 
             logger.debug("Publishing graph embeddings...")
 
@@ -547,10 +539,8 @@ class KnowledgeManager:
 
         finally:
 
-            logger.debug("Stopping publishers...")
-
-            if t_pub: await t_pub.stop()
-            if ge_pub: await ge_pub.stop()
+            if t_prod: await t_prod.close()
+            if ge_prod: await ge_prod.close()
 
     async def _load_de_core(self, request, respond, workspace, flow):
 
@@ -569,25 +559,20 @@ class KnowledgeManager:
             )
         )
 
-        de_pub = None
+        de_prod = None
 
         try:
 
             logger.debug(f"Document embeddings queue: {de_q}")
 
-            de_pub = Publisher(
-                self.flow_config.pubsub, de_q,
-                schema=DocumentEmbeddings,
+            de_prod = await self.flow_config.pubsub.create_producer(
+                topic=de_q, schema=DocumentEmbeddings,
             )
 
-            logger.debug("Starting publisher...")
-
-            await de_pub.start()
-
             async def publish_de(de):
                 if hasattr(de, 'metadata') and hasattr(de.metadata, 'collection'):
                     de.metadata.collection = request.collection or "default"
-                await de_pub.send(None, de)
+                await de_prod.send(de)
 
             logger.debug("Publishing document embeddings...")
 
@@ -605,6 +590,4 @@ class KnowledgeManager:
 
         finally:
 
-            logger.debug("Stopping publisher...")
-
-            if de_pub: await de_pub.stop()
+            if de_prod: await de_prod.close()

{trustgraph_flow-2.8.0 → trustgraph_flow-2.8.2}/trustgraph/flow/service/service.py

@@ -164,6 +164,7 @@ class Processor(WorkspaceProcessor):
             response_topic=config_response_queue,
             request_schema=ConfigRequest,
             response_schema=ConfigResponse,
+            default_timeout=self.config_timeout,
         )
 
         self.config_client = AsyncConfigClient(rr_client)

trustgraph_flow-2.8.2/trustgraph/flow_version.py

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

{trustgraph_flow-2.8.0 → trustgraph_flow-2.8.2}/trustgraph/gateway/auth.py

@@ -28,7 +28,8 @@ from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import ed25519
 
 from ..base.iam_client import IamClient
-from ..base.metrics import ProducerMetrics, SubscriberMetrics
+from ..base.request_response_client import RequestResponseClient
+from ..base.request_response_spec import _make_impl_wrapper
 from ..schema import (
     IamRequest, IamResponse,
     iam_request_queue, iam_response_queue,
@@ -155,33 +156,24 @@ class IamAuth:
     # ghosts from prior calls.
     # ------------------------------------------------------------------
 
-    def _make_client(self):
-        rr_id = str(uuid.uuid4())
-        return IamClient(
+    async def _make_client(self):
+        rr_client = await RequestResponseClient.create(
             backend=self.backend,
-            subscription=f"{self.id}--iam--{rr_id}",
-            consumer_name=self.id,
             request_topic=iam_request_queue,
-            request_schema=IamRequest,
-            request_metrics=ProducerMetrics(
-                processor=self.id, producer="iam-request",
-            ),
             response_topic=iam_response_queue,
+            request_schema=IamRequest,
             response_schema=IamResponse,
-            response_metrics=SubscriberMetrics(
-                processor=self.id, subscriber="iam-response",
-            ),
         )
+        return _make_impl_wrapper(rr_client, IamClient)
 
     async def _with_client(self, op):
         """Open a short-lived IamClient, run ``op(client)``, close."""
-        client = self._make_client()
-        await client.start()
+        client = await self._make_client()
         try:
             return await op(client)
         finally:
             try:
-                await client.stop()
+                await client.close()
             except Exception:
                 pass
 

{trustgraph_flow-2.8.0 → trustgraph_flow-2.8.2}/trustgraph/gateway/config/receiver.py

@@ -13,10 +13,7 @@ import json
 from ... schema import ConfigPush, ConfigRequest, ConfigResponse
 from ... schema import config_push_queue, config_request_queue
 from ... schema import config_response_queue
-from ... base import Consumer, Producer
-from ... base.subscriber import Subscriber
-from ... base.request_response_spec import RequestResponse
-from ... base.metrics import ProducerMetrics, SubscriberMetrics
+from ... base.request_response_client import RequestResponseClient
 
 logger = logging.getLogger("config.receiver")
 logger.setLevel(logging.INFO)
@@ -39,7 +36,7 @@ class ConfigReceiver:
     def add_handler(self, h):
         self.flow_handlers.append(h)
 
-    async def on_config_notify(self, msg, proc, flow):
+    async def on_config_notify(self, msg):
 
         try:
 
@@ -47,7 +44,6 @@ class ConfigReceiver:
             notify_version = v.version
             changes = v.changes
 
-            # Skip if we already have this version or newer
             if notify_version <= self.config_version:
                 logger.debug(
                     f"Ignoring config notify v{notify_version}, "
@@ -55,7 +51,6 @@ class ConfigReceiver:
                 )
                 return
 
-            # Track workspace lifecycle
             if v.workspace_changes and self.auth:
                 for ws in (v.workspace_changes.created or []):
                     self.auth.known_workspaces.add(ws)
@@ -64,8 +59,6 @@ class ConfigReceiver:
                     self.auth.known_workspaces.discard(ws)
                     logger.info(f"Workspace deregistered: {ws}")
 
-            # Gateway cares about flow config — check if any flow
-            # types changed in any workspace
             flow_workspaces = changes.get("flow", [])
             if changes and not flow_workspaces:
                 logger.debug(
@@ -80,7 +73,6 @@ class ConfigReceiver:
                 f"types={list(changes.keys())}, fetching config..."
             )
 
-            # Refresh config for each affected workspace
             for workspace in flow_workspaces:
                 await self.fetch_and_apply_workspace(workspace)
 
@@ -91,32 +83,16 @@ class ConfigReceiver:
                 f"Config notify processing exception: {e}", exc_info=True
             )
 
-    def _create_config_client(self):
-        """Create a short-lived config request/response client."""
-        id = str(uuid.uuid4())
-
-        config_req_metrics = ProducerMetrics(
-            processor="api-gateway", producer="config-request",
-        )
-        config_resp_metrics = SubscriberMetrics(
-            processor="api-gateway", subscriber="config-response",
-        )
-
-        return RequestResponse(
+    async def _create_config_client(self):
+        return await RequestResponseClient.create(
             backend=self.backend,
-            subscription=f"api-gateway--config--{id}",
-            consumer_name="api-gateway",
             request_topic=config_request_queue,
-            request_schema=ConfigRequest,
-            request_metrics=config_req_metrics,
             response_topic=config_response_queue,
+            request_schema=ConfigRequest,
             response_schema=ConfigResponse,
-            response_metrics=config_resp_metrics,
         )
 
     async def fetch_and_apply_workspace(self, workspace, retry=False):
-        """Fetch config for a single workspace and apply flow changes.
-        If retry=True, keeps retrying until successful."""
 
         while True:
 
@@ -125,9 +101,8 @@ class ConfigReceiver:
                     f"Fetching config for workspace {workspace}..."
                 )
 
-                client = self._create_config_client()
+                client = await self._create_config_client()
                 try:
-                    await client.start()
                     resp = await client.request(
                         ConfigRequest(
                             operation="config",
@@ -136,7 +111,7 @@ class ConfigReceiver:
                         timeout=10,
                     )
                 finally:
-                    await client.stop()
+                    await client.close()
 
                 logger.info(f"Config response received")
 
@@ -190,20 +165,15 @@ class ConfigReceiver:
                 return
 
     async def fetch_all_workspaces(self, retry=False):
-        """Fetch config for all workspaces at startup.
-        Discovers workspaces via the config service getvalues-all-ws
-        operation on the flow type."""
 
         while True:
 
             try:
                 logger.info("Discovering workspaces with flows...")
 
-                client = self._create_config_client()
+                client = await self._create_config_client()
                 try:
-                    await client.start()
 
-                    # Discover all known workspaces
                     ws_resp = await client.request(
                         ConfigRequest(
                             operation="getvalues",
@@ -230,7 +200,6 @@ class ConfigReceiver:
                         f"Known workspaces: {discovered}"
                     )
 
-                    # Discover workspaces that have any flow config
                     resp = await client.request(
                         ConfigRequest(
                             operation="getvalues-all-ws",
@@ -248,9 +217,6 @@ class ConfigReceiver:
                         v.workspace for v in resp.values if v.workspace
                     }
 
-                    # Always include the default workspace, even if
-                    # empty, so that newly-created flows in it can be
-                    # picked up by subsequent notifications.
                     workspaces.add("default")
 
                     logger.info(
@@ -258,9 +224,8 @@ class ConfigReceiver:
                     )
 
                 finally:
-                    await client.stop()
+                    await client.close()
 
-                # Fetch and apply config for each workspace
                 for workspace in workspaces:
                     await self.fetch_and_apply_workspace(
                         workspace, retry=retry
@@ -310,37 +275,35 @@ class ConfigReceiver:
 
         while True:
 
-            try:
-
-                async with asyncio.TaskGroup() as tg:
+            consumer = None
 
-                    id = str(uuid.uuid4())
+            try:
 
-                    # Subscribe to notify queue
-                    self.config_cons = Consumer(
-                        taskgroup=tg,
-                        flow=None,
-                        backend=self.backend,
-                        subscriber=f"gateway-{id}",
-                        topic=config_push_queue,
-                        schema=ConfigPush,
-                        handler=self.on_config_notify,
-                        start_of_messages=False,
-                    )
+                consumer = await self.backend.create_consumer(
+                    topic=config_push_queue,
+                    subscription=f"gateway-{uuid.uuid4()}",
+                    schema=ConfigPush,
+                    initial_position='latest',
+                )
 
-                    logger.info("Starting config notify consumer...")
-                    await self.config_cons.start()
-                    logger.info("Config notify consumer started")
+                logger.info("Config notify consumer started")
 
-                    # Fetch current config (subscribe-then-fetch pattern)
-                    # Retry until config service is available
-                    await self.fetch_all_workspaces(retry=True)
+                await self.fetch_all_workspaces(retry=True)
 
-                    logger.info(
-                        "Config loader initialised, waiting for notifys..."
-                    )
+                logger.info(
+                    "Config loader initialised, waiting for notifys..."
+                )
 
-                logger.warning("Config consumer exited, restarting...")
+                while True:
+                    msg = await consumer.receive()
+                    try:
+                        await self.on_config_notify(msg)
+                        await consumer.acknowledge(msg)
+                    except Exception as e:
+                        logger.error(
+                            f"Config notify error: {e}", exc_info=True
+                        )
+                        await consumer.negative_acknowledge(msg)
 
             except Exception as e:
                 logger.error(
@@ -348,6 +311,13 @@ class ConfigReceiver:
                     exc_info=True
                 )
 
+            finally:
+                if consumer:
+                    try:
+                        await consumer.close()
+                    except BaseException:
+                        pass
+
             await asyncio.sleep(4)
 
     async def start(self):

trustgraph_flow-2.8.2/trustgraph/gateway/dispatch/document_embeddings_export.py

@@ -0,0 +1,63 @@
+
+import asyncio
+import logging
+
+from ... schema import DocumentEmbeddings
+
+from . serialize import serialize_document_embeddings
+
+logger = logging.getLogger(__name__)
+
+class DocumentEmbeddingsExport:
+
+    def __init__(self, ws, running, backend, queue, consumer, subscriber):
+        self.ws = ws
+        self.running = running
+        self.backend = backend
+        self.queue = queue
+        self.subscription = subscriber
+        self._consumer = None
+
+    async def destroy(self):
+        self.running.stop()
+        if self._consumer:
+            try:
+                await self._consumer.close()
+            except BaseException:
+                pass
+            self._consumer = None
+        if self.ws and not self.ws.closed:
+            await self.ws.close()
+
+    async def receive(self, msg):
+        pass
+
+    async def run(self):
+        self._consumer = await self.backend.create_consumer(
+            topic=self.queue,
+            subscription=self.subscription,
+            schema=DocumentEmbeddings,
+        )
+
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+
+        while self.running.get():
+            try:
+                msg = await asyncio.wait_for(
+                    self._consumer.receive(), timeout=0.5,
+                )
+                await self.ws.send_json(serialize_document_embeddings(msg.value()))
+                await self._consumer.acknowledge(msg)
+                consecutive_errors = 0
+
+            except asyncio.TimeoutError:
+                continue
+
+            except Exception as e:
+                logger.error(f"Exception sending to websocket: {e}")
+                consecutive_errors += 1
+                if consecutive_errors >= max_consecutive_errors:
+                    logger.error("Too many consecutive errors, shutting down")
+                    break
+                await asyncio.sleep(0.1)

{trustgraph_flow-2.8.0 → trustgraph_flow-2.8.2}/trustgraph/gateway/dispatch/document_embeddings_import.py

@@ -1,59 +1,44 @@
 
 import asyncio
-import uuid
 import logging
-from aiohttp import WSMsgType
 
-from ... schema import Metadata
-from ... schema import DocumentEmbeddings, ChunkEmbeddings
-from ... base import Publisher
+from ... schema import DocumentEmbeddings
 from ... messaging.translators.document_loading import DocumentEmbeddingsTranslator
 
-# Module logger
 logger = logging.getLogger(__name__)
 
 class DocumentEmbeddingsImport:
 
-    def __init__(
-            self, ws, running, backend, queue
-    ):
-
+    def __init__(self, ws, running, backend, queue):
         self.ws = ws
         self.running = running
+        self.backend = backend
+        self.queue = queue
         self.translator = DocumentEmbeddingsTranslator()
-        
-        self.publisher = Publisher(
-            backend, topic = queue, schema = DocumentEmbeddings
-        )
+        self.producer = None
 
     async def start(self):
-        await self.publisher.start()
+        self.producer = await self.backend.create_producer(
+            topic=self.queue, schema=DocumentEmbeddings,
+        )
 
     async def destroy(self):
-        # Step 1: Stop accepting new messages
         self.running.stop()
-
-        # Step 2: Wait for publisher to drain its queue
-        logger.info("Draining publisher queue...")
-        await self.publisher.stop()
-
-        # Step 3: Close websocket only after queue is drained
+        if self.producer:
+            await self.producer.close()
+            self.producer = None
         if self.ws:
             await self.ws.close()
 
     async def receive(self, msg):
-
         data = msg.json()
         elt = self.translator.decode(data)
-        await self.publisher.send(None, elt)
+        await self.producer.send(elt)
 
     async def run(self):
-
         while self.running.get():
             await asyncio.sleep(0.5)
 
         if self.ws:
             await self.ws.close()
-
         self.ws = None
-

trustgraph_flow-2.8.2/trustgraph/gateway/dispatch/entity_contexts_export.py

@@ -0,0 +1,63 @@
+
+import asyncio
+import logging
+
+from ... schema import EntityContexts
+
+from . serialize import serialize_entity_contexts
+
+logger = logging.getLogger(__name__)
+
+class EntityContextsExport:
+
+    def __init__(self, ws, running, backend, queue, consumer, subscriber):
+        self.ws = ws
+        self.running = running
+        self.backend = backend
+        self.queue = queue
+        self.subscription = subscriber
+        self._consumer = None
+
+    async def destroy(self):
+        self.running.stop()
+        if self._consumer:
+            try:
+                await self._consumer.close()
+            except BaseException:
+                pass
+            self._consumer = None
+        if self.ws and not self.ws.closed:
+            await self.ws.close()
+
+    async def receive(self, msg):
+        pass
+
+    async def run(self):
+        self._consumer = await self.backend.create_consumer(
+            topic=self.queue,
+            subscription=self.subscription,
+            schema=EntityContexts,
+        )
+
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+
+        while self.running.get():
+            try:
+                msg = await asyncio.wait_for(
+                    self._consumer.receive(), timeout=0.5,
+                )
+                await self.ws.send_json(serialize_entity_contexts(msg.value()))
+                await self._consumer.acknowledge(msg)
+                consecutive_errors = 0
+
+            except asyncio.TimeoutError:
+                continue
+
+            except Exception as e:
+                logger.error(f"Exception sending to websocket: {e}")
+                consecutive_errors += 1
+                if consecutive_errors >= max_consecutive_errors:
+                    logger.error("Too many consecutive errors, shutting down")
+                    break
+                await asyncio.sleep(0.1)