mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
183 lines
4.2 KiB
Python
Executable file
183 lines
4.2 KiB
Python
Executable file
#!/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()
|
|
|