apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-22 11:41:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Merge branch 'release/v0.21' into pulsar-api-support

Browse source
This commit is contained in:
Tyler O 2025-02-10 17:29:00 +00:00
commit 5bdb9c1919
149 changed files with 3916 additions and 1823 deletions
Download patch file Download diff file Expand all files Collapse all files

3
trustgraph-cli/scripts/tg-dump-msgpack
View file

@ -9,6 +9,7 @@ diagnostic utility.
import msgpack
import sys
import argparse
import json
def dump(input_file, action):
@ -17,7 +18,7 @@ def dump(input_file, action):
unpacker = msgpack.Unpacker(f, raw=False)
for unpacked in unpacker:
print(unpacked)
print(json.dumps(unpacked))
def summary(input_file, action):

36
trustgraph-cli/scripts/tg-graph-show
View file

@ -6,24 +6,23 @@ Connects to the graph query service and dumps all graph edges.
import argparse
import os
from trustgraph.clients.triples_query_client import TriplesQueryClient
from trustgraph.api import Api
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
def show_graph(pulsar, user, collection, pulsar_api_key=None):
def show_graph(url, user, collection):
tq = TriplesQueryClient(pulsar_host=pulsar, pulsar_api_key=pulsar_api_key)
api = Api(url)
rows = tq.request(
user=user, collection=collection,
s=None, p=None, o=None, limit=10_000_000
rows = api.triples_query(
# user=user, collection=collection,
s=None, p=None, o=None, limit=10_000,
)
for row in rows:
print(row.s.value, row.p.value, row.o.value)
print(row.s, row.p, row.o)
def main():
@ -33,35 +32,30 @@ def main():
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
'-u', '--api-url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-u', '--user',
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-c', '--collection',
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
args = parser.parse_args()
try:
show_graph(
pulsar=args.pulsar_host, user=args.user,
url=args.api_url,
user=args.user,
collection=args.collection,
pulsar_api_key=args.pulsar_api_key,
)

61
trustgraph-cli/scripts/tg-graph-to-turtle
View file

@ -5,38 +5,45 @@ Connects to the graph query service and dumps all graph edges in Turtle
format.
"""
import argparse
import os
from trustgraph.clients.triples_query_client import TriplesQueryClient
import rdflib
import io
import sys
import argparse
import os
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
from trustgraph.api import Api, Uri
def show_graph(pulsar, pulsar_api_key=None):
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
tq = TriplesQueryClient(pulsar_host=pulsar, pulsar_api_key=pulsar_api_key)
def show_graph(url, user, collection):
rows = tq.request(None, None, None, limit=10_000_000)
api = Api(url)
rows = api.triples_query(
s=None, p=None, o=None,
limit=10_000)
# user=user, collection=collection,
g = rdflib.Graph()
for row in rows:
sv = rdflib.term.URIRef(row.s.value)
pv = rdflib.term.URIRef(row.p.value)
sv = rdflib.term.URIRef(row.s)
pv = rdflib.term.URIRef(row.p)
if row.o.is_uri:
if isinstance(row.o, Uri):
# Skip malformed URLs with spaces in
if " " in row.o.value:
if " " in row.o:
continue
ov = rdflib.term.URIRef(row.o.value)
ov = rdflib.term.URIRef(row.o)
else:
ov = rdflib.term.Literal(row.o.value)
ov = rdflib.term.Literal(row.o)
g.add((sv, pv, ov))
@ -57,22 +64,32 @@ def main():
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
'-u', '--api-url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
args = parser.parse_args()
try:
show_graph(args.pulsar_host, pulsar_api_key=args.pulsar_api_key)
show_graph(
url=args.api_url,
user=args.user,
collection=args.collection
)
except Exception as e:

103
trustgraph-cli/scripts/tg-invoke-agent
View file

@ -1,17 +1,18 @@
#!/usr/bin/env python3
"""
Uses the GraphRAG service to answer a query
Uses the GraphRAG service to answer a question
"""
import argparse
import os
import textwrap
import uuid
import asyncio
import json
from websockets.asyncio.client import connect
from trustgraph.clients.agent_client import AgentClient
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
@ -28,15 +29,18 @@ def output(text, prefix="> ", width=78):
)
print(out)
def query(
pulsar_host, query, user, collection,
plan=None, state=None, verbose=False, pulsar_api_key=None
async def question(
url, question, user, collection,
plan=None, state=None, verbose=False
):
am = AgentClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
if not url.endswith("/"):
url += "/"
url = url + "api/v1/socket"
if verbose:
output(wrap(query), "\U00002753 ")
output(wrap(question), "\U00002753 ")
print()
def think(x):
@ -49,11 +53,43 @@ def query(
output(wrap(x), "\U0001f4a1 ")
print()
resp = am.request(
question=query, think=think, observe=observe,
)
mid = str(uuid.uuid4())
print(resp)
async with connect(url) as ws:
req = json.dumps({
"id": mid,
"service": "agent",
"request": {
"question": question,
}
})
await ws.send(req)
while True:
msg = await ws.recv()
obj = json.loads(msg)
if obj["id"] != mid:
print("Ignore message")
continue
if "thought" in obj["response"]:
think(obj["response"]["thought"])
if "observation" in obj["response"]:
observe(obj["response"]["observation"])
if "answer" in obj["response"]:
print(obj["response"]["answer"])
if obj["complete"]: break
await ws.close()
def main():
@ -63,25 +99,25 @@ def main():
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-q', '--query',
'-q', '--question',
required=True,
help=f'Query to execute',
help=f'Question to answer',
)
parser.add_argument(
'-u', '--user',
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-c', '--collection',
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
@ -101,26 +137,21 @@ def main():
action="store_true",
help=f'Output thinking/observations'
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
args = parser.parse_args()
try:
query(
pulsar_host=args.pulsar_host,
query=args.query,
user=args.user,
collection=args.collection,
plan=args.plan,
state=args.state,
verbose=args.verbose,
pulsar_api_key=args.pulsar_api_key,
asyncio.run(
question(
url=args.url,
question=args.question,
user=args.user,
collection=args.collection,
plan=args.plan,
state=args.state,
verbose=args.verbose,
)
)
except Exception as e:

71
trustgraph-cli/scripts/tg-invoke-document-rag Executable file
View file

@ -0,0 +1,71 @@
#!/usr/bin/env python3
"""
Uses the GraphRAG service to answer a question
"""
import argparse
import os
from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
def question(url, question, user, collection):
rag = Api(url)
# user=user, collection=collection,
resp = rag.document_rag(question=question)
print(resp)
def main():
parser = argparse.ArgumentParser(
prog='tg-invoke-document-rag',
description=__doc__,
)
parser.add_argument(
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-q', '--question',
required=True,
help=f'Question to answer',
)
parser.add_argument(
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
args = parser.parse_args()
try:
question(
url=args.url,
question=args.question,
user=args.user,
collection=args.collection,
)
except Exception as e:
print("Exception:", e, flush=True)
main()

71
trustgraph-cli/scripts/tg-invoke-graph-rag Executable file
View file

@ -0,0 +1,71 @@
#!/usr/bin/env python3
"""
Uses the GraphRAG service to answer a question
"""
import argparse
import os
from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
def question(url, question, user, collection):
rag = Api(url)
# user=user, collection=collection,
resp = rag.graph_rag(question=question)
print(resp)
def main():
parser = argparse.ArgumentParser(
prog='tg-invoke-graph-rag',
description=__doc__,
)
parser.add_argument(
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-q', '--question',
required=True,
help=f'Question to answer',
)
parser.add_argument(
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
args = parser.parse_args()
try:
question(
url=args.url,
question=args.question,
user=args.user,
collection=args.collection,
)
except Exception as e:
print("Exception:", e, flush=True)
main()

28
trustgraph-cli/scripts/tg-invoke-llm
View file

@ -8,17 +8,15 @@ and user prompt. Both arguments are required.
import argparse
import os
import json
from trustgraph.clients.llm_client import LlmClient
from trustgraph.api import Api
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def query(url, system, prompt):
def query(pulsar_host, system, prompt, pulsar_api_key=None):
api = Api(url)
cli = LlmClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
resp = cli.request(system=system, prompt=prompt)
resp = api.text_completion(system=system, prompt=prompt)
print(resp)
@ -30,9 +28,9 @@ def main():
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
@ -46,23 +44,15 @@ def main():
nargs=1,
help='LLM prompt e.g. What is 2 + 2?',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
args = parser.parse_args()
try:
query(
pulsar_host=args.pulsar_host,
url=args.url,
system=args.system[0],
prompt=args.prompt[0],
pulsar_api_key=args.pulsar_api_key,
)
except Exception as e:

28
trustgraph-cli/scripts/tg-invoke-prompt
View file

@ -12,17 +12,15 @@ using key=value arguments on the command line, and these replace
import argparse
import os
import json
from trustgraph.clients.prompt_client import PromptClient
from trustgraph.api import Api
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def query(url, template_id, variables):
def query(pulsar_host, template_id, variables, pulsar_api_key=None):
api = Api(url)
cli = PromptClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
resp = cli.request(id=template_id, variables=variables)
resp = api.prompt(id=template_id, variables=variables)
if isinstance(resp, str):
print(resp)
@ -37,9 +35,9 @@ def main():
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
@ -57,13 +55,6 @@ def main():
specified multiple times''',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
args = parser.parse_args()
variables = {}
@ -79,10 +70,9 @@ specified multiple times''',
try:
query(
pulsar_host=args.pulsar_host,
url=args.url,
template_id=args.id[0],
variables=variables,
pulsar_api_key=args.pulsar_api_key,
)
except Exception as e:

224
trustgraph-cli/scripts/tg-load-doc-embeds Executable file
View file

@ -0,0 +1,224 @@
#!/usr/bin/env python3
"""This utility takes a knowledge core and loads it into a running TrustGraph
through the API. The knowledge core should be in msgpack format, which is the
default format produce by tg-save-kg-core.
"""
import aiohttp
import asyncio
import msgpack
import json
import sys
import argparse
import os
import signal
class Running:
def __init__(self): self.running = True
def get(self): return self.running
def stop(self): self.running = False
de_counts = 0
async def load_de(running, queue, url):
global de_counts
async with aiohttp.ClientSession() as session:
async with session.ws_connect(f"{url}load/document-embeddings") as ws:
while running.get():
try:
msg = await asyncio.wait_for(queue.get(), 1)
# End of load
if msg is None:
break
except:
# Hopefully it's TimeoutError. Annoying to match since
# it changed in 3.11.
continue
msg = {
"metadata": {
"id": msg["m"]["i"],
"metadata": msg["m"]["m"],
"user": msg["m"]["u"],
"collection": msg["m"]["c"],
},
"chunks": [
{
"chunk": chunk["c"],
"vectors": chunk["v"],
}
for chunk in msg["c"]
],
}
try:
await ws.send_json(msg)
except Exception as e:
print(e)
de_counts += 1
async def stats(running):
global de_counts
while running.get():
await asyncio.sleep(2)
print(
f"Graph embeddings: {de_counts:10d}"
)
async def loader(running, de_queue, path, format, user, collection):
if format == "json":
raise RuntimeError("Not implemented")
else:
with open(path, "rb") as f:
unpacker = msgpack.Unpacker(f, raw=False)
while running.get():
try:
unpacked = unpacker.unpack()
except:
break
if user:
unpacked["metadata"]["user"] = user
if collection:
unpacked["metadata"]["collection"] = collection
if unpacked[0] == "de":
qtype = de_queue
while running.get():
try:
await asyncio.wait_for(qtype.put(unpacked[1]), 0.5)
# Successful put message, move on
break
except:
# Hopefully it's TimeoutError. Annoying to match since
# it changed in 3.11.
continue
if not running.get(): break
# Put 'None' on end of queue to finish
while running.get():
try:
await asyncio.wait_for(de_queue.put(None), 1)
# Successful put message, move on
break
except:
# Hopefully it's TimeoutError. Annoying to match since
# it changed in 3.11.
continue
async def run(running, **args):
# Maxsize on queues reduces back-pressure so tg-load-kg-core doesn't
# grow to eat all memory
de_q = asyncio.Queue(maxsize=10)
load_task = asyncio.create_task(
loader(
running=running,
de_queue=de_q,
path=args["input_file"], format=args["format"],
user=args["user"], collection=args["collection"],
)
)
de_task = asyncio.create_task(
load_de(
running=running,
queue=de_q, url=args["url"] + "api/v1/"
)
)
stats_task = asyncio.create_task(stats(running))
await de_task
running.stop()
await load_task
await stats_task
async def main(running):
parser = argparse.ArgumentParser(
prog='tg-load-kg-core',
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(running, **vars(args))
running = Running()
def interrupt(sig, frame):
running.stop()
print('Interrupt')
signal.signal(signal.SIGINT, interrupt)
asyncio.run(main(running))

9
trustgraph-cli/scripts/tg-load-kg-core
View file

@ -51,8 +51,13 @@ async def load_ge(running, queue, url):
"user": msg["m"]["u"],
"collection": msg["m"]["c"],
},
"vectors": msg["v"],
"entity": msg["e"],
"entities": [
{
"entity": ent["e"],
"vectors": ent["v"],
}
for ent in msg["e"]
],
}
try:

116
trustgraph-cli/scripts/tg-load-pdf
View file

@ -6,21 +6,19 @@ Loads a PDF document into TrustGraph processing.
import pulsar
from pulsar.schema import JsonSchema
import base64
import hashlib
import argparse
import os
import time
import uuid
from trustgraph.schema import Document, document_ingest_queue
from trustgraph.schema import Metadata, Triple, Value
from trustgraph.log_level import LogLevel
from trustgraph.knowledge import hash, to_uri, Uri
from trustgraph.api import Api
from trustgraph.knowledge import hash, to_uri
from trustgraph.knowledge import PREF_PUBEV, PREF_DOC, PREF_ORG
from trustgraph.knowledge import Organization, PublicationEvent
from trustgraph.knowledge import DigitalDocument
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
@ -28,33 +26,13 @@ class Loader:
def __init__(
self,
pulsar_host,
output_queue,
url,
user,
collection,
log_level,
metadata,
pulsar_api_key=None,
):
if pulsar_api_key:
auth = pulsar.AuthenticationToken(pulsar_api_key)
self.client = pulsar.Client(
pulsar_host,
authentication=auth,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
else:
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.producer = self.client.create_producer(
topic=output_queue,
schema=JsonSchema(Document),
chunking_enabled=True,
)
self.api = Api(url)
self.user = user
self.collection = collection
@ -77,49 +55,18 @@ class Loader:
id = to_uri(PREF_DOC, id)
triples = []
def emit(t):
triples.append(t)
self.metadata.id = id
self.metadata.emit(emit)
r = Document(
metadata=Metadata(
id=id,
metadata=[
Triple(
s=Value(
value=t["s"],
is_uri=isinstance(t["s"], Uri)
),
p=Value(
value=t["p"],
is_uri=isinstance(t["p"], Uri)
),
o=Value(
value=t["o"],
is_uri=isinstance(t["o"], Uri)
),
)
for t in triples
],
user=self.user,
collection=self.collection,
),
data=base64.b64encode(data),
self.api.load_document(
document=data, id=id, metadata=self.metadata,
# user=self.user,
# collection=self.collection,
)
self.producer.send(r)
print(f"{file}: Loaded successfully.")
except Exception as e:
print(f"{file}: Failed: {str(e)}", flush=True)
def __del__(self):
self.client.close()
def main():
@ -128,36 +75,20 @@ def main():
description=__doc__,
)
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_output_queue = document_ingest_queue
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
parser.add_argument(
'-o', '--output-queue',
default=default_output_queue,
help=f'Output queue (default: {default_output_queue})'
)
parser.add_argument(
'-u', '--user',
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-c', '--collection',
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
@ -199,7 +130,7 @@ def main():
)
parser.add_argument(
'--url', help=f'Document URL'
'--document-url', help=f'Document URL'
)
parser.add_argument(
@ -210,14 +141,6 @@ def main():
'--identifier', '--id', help=f'Document ID'
)
parser.add_argument(
'-l', '--log-level',
type=LogLevel,
default=LogLevel.ERROR,
choices=list(LogLevel),
help=f'Output queue (default: info)'
)
parser.add_argument(
'files', nargs='+',
help=f'File to load'
@ -237,7 +160,7 @@ def main():
copyright_holder=args.copyright_holder,
copyright_year=args.copyright_year,
license=args.license,
url=args.url,
url=args.document_url,
keywords=args.keyword,
)
@ -255,12 +178,9 @@ def main():
)
p = Loader(
pulsar_host=args.pulsar_host,
pulsar_api_key=args.pulsar_api_key,
output_queue=args.output_queue,
url=args.url,
user=args.user,
collection=args.collection,
log_level=args.log_level,
metadata=document,
)

112
trustgraph-cli/scripts/tg-load-text
View file

@ -12,14 +12,13 @@ import os
import time
import uuid
from trustgraph.schema import TextDocument, text_ingest_queue
from trustgraph.schema import Metadata, Triple, Value
from trustgraph.log_level import LogLevel
from trustgraph.knowledge import hash, to_uri, Literal, Uri
from trustgraph.api import Api
from trustgraph.knowledge import hash, to_uri
from trustgraph.knowledge import PREF_PUBEV, PREF_DOC, PREF_ORG
from trustgraph.knowledge import Organization, PublicationEvent
from trustgraph.knowledge import DigitalDocument
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
@ -27,32 +26,14 @@ class Loader:
def __init__(
self,
pulsar_host,
output_queue,
url,
user,
collection,
log_level,
metadata,
pulsar_api_key=None,
):
if pulsar_api_key:
auth = pulsar.AuthenticationToken(pulsar_api_key)
self.client = pulsar.Client(
pulsar_host,
authentication=auth,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
else:
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.producer = self.client.create_producer(
topic=output_queue,
schema=JsonSchema(TextDocument),
chunking_enabled=True,
)
self.api = Api(url)
self.user = user
self.collection = collection
@ -75,49 +56,18 @@ class Loader:
id = to_uri(PREF_DOC, id)
triples = []
def emit(t):
triples.append(t)
self.metadata.id = id
self.metadata.emit(emit)
r = TextDocument(
metadata=Metadata(
id=id,
metadata=[
Triple(
s=Value(
value=t["s"],
is_uri=isinstance(t["s"], Uri)
),
p=Value(
value=t["p"],
is_uri=isinstance(t["p"], Uri)
),
o=Value(
value=t["o"],
is_uri=isinstance(t["o"], Uri)
),
)
for t in triples
],
user=self.user,
collection=self.collection,
),
text=data,
self.api.load_text(
text=data, id=id, metadata=self.metadata,
# user=self.user,
# collection=self.collection,
)
self.producer.send(r)
print(f"{file}: Loaded successfully.")
except Exception as e:
print(f"{file}: Failed: {str(e)}", flush=True)
def __del__(self):
self.client.close()
def main():
@ -126,37 +76,20 @@ def main():
description=__doc__,
)
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_output_queue = text_ingest_queue
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-o', '--output-queue',
default=default_output_queue,
help=f'Output queue (default: {default_output_queue})'
)
parser.add_argument(
'-u', '--user',
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-c', '--collection',
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
@ -198,7 +131,7 @@ def main():
)
parser.add_argument(
'--url', help=f'Document URL'
'--document-url', help=f'Document URL'
)
parser.add_argument(
@ -209,14 +142,6 @@ def main():
'--identifier', '--id', help=f'Document ID'
)
parser.add_argument(
'-l', '--log-level',
type=LogLevel,
default=LogLevel.ERROR,
choices=list(LogLevel),
help=f'Output queue (default: info)'
)
parser.add_argument(
'files', nargs='+',
help=f'File to load'
@ -236,7 +161,7 @@ def main():
copyright_holder=args.copyright_holder,
copyright_year=args.copyright_year,
license=args.license,
url=args.url,
url=args.document_url,
keywords=args.keyword,
)
@ -254,12 +179,9 @@ def main():
)
p = Loader(
pulsar_host=args.pulsar_host,
pulsar_api_key=args.pulsar_api_key,
output_queue=args.output_queue,
url=args.url,
user=args.user,
collection=args.collection,
log_level=args.log_level,
metadata=document,
)

76
trustgraph-cli/scripts/tg-query-document-rag
View file

@ -1,76 +0,0 @@
#!/usr/bin/env python3
"""
Uses the Document RAG service to answer a query
"""
import argparse
import os
from trustgraph.clients.document_rag_client import DocumentRagClient
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_user = 'trustgraph'
default_collection = 'default'
def query(pulsar_host, query, user, collection, pulsar_api_key=None):
rag = DocumentRagClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
resp = rag.request(user=user, collection=collection, query=query)
print(resp)
def main():
parser = argparse.ArgumentParser(
prog='tg-query-document-rag',
description=__doc__,
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
parser.add_argument(
'-q', '--query',
required=True,
help=f'Query to execute',
)
parser.add_argument(
'-u', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-c', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
args = parser.parse_args()
try:
query(
pulsar_host=args.pulsar_host,
pulsar_api_key=args.pulsar_api_key,
query=args.query,
user=args.user,
collection=args.collection,
)
except Exception as e:
print("Exception:", e, flush=True)
main()

77
trustgraph-cli/scripts/tg-query-graph-rag
View file

@ -1,77 +0,0 @@
#!/usr/bin/env python3
"""
Uses the GraphRAG service to answer a query
"""
import argparse
import os
from trustgraph.clients.graph_rag_client import GraphRagClient
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
default_user = 'trustgraph'
default_collection = 'default'
def query(pulsar_host, query, user, collection, pulsar_api_key=None):
rag = GraphRagClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
resp = rag.request(user=user, collection=collection, query=query)
print(resp)
def main():
parser = argparse.ArgumentParser(
prog='tg-graph-query-rag',
description=__doc__,
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
parser.add_argument(
'--pulsar-api-key',
default=default_pulsar_api_key,
help=f'Pulsar API key',
)
parser.add_argument(
'-q', '--query',
required=True,
help=f'Query to execute',
)
parser.add_argument(
'-u', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-c', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
args = parser.parse_args()
try:
query(
pulsar_host=args.pulsar_host,
pulsar_api_key=args.pulsar_api_key,
query=args.query,
user=args.user,
collection=args.collection,
)
except Exception as e:
print("Exception:", e, flush=True)
main()

198
trustgraph-cli/scripts/tg-save-doc-embeds Executable file
View file

@ -0,0 +1,198 @@
#!/usr/bin/env python3
"""
This utility connects to a running TrustGraph through the API and creates
a knowledge core from the data streaming through the processing queues.
For completeness of data, tg-save-kg-core should be initiated before data
loading takes place. The default output format, msgpack should be used.
JSON output format is also available - msgpack produces a more compact
representation, which is also more performant to load.
"""
import aiohttp
import asyncio
import msgpack
import json
import sys
import argparse
import os
import signal
class Running:
def __init__(self): self.running = True
def get(self): return self.running
def stop(self): self.running = False
async def fetch_de(running, queue, user, collection, url):
async with aiohttp.ClientSession() as session:
de_url = f"{url}stream/document-embeddings"
async with session.ws_connect(de_url) as ws:
while running.get():
try:
msg = await asyncio.wait_for(ws.receive(), 1)
except:
continue
if msg.type == aiohttp.WSMsgType.TEXT:
data = msg.json()
if user:
if data["metadata"]["user"] != user:
continue
if collection:
if data["metadata"]["collection"] != collection:
continue
await queue.put([
"de",
{
"m": {
"i": data["metadata"]["id"],
"m": data["metadata"]["metadata"],
"u": data["metadata"]["user"],
"c": data["metadata"]["collection"],
},
"c": [
{
"c": chunk["chunk"],
"v": chunk["vectors"],
}
for chunk in data["chunks"]
]
}
])
if msg.type == aiohttp.WSMsgType.ERROR:
print("Error")
break
de_counts = 0
async def stats(running):
global t_counts
global de_counts
while running.get():
await asyncio.sleep(2)
print(
f"Document embeddings: {de_counts:10d}"
)
async def output(running, queue, path, format):
global t_counts
global de_counts
with open(path, "wb") as f:
while running.get():
try:
msg = await asyncio.wait_for(queue.get(), 0.5)
except:
# Hopefully it's TimeoutError. Annoying to match since
# it changed in 3.11.
continue
if format == "msgpack":
f.write(msgpack.packb(msg, use_bin_type=True))
else:
f.write(json.dumps(msg).encode("utf-8"))
if msg[0] == "de":
de_counts += 1
print("Output file closed")
async def run(running, **args):
q = asyncio.Queue()
de_task = asyncio.create_task(
fetch_de(
running=running,
queue=q, user=args["user"], collection=args["collection"],
url=args["url"] + "api/v1/"
)
)
output_task = asyncio.create_task(
output(
running=running, queue=q,
path=args["output_file"], format=args["format"],
)
)
stats_task = asyncio.create_task(stats(running))
await output_task
await de_task
await stats_task
print("Exiting")
async def main(running):
parser = argparse.ArgumentParser(
prog='tg-save-kg-core',
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(
'-o', '--output-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 filter on (default: no filter)'
)
parser.add_argument(
'--collection',
help=f'Collection ID to filter on (default: no filter)'
)
args = parser.parse_args()
await run(running, **vars(args))
running = Running()
def interrupt(sig, frame):
running.stop()
print('Interrupt')
signal.signal(signal.SIGINT, interrupt)
asyncio.run(main(running))

9
trustgraph-cli/scripts/tg-save-kg-core
View file

@ -57,8 +57,13 @@ async def fetch_ge(running, queue, user, collection, url):
"u": data["metadata"]["user"],
"c": data["metadata"]["collection"],
},
"v": data["vectors"],
"e": data["entity"],
"e": [
{
"e": ent["entity"],
"v": ent["vectors"],
}
for ent in data["entities"]
]
}
])
if msg.type == aiohttp.WSMsgType.ERROR: