Tables for data from queue outputs

trustgraph-flow/trustgraph/librarian/librarian.py

@@ -53,4 +53,22 @@ class Librarian:
             info = None,
         )
 
+    def handle_triples(self, m):
+        self.table_store.add_triples(m)
+
+    def handle_graph_embeddings(self, m):
+        self.table_store.add_graph_embeddings(m)
+
+    def handle_document_embeddings(self, m):
+        self.table_store.add_document_embeddings(m)
+        
+
+    def handle_triples(self, m):
+        self.table_store.add_triples(m)
+
+    def handle_graph_embeddings(self, m):
+        self.table_store.add_graph_embeddings(m)
+
+    def handle_document_embeddings(self, m):
+        self.table_store.add_document_embeddings(m)
 

trustgraph-flow/trustgraph/librarian/service.py

@@ -103,14 +103,34 @@ class Processor(ConsumerProducer):
             listener=self.pulsar_listener,
         )
 
-        self.triples_load = Subscriber(
+        self.triples_brk = Subscriber(
             self.pulsar_host, triples_store_queue,
             "librarian", "librarian",
             schema=JsonSchema(Triples),
             listener=self.pulsar_listener,
         )
+        self.graph_embeddings_brk = Subscriber(
+            self.pulsar_host, graph_embeddings_store_queue,
+            "librarian", "librarian",
+            schema=JsonSchema(GraphEmbeddings),
+            listener=self.pulsar_listener,
+        )
+        self.document_embeddings_brk = Subscriber(
+            self.pulsar_host, document_embeddings_store_queue,
+            "librarian", "librarian",
+            schema=JsonSchema(DocumentEmbeddings),
+            listener=self.pulsar_listener,
+        )
 
-        self.triples_reader = threading.Thread(target=self.receive_triples)
+        self.triples_reader = threading.Thread(
+            target=self.receive_triples
+        )
+        self.graph_embeddings_reader = threading.Thread(
+            target=self.receive_graph_embeddings
+        )
+        self.document_embeddings_reader = threading.Thread(
+            target=self.receive_document_embeddings
+        )
 
         self.librarian = Librarian(
             cassandra_host = cassandra_host.split(","),
@@ -131,34 +151,23 @@ class Processor(ConsumerProducer):
         
         self.document_load.start()
         self.text_load.start()
-        self.triples_load.start()
 
-        self.triples_sub = self.triples_load.subscribe_all("x")
+        self.triples_brk.start()
+        self.graph_embeddings_brk.start()
+        self.document_embeddings_brk.start()
+
+        self.triples_sub = self.triples_brk.subscribe_all("x")
+        self.graph_embeddings_sub = self.graph_embeddings_brk.subscribe_all("x")
+        self.document_embeddings_sub = self.document_embeddings_brk.subscribe_all("x")
 
         self.triples_reader.start()
-
+        self.graph_embeddings_reader.start()
-    def receive_triples(self):
+        self.document_embeddings_reader.start()
-
-        print("Receive triples!")
-
-        while self.running:
-            try:
-                msg = self.triples_sub.get(timeout=1)
-            except queue.Empty:
-                print("Tick")
-                continue
-
-            print(msg)
-
-        print("BYE")
 
     def __del__(self):
 
         self.running = False
 
-        if hasattr(self, "triples_sub"):
-            self.triples_sub.unsubscribe_all("x")
-
         if hasattr(self, "document_load"):
             self.document_load.stop()
             self.document_load.join()
@@ -167,9 +176,56 @@ class Processor(ConsumerProducer):
             self.text_load.stop()
             self.text_load.join()
 
-        if hasattr(self, "triples_load"):
+        if hasattr(self, "triples_sub"):
-            self.triples_load.stop()
+            self.triples_sub.unsubscribe_all("x")
-            self.triples_load.join()
+
+        if hasattr(self, "graph_embeddings_sub"):
+            self.graph_embeddings_sub.unsubscribe_all("x")
+
+        if hasattr(self, "document_embeddings_sub"):
+            self.document_embeddings_sub.unsubscribe_all("x")
+
+        if hasattr(self, "triples_brk"):
+            self.triples_brk.stop()
+            self.triples_brk.join()
+
+        if hasattr(self, "graph_embeddings_brk"):
+            self.graph_embeddings_brk.stop()
+            self.graph_embeddings_brk.join()
+
+        if hasattr(self, "document_embeddings_brk"):
+            self.document_embeddings_brk.stop()
+            self.document_embeddings_brk.join()
+
+    def receive_triples(self):
+
+        while self.running:
+            try:
+                msg = self.triples_sub.get(timeout=1)
+            except queue.Empty:
+                continue
+
+            self.librarian.handle_triples(msg)
+
+    def receive_graph_embeddings(self):
+
+        while self.running:
+            try:
+                msg = self.graph_embeddings_sub.get(timeout=1)
+            except queue.Empty:
+                continue
+
+            self.librarian.handle_graph_embeddings(msg)
+
+    def receive_document_embeddings(self):
+
+        while self.running:
+            try:
+                msg = self.document_embeddings_sub.get(timeout=1)
+            except queue.Empty:
+                continue
+
+            self.librarian.handle_document_embeddings(msg)
 
     async def load_document(self, id, document):
 

trustgraph-flow/trustgraph/librarian/table_store.py

@@ -36,11 +36,7 @@ class TableStore:
 
         self.ensure_cassandra_schema()
 
-        self.insert_document_stmt = self.cassandra.prepare("""
+        self.prepare_statements()
-            insert into document
-            (id, user, collection, kind, object_id, metadata)
-            values (?, ?, ?, ?, ?, ?)
-        """)
 
     def ensure_cassandra_schema(self):
 
@@ -82,8 +78,94 @@ class TableStore:
             on document ( object_id)
         """);
 
+        print("triples table...", flush=True)
+
+        self.cassandra.execute("""
+            create table if not exists triples (
+                user text,
+                collection text,
+                document_id text,
+                id uuid,
+                metadata list<tuple<
+                    text, boolean, text, boolean, text, boolean
+                >>,
+                triples list<tuple<
+                    text, boolean, text, boolean, text, boolean
+                >>,
+                PRIMARY KEY (user, collection, document_id, id)
+            );
+        """);
+
+        print("graph_embeddings table...", flush=True)
+
+        self.cassandra.execute("""
+            create table if not exists graph_embeddings (
+                user text,
+                collection text,
+                document_id text,
+                id uuid,
+                metadata list<tuple<
+                    text, boolean, text, boolean, text, boolean
+                >>,
+                entity_embeddings list<
+                    tuple<
+                        tuple<text, boolean>,
+                        list<list<double>>
+                    >
+                >,
+                PRIMARY KEY (user, collection, document_id, id)
+            );
+        """);
+
+        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,
+                metadata list<tuple<
+                    text, boolean, text, boolean, text, boolean
+                >>,
+                chunks list<
+                    tuple<
+                        blob,
+                        list<list<double>>
+                    >
+                >,
+                PRIMARY KEY (user, collection, document_id, id)
+            );
+        """);
+
         print("Cassandra schema OK.", flush=True)
 
+    def prepare_statements(self):
+
+        self.insert_document_stmt = self.cassandra.prepare("""
+            insert into document
+            (id, user, collection, kind, object_id, metadata)
+            values (?, ?, ?, ?, ?, ?)
+        """)
+
+        self.insert_triples_stmt = self.cassandra.prepare("""
+            insert into triples
+            (id, user, collection, document_id, metadata, triples)
+            values (?, ?, ?, ?, ?, ?)
+        """)
+
+        self.insert_graph_embeddings_stmt = self.cassandra.prepare("""
+            insert into graph_embeddings
+            (id, user, collection, document_id, metadata, entity_embeddings)
+            values (?, ?, ?, ?, ?, ?)
+        """)
+
+        self.insert_document_embeddings_stmt = self.cassandra.prepare("""
+            insert into document_embeddings
+            (id, user, collection, document_id, metadata, chunks)
+            values (?, ?, ?, ?, ?, ?)
+        """)
+
     def add(self, object_id, document):
 
         if document.kind not in (
@@ -126,6 +208,130 @@ class TableStore:
 
         print("Add complete", flush=True)
 
+    def add_triples(self, m):
 
+        if m.metadata.metadata:
+            metadata = [
+                (
+                    v.s.value, v.s.is_uri, v.p.value, v.p.is_uri,
+                    v.o.value, v.o.is_uri
+                )
+                for v in m.metadata.metadata
+            ]
+        else:
+            metadata = []
 
+        triples = [
+            (
+                v.s.value, v.s.is_uri, v.p.value, v.p.is_uri,
+                v.o.value, v.o.is_uri
+            )
+            for v in m.triples
+        ]
 
+        while True:
+
+            try:
+
+                resp = self.cassandra.execute(
+                    self.insert_triples_stmt,
+                    (
+                        uuid.uuid4(), m.metadata.user,
+                        m.metadata.collection, m.metadata.id, 
+                        metadata, triples,
+                    )
+                )
+
+                break
+
+            except Exception as e:
+
+                print("Exception:", type(e))
+                print(f"{e}, retry...", flush=True)
+                time.sleep(1)
+
+    def add_graph_embeddings(self, m):
+
+        if m.metadata.metadata:
+            metadata = [
+                (
+                    v.s.value, v.s.is_uri, v.p.value, v.p.is_uri,
+                    v.o.value, v.o.is_uri
+                )
+                for v in m.metadata.metadata
+            ]
+        else:
+            metadata = []
+
+        entities = [
+            (
+                (v.entity.value, v.entity.is_uri),
+                v.vectors
+            )
+            for v in m.entities
+        ]
+
+        while True:
+
+            try:
+
+                resp = self.cassandra.execute(
+                    self.insert_graph_embeddings_stmt,
+                    (
+                        uuid.uuid4(), m.metadata.user,
+                        m.metadata.collection, m.metadata.id, 
+                        metadata, entities,
+                    )
+                )
+
+                break
+
+            except Exception as e:
+
+                print("Exception:", type(e))
+                print(f"{e}, retry...", flush=True)
+                time.sleep(1)
+
+    def add_document_embeddings(self, m):
+
+        if m.metadata.metadata:
+            metadata = [
+                (
+                    v.s.value, v.s.is_uri, v.p.value, v.p.is_uri,
+                    v.o.value, v.o.is_uri
+                )
+                for v in m.metadata.metadata
+            ]
+        else:
+            metadata = []
+
+        chunks = [
+            (
+                v.chunk,
+                v.vectors,
+            )
+            for v in m.chunks
+        ]
+
+        while True:
+
+            try:
+
+                resp = self.cassandra.execute(
+                    self.insert_document_embeddings_stmt,
+                    (
+                        uuid.uuid4(), m.metadata.user,
+                        m.metadata.collection, m.metadata.id, 
+                        metadata, chunks,
+                    )
+                )
+
+                break
+
+            except Exception as e:
+
+                print("Exception:", type(e))
+                print(f"{e}, retry...", flush=True)
+                time.sleep(1)
+
+