diff --git a/trustgraph-cli/scripts/tg-load-kg-core b/trustgraph-cli/scripts/tg-load-kg-core new file mode 100755 index 00000000..2469772d --- /dev/null +++ b/trustgraph-cli/scripts/tg-load-kg-core @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 + +import aiohttp +import asyncio +import msgpack +import json +import sys +import argparse +import os + +async def load_ge(queue, url): + + async with aiohttp.ClientSession() as session: + + async with session.ws_connect(f"{url}load/graph-embeddings") as ws: + + while True: + + msg = await queue.get() + + msg = { + "metadata": { + "id": msg["m"]["i"], + "metadata": msg["m"]["m"], + "user": msg["m"]["u"], + "collection": msg["m"]["c"], + }, + "vectors": msg["v"], + "entity": msg["e"], + } + + await ws.send_json(msg) + +async def load_triples(queue, url): + async with aiohttp.ClientSession() as session: + async with session.ws_connect(f"{url}load/triples") as ws: + + while True: + + msg = await queue.get() + + msg ={ + "metadata": { + "id": msg["m"]["i"], + "metadata": msg["m"]["m"], + "user": msg["m"]["u"], + "collection": msg["m"]["c"], + }, + "triples": msg["t"], + } + + await ws.send_json(msg) + +ge_counts = 0 +t_counts = 0 + +async def stats(): + + global t_counts + global ge_counts + + while True: + await asyncio.sleep(5) + print( + f"Graph embeddings: {ge_counts:10d} Triples: {t_counts:10d}" + ) + +async def loader(ge_queue, t_queue, path, format, user, collection): + + global t_counts + global ge_counts + + if format == "json": + + raise RuntimeError("Not implemented") + + else: + + with open(path, "rb") as f: + + unpacker = msgpack.Unpacker(f, raw=False) + + for unpacked in unpacker: + + if user: + unpacked["metadata"]["user"] = user + + if collection: + unpacked["metadata"]["collection"] = collection + + + if unpacked[0] == "t": + await t_queue.put(unpacked[1]) + t_counts += 1 + else: + if unpacked[0] == "ge": + await ge_queue.put(unpacked[1]) + ge_counts += 1 + +async def run(**args): + + ge_q = asyncio.Queue() + t_q = asyncio.Queue() + + load_task = asyncio.create_task( + loader( + ge_queue=ge_q, t_queue=t_q, + path=args["input_file"], format=args["format"], + user=args["user"], collection=args["collection"], + ) + + ) + + ge_task = asyncio.create_task( + load_ge( + queue=ge_q, url=args["url"] + "api/v1/" + ) + ) + + triples_task = asyncio.create_task( + load_triples( + queue=t_q, url=args["url"] + "api/v1/" + ) + ) + + stats_task = asyncio.create_task(stats()) + + await load_task + await triples_task + await ge_task + await stats_task + +async def main(): + + parser = argparse.ArgumentParser( + prog='tg-load-pdf', + description=__doc__, + ) + + default_url = os.getenv("TRUSTGRAPH_API", "http://localhost:8088/") + default_user = "trustgraph" + collection = "default" + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'TrustGraph API URL (default: {default_url})', + ) + + parser.add_argument( + '-i', '--input-file', + # Make it mandatory, difficult to over-write an existing file + required=True, + help=f'Output file' + ) + + parser.add_argument( + '--format', + default="msgpack", + choices=["msgpack", "json"], + help=f'Output format (default: msgpack)', + ) + + parser.add_argument( + '--user', + help=f'User ID to load as (default: from input)' + ) + + parser.add_argument( + '--collection', + help=f'Collection ID to load as (default: from input)' + ) + + args = parser.parse_args() + + await run(**vars(args)) + +asyncio.run(main()) + diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index f4c4e92d..1608cfdb 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -56,6 +56,7 @@ setuptools.setup( "scripts/tg-invoke-prompt", "scripts/tg-invoke-llm", "scripts/tg-save-kg-core", + "scripts/tg-load-kg-core", "scripts/tg-dump-msgpack", ] ) diff --git a/trustgraph-flow/trustgraph/api/gateway/service.py b/trustgraph-flow/trustgraph/api/gateway/service.py index 6bad1649..6d5f70ce 100755 --- a/trustgraph-flow/trustgraph/api/gateway/service.py +++ b/trustgraph-flow/trustgraph/api/gateway/service.py @@ -321,12 +321,22 @@ class Api: schema=JsonSchema(Triples) ) + self.triples_pub = Publisher( + self.pulsar_host, triples_store_queue, + schema=JsonSchema(Triples) + ) + self.graph_embeddings_tap = Subscriber( self.pulsar_host, graph_embeddings_store_queue, "api-gateway", "api-gateway", schema=JsonSchema(GraphEmbeddings) ) + self.graph_embeddings_pub = Publisher( + self.pulsar_host, graph_embeddings_store_queue, + schema=JsonSchema(GraphEmbeddings) + ) + self.document_out = Publisher( self.pulsar_host, document_ingest_queue, schema=JsonSchema(Document), @@ -349,11 +359,19 @@ class Api: web.post("/api/v1/load/document", self.load_document), web.post("/api/v1/load/text", self.load_text), web.get("/api/v1/ws", self.socket), + web.get("/api/v1/stream/triples", self.stream_triples), web.get( "/api/v1/stream/graph-embeddings", self.stream_graph_embeddings ), + + web.get("/api/v1/load/triples", self.load_triples), + web.get( + "/api/v1/load/graph-embeddings", + self.load_graph_embeddings + ), + ]) async def llm(self, request): @@ -844,6 +862,75 @@ class Api: return ws + async def load_triples(self, request): + + ws = web.WebSocketResponse() + await ws.prepare(request) + + async for msg in ws: + + try: + + if msg.type == WSMsgType.TEXT: + + 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.triples_pub.send(None, elt) + + elif msg.type == WSMsgType.ERROR: + break + + except Exception as e: + + print("Exception:", e) + + return ws + + async def load_graph_embeddings(self, request): + + ws = web.WebSocketResponse() + await ws.prepare(request) + + async for msg in ws: + + try: + + if msg.type == WSMsgType.TEXT: + + 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"], + ), + entity=to_value(data["entity"]), + vectors=data["vectors"], + ) + + await self.graph_embeddings_pub.send(None, elt) + + elif msg.type == WSMsgType.ERROR: + break + + except Exception as e: + + print("Exception:", e) + + return ws + async def app_factory(self): self.llm_pub_task = asyncio.create_task(self.llm_in.run()) @@ -876,10 +963,18 @@ class Api: self.triples_tap.run() ) + self.triples_pub_task = asyncio.create_task( + self.triples_pub.run() + ) + self.graph_embeddings_tap_task = asyncio.create_task( self.graph_embeddings_tap.run() ) + self.graph_embeddings_pub_task = asyncio.create_task( + self.graph_embeddings_pub.run() + ) + self.doc_ingest_pub_task = asyncio.create_task(self.document_out.run()) self.text_ingest_pub_task = asyncio.create_task(self.text_out.run()) @@ -891,7 +986,6 @@ class Api: def run(): - parser = argparse.ArgumentParser( prog="api-gateway", description=__doc__