mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
123 lines
2.6 KiB
Python
Executable file
123 lines
2.6 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'
|
|
|
|
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": "fetch-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 "error" in obj:
|
|
raise RuntimeError(obj["error"])
|
|
|
|
if "response" not in obj:
|
|
continue
|
|
|
|
if "eos" in obj["response"]:
|
|
if obj["response"]["eos"]: break
|
|
|
|
if "triples" in obj["response"]:
|
|
t += 1
|
|
msg = obj["response"]["triples"]
|
|
f.write(msgpack.packb(msg, use_bin_type=True))
|
|
|
|
if "graph-embeddings" in obj["response"]:
|
|
ge += 1
|
|
msg = obj["response"]["graph-embeddings"]
|
|
f.write(msgpack.packb(msg, use_bin_type=True))
|
|
|
|
print(f"Wrote: {t} triple, {ge} GE messages.")
|
|
await ws.close()
|
|
|
|
def main():
|
|
|
|
parser = argparse.ArgumentParser(
|
|
prog='tg-invoke-agent',
|
|
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()
|
|
|