From 0393c7636a04fe4d1d89b7170cddf6bdb0d108e4 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 6 May 2025 17:49:28 +0100 Subject: [PATCH] Write knowledge core elements to Cassandra --- trustgraph-flow/scripts/kg-store | 6 + trustgraph-flow/setup.py | 1 + .../trustgraph/librarian/table_store.py | 227 -------------- .../trustgraph/storage/knowledge/__init__.py | 3 + .../trustgraph/storage/knowledge/__main__.py | 5 + .../trustgraph/storage/knowledge/store.py | 78 +++++ .../storage/knowledge/table_store.py | 285 ++++++++++++++++++ 7 files changed, 378 insertions(+), 227 deletions(-) create mode 100644 trustgraph-flow/scripts/kg-store create mode 100644 trustgraph-flow/trustgraph/storage/knowledge/__init__.py create mode 100644 trustgraph-flow/trustgraph/storage/knowledge/__main__.py create mode 100644 trustgraph-flow/trustgraph/storage/knowledge/store.py create mode 100644 trustgraph-flow/trustgraph/storage/knowledge/table_store.py diff --git a/trustgraph-flow/scripts/kg-store b/trustgraph-flow/scripts/kg-store new file mode 100644 index 00000000..1a5ba9ef --- /dev/null +++ b/trustgraph-flow/scripts/kg-store @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.storage.knowledge import run + +run() + diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index de4bf95c..05a72a55 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -95,6 +95,7 @@ setuptools.setup( "scripts/kg-extract-definitions", "scripts/kg-extract-relationships", "scripts/kg-extract-topics", + "scripts/kg-store", "scripts/librarian", "scripts/metering", "scripts/object-extract-row", diff --git a/trustgraph-flow/trustgraph/librarian/table_store.py b/trustgraph-flow/trustgraph/librarian/table_store.py index 0646f50b..f6de0a74 100644 --- a/trustgraph-flow/trustgraph/librarian/table_store.py +++ b/trustgraph-flow/trustgraph/librarian/table_store.py @@ -104,71 +104,6 @@ class TableStore: ); """); - return - - print("triples table...", flush=True) - - self.cassandra.execute(""" - CREATE TABLE IF NOT EXISTS triples ( - user text, - collection text, - document_id text, - id uuid, - time timestamp, - 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, - time timestamp, - 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, - time timestamp, - metadata list>, - chunks list< - tuple< - blob, - list> - > - >, - PRIMARY KEY (user, collection, document_id, id) - ); - """); - print("Cassandra schema OK.", flush=True) def prepare_statements(self): @@ -252,35 +187,6 @@ class TableStore: WHERE user = ? """) - return - - self.insert_triples_stmt = self.cassandra.prepare(""" - INSERT INTO triples - ( - id, user, collection, document_id, time, - metadata, triples - ) - VALUES (?, ?, ?, ?, ?, ?, ?) - """) - - self.insert_graph_embeddings_stmt = self.cassandra.prepare(""" - INSERT INTO graph_embeddings - ( - id, user, collection, document_id, time, - metadata, entity_embeddings - ) - VALUES (?, ?, ?, ?, ?, ?, ?) - """) - - self.insert_document_embeddings_stmt = self.cassandra.prepare(""" - INSERT INTO document_embeddings - ( - id, user, collection, document_id, time, - metadata, chunks - ) - VALUES (?, ?, ?, ?, ?, ?, ?) - """) - async def document_exists(self, user, id): resp = self.cassandra.execute( @@ -391,50 +297,6 @@ class TableStore: print("Delete complete", flush=True) - async def add_triples(self, m): - - when = int(time.time() * 1000) - - 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, when, - metadata, triples, - ) - ) - - break - - except Exception as e: - - print("Exception:", type(e)) - print(f"{e}, retry...", flush=True) - await asyncio.sleep(1) - async def list_documents(self, user): print("List documents...") @@ -661,92 +523,3 @@ class TableStore: return lst - async def add_graph_embeddings(self, m): - - when = int(time.time() * 1000) - - 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, when, - metadata, entities, - ) - ) - - break - - except Exception as e: - - print("Exception:", type(e)) - print(f"{e}, retry...", flush=True) - await asyncio.sleep(1) - - async def add_document_embeddings(self, m): - - when = int(time.time() * 1000) - - 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, when, - metadata, chunks, - ) - ) - - break - - except Exception as e: - - print("Exception:", type(e)) - print(f"{e}, retry...", flush=True) - await asyncio.sleep(1) - - diff --git a/trustgraph-flow/trustgraph/storage/knowledge/__init__.py b/trustgraph-flow/trustgraph/storage/knowledge/__init__.py new file mode 100644 index 00000000..ff60c5fa --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/knowledge/__init__.py @@ -0,0 +1,3 @@ + +from . store import run + diff --git a/trustgraph-flow/trustgraph/storage/knowledge/__main__.py b/trustgraph-flow/trustgraph/storage/knowledge/__main__.py new file mode 100644 index 00000000..92825a02 --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/knowledge/__main__.py @@ -0,0 +1,5 @@ + +from . store import run + +if __name__ == '__main__': + run() diff --git a/trustgraph-flow/trustgraph/storage/knowledge/store.py b/trustgraph-flow/trustgraph/storage/knowledge/store.py new file mode 100644 index 00000000..4470f5a5 --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/knowledge/store.py @@ -0,0 +1,78 @@ + +""" +Stores knowledge-cores in Cassandra +""" + +import json +import urllib.parse + +from ... schema import Triples, GraphEmbeddings +from ... base import FlowProcessor, ConsumerSpec + +from . table_store import TableStore + +default_ident = "kg-store" + +default_cassandra_host = "cassandra" +keyspace = "knowledge" + +class Processor(FlowProcessor): + + def __init__(self, **params): + + id = params.get("id") + + cassandra_host = params.get("cassandra_host", default_cassandra_host) + cassandra_user = params.get("cassandra_user") + cassandra_password = params.get("cassandra_password") + + super(Processor, self).__init__( + **params | { + "id": id, + "cassandra_host": cassandra_host, + "cassandra_user": cassandra_user, + } + ) + + self.register_specification( + ConsumerSpec( + name = "triples-input", + schema = Triples, + handler = self.on_triples + ) + ) + + self.register_specification( + ConsumerSpec( + name = "graph-embeddings-input", + schema = GraphEmbeddings, + handler = self.on_graph_embeddings + ) + ) + + self.table_store = TableStore( + cassandra_host = cassandra_host.split(","), + cassandra_user = cassandra_user, + cassandra_password = cassandra_password, + keyspace = keyspace, + ) + + async def on_triples(self, msg, consumer, flow): + + v = msg.value() + await self.table_store.add_triples(v) + + async def on_graph_embeddings(self, msg, consumer, flow): + + v = msg.value() + await self.table_store.add_graph_embeddings(v) + + @staticmethod + def add_args(parser): + + FlowProcessor.add_args(parser) + +def run(): + + Processor.launch(default_ident, __doc__) + diff --git a/trustgraph-flow/trustgraph/storage/knowledge/table_store.py b/trustgraph-flow/trustgraph/storage/knowledge/table_store.py new file mode 100644 index 00000000..8a75e622 --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/knowledge/table_store.py @@ -0,0 +1,285 @@ + +from cassandra.cluster import Cluster +from cassandra.auth import PlainTextAuthProvider +from ssl import SSLContext, PROTOCOL_TLSv1_2 + +import uuid +import time +import asyncio + +class TableStore: + + def __init__( + self, + cassandra_host, cassandra_user, cassandra_password, keyspace, + ): + + self.keyspace = keyspace + + print("Connecting to Cassandra...", flush=True) + + if cassandra_user and cassandra_password: + ssl_context = SSLContext(PROTOCOL_TLSv1_2) + auth_provider = PlainTextAuthProvider( + username=cassandra_user, password=cassandra_password + ) + self.cluster = Cluster( + cassandra_host, + auth_provider=auth_provider, + ssl_context=ssl_context + ) + else: + self.cluster = Cluster(cassandra_host) + + self.cassandra = self.cluster.connect() + + print("Connected.", flush=True) + + self.ensure_cassandra_schema() + + self.prepare_statements() + + def ensure_cassandra_schema(self): + + print("Ensure Cassandra schema...", flush=True) + + print("Keyspace...", flush=True) + + # FIXME: Replication factor should be configurable + self.cassandra.execute(f""" + create keyspace if not exists {self.keyspace} + with replication = {{ + 'class' : 'SimpleStrategy', + 'replication_factor' : 1 + }}; + """); + + self.cassandra.set_keyspace(self.keyspace) + + print("triples table...", flush=True) + + self.cassandra.execute(""" + CREATE TABLE IF NOT EXISTS triples ( + user text, + collection text, + document_id text, + id uuid, + time timestamp, + 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, + time timestamp, + 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, + time timestamp, + 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_triples_stmt = self.cassandra.prepare(""" + INSERT INTO triples + ( + id, user, collection, document_id, time, + metadata, triples + ) + VALUES (?, ?, ?, ?, ?, ?, ?) + """) + + self.insert_graph_embeddings_stmt = self.cassandra.prepare(""" + INSERT INTO graph_embeddings + ( + id, user, collection, document_id, time, + metadata, entity_embeddings + ) + VALUES (?, ?, ?, ?, ?, ?, ?) + """) + + self.insert_document_embeddings_stmt = self.cassandra.prepare(""" + INSERT INTO document_embeddings + ( + id, user, collection, document_id, time, + metadata, chunks + ) + VALUES (?, ?, ?, ?, ?, ?, ?) + """) + + async def add_triples(self, m): + + when = int(time.time() * 1000) + + 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, when, + metadata, triples, + ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + async def add_graph_embeddings(self, m): + + when = int(time.time() * 1000) + + 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, when, + metadata, entities, + ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + + async def add_document_embeddings(self, m): + + when = int(time.time() * 1000) + + 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, when, + metadata, chunks, + ) + ) + + break + + except Exception as e: + + print("Exception:", type(e)) + print(f"{e}, retry...", flush=True) + await asyncio.sleep(1) + +