tg-put-kg-core command (#369)

trustgraph-cli/scripts/tg-get-kg-core

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 
 """
-Uses the agent service to answer a question
+Uses the knowledge service to fetch a knowledge core which is saved
+to a local file in msgpack format.
 """
 
 import argparse
@@ -16,6 +17,42 @@ import msgpack
 default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
 default_user = 'trustgraph'
 
+def write_triple(f, data):
+    msg = (
+        "t",
+        {
+            "m": {
+                "i": data["metadata"]["id"], 
+                "m": data["metadata"]["metadata"],
+                "u": data["metadata"]["user"],
+                "c": data["metadata"]["collection"],
+            },
+            "t": data["triples"],
+        }
+    )
+    f.write(msgpack.packb(msg, use_bin_type=True))
+
+def write_ge(f, data):
+    msg = (
+        "ge",
+        {
+            "m": {
+                "i": data["metadata"]["id"], 
+                "m": data["metadata"]["metadata"],
+                "u": data["metadata"]["user"],
+                "c": data["metadata"]["collection"],
+            },
+            "e": [
+                {
+                    "e": ent["entity"],
+                    "v": ent["vectors"],
+                }
+                for ent in data["entities"]
+            ]
+        }
+    )
+    f.write(msgpack.packb(msg, use_bin_type=True))
+
 async def fetch(url, user, id, output):
 
     if not url.endswith("/"):
@@ -61,21 +98,20 @@ async def fetch(url, user, id, output):
 
                 if "triples" in obj["response"]:
                     t += 1
-                    msg = obj["response"]["triples"]
-                    f.write(msgpack.packb(msg, use_bin_type=True))
+                    write_triple(f, obj["response"]["triples"])
 
                 if "graph-embeddings" in obj["response"]:
                     ge += 1
-                    msg = obj["response"]["graph-embeddings"]
-                    f.write(msgpack.packb(msg, use_bin_type=True))
+                    write_ge(f, obj["response"]["graph-embeddings"])
+
+        print(f"Got: {t} triple, {ge} GE messages.")
 
-        print(f"Wrote: {t} triple, {ge} GE messages.")
         await ws.close()
 
 def main():
 
     parser = argparse.ArgumentParser(
-        prog='tg-invoke-agent',
+        prog='tg-get-kg-core',
         description=__doc__,
     )
 

trustgraph-cli/scripts/tg-put-kg-core

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+
+"""
+Uses the agent service to answer a question
+"""
+
+import argparse
+import os
+import textwrap
+import uuid
+import asyncio
+import json
+from websockets.asyncio.client import connect
+import msgpack
+
+default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
+default_user = 'trustgraph'
+
+def read_message(unpacked, id, user):
+    
+    if unpacked[0] == "ge":
+        msg = unpacked[1]
+        return "ge", {
+            "metadata": {
+                "id": id,
+                "metadata": msg["m"]["m"],
+                "user": user,
+                "collection": "default",  # Not used?
+            },
+            "entities": [
+                {
+                    "entity": ent["e"],
+                    "vectors": ent["v"],
+                }
+                for ent in msg["e"]
+            ],
+        }
+    elif unpacked[0] == "t":
+        msg = unpacked[1]
+        return "t", {
+            "metadata": {
+                "id": id,
+                "metadata": msg["m"]["m"],
+                "user": user,
+                "collection": "default",   # Not used by receiver?
+            },
+            "triples": msg["t"],
+        }
+    else:
+        raise RuntimeError("Unpacked unexpected messsage type", unpacked[0])
+
+async def put(url, user, id, input):
+
+    if not url.endswith("/"):
+        url += "/"
+
+    url = url + "api/v1/socket"
+
+    async with connect(url) as ws:
+
+
+        ge = 0
+        t = 0
+
+        with open(input, "rb") as f:
+
+            unpacker = msgpack.Unpacker(f, raw=False)
+
+            while True:
+
+                try:
+                    unpacked = unpacker.unpack()
+                except:
+                    break
+
+                kind, msg = read_message(unpacked, id, user)
+
+                mid = str(uuid.uuid4())
+
+                if kind == "ge":
+
+                    ge += 1
+
+                    req = json.dumps({
+                        "id": mid,
+                        "service": "knowledge",
+                        "request": {
+                            "operation": "put-kg-core",
+                            "user": user,
+                            "id": id,
+                            "graph-embeddings": msg
+                        }
+                    })
+
+                elif kind == "t":
+
+                    t += 1
+
+                    req = json.dumps({
+                        "id": mid,
+                        "service": "knowledge",
+                        "request": {
+                            "operation": "put-kg-core",
+                            "user": user,
+                            "id": id,
+                            "triples": msg
+                        }
+                    })
+
+                else:
+
+                    raise RuntimeError("Unexpected message kind", kind)
+
+                await ws.send(req)
+
+                # Retry loop, wait for right response to come back
+                while True:
+                    
+                    msg = await ws.recv()
+                    msg = json.loads(msg)
+
+                    if msg["id"] != mid:
+                        continue
+
+                    if "response" in msg:
+                        if "error" in msg["response"]:
+                            raise RuntimeError(msg["response"]["error"])
+
+                    break
+
+        print(f"Put: {t} triple, {ge} GE messages.")
+
+        await ws.close()
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        prog='tg-put-kg-core',
+        description=__doc__,
+    )
+
+    parser.add_argument(
+        '-u', '--url',
+        default=default_url,
+        help=f'API URL (default: {default_url})',
+    )
+    parser.add_argument(
+        '-U', '--user',
+        default=default_user,
+        help=f'User ID (default: {default_user})'
+    )
+
+    parser.add_argument(
+        '--id', '--identifier',
+        required=True,
+        help=f'Knowledge core ID',
+    )
+
+    parser.add_argument(
+        '-i', '--input',
+        required=True,
+        help=f'Input file'
+    )
+
+    args = parser.parse_args()
+
+    try:
+
+        asyncio.run(
+            put(
+                url = args.url,
+                user = args.user,
+                id = args.id,
+                input = args.input,
+            )
+        )
+
+    except Exception as e:
+
+        print("Exception:", e, flush=True)
+
+main()
+