mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-20 18:51:03 +02:00
Tables for data from queue outputs
This commit is contained in:
Cyber MacGeddon
parent
7fdbfd2d44
commit
b5f51731bd
3 changed files with 310 additions and 30 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
||||
|
|
|
|||
|
|
@ -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<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)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue