mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 16:36:21 +02:00
161 lines
3.5 KiB
Python
Executable file
161 lines
3.5 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
"""
|
|
Uses the knowledge service to fetch a knowledge core which is saved
|
|
to a local file in msgpack format.
|
|
"""
|
|
|
|
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 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("/"):
|
|
url += "/"
|
|
|
|
url = url + "api/v1/socket"
|
|
|
|
mid = str(uuid.uuid4())
|
|
|
|
async with connect(url) as ws:
|
|
|
|
req = json.dumps({
|
|
"id": mid,
|
|
"service": "knowledge",
|
|
"request": {
|
|
"operation": "get-kg-core",
|
|
"user": user,
|
|
"id": id,
|
|
}
|
|
})
|
|
|
|
await ws.send(req)
|
|
|
|
ge = 0
|
|
t = 0
|
|
|
|
with open(output, "wb") as f:
|
|
|
|
while True:
|
|
|
|
msg = await ws.recv()
|
|
|
|
obj = json.loads(msg)
|
|
|
|
if "response" not in obj:
|
|
raise RuntimeError("No response?")
|
|
|
|
response = obj["response"]
|
|
|
|
if "error" in response:
|
|
raise RuntimeError(obj["error"])
|
|
|
|
if "eos" in response:
|
|
if response["eos"]: break
|
|
|
|
if "triples" in response:
|
|
t += 1
|
|
write_triple(f, response["triples"])
|
|
|
|
if "graph-embeddings" in response:
|
|
ge += 1
|
|
write_ge(f, response["graph-embeddings"])
|
|
|
|
print(f"Got: {t} triple, {ge} GE messages.")
|
|
|
|
await ws.close()
|
|
|
|
def main():
|
|
|
|
parser = argparse.ArgumentParser(
|
|
prog='tg-get-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(
|
|
'-o', '--output',
|
|
required=True,
|
|
help=f'Output file'
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
|
|
asyncio.run(
|
|
fetch(
|
|
url = args.url,
|
|
user = args.user,
|
|
id = args.id,
|
|
output = args.output,
|
|
)
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
print("Exception:", e, flush=True)
|
|
|
|
main()
|
|
|