From cb529db5231d809fdb52fe207cc0ef0c95db1117 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 6 May 2025 22:37:15 +0100 Subject: [PATCH] Fetch and list mostly there --- .../trustgraph/schema/knowledge.py | 7 +- trustgraph-flow/trustgraph/cores/knowledge.py | 218 +++-------------- trustgraph-flow/trustgraph/cores/service.py | 12 +- .../trustgraph/gateway/dispatch/knowledge.py | 55 +++++ .../trustgraph/gateway/dispatch/manager.py | 2 + .../trustgraph/librarian/librarian.py | 2 +- .../trustgraph/tables/knowledge.py | 222 ++++++++++++++++-- 7 files changed, 300 insertions(+), 218 deletions(-) create mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/knowledge.py diff --git a/trustgraph-base/trustgraph/schema/knowledge.py b/trustgraph-base/trustgraph/schema/knowledge.py index b663ac35..88892f8c 100644 --- a/trustgraph-base/trustgraph/schema/knowledge.py +++ b/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() - - # fetch-kg-core, list-kg-cores, delete-kg-core - document_id = String() + 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() diff --git a/trustgraph-flow/trustgraph/cores/knowledge.py b/trustgraph-flow/trustgraph/cores/knowledge.py index 05a25605..82224d71 100644 --- a/trustgraph-flow/trustgraph/cores/knowledge.py +++ b/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): - - 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): + async def delete_kg_core(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( - request.user, - request.document_id - ) - - 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") + async def publish(obj): + print(obj) +# await publisher.send(triples) # 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( - error = None, - document_metadata = None, - content = None, - document_metadatas = None, - processing_metadatas = None, + # Remove doc table row + await self.table_store.get_graph_embeddings( + request.user, + request.id, + publish ) - async def list_cores(self, request): - - ids = await self.table_store.list_cores(request.user) + print("Fetch complete", flush=True) 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 diff --git a/trustgraph-flow/trustgraph/cores/service.py b/trustgraph-flow/trustgraph/cores/service.py index a2ac38bc..eca6ceb3 100755 --- a/trustgraph-flow/trustgraph/cores/service.py +++ b/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, - "fetch-core": self.knowledge.fetch_core, - "delet-core": self.knowledge.delete_core, + "list-kg-cores": self.knowledge.list_kg_cores, + "fetch-kg-core": self.knowledge.fetch_kg_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} diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/knowledge.py b/trustgraph-flow/trustgraph/gateway/dispatch/knowledge.py new file mode 100644 index 00000000..3369b4e3 --- /dev/null +++ b/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 + diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index ddd396a1..7896d588 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/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 = { diff --git a/trustgraph-flow/trustgraph/librarian/librarian.py b/trustgraph-flow/trustgraph/librarian/librarian.py index 3247c38e..89750c42 100644 --- a/trustgraph-flow/trustgraph/librarian/librarian.py +++ b/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 ) diff --git a/trustgraph-flow/trustgraph/tables/knowledge.py b/trustgraph-flow/trustgraph/tables/knowledge.py index 628c3fae..1dbb4fd5 100644 --- a/trustgraph-flow/trustgraph/tables/knowledge.py +++ b/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>, - 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> > >, - 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> > >, - 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, - metadata, triples + id, user, document_id, + time, metadata, triples ) - VALUES (?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?) """) self.insert_graph_embeddings_stmt = self.cassandra.prepare(""" INSERT INTO graph_embeddings ( - id, user, collection, document_id, time, - metadata, entity_embeddings + id, user, document_id, time, metadata, entity_embeddings ) - VALUES (?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?) """) self.insert_document_embeddings_stmt = self.cassandra.prepare(""" INSERT INTO document_embeddings ( - id, user, collection, document_id, time, - metadata, chunks + id, user, document_id, time, 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") +