diff --git a/trustgraph-cli/scripts/tg-load-kg-core b/trustgraph-cli/scripts/tg-load-kg-core index b79ec992..46811bd3 100755 --- a/trustgraph-cli/scripts/tg-load-kg-core +++ b/trustgraph-cli/scripts/tg-load-kg-core @@ -28,7 +28,7 @@ async def load_ge(running, queue, url): async with aiohttp.ClientSession() as session: - async with session.ws_connect(f"{url}load/graph-embeddings") as ws: + async with session.ws_connect(url) as ws: while running.get(): @@ -73,7 +73,7 @@ async def load_triples(running, queue, url): async with aiohttp.ClientSession() as session: - async with session.ws_connect(f"{url}load/triples") as ws: + async with session.ws_connect(url) as ws: while running.get(): @@ -200,6 +200,9 @@ async def run(running, **args): ge_q = asyncio.Queue(maxsize=10) t_q = asyncio.Queue(maxsize=10) + flow_id = args["flow"] + url = args["url"] + load_task = asyncio.create_task( loader( running=running, @@ -212,15 +215,17 @@ async def run(running, **args): ge_task = asyncio.create_task( load_ge( - running=running, - queue=ge_q, url=args["url"] + "api/v1/" + running = running, + queue = ge_q, + url = f"{url}api/v1/{flow_id}/import/graph-embeddings" ) ) triples_task = asyncio.create_task( load_triples( - running=running, - queue=t_q, url=args["url"] + "api/v1/" + running = running, + queue = t_q, + url = f"{url}api/v1/{flow_id}/import/triples" ) ) @@ -258,6 +263,13 @@ async def main(running): help=f'Output file' ) + parser.add_argument( + '-f', '--flow-id', + # Make it mandatory, difficult to over-write an existing file + default="0000", + help=f'Flow ID' + ) + parser.add_argument( '--format', default="msgpack", diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py index 306938d3..46ef805e 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/document_embeddings_import.py @@ -3,11 +3,10 @@ import asyncio import uuid from aiohttp import WSMsgType -from .. schema import Metadata -from .. schema import DocumentEmbeddings, ChunkEmbeddings -from .. base import Publisher +from ... schema import Metadata +from ... schema import DocumentEmbeddings, ChunkEmbeddings +from ... base import Publisher -from . socket import SocketEndpoint from . serialize import to_subgraph class DocumentEmbeddingsImport: @@ -20,7 +19,7 @@ class DocumentEmbeddingsImport: self.running = running self.publisher = Publisher( - pulsar_client, queue = queue, schema = DocumentEmbeddings + pulsar_client, topic = queue, schema = DocumentEmbeddings ) async def destroy(self): diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py index ca72b822..c27f0c60 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/graph_embeddings_import.py @@ -3,11 +3,10 @@ import asyncio import uuid from aiohttp import WSMsgType -from .. schema import Metadata -from .. schema import GraphEmbeddings, EntityEmbeddings -from .. base import Publisher +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: @@ -20,7 +19,7 @@ class GraphEmbeddingsImport: self.running = running self.publisher = Publisher( - pulsar_client, queue = queue, schema = GraphEmbeddings + pulsar_client, topic = queue, schema = GraphEmbeddings ) async def destroy(self): diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 3adee9cb..f277e255 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -44,9 +44,9 @@ sender_dispatchers = { } export_dispatchers = { - "triples": TriplesStream, - "graph-embeddings": GraphEmbeddingsStream, - "document-embeddings": DocumentEmbeddingsStream, + "triples": TriplesExport, + "graph-embeddings": GraphEmbeddingsExport, + "document-embeddings": DocumentEmbeddingsExport, } import_dispatchers = { diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py b/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py index c5346091..84aa927c 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/triples_import.py @@ -3,11 +3,10 @@ import asyncio import uuid from aiohttp import WSMsgType -from .. schema import Metadata -from .. schema import Triples -from .. base import Publisher +from ... schema import Metadata +from ... schema import Triples +from ... base import Publisher -from . socket import SocketEndpoint from . serialize import to_subgraph class TriplesImport: @@ -20,7 +19,7 @@ class TriplesImport: self.running = running self.publisher = Publisher( - pulsar_client, queue = queue, schema = Triples + pulsar_client, topic = queue, schema = Triples ) async def destroy(self):