Fetch and list mostly there

trustgraph-base/trustgraph/schema/knowledge.py

@@ -31,14 +31,11 @@ class KnowledgeRequest(Record):
     user = String()
 
     # fetch-kg-core, list-kg-cores, delete-kg-core
-    collection = String()
+    id = String()
-
-    # fetch-kg-core, list-kg-cores, delete-kg-core
-    document_id = String()
 
 class KnowledgeResponse(Record):
     error = Error()
-    document_ids = Array(String())
+    ids = Array(String())
     eos = Boolean()     # Indicates end of knowledge core stream
     triples = Triples()
     graph_embeddings = GraphEmbeddings()

trustgraph-flow/trustgraph/cores/knowledge.py

@@ -1,5 +1,5 @@
 
-from .. schema import KnowledgeResponse, Error, Triple
+from .. schema import KnowledgeResponse, Error
 from .. knowledge import hash
 from .. exceptions import RequestError
 from .. tables.knowledge import KnowledgeTableStore
@@ -18,82 +18,7 @@ class KnowledgeManager:
             cassandra_host, cassandra_user, cassandra_password, keyspace
         )
 
-    async def ASDlist_cores(self, request):
+    async def delete_kg_core(self, request):
-
-        if request.document_metadata.kind not in (
-                "text/plain", "application/pdf"
-        ):
-            raise RequestError(
-                "Invalid document kind: " + request.document_metadata.kind
-            )
-
-        if await self.table_store.document_exists(
-                request.document_metadata.user,
-                request.document_metadata.id
-        ):
-            raise RuntimeError("Document already exists")
-
-        # Create object ID for blob
-        object_id = uuid.uuid4()
-
-        print("Add blob...")
-
-        await self.blob_store.add(
-            object_id, base64.b64decode(request.content),
-            request.document_metadata.kind
-        )
-
-        print("Add table...")
-
-        await self.table_store.add_document(
-            request.document_metadata, object_id
-        )
-
-        print("Add complete", flush=True)
-
-        return LibrarianResponse(
-            error = None,
-            document_metadata = None,
-            content = None,
-            document_metadatas = None,
-            processing_metadatas = None,
-        )
-
-    async def fetch_core(self, request):
-
-        print("Removing doc...")
-
-        if not await self.table_store.document_exists(
-                request.user,
-                request.document_id,
-        ):
-            raise RuntimeError("Document does not exist")
-
-        object_id = await self.table_store.get_document_object_id(
-            request.user,
-            request.document_id
-        )
-
-        # Remove blob...
-        await self.blob_store.remove(object_id)
-
-        # Remove doc table row
-        await self.table_store.remove_document(
-            request.user,
-            request.document_id
-        )
-
-        print("Remove complete", flush=True)
-
-        return LibrarianResponse(
-            error = None,
-            document_metadata = None,
-            content = None,
-            document_metadatas = None,
-            processing_metadatas = None,
-        )
-
-    async def update_document(self, request):
 
         print("Updating doc...")
 
@@ -117,129 +42,48 @@ class KnowledgeManager:
             processing_metadatas = None,
         )
 
-    async def get_document_metadata(self, request):
+    async def fetch_kg_core(self, request, flow):
 
-        print("Get doc...")
+        print("Fetch core...")
 
-        doc = await self.table_store.get_document(
+        async def publish(obj):
-            request.user,
+            print(obj)
-            request.document_id
+#            await publisher.send(triples)
-        )
-
-        print("Get complete", flush=True)
-
-        return LibrarianResponse(
-            error = None,
-            document_metadata = doc,
-            content = None,
-            document_metadatas = None,
-            processing_metadatas = None,
-        )
-
-    async def get_document_content(self, request):
-
-        print("Get doc content...")
-
-        object_id = await self.table_store.get_document_object_id(
-            request.user,
-            request.document_id
-        )
-
-        content = await self.blob_store.get(
-            object_id
-        )
-
-        print("Get complete", flush=True)
-
-        return LibrarianResponse(
-            error = None,
-            document_metadata = None,
-            content = base64.b64encode(content),
-            document_metadatas = None,
-            processing_metadatas = None,
-        )
-
-    async def add_processing(self, request):
-
-        print("Add processing")
-
-        if await self.table_store.processing_exists(
-                request.processing_metadata.user,
-                request.processing_metadata.id
-        ):
-            raise RuntimeError("Processing already exists")
-
-        doc = await self.table_store.get_document(
-            request.processing_metadata.user,
-            request.processing_metadata.document_id
-        )
-
-        object_id = await self.table_store.get_document_object_id(
-            request.processing_metadata.user,
-            request.processing_metadata.document_id
-        )
-
-        content = await self.blob_store.get(
-            object_id
-        )
-
-        print("Got content")
-
-        print("Add processing...")
-
-        await self.table_store.add_processing(request.processing_metadata)
-
-        print("Invoke document processing...")
-
-        await self.load_document(
-            document = doc,
-            processing = request.processing_metadata,
-            content = content,
-        )
-
-        print("Add complete", flush=True)
-
-        return LibrarianResponse(
-            error = None,
-            document_metadata = None,
-            content = None,
-            document_metadatas = None,
-            processing_metadatas = None,
-        )
-
-    async def remove_processing(self, request):
-
-        print("Removing processing...")
-
-        if not await self.table_store.processing_exists(
-                request.user,
-                request.processing_id,
-        ):
-            raise RuntimeError("Processing object does not exist")
 
         # Remove doc table row
-        await self.table_store.remove_processing(
+        await self.table_store.get_triples(
             request.user,
-            request.processing_id
+            request.id,
+            publish
         )
 
-        print("Remove complete", flush=True)
+        async def publish_ge(obj):
+            print(obj)
 
-        return LibrarianResponse(
+        # Remove doc table row
-            error = None,
+        await self.table_store.get_graph_embeddings(
-            document_metadata = None,
+            request.user,
-            content = None,
+            request.id,
-            document_metadatas = None,
+            publish
-            processing_metadatas = None,
         )
 
-    async def list_cores(self, request):
+        print("Fetch complete", flush=True)
-
-        ids = await self.table_store.list_cores(request.user)
 
         return KnowledgeResponse(
             error = None,
-            document_ids = ids,
+            ids = None,
+            eos = True,
+            triples = None,
+            graph_embeddings = None,
+        )
+
+    async def list_kg_cores(self, request, flow):
+
+        ids = await self.table_store.list_kg_cores(request.user)
+
+        return KnowledgeResponse(
+            error = None,
+            ids = ids,
             eos = False,
             triples = None,
             graph_embeddings = None

trustgraph-flow/trustgraph/cores/service.py

@@ -181,7 +181,7 @@ class Processor(AsyncProcessor):
 
         print("Document submitted")
 
-    async def process_request(self, v):
+    async def process_request(self, v, flow):
 
         if v.operation is None:
             raise RequestError("Null operation")
@@ -189,15 +189,15 @@ class Processor(AsyncProcessor):
         print("request", v.operation)
 
         impls = {
-            "list-cores": self.knowledge.list_cores,
+            "list-kg-cores": self.knowledge.list_kg_cores,
-            "fetch-core": self.knowledge.fetch_core,
+            "fetch-kg-core": self.knowledge.fetch_kg_core,
-            "delet-core": self.knowledge.delete_core,
+            "delete-kg-core": self.knowledge.delete_kg_core,
         }
 
         if v.operation not in impls:
             raise RequestError(f"Invalid operation: {v.operation}")
 
-        return await impls[v.operation](v)
+        return await impls[v.operation](v, flow)
 
     async def on_knowledge_request(self, msg, consumer, flow):
 
@@ -211,7 +211,7 @@ class Processor(AsyncProcessor):
 
         try:
 
-            resp = await self.process_request(v)
+            resp = await self.process_request(v, flow)
 
             await self.knowledge_response_producer.send(
                 resp, properties={"id": id}

trustgraph-flow/trustgraph/gateway/dispatch/knowledge.py

@@ -0,0 +1,55 @@
+
+import base64
+
+from ... schema import KnowledgeRequest, KnowledgeResponse
+from ... schema import knowledge_request_queue
+from ... schema import knowledge_response_queue
+
+from . requestor import ServiceRequestor
+from . serialize import serialize_graph_embeddings
+from . serialize import serialize_triples
+from . serialize import to_document_metadata, to_processing_metadata
+
+class KnowledgeRequestor(ServiceRequestor):
+    def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
+
+        super(KnowledgeRequestor, self).__init__(
+            pulsar_client=pulsar_client,
+            consumer_name = consumer,
+            subscription = subscriber,
+            request_queue=knowledge_request_queue,
+            response_queue=knowledge_response_queue,
+            request_schema=KnowledgeRequest,
+            response_schema=KnowledgeResponse,
+            timeout=timeout,
+        )
+
+    def to_request(self, body):
+
+        return KnowledgeRequest(
+            operation = body.get("operation", None),
+            user = body.get("user", None),
+            id = body.get("id", None),
+        )
+
+    def from_response(self, message):
+
+        print(message)
+
+        response = {
+            "eos": message.eos
+        }
+
+        if message.ids is not None:
+            response["ids"] = message.ids
+
+        if message.triples:
+            response["triples"] = serialize_triples(message.triples)
+
+        if message.graph_embeddings:
+            response["graph-embeddings"] = serialize_graph_embeddings(
+                message.graph_embeddings
+            )
+        
+        return response, True
+

trustgraph-flow/trustgraph/gateway/dispatch/manager.py

@@ -5,6 +5,7 @@ import uuid
 from . config import ConfigRequestor
 from . flow import FlowRequestor
 from . librarian import LibrarianRequestor
+from . knowledge import KnowledgeRequestor
 
 from . embeddings import EmbeddingsRequestor
 from . agent import AgentRequestor
@@ -44,6 +45,7 @@ global_dispatchers = {
     "config": ConfigRequestor,
     "flow": FlowRequestor,
     "librarian": LibrarianRequestor,
+    "knowledge": KnowledgeRequestor,
 }
 
 sender_dispatchers = {

trustgraph-flow/trustgraph/librarian/librarian.py

@@ -21,7 +21,7 @@ class Librarian:
             minio_host, minio_access_key, minio_secret_key, bucket_name
         )
 
-        self.table_store = TableStore(
+        self.table_store = LibraryTableStore(
             cassandra_host, cassandra_user, cassandra_password, keyspace
         )
 

trustgraph-flow/trustgraph/tables/knowledge.py

@@ -1,4 +1,7 @@
 
+from .. schema import KnowledgeResponse, Triple, Triples, EntityEmbeddings
+from .. schema import Metadata, Value, GraphEmbeddings
+
 from cassandra.cluster import Cluster
 from cassandra.auth import PlainTextAuthProvider
 from ssl import SSLContext, PROTOCOL_TLSv1_2
@@ -61,7 +64,6 @@ class KnowledgeTableStore:
         self.cassandra.execute("""
             CREATE TABLE IF NOT EXISTS triples (
                 user text,
-                collection text,
                 document_id text,
                 id uuid,
                 time timestamp,
@@ -71,7 +73,7 @@ class KnowledgeTableStore:
                 triples list<tuple<
                     text, boolean, text, boolean, text, boolean
                 >>,
-                PRIMARY KEY (user, collection, document_id, id)
+                PRIMARY KEY ((user, document_id), id)
             );
         """);
 
@@ -80,7 +82,6 @@ class KnowledgeTableStore:
         self.cassandra.execute("""
             create table if not exists graph_embeddings (
                 user text,
-                collection text,
                 document_id text,
                 id uuid,
                 time timestamp,
@@ -93,16 +94,20 @@ class KnowledgeTableStore:
                         list<list<double>>
                     >
                 >,
-                PRIMARY KEY (user, collection, document_id, id)
+                PRIMARY KEY ((user, document_id), id)
             );
         """);
 
+        self.cassandra.execute("""
+            CREATE INDEX IF NOT EXISTS graph_embeddings_user ON
+            graph_embeddings ( user );
+        """);
+
         print("document_embeddings table...", flush=True)
 
         self.cassandra.execute("""
             create table if not exists document_embeddings (
                 user text,
-                collection text,
                 document_id text,
                 id uuid,
                 time timestamp,
@@ -115,10 +120,15 @@ class KnowledgeTableStore:
                         list<list<double>>
                     >
                 >,
-                PRIMARY KEY (user, collection, document_id, id)
+                PRIMARY KEY ((user, document_id), id)
             );
         """);
 
+        self.cassandra.execute("""
+            CREATE INDEX IF NOT EXISTS document_embeddings_user ON
+            document_embeddings ( user );
+        """);
+
         print("Cassandra schema OK.", flush=True)
 
     def prepare_statements(self):
@@ -126,28 +136,49 @@ class KnowledgeTableStore:
         self.insert_triples_stmt = self.cassandra.prepare("""
             INSERT INTO triples
             (
-                id, user, collection, document_id, time,
+                id, user, document_id,
-                metadata, triples
+                time, metadata, triples
             )
-            VALUES (?, ?, ?, ?, ?, ?, ?)
+            VALUES (?, ?, ?, ?, ?, ?)
         """)
 
         self.insert_graph_embeddings_stmt = self.cassandra.prepare("""
             INSERT INTO graph_embeddings
             (
-                id, user, collection, document_id, time,
+                id, user, document_id, time, metadata, entity_embeddings
-                metadata, entity_embeddings
             )
-            VALUES (?, ?, ?, ?, ?, ?, ?)
+            VALUES (?, ?, ?, ?, ?, ?)
         """)
 
         self.insert_document_embeddings_stmt = self.cassandra.prepare("""
             INSERT INTO document_embeddings
             (
-                id, user, collection, document_id, time,
+                id, user, document_id, time, metadata, chunks
-                metadata, chunks
             )
-            VALUES (?, ?, ?, ?, ?, ?, ?)
+            VALUES (?, ?, ?, ?, ?, ?)
+        """)
+
+        self.list_cores_stmt = self.cassandra.prepare("""
+            SELECT DISTINCT user, document_id FROM graph_embeddings
+            WHERE user = ?
+        """)
+
+        self.get_triples_stmt = self.cassandra.prepare("""
+            SELECT id, time, metadata, triples
+            FROM triples
+            WHERE user = ? AND document_id = ?
+        """)
+
+        self.get_graph_embeddings_stmt = self.cassandra.prepare("""
+            SELECT id, time, metadata, entity_embeddings
+            FROM graph_embeddings
+            WHERE user = ? AND document_id = ?
+        """)
+
+        self.get_document_embeddings_stmt = self.cassandra.prepare("""
+            SELECT id, time, metadata, chunks
+            FROM document_embeddings
+            WHERE user = ? AND document_id = ?
         """)
 
     async def add_triples(self, m):
@@ -181,7 +212,7 @@ class KnowledgeTableStore:
                     self.insert_triples_stmt,
                     (
                         uuid.uuid4(), m.metadata.user,
-                        m.metadata.collection, m.metadata.id, when,
+                        m.metadata.id, when,
                         metadata, triples,
                     )
                 )
@@ -225,7 +256,7 @@ class KnowledgeTableStore:
                     self.insert_graph_embeddings_stmt,
                     (
                         uuid.uuid4(), m.metadata.user,
-                        m.metadata.collection, m.metadata.id, when,
+                        m.metadata.id, when,
                         metadata, entities,
                     )
                 )
@@ -269,7 +300,7 @@ class KnowledgeTableStore:
                     self.insert_document_embeddings_stmt,
                     (
                         uuid.uuid4(), m.metadata.user,
-                        m.metadata.collection, m.metadata.id, when,
+                        m.metadata.id, when,
                         metadata, chunks,
                     )
                 )
@@ -282,4 +313,157 @@ class KnowledgeTableStore:
                 print(f"{e}, retry...", flush=True)
                 await asyncio.sleep(1)
 
-        
+    async def list_kg_cores(self, user):
+
+        print("List kg cores...")
+
+        while True:
+
+            try:
+
+                resp = self.cassandra.execute(
+                    self.list_cores_stmt,
+                    (user,)
+                )
+
+                break
+
+            except Exception as e:
+                print("Exception:", type(e))
+                print(f"{e}, retry...", flush=True)
+                await asyncio.sleep(1)
+
+
+        lst = [
+            row[1]
+            for row in resp
+        ]
+
+        print("Done")
+
+        return lst
+
+    async def get_triples(self, user, document_id, receiver):
+
+        print("Get triples...")
+
+        while True:
+
+            try:
+
+                print("Executing!")
+                print(self.get_triples_stmt)
+                print(user, document_id)
+
+                resp = self.cassandra.execute(
+                    self.get_triples_stmt,
+                    (user, document_id)
+                )
+                print("HERE")
+
+                print("Executed")
+
+                break
+
+            except Exception as e:
+                print("Exception:", type(e))
+                print(f"{e}, retry...", flush=True)
+                await asyncio.sleep(1)
+
+        print("Unpack...")
+        for row in resp:
+
+            print("ROW")
+            print(row)
+
+            if row[2]:
+                metadata = [
+                    Triple(
+                        s = Value(value = elt[0], is_uri = elt[1]),
+                        p = Value(value = elt[2], is_uri = elt[3]),
+                        o = Value(value = elt[4], is_uri = elt[5]),
+                    )
+                    for elt in row[2]
+                ]
+            else:
+                metadata = []
+
+            triples = [
+                Triple(
+                    s = Value(value = elt[0], is_uri = elt[1]),
+                    p = Value(value = elt[2], is_uri = elt[3]),
+                    o = Value(value = elt[4], is_uri = elt[5]),
+                )
+                for elt in row[3]
+            ]
+
+            await receiver(
+                Triples(
+                    metadata = Metadata(
+                        id = document_id,
+                        user = user,
+                        metadata = metadata,
+                    ),
+                    triples = triples
+                )
+            )
+
+        print("Done")
+
+    async def get_graph_embeddings(self, user, document_id, receiver):
+
+        print("Get GE...")
+
+        while True:
+
+            try:
+
+                resp = self.cassandra.execute(
+                    self.get_graph_embeddings_stmt,
+                    (user, document_id)
+                )
+
+                break
+
+            except Exception as e:
+                print("Exception:", type(e))
+                print(f"{e}, retry...", flush=True)
+                await asyncio.sleep(1)
+
+        for row in resp:
+
+            print("MD...")
+            if row[2]:
+                metadata = [
+                    Triple(
+                        s = Value(value = elt[0], is_uri = elt[1]),
+                        p = Value(value = elt[2], is_uri = elt[3]),
+                        o = Value(value = elt[4], is_uri = elt[5]),
+                    )
+                    for elt in row[2]
+                ]
+            else:
+                metadata = []
+
+            print("EE...")
+            entities = [
+                EntityEmbeddings(
+                    entity = Value(value = ent[0][0], is_uri = ent[0][1]),
+                    vectors = ent[1]
+                )
+                for ent in row[3]
+            ]
+
+            await receiver(
+                GraphEmbeddings(
+                    metadata = Metadata(
+                        id = document_id,
+                        user = user,
+                        metadata = metadata,
+                    ),
+                    entities = entities
+                )
+            )                    
+
+        print("Done")
+