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

trustgraph-cli/scripts/tg-dump-msgpack

@@ -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):
 

trustgraph-cli/scripts/tg-graph-show

@@ -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,
         )

trustgraph-cli/scripts/tg-graph-to-turtle

@@ -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:
 

trustgraph-cli/scripts/tg-invoke-agent

@@ -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:

trustgraph-cli/scripts/tg-invoke-document-rag

@@ -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()
+

trustgraph-cli/scripts/tg-invoke-graph-rag

@@ -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()
+

trustgraph-cli/scripts/tg-invoke-llm

@@ -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:

trustgraph-cli/scripts/tg-invoke-prompt

@@ -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:

trustgraph-cli/scripts/tg-load-doc-embeds

@@ -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))
+

trustgraph-cli/scripts/tg-load-kg-core

@@ -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:

trustgraph-cli/scripts/tg-load-pdf

@@ -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,
             )
 

trustgraph-cli/scripts/tg-load-text

@@ -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,
             )
 

trustgraph-cli/scripts/tg-query-document-rag

@@ -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()
-

trustgraph-cli/scripts/tg-query-graph-rag

@@ -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()
-

trustgraph-cli/scripts/tg-save-doc-embeds

@@ -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))
+

trustgraph-cli/scripts/tg-save-kg-core

@@ -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:

trustgraph-cli/setup.py

@@ -34,29 +34,33 @@ setuptools.setup(
     python_requires='>=3.8',
     download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz",
     install_requires=[
-        "trustgraph-base>=0.18,<0.19",
+        "trustgraph-base>=0.21,<0.22",
         "requests",
         "pulsar-client",
+        "aiohttp",
         "rdflib",
         "tabulate",
         "msgpack",
+        "websockets",
     ],
     scripts=[
+        "scripts/tg-dump-msgpack",
         "scripts/tg-graph-show",
         "scripts/tg-graph-to-turtle",
+        "scripts/tg-init-pulsar",
         "scripts/tg-init-pulsar-manager",
+        "scripts/tg-invoke-agent",
+        "scripts/tg-invoke-document-rag",
+        "scripts/tg-invoke-graph-rag",
+        "scripts/tg-invoke-llm",
+        "scripts/tg-invoke-prompt",
+        "scripts/tg-load-kg-core",
+        "scripts/tg-load-doc-embeds",
         "scripts/tg-load-pdf",
         "scripts/tg-load-text",
         "scripts/tg-load-turtle",
-        "scripts/tg-query-document-rag",
-        "scripts/tg-query-graph-rag",
-        "scripts/tg-init-pulsar",
         "scripts/tg-processor-state",
-        "scripts/tg-invoke-agent",
-        "scripts/tg-invoke-prompt",
-        "scripts/tg-invoke-llm",
         "scripts/tg-save-kg-core",
-        "scripts/tg-load-kg-core",
-        "scripts/tg-dump-msgpack",
+        "scripts/tg-save-doc-embeds",
     ]
 )