trustgraph-flow 0.23.4__tar.gz → 0.23.6__tar.gz

{trustgraph-flow-0.23.4 → trustgraph-flow-0.23.6}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 0.23.4
+Version: 0.23.6
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Home-page: https://github.com/trustgraph-ai/trustgraph
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.23.4.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.23.6.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

trustgraph-flow-0.23.6/trustgraph/flow_version.py

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

trustgraph-flow-0.23.6/trustgraph/gateway/config/receiver.py

@@ -0,0 +1,121 @@
+"""
+API gateway.  Offers HTTP services which are translated to interaction on the
+Pulsar bus.
+"""
+
+module = "api-gateway"
+
+# FIXME: Subscribes to Pulsar unnecessarily, should only do it when there
+# are active listeners
+
+# FIXME: Connection errors in publishers / subscribers cause those threads
+# to fail and are not failed or retried
+
+import asyncio
+import argparse
+from aiohttp import web
+import logging
+import os
+import base64
+import uuid
+import json
+
+import pulsar
+from prometheus_client import start_http_server
+
+from ... schema import ConfigPush, config_push_queue
+from ... base import Consumer
+
+logger = logging.getLogger("config.receiver")
+logger.setLevel(logging.INFO)
+
+class ConfigReceiver:
+
+    def __init__(self, pulsar_client):
+
+        self.pulsar_client = pulsar_client
+
+        self.flow_handlers = []
+
+        self.flows = {}
+
+    def add_handler(self, h):
+        self.flow_handlers.append(h)
+
+    async def on_config(self, msg, proc, flow):
+
+        try:
+
+            v = msg.value()
+
+            print(f"Config version", v.version)
+
+            if "flows" in v.config:
+
+                flows = v.config["flows"]
+
+                wanted = list(flows.keys())
+                current = list(self.flows.keys())
+
+                for k in wanted:
+                    if k not in current:
+                        self.flows[k] = json.loads(flows[k])
+                        await self.start_flow(k, self.flows[k])
+
+                for k in current:
+                    if k not in wanted:
+                        await self.stop_flow(k, self.flows[k])
+                        del self.flows[k]
+
+        except Exception as e:
+            print(f"Exception: {e}", flush=True)
+
+    async def start_flow(self, id, flow):
+
+        print("Start flow", id)
+
+        for handler in self.flow_handlers:
+
+            try:
+                await handler.start_flow(id, flow)
+            except Exception as e:
+                print(f"Exception: {e}", flush=True)
+
+    async def stop_flow(self, id, flow):
+
+        print("Stop flow", id)
+
+        for handler in self.flow_handlers:
+
+            try:
+                await handler.stop_flow(id, flow)
+            except Exception as e:
+                print(f"Exception: {e}", flush=True)
+
+    async def config_loader(self):
+
+        async with asyncio.TaskGroup() as tg:
+
+            id = str(uuid.uuid4())
+
+            self.config_cons = Consumer(
+                taskgroup = tg,
+                flow = None,
+                client = self.pulsar_client,
+                subscriber = f"gateway-{id}",                
+                topic = config_push_queue,
+                schema = ConfigPush,
+                handler = self.on_config,
+                start_of_messages = True,
+            )
+
+            await self.config_cons.start()
+
+            print("Waiting...")
+
+        print("Config consumer done. :/")
+
+    async def start(self):
+        
+        asyncio.create_task(self.config_loader())
+

{trustgraph-flow-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/agent.py

@@ -1,12 +1,11 @@
 
-from .. schema import AgentRequest, AgentResponse
+from ... schema import AgentRequest, AgentResponse
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class AgentRequestor(ServiceRequestor):
     def __init__(
-            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            self, pulsar_client, request_queue, response_queue, timeout,
             consumer, subscriber,
     ):
 

{trustgraph-flow-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/config.py

@@ -1,16 +1,17 @@
 
-from .. schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
-from .. schema import config_request_queue
-from .. schema import config_response_queue
+from ... schema import ConfigRequest, ConfigResponse, ConfigKey, ConfigValue
+from ... schema import config_request_queue
+from ... schema import config_response_queue
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class ConfigRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
 
         super(ConfigRequestor, self).__init__(
             pulsar_client=pulsar_client,
+            consumer_name = consumer,
+            subscription = subscriber,
             request_queue=config_request_queue,
             response_queue=config_response_queue,
             request_schema=ConfigRequest,

trustgraph-flow-0.23.6/trustgraph/gateway/dispatch/document_embeddings_export.py

@@ -0,0 +1,67 @@
+
+import asyncio
+import queue
+import uuid
+
+from ... schema import DocumentEmbeddings
+from ... base import Subscriber
+
+from . serialize import serialize_document_embeddings
+
+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):
+        self.running.stop()
+        await self.ws.close()
+
+    async def receive(self, msg):
+        # Ignore incoming info from websocket
+        pass
+
+    async def run(self):
+
+        subs = Subscriber(
+            client = self.pulsar_client, topic = self.queue,
+            consumer_name = self.consumer, subscription = self.subscriber,
+            schema = DocumentEmbeddings
+        )
+
+        await subs.start()
+
+        id = str(uuid.uuid4())
+        q = await subs.subscribe_all(id)
+
+        while self.running.get():
+            try:
+
+                resp = await asyncio.wait_for(q.get(), timeout=0.5)
+                await self.ws.send_json(serialize_document_embeddings(resp))
+
+            except TimeoutError:
+                continue
+
+            except queue.Empty:
+                continue
+
+            except Exception as e:
+                print(f"Exception: {str(e)}", flush=True)
+                break
+
+        await subs.unsubscribe_all(id)
+
+        await subs.stop()
+
+        await self.ws.close()
+        self.running.stop()
+

trustgraph-flow-0.23.6/trustgraph/gateway/dispatch/document_embeddings_import.py

@@ -0,0 +1,64 @@
+
+import asyncio
+import uuid
+from aiohttp import WSMsgType
+
+from ... schema import Metadata
+from ... schema import DocumentEmbeddings, ChunkEmbeddings
+from ... base import Publisher
+
+from . serialize import to_subgraph
+
+class DocumentEmbeddingsImport:
+
+    def __init__(
+            self, ws, running, pulsar_client, queue
+    ):
+
+        self.ws = ws
+        self.running = running
+        
+        self.publisher = Publisher(
+            pulsar_client, topic = queue, schema = DocumentEmbeddings
+        )
+
+    async def destroy(self):
+        self.running.stop()
+
+        if self.ws:
+            await self.ws.close()
+
+        await self.publisher.stop()
+
+    async def receive(self, msg):
+
+        data = msg.json()
+
+        elt = DocumentEmbeddings(
+            metadata=Metadata(
+                id=data["metadata"]["id"],
+                metadata=to_subgraph(data["metadata"]["metadata"]),
+                user=data["metadata"]["user"],
+                collection=data["metadata"]["collection"],
+            ),
+            chunks=[
+                ChunkEmbeddings(
+                    chunk=de["chunk"].encode("utf-8"),
+                    vectors=de["vectors"],
+                )
+                for de in data["chunks"]
+            ],
+        )
+
+        await self.publisher.send(None, 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-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/document_load.py

@@ -1,19 +1,18 @@
 
 import base64
 
-from .. schema import Document, Metadata
-from .. schema import document_ingest_queue
+from ... schema import Document, Metadata
 
 from . sender import ServiceSender
 from . serialize import to_subgraph
 
-class DocumentLoadSender(ServiceSender):
-    def __init__(self, pulsar_client):
+class DocumentLoad(ServiceSender):
+    def __init__(self, pulsar_client, queue):
 
-        super(DocumentLoadSender, self).__init__(
-            pulsar_client=pulsar_client,
-            request_queue=document_ingest_queue,
-            request_schema=Document,
+        super(DocumentLoad, self).__init__(
+            pulsar_client = pulsar_client,
+            queue = queue,
+            schema = Document,
         )
 
     def to_request(self, body):

{trustgraph-flow-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/document_rag.py

@@ -1,20 +1,22 @@
 
-from .. schema import DocumentRagQuery, DocumentRagResponse
-from .. schema import document_rag_request_queue
-from .. schema import document_rag_response_queue
+from ... schema import DocumentRagQuery, DocumentRagResponse
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class DocumentRagRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout,
+            consumer, subscriber,
+    ):
 
         super(DocumentRagRequestor, self).__init__(
             pulsar_client=pulsar_client,
-            request_queue=document_rag_request_queue,
-            response_queue=document_rag_response_queue,
+            request_queue=request_queue,
+            response_queue=response_queue,
             request_schema=DocumentRagQuery,
             response_schema=DocumentRagResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
             timeout=timeout,
         )
 

{trustgraph-flow-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/embeddings.py

@@ -1,12 +1,11 @@
 
-from .. schema import EmbeddingsRequest, EmbeddingsResponse
+from ... schema import EmbeddingsRequest, EmbeddingsResponse
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class EmbeddingsRequestor(ServiceRequestor):
     def __init__(
-            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            self, pulsar_client, request_queue, response_queue, timeout,
             consumer, subscriber,
     ):
 

{trustgraph-flow-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/flow.py

@@ -1,16 +1,17 @@
 
-from .. schema import FlowRequest, FlowResponse, ConfigKey, ConfigValue
-from .. schema import flow_request_queue
-from .. schema import flow_response_queue
+from ... schema import FlowRequest, FlowResponse
+from ... schema import flow_request_queue
+from ... schema import flow_response_queue
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class FlowRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
 
         super(FlowRequestor, self).__init__(
             pulsar_client=pulsar_client,
+            consumer_name = consumer,
+            subscription = subscriber,
             request_queue=flow_request_queue,
             response_queue=flow_response_queue,
             request_schema=FlowRequest,

trustgraph-flow-0.23.6/trustgraph/gateway/dispatch/graph_embeddings_export.py

@@ -0,0 +1,67 @@
+
+import asyncio
+import queue
+import uuid
+
+from ... schema import GraphEmbeddings
+from ... base import Subscriber
+
+from . serialize import serialize_graph_embeddings
+
+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):
+        self.running.stop()
+        await self.ws.close()
+
+    async def receive(self, msg):
+        # Ignore incoming info from websocket
+        pass
+
+    async def run(self):
+
+        subs = Subscriber(
+            client = self.pulsar_client, topic = self.queue,
+            consumer_name = self.consumer, subscription = self.subscriber,
+            schema = GraphEmbeddings
+        )
+
+        await subs.start()
+
+        id = str(uuid.uuid4())
+        q = await subs.subscribe_all(id)
+
+        while self.running.get():
+            try:
+
+                resp = await asyncio.wait_for(q.get(), timeout=0.5)
+                await self.ws.send_json(serialize_graph_embeddings(resp))
+
+            except TimeoutError:
+                continue
+
+            except queue.Empty:
+                continue
+
+            except Exception as e:
+                print(f"Exception: {str(e)}", flush=True)
+                break
+
+        await subs.unsubscribe_all(id)
+
+        await subs.stop()
+
+        await self.ws.close()
+        self.running.stop()
+

trustgraph-flow-0.23.6/trustgraph/gateway/dispatch/graph_embeddings_import.py

@@ -0,0 +1,64 @@
+
+import asyncio
+import uuid
+from aiohttp import WSMsgType
+
+from ... schema import Metadata
+from ... schema import GraphEmbeddings, EntityEmbeddings
+from ... base import Publisher
+
+from . serialize import to_subgraph, to_value
+
+class GraphEmbeddingsImport:
+
+    def __init__(
+            self, ws, running, pulsar_client, queue
+    ):
+
+        self.ws = ws
+        self.running = running
+        
+        self.publisher = Publisher(
+            pulsar_client, topic = queue, schema = GraphEmbeddings
+        )
+
+    async def destroy(self):
+        self.running.stop()
+
+        if self.ws:
+            await self.ws.close()
+
+        await self.publisher.stop()
+
+    async def receive(self, msg):
+
+        data = msg.json()
+
+        elt = GraphEmbeddings(
+            metadata=Metadata(
+                id=data["metadata"]["id"],
+                metadata=to_subgraph(data["metadata"]["metadata"]),
+                user=data["metadata"]["user"],
+                collection=data["metadata"]["collection"],
+            ),
+            entities=[
+                EntityEmbeddings(
+                    entity=to_value(ent["entity"]),
+                    vectors=ent["vectors"],
+                )
+                for ent in data["entities"]
+            ]
+        )
+
+        await self.publisher.send(None, 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-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/graph_embeddings_query.py

@@ -1,13 +1,12 @@
 
-from .. schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
+from ... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 from . serialize import serialize_value
 
 class GraphEmbeddingsQueryRequestor(ServiceRequestor):
     def __init__(
-            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            self, pulsar_client, request_queue, response_queue, timeout,
             consumer, subscriber,
     ):
 

{trustgraph-flow-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/graph_rag.py

@@ -1,12 +1,11 @@
 
-from .. schema import GraphRagQuery, GraphRagResponse
+from ... schema import GraphRagQuery, GraphRagResponse
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class GraphRagRequestor(ServiceRequestor):
     def __init__(
-            self, pulsar_client, request_queue, response_queue, timeout, auth,
+            self, pulsar_client, request_queue, response_queue, timeout,
             consumer, subscriber,
     ):
 

{trustgraph-flow-0.23.4/trustgraph/gateway → trustgraph-flow-0.23.6/trustgraph/gateway/dispatch}/librarian.py

@@ -1,18 +1,19 @@
 
-from .. schema import LibrarianRequest, LibrarianResponse, Triples
-from .. schema import librarian_request_queue
-from .. schema import librarian_response_queue
+from ... schema import LibrarianRequest, LibrarianResponse
+from ... schema import librarian_request_queue
+from ... schema import librarian_response_queue
 
-from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 from . serialize import serialize_document_package, serialize_document_info
 from . serialize import to_document_package, to_document_info, to_criteria
 
 class LibrarianRequestor(ServiceRequestor):
-    def __init__(self, pulsar_client, timeout, auth):
+    def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
 
         super(LibrarianRequestor, self).__init__(
             pulsar_client=pulsar_client,
+            consumer_name = consumer,
+            subscription = subscriber,
             request_queue=librarian_request_queue,
             response_queue=librarian_response_queue,
             request_schema=LibrarianRequest,
@@ -22,20 +23,16 @@ class LibrarianRequestor(ServiceRequestor):
 
     def to_request(self, body):
 
-        print("TRR")
         if "document" in body:
             dp = to_document_package(body["document"])
         else:
             dp = None
 
-        print("GOT")
         if "criteria" in body:
             criteria = to_criteria(body["criteria"])
         else:
             criteria = None
 
-        print("ASLDKJ")
-
         return LibrarianRequest(
             operation = body.get("operation", None),
             id = body.get("id", None),