diff --git a/trustgraph-flow/trustgraph/librarian/librarian.py b/trustgraph-flow/trustgraph/librarian/librarian.py index 08010f3f..84ab3793 100644 --- a/trustgraph-flow/trustgraph/librarian/librarian.py +++ b/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) diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index ade4ca38..4ec0dc0b 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/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() - - def receive_triples(self): - - print("Receive triples!") - - while self.running: - try: - msg = self.triples_sub.get(timeout=1) - except queue.Empty: - print("Tick") - continue - - print(msg) - - print("BYE") + self.graph_embeddings_reader.start() + self.document_embeddings_reader.start() 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"): - self.triples_load.stop() - self.triples_load.join() + if hasattr(self, "triples_sub"): + self.triples_sub.unsubscribe_all("x") + + 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): diff --git a/trustgraph-flow/trustgraph/librarian/table_store.py b/trustgraph-flow/trustgraph/librarian/table_store.py index 99e0d845..7c033789 100644 --- a/trustgraph-flow/trustgraph/librarian/table_store.py +++ b/trustgraph-flow/trustgraph/librarian/table_store.py @@ -36,11 +36,7 @@ class TableStore: self.ensure_cassandra_schema() - self.insert_document_stmt = self.cassandra.prepare(""" - insert into document - (id, user, collection, kind, object_id, metadata) - values (?, ?, ?, ?, ?, ?) - """) + self.prepare_statements() 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>, + triples list>, + 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>, + entity_embeddings list< + tuple< + tuple, + list> + > + >, + 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>, + chunks list< + tuple< + blob, + list> + > + >, + 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) + +