mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-16 16:51:02 +02:00
Created a loader
This commit is contained in:
Cyber MacGeddon
parent
f3d090cd52
commit
608fb2d5a9
3 changed files with 275 additions and 1 deletions
179
trustgraph-cli/scripts/tg-load-kg-core
Executable file
179
trustgraph-cli/scripts/tg-load-kg-core
Executable file
|
|
@ -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())
|
||||
|
||||
|
|
@ -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",
|
||||
]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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__
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue