From 03d59b62af48a75df4b0fb1811ef9ac1149c7904 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 2 May 2025 20:03:17 +0100 Subject: [PATCH] Rework loaders, refactored stream/load to be export/import --- ...tream.py => document_embeddings_export.py} | 2 +- .../dispatch/document_embeddings_import.py | 60 +++++++++++++++++ .../dispatch/document_embeddings_load.py | 63 ------------------ ...s_stream.py => graph_embeddings_export.py} | 2 +- .../dispatch/graph_embeddings_import.py | 60 +++++++++++++++++ .../gateway/dispatch/graph_embeddings_load.py | 64 ------------------- .../trustgraph/gateway/dispatch/manager.py | 64 ++++++++++++++++--- .../{triples_stream.py => triples_export.py} | 2 +- .../gateway/dispatch/triples_import.py | 54 ++++++++++++++++ .../gateway/dispatch/triples_load.py | 56 ---------------- .../trustgraph/gateway/endpoint/manager.py | 13 ++-- 11 files changed, 240 insertions(+), 200 deletions(-) rename trustgraph-flow/trustgraph/gateway/dispatch/{document_embeddings_stream.py => document_embeddings_export.py} (97%) create mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py delete mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_load.py rename trustgraph-flow/trustgraph/gateway/dispatch/{graph_embeddings_stream.py => graph_embeddings_export.py} (98%) create mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py delete mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_load.py rename trustgraph-flow/trustgraph/gateway/dispatch/{triples_stream.py => triples_export.py} (98%) create mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py delete mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/triples_load.py diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_export.py similarity index 97% rename from trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_stream.py rename to trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_export.py index 2d61763b..2587132d 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_stream.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py new file mode 100644 index 00000000..306938d3 --- /dev/null +++ b/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 + diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_load.py deleted file mode 100644 index bbfb51a3..00000000 --- a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_load.py +++ /dev/null @@ -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() diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_export.py similarity index 98% rename from trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_stream.py rename to trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_export.py index b46c9f9f..07f72550 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_stream.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py new file mode 100644 index 00000000..ca72b822 --- /dev/null +++ b/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 + diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_load.py deleted file mode 100644 index 27e92a30..00000000 --- a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_load.py +++ /dev/null @@ -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() diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 34fc2627..3adee9cb 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/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 . graph_embeddings_stream import GraphEmbeddingsStream -from . document_embeddings_stream import DocumentEmbeddingsStream +from . triples_export import TriplesExport +from . graph_embeddings_export import GraphEmbeddingsExport +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): - return self.dispatch_receive + def dispatch_import(self): + 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, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py similarity index 98% rename from trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py rename to trustgraph-flow/trustgraph/gateway/dispatch/triples_export.py index 31b08069..d065550e 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/triples_stream.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py new file mode 100644 index 00000000..c5346091 --- /dev/null +++ b/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 + diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/triples_load.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_load.py deleted file mode 100644 index 81c8ea82..00000000 --- a/trustgraph-flow/trustgraph/gateway/dispatch/triples_load.py +++ /dev/null @@ -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() diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py b/trustgraph-flow/trustgraph/gateway/endpoint/manager.py index eef7059c..cee904a5 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py +++ b/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() ), ]