Created a loader

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())
+

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",
     ]
 )

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__