Feature / collections (#96)

* Update schema defs for source -> metadata
* Migrate to use metadata part of schema, also add metadata to triples & vecs
* Add user/collection metadata to query
* Use user/collection in RAG
* Write and query working on triples

trustgraph-cli/scripts/tg-graph-show

@@ -9,12 +9,17 @@ import os
 from trustgraph.clients.triples_query_client import TriplesQueryClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_user = 'trustgraph'
+default_collection = 'default'
 
-def show_graph(pulsar):
+def show_graph(pulsar, user, collection):
 
     tq = TriplesQueryClient(pulsar_host=pulsar)
 
-    rows = tq.request(None, None, None, limit=10_000_000)
+    rows = tq.request(
+        user=user, collection=collection,
+        s=None, p=None, o=None, limit=10_000_000
+    )
 
     for row in rows:
         print(row.s.value, row.p.value, row.o.value)
@@ -32,11 +37,26 @@ def main():
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
 
+    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:
 
-        show_graph(args.pulsar_host)
+        show_graph(
+            pulsar=args.pulsar_host, user=args.user,
+            collection=args.collection,
+        )
 
     except Exception as e:
 

trustgraph-cli/scripts/tg-load-pdf

@@ -6,7 +6,7 @@ Loads a PDF document into TrustGraph processing.
 
 import pulsar
 from pulsar.schema import JsonSchema
-from trustgraph.schema import Document, Source, document_ingest_queue
+from trustgraph.schema import Document, document_ingest_queue, Metadata
 import base64
 import hashlib
 import argparse
@@ -15,12 +15,17 @@ import time
 
 from trustgraph.log_level import LogLevel
 
+default_user = 'trustgraph'
+default_collection = 'default'
+
 class Loader:
 
     def __init__(
             self,
             pulsar_host,
             output_queue,
+            user,
+            collection,
             log_level,
     ):
 
@@ -35,6 +40,9 @@ class Loader:
             chunking_enabled=True,
         )
 
+        self.user = user
+        self.collection = collection
+
     def load(self, files):
 
         for file in files:
@@ -50,10 +58,12 @@ class Loader:
             id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8]
 
             r = Document(
-                source=Source(
+                metadata=Metadata(
                     source=path,
                     title=path,
                     id=id,
+                    user=self.user,
+                    collection=self.collection,
                 ),
                 data=base64.b64encode(data),
             )
@@ -90,6 +100,18 @@ def main():
         help=f'Output queue (default: {default_output_queue})'
     )
 
+    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})'
+    )
+
     parser.add_argument(
         '-l', '--log-level',
         type=LogLevel,
@@ -112,6 +134,8 @@ def main():
             p = Loader(
                 pulsar_host=args.pulsar_host,
                 output_queue=args.output_queue,
+                user=args.user,
+                collection=args.collection,
                 log_level=args.log_level,
             )
 

trustgraph-cli/scripts/tg-load-text

@@ -6,7 +6,7 @@ Loads a text document into TrustGraph processing.
 
 import pulsar
 from pulsar.schema import JsonSchema
-from trustgraph.schema import TextDocument, Source, text_ingest_queue
+from trustgraph.schema import TextDocument, text_ingest_queue, Metadata
 import base64
 import hashlib
 import argparse
@@ -15,12 +15,17 @@ import time
 
 from trustgraph.log_level import LogLevel
 
+default_user = 'trustgraph'
+default_collection = 'default'
+
 class Loader:
 
     def __init__(
             self,
             pulsar_host,
             output_queue,
+            user,
+            collection,
             log_level,
     ):
 
@@ -35,6 +40,9 @@ class Loader:
             chunking_enabled=True,
         )
 
+        self.user = user
+        self.collection = collection
+
     def load(self, files):
 
         for file in files:
@@ -50,10 +58,12 @@ class Loader:
             id = hashlib.sha256(path.encode("utf-8")).hexdigest()[0:8]
 
             r = TextDocument(
-                source=Source(
+                metadata=Metadata(
                     source=path,
                     title=path,
                     id=id,
+                    user=self.user,
+                    collection=self.collection,
                 ),
                 text=data,
             )
@@ -90,6 +100,18 @@ def main():
         help=f'Output queue (default: {default_output_queue})'
     )
 
+    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})'
+    )
+
     parser.add_argument(
         '-l', '--log-level',
         type=LogLevel,
@@ -112,6 +134,8 @@ def main():
             p = Loader(
                 pulsar_host=args.pulsar_host,
                 output_queue=args.output_queue,
+                user=args.user,
+                collection=args.collection,
                 log_level=args.log_level,
             )
 

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

@@ -9,17 +9,19 @@ import os
 from trustgraph.clients.document_rag_client import DocumentRagClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_user = 'trustgraph'
+default_collection = 'default'
 
-def query(pulsar, query):
+def query(pulsar_host, query, user, collection):
 
     rag = DocumentRagClient(pulsar_host=pulsar)
-    resp = rag.request(query)
+    resp = rag.request(user=user, collection=collection, query=query)
     print(resp)
 
 def main():
 
     parser = argparse.ArgumentParser(
-        prog='graph-show',
+        prog='tg-query-document-rag',
         description=__doc__,
     )
 
@@ -35,11 +37,28 @@ def main():
         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(args.pulsar_host, args.query)
+        query(
+            pulsar_host=args.pulsar_host,
+            query=args.query,
+            user=args.user,
+            collection=args.collection,
+        )
 
     except Exception as e:
 

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

@@ -9,17 +9,19 @@ import os
 from trustgraph.clients.graph_rag_client import GraphRagClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_user = 'trustgraph'
+default_collection = 'default'
 
-def query(pulsar, query):
+def query(pulsar_host, query, user, collection):
 
-    rag = GraphRagClient(pulsar_host=pulsar)
-    resp = rag.request(query)
+    rag = GraphRagClient(pulsar_host=pulsar_host)
+    resp = rag.request(user=user, collection=collection, query=query)
     print(resp)
 
 def main():
 
     parser = argparse.ArgumentParser(
-        prog='graph-show',
+        prog='tg-graph-query-rag',
         description=__doc__,
     )
 
@@ -35,11 +37,28 @@ def main():
         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(args.pulsar_host, args.query)
+        query(
+            pulsar_host=args.pulsar_host,
+            query=args.query,
+            user=args.user,
+            collection=args.collection,
+        )
 
     except Exception as e: