trustgraph-flow 1.2.20__tar.gz → 1.3.1__tar.gz

{trustgraph_flow-1.2.20 → trustgraph_flow-1.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trustgraph-flow
-Version: 1.2.20
+Version: 1.3.1
 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
@@ -8,7 +8,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Requires-Dist: trustgraph-base<1.3,>=1.2
+Requires-Dist: trustgraph-base<1.4,>=1.3
 Requires-Dist: aiohttp
 Requires-Dist: anthropic
 Requires-Dist: cassandra-driver

{trustgraph_flow-1.2.20 → trustgraph_flow-1.3.1}/pyproject.toml

@@ -10,7 +10,7 @@ description = "TrustGraph provides a means to run a pipeline of flexible AI proc
 readme = "README.md"
 requires-python = ">=3.8"
 dependencies = [
-    "trustgraph-base>=1.2,<1.3",
+    "trustgraph-base>=1.3,<1.4",
     "aiohttp",
     "anthropic",
     "cassandra-driver",

trustgraph_flow-1.3.1/trustgraph/flow_version.py

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

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

@@ -0,0 +1,91 @@
+
+import asyncio
+import queue
+import uuid
+import logging
+
+from ... schema import DocumentEmbeddings
+from ... base import Subscriber
+
+from . serialize import serialize_document_embeddings
+
+# Module logger
+logger = logging.getLogger(__name__)
+
+class DocumentEmbeddingsExport:
+
+    def __init__(
+            self, ws, running, pulsar_client, queue, consumer, subscriber
+    ):
+
+        self.ws = ws
+        self.running = running
+        self.pulsar_client = pulsar_client
+        self.queue = queue
+        self.consumer = consumer
+        self.subscriber = subscriber
+
+    async def destroy(self):
+        # Step 1: Signal stop to prevent new messages
+        self.running.stop()
+        
+        # Step 2: Wait briefly for in-flight messages
+        await asyncio.sleep(0.5)
+        
+        # Step 3: Unsubscribe and stop subscriber (triggers queue drain)
+        if hasattr(self, 'subs'):
+            await self.subs.unsubscribe_all(self.id)
+            await self.subs.stop()
+        
+        # Step 4: Close websocket last
+        if self.ws and not self.ws.closed:
+            await self.ws.close()
+
+    async def receive(self, msg):
+        # Ignore incoming info from websocket
+        pass
+
+    async def run(self):
+        """Enhanced run with better error handling"""
+        self.subs = Subscriber(
+            client = self.pulsar_client, 
+            topic = self.queue,
+            consumer_name = self.consumer, 
+            subscription = self.subscriber,
+            schema = DocumentEmbeddings,
+            backpressure_strategy = "block"  # Configurable
+        )
+        
+        await self.subs.start()
+        
+        self.id = str(uuid.uuid4())
+        q = await self.subs.subscribe_all(self.id)
+        
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+        
+        while self.running.get():
+            try:
+                resp = await asyncio.wait_for(q.get(), timeout=0.5)
+                await self.ws.send_json(serialize_document_embeddings(resp))
+                consecutive_errors = 0  # Reset on success
+                
+            except asyncio.TimeoutError:
+                continue
+                
+            except queue.Empty:
+                continue
+                
+            except Exception as e:
+                logger.error(f"Exception sending to websocket: {str(e)}")
+                consecutive_errors += 1
+                
+                if consecutive_errors >= max_consecutive_errors:
+                    logger.error("Too many consecutive errors, shutting down")
+                    break
+                    
+                # Brief pause before retry
+                await asyncio.sleep(0.1)
+        
+        # Graceful cleanup handled in destroy()
+

{trustgraph_flow-1.2.20 → trustgraph_flow-1.3.1}/trustgraph/gateway/dispatch/document_embeddings_import.py

@@ -1,6 +1,7 @@
 
 import asyncio
 import uuid
+import logging
 from aiohttp import WSMsgType
 
 from ... schema import Metadata
@@ -8,6 +9,9 @@ from ... schema import DocumentEmbeddings, ChunkEmbeddings
 from ... base import Publisher
 from ... messaging.translators.document_loading import DocumentEmbeddingsTranslator
 
+# Module logger
+logger = logging.getLogger(__name__)
+
 class DocumentEmbeddingsImport:
 
     def __init__(
@@ -26,13 +30,17 @@ class DocumentEmbeddingsImport:
         await self.publisher.start()
 
     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.ws:
             await self.ws.close()
 
-        await self.publisher.stop()
-
     async def receive(self, msg):
 
         data = msg.json()

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

@@ -0,0 +1,91 @@
+
+import asyncio
+import queue
+import uuid
+import logging
+
+from ... schema import EntityContexts
+from ... base import Subscriber
+
+from . serialize import serialize_entity_contexts
+
+# Module logger
+logger = logging.getLogger(__name__)
+
+class EntityContextsExport:
+
+    def __init__(
+            self, ws, running, pulsar_client, queue, consumer, subscriber
+    ):
+
+        self.ws = ws
+        self.running = running
+        self.pulsar_client = pulsar_client
+        self.queue = queue
+        self.consumer = consumer
+        self.subscriber = subscriber
+
+    async def destroy(self):
+        # Step 1: Signal stop to prevent new messages
+        self.running.stop()
+        
+        # Step 2: Wait briefly for in-flight messages
+        await asyncio.sleep(0.5)
+        
+        # Step 3: Unsubscribe and stop subscriber (triggers queue drain)
+        if hasattr(self, 'subs'):
+            await self.subs.unsubscribe_all(self.id)
+            await self.subs.stop()
+        
+        # Step 4: Close websocket last
+        if self.ws and not self.ws.closed:
+            await self.ws.close()
+
+    async def receive(self, msg):
+        # Ignore incoming info from websocket
+        pass
+
+    async def run(self):
+        """Enhanced run with better error handling"""
+        self.subs = Subscriber(
+            client = self.pulsar_client, 
+            topic = self.queue,
+            consumer_name = self.consumer, 
+            subscription = self.subscriber,
+            schema = EntityContexts,
+            backpressure_strategy = "block"  # Configurable
+        )
+        
+        await self.subs.start()
+        
+        self.id = str(uuid.uuid4())
+        q = await self.subs.subscribe_all(self.id)
+        
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+        
+        while self.running.get():
+            try:
+                resp = await asyncio.wait_for(q.get(), timeout=0.5)
+                await self.ws.send_json(serialize_entity_contexts(resp))
+                consecutive_errors = 0  # Reset on success
+                
+            except asyncio.TimeoutError:
+                continue
+                
+            except queue.Empty:
+                continue
+                
+            except Exception as e:
+                logger.error(f"Exception sending to websocket: {str(e)}")
+                consecutive_errors += 1
+                
+                if consecutive_errors >= max_consecutive_errors:
+                    logger.error("Too many consecutive errors, shutting down")
+                    break
+                    
+                # Brief pause before retry
+                await asyncio.sleep(0.1)
+        
+        # Graceful cleanup handled in destroy()
+

{trustgraph_flow-1.2.20 → trustgraph_flow-1.3.1}/trustgraph/gateway/dispatch/entity_contexts_import.py

@@ -1,6 +1,7 @@
 
 import asyncio
 import uuid
+import logging
 from aiohttp import WSMsgType
 
 from ... schema import Metadata
@@ -9,6 +10,9 @@ from ... base import Publisher
 
 from . serialize import to_subgraph, to_value
 
+# Module logger
+logger = logging.getLogger(__name__)
+
 class EntityContextsImport:
 
     def __init__(
@@ -26,13 +30,17 @@ class EntityContextsImport:
         await self.publisher.start()
 
     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.ws:
             await self.ws.close()
 
-        await self.publisher.stop()
-
     async def receive(self, msg):
 
         data = msg.json()

trustgraph_flow-1.3.1/trustgraph/gateway/dispatch/graph_embeddings_export.py

@@ -0,0 +1,91 @@
+
+import asyncio
+import queue
+import uuid
+import logging
+
+from ... schema import GraphEmbeddings
+from ... base import Subscriber
+
+from . serialize import serialize_graph_embeddings
+
+# Module logger
+logger = logging.getLogger(__name__)
+
+class GraphEmbeddingsExport:
+
+    def __init__(
+            self, ws, running, pulsar_client, queue, consumer, subscriber
+    ):
+
+        self.ws = ws
+        self.running = running
+        self.pulsar_client = pulsar_client
+        self.queue = queue
+        self.consumer = consumer
+        self.subscriber = subscriber
+
+    async def destroy(self):
+        # Step 1: Signal stop to prevent new messages
+        self.running.stop()
+        
+        # Step 2: Wait briefly for in-flight messages
+        await asyncio.sleep(0.5)
+        
+        # Step 3: Unsubscribe and stop subscriber (triggers queue drain)
+        if hasattr(self, 'subs'):
+            await self.subs.unsubscribe_all(self.id)
+            await self.subs.stop()
+        
+        # Step 4: Close websocket last
+        if self.ws and not self.ws.closed:
+            await self.ws.close()
+
+    async def receive(self, msg):
+        # Ignore incoming info from websocket
+        pass
+
+    async def run(self):
+        """Enhanced run with better error handling"""
+        self.subs = Subscriber(
+            client = self.pulsar_client, 
+            topic = self.queue,
+            consumer_name = self.consumer, 
+            subscription = self.subscriber,
+            schema = GraphEmbeddings,
+            backpressure_strategy = "block"  # Configurable
+        )
+        
+        await self.subs.start()
+        
+        self.id = str(uuid.uuid4())
+        q = await self.subs.subscribe_all(self.id)
+        
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+        
+        while self.running.get():
+            try:
+                resp = await asyncio.wait_for(q.get(), timeout=0.5)
+                await self.ws.send_json(serialize_graph_embeddings(resp))
+                consecutive_errors = 0  # Reset on success
+                
+            except asyncio.TimeoutError:
+                continue
+                
+            except queue.Empty:
+                continue
+                
+            except Exception as e:
+                logger.error(f"Exception sending to websocket: {str(e)}")
+                consecutive_errors += 1
+                
+                if consecutive_errors >= max_consecutive_errors:
+                    logger.error("Too many consecutive errors, shutting down")
+                    break
+                    
+                # Brief pause before retry
+                await asyncio.sleep(0.1)
+        
+        # Graceful cleanup handled in destroy()
+

{trustgraph_flow-1.2.20 → trustgraph_flow-1.3.1}/trustgraph/gateway/dispatch/graph_embeddings_import.py

@@ -1,6 +1,7 @@
 
 import asyncio
 import uuid
+import logging
 from aiohttp import WSMsgType
 
 from ... schema import Metadata
@@ -9,6 +10,9 @@ from ... base import Publisher
 
 from . serialize import to_subgraph, to_value
 
+# Module logger
+logger = logging.getLogger(__name__)
+
 class GraphEmbeddingsImport:
 
     def __init__(
@@ -26,13 +30,17 @@ class GraphEmbeddingsImport:
         await self.publisher.start()
 
     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.ws:
             await self.ws.close()
 
-        await self.publisher.stop()
-
     async def receive(self, msg):
 
         data = msg.json()

trustgraph_flow-1.3.1/trustgraph/gateway/dispatch/triples_export.py

@@ -0,0 +1,91 @@
+
+import asyncio
+import queue
+import uuid
+import logging
+
+from ... schema import Triples
+from ... base import Subscriber
+
+from . serialize import serialize_triples
+
+# Module logger
+logger = logging.getLogger(__name__)
+
+class TriplesExport:
+
+    def __init__(
+            self, ws, running, pulsar_client, queue, consumer, subscriber
+    ):
+
+        self.ws = ws
+        self.running = running
+        self.pulsar_client = pulsar_client
+        self.queue = queue
+        self.consumer = consumer
+        self.subscriber = subscriber
+
+    async def destroy(self):
+        # Step 1: Signal stop to prevent new messages
+        self.running.stop()
+        
+        # Step 2: Wait briefly for in-flight messages
+        await asyncio.sleep(0.5)
+        
+        # Step 3: Unsubscribe and stop subscriber (triggers queue drain)
+        if hasattr(self, 'subs'):
+            await self.subs.unsubscribe_all(self.id)
+            await self.subs.stop()
+        
+        # Step 4: Close websocket last
+        if self.ws and not self.ws.closed:
+            await self.ws.close()
+
+    async def receive(self, msg):
+        # Ignore incoming info from websocket
+        pass
+
+    async def run(self):
+        """Enhanced run with better error handling"""
+        self.subs = Subscriber(
+            client = self.pulsar_client, 
+            topic = self.queue,
+            consumer_name = self.consumer, 
+            subscription = self.subscriber,
+            schema = Triples,
+            backpressure_strategy = "block"  # Configurable
+        )
+        
+        await self.subs.start()
+        
+        self.id = str(uuid.uuid4())
+        q = await self.subs.subscribe_all(self.id)
+        
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+        
+        while self.running.get():
+            try:
+                resp = await asyncio.wait_for(q.get(), timeout=0.5)
+                await self.ws.send_json(serialize_triples(resp))
+                consecutive_errors = 0  # Reset on success
+                
+            except asyncio.TimeoutError:
+                continue
+                
+            except queue.Empty:
+                continue
+                
+            except Exception as e:
+                logger.error(f"Exception sending to websocket: {str(e)}")
+                consecutive_errors += 1
+                
+                if consecutive_errors >= max_consecutive_errors:
+                    logger.error("Too many consecutive errors, shutting down")
+                    break
+                    
+                # Brief pause before retry
+                await asyncio.sleep(0.1)
+        
+        # Graceful cleanup handled in destroy()
+

{trustgraph_flow-1.2.20 → trustgraph_flow-1.3.1}/trustgraph/gateway/dispatch/triples_import.py

@@ -1,6 +1,7 @@
 
 import asyncio
 import uuid
+import logging
 from aiohttp import WSMsgType
 
 from ... schema import Metadata
@@ -9,6 +10,9 @@ from ... base import Publisher
 
 from . serialize import to_subgraph
 
+# Module logger
+logger = logging.getLogger(__name__)
+
 class TriplesImport:
 
     def __init__(
@@ -26,13 +30,17 @@ class TriplesImport:
         await self.publisher.start()
 
     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.ws:
             await self.ws.close()
 
-        await self.publisher.stop()
-
     async def receive(self, msg):
 
         data = msg.json()

trustgraph_flow-1.3.1/trustgraph/gateway/endpoint/socket.py

@@ -0,0 +1,152 @@
+
+import asyncio
+from aiohttp import web, WSMsgType
+import logging
+
+from .. running import Running
+
+logger = logging.getLogger("socket")
+logger.setLevel(logging.INFO)
+
+class SocketEndpoint:
+
+    def __init__(
+            self, endpoint_path, auth, dispatcher,
+    ):
+
+        self.path = endpoint_path
+        self.auth = auth
+        self.operation = "socket"
+
+        self.dispatcher = dispatcher
+
+    async def worker(self, ws, dispatcher, running):
+
+        await dispatcher.run()
+
+    async def listener(self, ws, dispatcher, running):
+        """Enhanced listener with graceful shutdown"""
+        try:
+            async for msg in ws:
+                # On error, finish
+                if msg.type == WSMsgType.TEXT:
+                    await dispatcher.receive(msg)
+                    continue
+                elif msg.type == WSMsgType.BINARY:
+                    await dispatcher.receive(msg)
+                    continue
+                else:
+                    # Graceful shutdown on close
+                    logger.info("Websocket closing, initiating graceful shutdown")
+                    running.stop()
+                    
+                    # Allow time for dispatcher cleanup
+                    await asyncio.sleep(1.0)
+                    
+                    # Close websocket if not already closed
+                    if not ws.closed:
+                        await ws.close()
+                    break
+            else:
+                # This executes when the async for loop completes normally (no break)
+                logger.debug("Websocket iteration completed, performing cleanup")
+                running.stop()
+                if not ws.closed:
+                    await ws.close()
+        except Exception:
+            # Handle exceptions and cleanup
+            running.stop()
+            if not ws.closed:
+                await ws.close()
+            raise
+        
+    async def handle(self, request):
+        """Enhanced handler with better cleanup"""
+        try:
+            token = request.query['token']
+        except:
+            token = ""
+
+        if not self.auth.permitted(token, self.operation):
+            return web.HTTPUnauthorized()
+        
+        # 50MB max message size
+        ws = web.WebSocketResponse(max_msg_size=52428800)
+
+        await ws.prepare(request)
+        
+        dispatcher = None
+        
+        try:
+
+            async with asyncio.TaskGroup() as tg:
+
+                running = Running()
+
+                dispatcher = await self.dispatcher(
+                    ws, running, request.match_info
+                )
+
+                worker_task = tg.create_task(
+                    self.worker(ws, dispatcher, running)
+                )
+
+                lsnr_task = tg.create_task(
+                    self.listener(ws, dispatcher, running)
+                )
+
+                logger.debug("Created task group, waiting for completion...")
+
+                # Wait for threads to complete
+
+            logger.debug("Task group closed")
+
+        except ExceptionGroup as e:
+
+            logger.error("Exception group occurred:", exc_info=True)
+
+            for se in e.exceptions:
+                logger.error(f"  Exception type: {type(se)}")
+                logger.error(f"  Exception: {se}")
+                
+            # Attempt graceful dispatcher shutdown
+            if dispatcher and hasattr(dispatcher, 'destroy'):
+                try:
+                    await asyncio.wait_for(
+                        dispatcher.destroy(), 
+                        timeout=5.0
+                    )
+                except asyncio.TimeoutError:
+                    logger.warning("Dispatcher shutdown timed out")
+                except Exception as de:
+                    logger.error(f"Error during dispatcher cleanup: {de}")
+                    
+        except Exception as e:
+            logger.error(f"Socket exception: {e}", exc_info=True)
+            
+        finally:
+            # Ensure dispatcher cleanup
+            if dispatcher and hasattr(dispatcher, 'destroy'):
+                try:
+                    await dispatcher.destroy()
+                except Exception as de:
+                    logger.error(f"Error in final dispatcher cleanup: {de}")
+                    
+            # Ensure websocket is closed
+            if ws and not ws.closed:
+                await ws.close()
+                
+        return ws
+
+    async def start(self):
+        pass
+
+    async def stop(self):
+        self.running.stop()
+
+    def add_routes(self, app):
+
+        app.add_routes([
+            web.get(self.path, self.handle),
+        ])
+