Rework loaders, refactored stream/load to be export/import

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

@@ -8,7 +8,7 @@ from ... base import Subscriber
 
 from . serialize import serialize_document_embeddings
 
-class DocumentEmbeddingsStream:
+class DocumentEmbeddingsExport:
 
     def __init__(
             self, ws, running, pulsar_client, queue, consumer, subscriber

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

@@ -0,0 +1,60 @@
+
+import asyncio
+import uuid
+from aiohttp import WSMsgType
+
+from .. schema import Metadata
+from .. schema import DocumentEmbeddings, ChunkEmbeddings
+from .. base import Publisher
+
+from . socket import SocketEndpoint
+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, queue = queue, schema = DocumentEmbeddings
+        )
+
+    async def destroy(self):
+        self.running.stop()
+        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)
+
+        await self.ws.close()
+        self.ws = None
+

trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_load.py

@@ -1,63 +0,0 @@
-
-import asyncio
-import uuid
-from aiohttp import WSMsgType
-
-from .. schema import Metadata
-from .. schema import DocumentEmbeddings, ChunkEmbeddings
-from .. schema import document_embeddings_store_queue
-from .. base import Publisher
-
-from . socket import SocketEndpoint
-from . serialize import to_subgraph
-
-class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
-
-    def __init__(
-            self, pulsar_client, auth, path="/api/v1/load/document-embeddings",
-    ):
-
-        super(DocumentEmbeddingsLoadEndpoint, self).__init__(
-            endpoint_path=path, auth=auth,
-        )
-
-        self.pulsar_client=pulsar_client
-
-        self.publisher = Publisher(
-            self.pulsar_client, document_embeddings_store_queue,
-            schema=DocumentEmbeddings
-        )
-
-    async def start(self):
-
-        await self.publisher.start()
-
-    async def listener(self, ws, running):
-        
-        async for msg in ws:
-            # On error, finish
-            if msg.type == WSMsgType.ERROR:
-                break
-            else:
-
-                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)
-
-        running.stop()

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

@@ -8,7 +8,7 @@ from ... base import Subscriber
 
 from . serialize import serialize_graph_embeddings
 
-class GraphEmbeddingsStream:
+class GraphEmbeddingsExport:
 
     def __init__(
             self, ws, running, pulsar_client, queue, consumer, subscriber

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

@@ -0,0 +1,60 @@
+
+import asyncio
+import uuid
+from aiohttp import WSMsgType
+
+from .. schema import Metadata
+from .. schema import GraphEmbeddings, EntityEmbeddings
+from .. base import Publisher
+
+from . socket import SocketEndpoint
+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, queue = queue, schema = GraphEmbeddings
+        )
+
+    async def destroy(self):
+        self.running.stop()
+        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)
+
+        await self.ws.close()
+        self.ws = None
+

trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_load.py

@@ -1,64 +0,0 @@
-
-import asyncio
-import uuid
-from aiohttp import WSMsgType
-
-from .. schema import Metadata
-from .. schema import GraphEmbeddings, EntityEmbeddings
-from .. schema import graph_embeddings_store_queue
-from .. base import Publisher
-
-from . socket import SocketEndpoint
-from . serialize import to_subgraph, to_value
-
-class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
-
-    def __init__(
-            self, pulsar_client, auth, path="/api/v1/load/graph-embeddings",
-    ):
-
-        super(GraphEmbeddingsLoadEndpoint, self).__init__(
-            endpoint_path=path, auth=auth,
-        )
-
-        self.pulsar_client=pulsar_client
-
-        self.publisher = Publisher(
-            self.pulsar_client, graph_embeddings_store_queue,
-            schema=GraphEmbeddings
-        )
-
-    async def start(self):
-
-        await self.publisher.start()
-
-    async def listener(self, ws, running):
-        
-        async for msg in ws:
-
-            # On error, finish
-            if msg.type == WSMsgType.ERROR:
-                break
-            else:
-
-                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)
-
-        running.stop()

trustgraph-flow/trustgraph/gateway/dispatch/manager.py

@@ -19,9 +19,13 @@ from . prompt import PromptRequestor
 from . text_load import TextLoad
 from . document_load import DocumentLoad
 
-from . triples_stream import TriplesStream
+from . triples_export import TriplesExport
-from . graph_embeddings_stream import GraphEmbeddingsStream
+from . graph_embeddings_export import GraphEmbeddingsExport
-from . document_embeddings_stream import DocumentEmbeddingsStream
+from . document_embeddings_export import DocumentEmbeddingsExport
+
+from . triples_import import TriplesImport
+from . graph_embeddings_import import GraphEmbeddingsImport
+from . document_embeddings_import import DocumentEmbeddingsImport
 
 request_response_dispatchers = {
     "agent": AgentRequestor,
@@ -39,12 +43,18 @@ sender_dispatchers = {
     "document-load": DocumentLoad,
 }
 
-receive_dispatchers = {
+export_dispatchers = {
     "triples": TriplesStream,
     "graph-embeddings": GraphEmbeddingsStream,
     "document-embeddings": DocumentEmbeddingsStream,
 }
 
+import_dispatchers = {
+    "triples": TriplesImport,
+    "graph-embeddings": GraphEmbeddingsImport,
+    "document-embeddings": DocumentEmbeddingsImport,
+}
+
 class DispatcherWrapper:
     def __init__(self, mgr, name, impl):
         self.mgr = mgr
@@ -101,13 +111,16 @@ class DispatcherManager:
 
         return await dispatcher.process(data, responder)
 
-    def dispatch_flow_service(self):
+    def dispatch_service(self):
         return self
 
-    def dispatch_flow_receive(self):
+    def dispatch_import(self):
-        return self.dispatch_receive
+        return self.invoke_import
 
-    async def dispatch_receive(self, ws, running, params):
+    def dispatch_export(self):
+        return self.invoke_export
+
+    async def invoke_import(self, ws, running, params):
 
         flow = params.get("flow")
         kind = params.get("kind")
@@ -115,7 +128,7 @@ class DispatcherManager:
         if flow not in self.flows:
             raise RuntimeError("Invalid flow")
 
-        if kind not in receive_dispatchers:
+        if kind not in import_dispatchers:
             raise RuntimeError("Invalid kind")
 
         key = (flow, kind)
@@ -129,7 +142,38 @@ class DispatcherManager:
         qconfig = intf_defs[kind + "-store"]
 
         id = str(uuid.uuid4())
-        dispatcher = receive_dispatchers[kind](
+        dispatcher = import_dispatchers[kind](
+            pulsar_client = self.pulsar_client,
+            ws = ws,
+            running = running,
+            queue = qconfig,
+        )
+
+        return dispatcher
+
+    async def invoke_export(self, ws, running, params):
+
+        flow = params.get("flow")
+        kind = params.get("kind")
+
+        if flow not in self.flows:
+            raise RuntimeError("Invalid flow")
+
+        if kind not in export_dispatchers:
+            raise RuntimeError("Invalid kind")
+
+        key = (flow, kind)
+
+        intf_defs = self.flows[flow]["interfaces"]
+
+        if kind not in intf_defs:
+            raise RuntimeError("This kind not supported by flow")
+
+        # FIXME: The -store bit, does it make sense?
+        qconfig = intf_defs[kind + "-store"]
+
+        id = str(uuid.uuid4())
+        dispatcher = export_dispatchers[kind](
             pulsar_client = self.pulsar_client,
             ws = ws,
             running = running,

trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py

@@ -8,7 +8,7 @@ from ... base import Subscriber
 
 from . serialize import serialize_triples
 
-class TriplesStream:
+class TriplesExport:
 
     def __init__(
             self, ws, running, pulsar_client, queue, consumer, subscriber

trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py

@@ -0,0 +1,54 @@
+
+import asyncio
+import uuid
+from aiohttp import WSMsgType
+
+from .. schema import Metadata
+from .. schema import Triples
+from .. base import Publisher
+
+from . socket import SocketEndpoint
+from . serialize import to_subgraph
+
+class TriplesImport:
+
+    def __init__(
+            self, ws, running, pulsar_client, queue
+    ):
+
+        self.ws = ws
+        self.running = running
+        
+        self.publisher = Publisher(
+            pulsar_client, queue = queue, schema = Triples
+        )
+
+    async def destroy(self):
+        self.running.stop()
+        await self.ws.close()
+        await self.publisher.stop()
+
+    async def receive(self, msg):
+
+        data = msg.json()
+
+        elt = Triples(
+            metadata=Metadata(
+                id=data["metadata"]["id"],
+                metadata=to_subgraph(data["metadata"]["metadata"]),
+                user=data["metadata"]["user"],
+                collection=data["metadata"]["collection"],
+            ),
+            triples=to_subgraph(data["triples"]),
+        )
+
+        await self.publisher.send(None, elt)
+
+    async def run(self):
+
+        while self.running.get():
+            await asyncio.sleep(0.5)
+
+        await self.ws.close()
+        self.ws = None
+

trustgraph-flow/trustgraph/gateway/dispatch/triples_load.py

@@ -1,56 +0,0 @@
-
-import asyncio
-import uuid
-from aiohttp import WSMsgType
-
-from .. schema import Metadata
-from .. schema import Triples
-from .. schema import triples_store_queue
-from .. base import Publisher
-
-from . socket import SocketEndpoint
-from . serialize import to_subgraph
-
-class TriplesLoadEndpoint(SocketEndpoint):
-
-    def __init__(self, pulsar_client, auth, path="/api/v1/load/triples"):
-
-        super(TriplesLoadEndpoint, self).__init__(
-            endpoint_path=path, auth=auth,
-        )
-
-        self.pulsar_client=pulsar_client
-
-        self.publisher = Publisher(
-            self.pulsar_client, triples_store_queue,
-            schema=Triples
-        )
-
-    async def start(self):
-
-        await self.publisher.start()
-
-    async def listener(self, ws, running):
-        
-        async for msg in ws:
-            # On error, finish
-            if msg.type == WSMsgType.ERROR:
-                break
-            else:
-
-                data = msg.json()
-
-                elt = Triples(
-                    metadata=Metadata(
-                        id=data["metadata"]["id"],
-                        metadata=to_subgraph(data["metadata"]["metadata"]),
-                        user=data["metadata"]["user"],
-                        collection=data["metadata"]["collection"],
-                    ),
-                    triples=to_subgraph(data["triples"]),
-                )
-
-                await self.publisher.send(None, elt)
-
-
-        running.stop()

trustgraph-flow/trustgraph/gateway/endpoint/manager.py

@@ -41,14 +41,19 @@ class EndpointManager:
                 auth = auth,
             ),
             VariableEndpoint(
-                endpoint_path = "/api/v1/flow/{flow}/{kind}",
+                endpoint_path = "/api/v1/flow/{flow}/service/{kind}",
                 auth = auth,
-                dispatcher = dispatcher_manager.dispatch_flow_service(),
+                dispatcher = dispatcher_manager.dispatch_service(),
             ),
             SocketEndpoint(
-                endpoint_path = "/api/v1/flow/{flow}/receive/{kind}",
+                endpoint_path = "/api/v1/flow/{flow}/import/{kind}",
                 auth = auth,
-                dispatcher = dispatcher_manager.dispatch_flow_receive()
+                dispatcher = dispatcher_manager.dispatch_import()
+            ),
+            SocketEndpoint(
+                endpoint_path = "/api/v1/flow/{flow}/export/{kind}",
+                auth = auth,
+                dispatcher = dispatcher_manager.dispatch_export()
             ),
         ]