trustgraph-flow 0.21.3__tar.gz → 0.21.5__tar.gz

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 0.21.3
+Version: 0.21.5
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Home-page: https://github.com/trustgraph-ai/trustgraph
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.21.3.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.21.5.tar.gz
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 Classifier: Programming Language :: Python :: 3

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/document_rag.py

@@ -18,11 +18,15 @@ DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
 
 class Query:
 
-    def __init__(self, rag, user, collection, verbose):
+    def __init__(
+            self, rag, user, collection, verbose,
+            doc_limit=20
+    ):
         self.rag = rag
         self.user = user
         self.collection = collection
         self.verbose = verbose
+        self.doc_limit = doc_limit
 
     def get_vector(self, query):
 
@@ -44,7 +48,7 @@ class Query:
             print("Get entities...", flush=True)
 
         docs = self.rag.de_client.request(
-            vectors, limit=self.rag.doc_limit
+            vectors, limit=self.doc_limit
         )
 
         if self.verbose:
@@ -93,9 +97,6 @@ class DocumentRag:
         if self.verbose:
             print("Initialising...", flush=True)
 
-        # FIXME: Configurable
-        self.doc_limit = 20
-
         self.de_client = DocumentEmbeddingsClient(
             pulsar_host=pulsar_host,
             subscriber=module + "-de",
@@ -123,13 +124,17 @@ class DocumentRag:
         if self.verbose:
             print("Initialised", flush=True)
 
-    def query(self, query, user="trustgraph", collection="default"):
+    def query(
+            self, query, user="trustgraph", collection="default",
+            doc_limit=20,
+    ):
 
         if self.verbose:
             print("Construct prompt...", flush=True)
 
         q = Query(
-            rag=self, user=user, collection=collection, verbose=self.verbose
+            rag=self, user=user, collection=collection, verbose=self.verbose,
+            doc_limit=doc_limit
         )
 
         docs = q.get_docs(query)

trustgraph-flow-0.21.5/trustgraph/flow_version.py

@@ -0,0 +1 @@
+__version__ = "0.21.5"

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/gateway/document_rag.py

@@ -23,6 +23,7 @@ class DocumentRagRequestor(ServiceRequestor):
             query=body["query"],
             user=body.get("user", "trustgraph"),
             collection=body.get("collection", "default"),
+            doc_limit=int(body.get("doc-limit", 20)),
         )
 
     def from_response(self, message):

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/gateway/graph_rag.py

@@ -23,6 +23,10 @@ class GraphRagRequestor(ServiceRequestor):
             query=body["query"],
             user=body.get("user", "trustgraph"),
             collection=body.get("collection", "default"),
+            entity_limit=int(body.get("entity-limit", 50)),
+            triple_limit=int(body.get("triple-limit", 30)),
+            max_subgraph_size=int(body.get("max-subgraph-size", 1000)),
+            max_path_length=int(body.get("max-path-length", 2)),
         )
 
     def from_response(self, message):

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/gateway/librarian.py

@@ -53,7 +53,10 @@ class LibrarianRequestor(ServiceRequestor):
             response["document"] = serialize_document_package(message.document)
 
         if message.info:
-            response["info"] = serialize_document_info(message.info)
+            response["info"] = [
+                serialize_document_info(v)
+                for v in message.info
+            ]
         
         return response, True
 

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/gateway/serialize.py

@@ -84,6 +84,9 @@ def serialize_document_package(message):
 
     ret = {}
 
+    if message.id:
+        ret["id"] = message.id
+
     if message.metadata:
         ret["metadata"] = serialize_subgraph(message.metdata)
 
@@ -108,8 +111,8 @@ def serialize_document_info(message):
 
     ret = {}
 
-    if message.metadata:
-        ret["metadata"] = serialize_subgraph(message.metdata)
+    if message.id:
+        ret["id"] = message.id
 
     if message.kind:
         ret["kind"] = message.kind
@@ -120,25 +123,45 @@ def serialize_document_info(message):
     if message.collection:
         ret["collection"] = message.collection
 
+    if message.title:
+        ret["title"] = message.title
+
+    if message.comments:
+        ret["comments"] = message.comments
+
+    if message.time:
+        ret["time"] = message.time
+
+    if message.metadata:
+        ret["metadata"] = serialize_subgraph(message.metadata)
+
     return ret
 
 def to_document_package(x):
 
     return DocumentPackage(
-        metadata = to_subgraph(x["metadata"]),
-        document = x.get("document", None),
+        id = x.get("id", None),
         kind = x.get("kind", None),
         user = x.get("user", None),
         collection = x.get("collection", None),
+        title = x.get("title", None),
+        comments = x.get("comments", None),
+        time = x.get("time", None),
+        document = x.get("document", None),
+        metadata = to_subgraph(x["metadata"]),
     )
 
 def to_document_info(x):
 
     return DocumentInfo(
-        metadata = to_subgraph(x["metadata"]),
+        id = x.get("id", None),
         kind = x.get("kind", None),
         user = x.get("user", None),
         collection = x.get("collection", None),
+        title = x.get("title", None),
+        comments = x.get("comments", None),
+        time = x.get("time", None),
+        metadata = to_subgraph(x["metadata"]),
     )
 
 def to_criteria(x):

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/graph_rag.py

@@ -20,11 +20,19 @@ DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
 
 class Query:
 
-    def __init__(self, rag, user, collection, verbose):
+    def __init__(
+            self, rag, user, collection, verbose,
+            entity_limit=50, triple_limit=30, max_subgraph_size=1000,
+            max_path_length=2,
+    ):
         self.rag = rag
         self.user = user
         self.collection = collection
         self.verbose = verbose
+        self.entity_limit = entity_limit
+        self.triple_limit = triple_limit
+        self.max_subgraph_size = max_subgraph_size
+        self.max_path_length = max_path_length
 
     def get_vector(self, query):
 
@@ -47,7 +55,7 @@ class Query:
 
         entities = self.rag.ge_client.request(
             user=self.user, collection=self.collection,
-            vectors=vectors, limit=self.rag.entity_limit,
+            vectors=vectors, limit=self.entity_limit,
         )
 
         entities = [
@@ -79,61 +87,66 @@ class Query:
         self.rag.label_cache[e] = res[0].o.value
         return self.rag.label_cache[e]
 
-    def get_subgraph(self, query):
-
-        entities = self.get_entities(query)
+    def follow_edges(self, ent, subgraph, path_length):
 
-        subgraph = set()
+        # Not needed?
+        if path_length <= 0:
+            return
 
-        if self.verbose:
-            print("Get subgraph...", flush=True)
+        # Stop spanning around if the subgraph is already maxed out
+        if len(subgraph) >= self.max_subgraph_size:
+            return
 
-        for e in entities:
+        res = self.rag.triples_client.request(
+            user=self.user, collection=self.collection,
+            s=ent, p=None, o=None,
+            limit=self.triple_limit
+        )
 
-            res = self.rag.triples_client.request(
-                user=self.user, collection=self.collection,
-                s=e, p=None, o=None,
-                limit=self.rag.query_limit
+        for triple in res:
+            subgraph.add(
+                (triple.s.value, triple.p.value, triple.o.value)
             )
+            if path_length > 1:
+                self.follow_edges(triple.o.value, subgraph, path_length-1)
 
-            for triple in res:
-                subgraph.add(
-                    (triple.s.value, triple.p.value, triple.o.value)
-                )
+        res = self.rag.triples_client.request(
+            user=self.user, collection=self.collection,
+            s=None, p=ent, o=None,
+            limit=self.triple_limit
+        )
 
-            res = self.rag.triples_client.request(
-                user=self.user, collection=self.collection,
-                s=None, p=e, o=None,
-                limit=self.rag.query_limit
+        for triple in res:
+            subgraph.add(
+                (triple.s.value, triple.p.value, triple.o.value)
             )
 
-            for triple in res:
-                subgraph.add(
-                    (triple.s.value, triple.p.value, triple.o.value)
-                )
+        res = self.rag.triples_client.request(
+            user=self.user, collection=self.collection,
+            s=None, p=None, o=ent,
+            limit=self.triple_limit,
+        )
 
-            res = self.rag.triples_client.request(
-                user=self.user, collection=self.collection,
-                s=None, p=None, o=e,
-                limit=self.rag.query_limit,
+        for triple in res:
+            subgraph.add(
+                (triple.s.value, triple.p.value, triple.o.value)
             )
+            if path_length > 1:
+                self.follow_edges(triple.s.value, subgraph, path_length-1)
 
-            for triple in res:
-                subgraph.add(
-                    (triple.s.value, triple.p.value, triple.o.value)
-                )
-
-        subgraph = list(subgraph)
+    def get_subgraph(self, query):
 
-        subgraph = subgraph[0:self.rag.max_subgraph_size]
+        entities = self.get_entities(query)
 
         if self.verbose:
-            print("Subgraph:", flush=True)
-            for edge in subgraph:
-                print(" ", str(edge), flush=True)
+            print("Get subgraph...", flush=True)
 
-        if self.verbose:
-            print("Done.", flush=True)
+        subgraph = set()
+
+        for ent in entities:
+            self.follow_edges(ent, subgraph, self.max_path_length)
+
+        subgraph = list(subgraph)
 
         return subgraph
 
@@ -154,6 +167,16 @@ class Query:
 
             sg2.append((s, p, o))
 
+        sg2 = sg2[0:self.max_subgraph_size]
+
+        if self.verbose:
+            print("Subgraph:", flush=True)
+            for edge in sg2:
+                print(" ", str(edge), flush=True)
+
+        if self.verbose:
+            print("Done.", flush=True)
+
         return sg2
     
 class GraphRag:
@@ -171,9 +194,6 @@ class GraphRag:
             tpl_request_queue=None,
             tpl_response_queue=None,
             verbose=False,
-            entity_limit=50,
-            triple_limit=30,
-            max_subgraph_size=3000,
             module="test",
     ):
 
@@ -230,10 +250,6 @@ class GraphRag:
             subscriber=module + "-emb",
         )
 
-        self.entity_limit=entity_limit
-        self.query_limit=triple_limit
-        self.max_subgraph_size=max_subgraph_size
-
         self.label_cache = {}
 
         self.prompt = PromptClient(
@@ -247,13 +263,20 @@ class GraphRag:
         if self.verbose:
             print("Initialised", flush=True)
 
-    def query(self, query, user="trustgraph", collection="default"):
+    def query(
+            self, query, user="trustgraph", collection="default",
+            entity_limit=50, triple_limit=30, max_subgraph_size=1000,
+            max_path_length=2,
+    ):
 
         if self.verbose:
             print("Construct prompt...", flush=True)
 
         q = Query(
-            rag=self, user=user, collection=collection, verbose=self.verbose
+            rag=self, user=user, collection=collection, verbose=self.verbose,
+            entity_limit=entity_limit, triple_limit=triple_limit,
+            max_subgraph_size=max_subgraph_size,
+            max_path_length=max_path_length,
         )
 
         kg = q.get_labelgraph(query)

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/librarian/librarian.py

@@ -1,4 +1,4 @@
-from .. schema import LibrarianRequest, LibrarianResponse, Error
+from .. schema import LibrarianRequest, LibrarianResponse, Error, Triple
 from .. knowledge import hash
 from .. exceptions import RequestError
 from . table_store import TableStore
@@ -26,7 +26,7 @@ class Librarian:
         self.load_document = load_document
         self.load_text = load_text
 
-    async def add(self, id, document):
+    async def add(self, document):
 
         if document.kind not in (
                 "text/plain", "application/pdf"
@@ -41,9 +41,9 @@ class Librarian:
         self.table_store.add(object_id, document)
 
         if document.kind == "application/pdf":
-            await self.load_document(id, document)
+            await self.load_document(document)
         elif document.kind == "text/plain":
-            await self.load_text(id, document)
+            await self.load_text(document)
 
         print("Add complete", flush=True)
 
@@ -53,6 +53,20 @@ class Librarian:
             info = None,
         )
 
+    async def list(self, user, collection):
+
+        print("list")
+
+        info = self.table_store.list(user, collection)
+
+        print(">>", info)
+
+        return LibrarianResponse(
+            error = None,
+            document = None,
+            info = info,
+        )
+
     def handle_triples(self, m):
         self.table_store.add_triples(m)
 

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/librarian/service.py

@@ -223,11 +223,11 @@ class Processor(ConsumerProducer):
 
             self.librarian.handle_document_embeddings(msg)
 
-    async def load_document(self, id, document):
+    async def load_document(self, document):
 
         doc = Document(
             metadata = Metadata(
-                id = id,
+                id = document.id,
                 metadata = document.metadata,
                 user = document.user,
                 collection = document.collection
@@ -237,14 +237,14 @@ class Processor(ConsumerProducer):
 
         self.document_load.send(None, doc)
 
-    async def load_text(self, id, document):
+    async def load_text(self, document):
 
         text = base64.b64decode(document.document)
         text = text.decode("utf-8")
 
         doc = TextDocument(
             metadata = Metadata(
-                id = id,
+                id = document.id,
                 metadata = document.metadata,
                 user = document.user,
                 collection = document.collection
@@ -259,19 +259,33 @@ class Processor(ConsumerProducer):
         if v.operation is None:
             raise RequestError("Null operation")
 
+        print("op", v.operation)
+
         if v.operation == "add":
             if (
-                    v.id and v.document and v.document.metadata and
+                    v.document and v.document.id and v.document.metadata and
                     v.document.document and v.document.kind
             ):
                 return partial(
                     self.librarian.add,
-                    id = v.id,
                     document = v.document,
                 )
             else:
                 raise RequestError("Invalid call")
 
+        if v.operation == "list":
+            print("list", v)
+            print(v.user)
+            if v.user:
+                return partial(
+                    self.librarian.list,
+                    user = v.user,
+                    collection = v.collection,
+                )
+            else:
+                print("BROK")
+                raise RequestError("Invalid call")
+
         raise RequestError("Invalid operation: " + v.operation)
 
     async def handle(self, msg):
@@ -298,6 +312,7 @@ class Processor(ConsumerProducer):
 
         try:
             resp = await func()
+            print("->", resp)
         except RequestError as e:
             resp = LibrarianResponse(
                 error = Error(
@@ -318,7 +333,7 @@ class Processor(ConsumerProducer):
             await self.send(resp, properties={"id": id})
             return
 
-        print("Send response...", flush=True)
+        print("Send response..!.", flush=True)
 
         await self.send(resp, properties={"id": id})
 

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/librarian/table_store.py

@@ -1,4 +1,5 @@
-from .. schema import LibrarianRequest, LibrarianResponse, Error
+from .. schema import LibrarianRequest, LibrarianResponse
+from .. schema import DocumentInfo, Error, Triple, Value
 from .. knowledge import hash
 from .. exceptions import RequestError
 
@@ -61,7 +62,7 @@ class TableStore:
             CREATE TABLE IF NOT EXISTS document (
                 user text,
                 collection text,
-                id uuid,
+                id text,
                 time timestamp,
                 title text,
                 comments text,
@@ -157,6 +158,20 @@ class TableStore:
             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
         """)
 
+        self.list_document_stmt = self.cassandra.prepare("""
+            SELECT
+                id, kind, user, collection, title, comments, time, metadata
+            FROM document
+            WHERE user = ?
+        """)
+
+        self.list_document_by_collection_stmt = self.cassandra.prepare("""
+            SELECT
+                id, kind, user, collection, title, comments, time, metadata
+            FROM document
+            WHERE user = ? AND collection = ?
+        """)
+
         self.insert_triples_stmt = self.cassandra.prepare("""
             INSERT INTO triples
             (
@@ -192,10 +207,9 @@ class TableStore:
             raise RequestError("Invalid document kind: " + document.kind)
 
         # Create random doc ID
-        doc_id = uuid.uuid4()
         when = int(time.time() * 1000)
 
-        print("Adding", object_id, doc_id)
+        print("Adding", document.id, object_id)
 
         metadata = [
             (
@@ -205,8 +219,6 @@ class TableStore:
             for v in document.metadata
         ]
 
-        # FIXME: doc_id should be the user-supplied ID???
-
         while True:
 
             try:
@@ -214,7 +226,7 @@ class TableStore:
                 resp = self.cassandra.execute(
                     self.insert_document_stmt,
                     (
-                        doc_id, document.user, document.collection,
+                        document.id, document.user, document.collection,
                         document.kind, object_id, when,
                         document.title, document.comments,
                         metadata
@@ -275,6 +287,73 @@ class TableStore:
                 print(f"{e}, retry...", flush=True)
                 time.sleep(1)
 
+    def list(self, user, collection=None):
+
+        print("LIST")
+        while True:
+
+            print("TRY")
+
+            print(self.list_document_stmt)
+            try:
+
+                if collection:
+                    resp = self.cassandra.execute(
+                        self.list_document_by_collection_stmt,
+                        (user, collection)
+                    )
+                else:
+                    resp = self.cassandra.execute(
+                        self.list_document_stmt,
+                        (user,)
+                    )
+                break
+
+                print("OK")
+
+            except Exception as e:
+                print("Exception:", type(e))
+                print(f"{e}, retry...", flush=True)
+                time.sleep(1)
+
+        print("OK2")
+
+        info = [
+            DocumentInfo(
+                id = row[0],
+                kind = row[1],
+                user = row[2],
+                collection = row[3],
+                title = row[4],
+                comments = row[5],
+                time = int(1000 * row[6].timestamp()),
+                metadata = [
+                    Triple(
+                        s=Value(value=m[0], is_uri=m[1]),
+                        p=Value(value=m[2], is_uri=m[3]),
+                        o=Value(value=m[4], is_uri=m[5])
+                    )
+                    for m in row[7]
+                ],
+            )
+            for row in resp
+        ]
+
+        print("OK3")
+
+        print(info[0])
+
+        print(info[0].user)
+        print(info[0].time)
+        print(info[0].kind)
+        print(info[0].collection)
+        print(info[0].title)
+        print(info[0].comments)
+        print(info[0].metadata)
+        print(info[0].metadata)
+
+        return info
+
     def add_graph_embeddings(self, m):
 
         when = int(time.time() * 1000)

{trustgraph-flow-0.21.3 → trustgraph-flow-0.21.5}/trustgraph/retrieval/document_rag/rag.py

@@ -50,6 +50,8 @@ class Processor(ConsumerProducer):
             document_embeddings_response_queue
         )
 
+        doc_limit = params.get("doc_limit", 10)
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,
@@ -79,6 +81,8 @@ class Processor(ConsumerProducer):
             module=module,
         )
 
+        self.doc_limit = doc_limit
+
     async def handle(self, msg):
 
         try:
@@ -90,7 +94,12 @@ class Processor(ConsumerProducer):
 
             print(f"Handling input {id}...", flush=True)
 
-            response = self.rag.query(v.query)
+            if v.doc_limit:
+                doc_limit = v.doc_limit
+            else:
+                doc_limit = self.doc_limit
+
+            response = self.rag.query(v.query, doc_limit=doc_limit)
 
             print("Send response...", flush=True)
             r = DocumentRagResponse(response = response, error=None)
@@ -124,6 +133,13 @@ class Processor(ConsumerProducer):
             default_output_queue,
         )
 
+        parser.add_argument(
+            '-d', '--doc-limit',
+            type=int,
+            default=20,
+            help=f'Default document fetch limit (default: 10)'
+        )
+
         parser.add_argument(
             '--prompt-request-queue',
             default=prompt_request_queue,