diff --git a/trustgraph-base/trustgraph/schema/__init__.py b/trustgraph-base/trustgraph/schema/__init__.py index a9bb30a6..957ebcbd 100644 --- a/trustgraph-base/trustgraph/schema/__init__.py +++ b/trustgraph-base/trustgraph/schema/__init__.py @@ -13,4 +13,5 @@ from . lookup import * from . library import * from . config import * from . flows import * +from . knowledge import * diff --git a/trustgraph-base/trustgraph/schema/knowledge.py b/trustgraph-base/trustgraph/schema/knowledge.py new file mode 100644 index 00000000..89bf6036 --- /dev/null +++ b/trustgraph-base/trustgraph/schema/knowledge.py @@ -0,0 +1,52 @@ + +from pulsar.schema import Record, Bytes, String, Array, Long +from . types import Triple +from . topic import topic +from . types import Error +from . metadata import Metadata +from . documents import Document, TextDocument +from . graph import Triples, GraphEmbeddings + +# fetch-kg-core +# -> (???) +# <- () +# <- (error) + +# delete-kg-core +# -> (???) +# <- () +# <- (error) + +# list-kg-cores +# -> (user) +# <- () +# <- (error) + +class KnowledgeRequest(Record): + + # fetch-kg-core, delete-kg-core, list-kg-cores + operation = String() + + # list-kg-cores, delete-kg-core + 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() + +class KnowledgeResponse(Record): + error = Error() + document_ids = Array(String()) + eos = Boolean() # Indicates end of knowledge core stream + triples = Triples() + graph_embeddings = GraphEmbeddings() + +knowledge_request_queue = topic( + 'knowledge', kind='non-persistent', namespace='request' +) +knowledge_response_queue = topic( + 'knowledge', kind='non-persistent', namespace='response', +) + diff --git a/trustgraph-base/trustgraph/schema/library.py b/trustgraph-base/trustgraph/schema/library.py index e6854987..6504fa78 100644 --- a/trustgraph-base/trustgraph/schema/library.py +++ b/trustgraph-base/trustgraph/schema/library.py @@ -51,17 +51,6 @@ from . documents import Document, TextDocument # <- (processing_metadata[]) # <- (error) -# OLD: -# add(Metadata, Bytes) : error? -# copy(id, user, collection) -# move(id, user, collection) -# delete(id) -# get(id) : Bytes -# reindex(id) -# list(user, collection) : id[] -# info(id[]) : DocumentInfo[] -# search([]) : id[] - class DocumentMetadata(Record): id = String() time = Long() diff --git a/trustgraph-flow/trustgraph/cores/__init__.py b/trustgraph-flow/trustgraph/cores/__init__.py new file mode 100644 index 00000000..9843ccfb --- /dev/null +++ b/trustgraph-flow/trustgraph/cores/__init__.py @@ -0,0 +1,3 @@ + +from . service import run + diff --git a/trustgraph-flow/trustgraph/cores/__main__.py b/trustgraph-flow/trustgraph/cores/__main__.py new file mode 100644 index 00000000..1729017e --- /dev/null +++ b/trustgraph-flow/trustgraph/cores/__main__.py @@ -0,0 +1,5 @@ + +from . service import run + +if __name__ == '__main__': + run() diff --git a/trustgraph-flow/trustgraph/cores/knowledge.py b/trustgraph-flow/trustgraph/cores/knowledge.py new file mode 100644 index 00000000..5c5a83fe --- /dev/null +++ b/trustgraph-flow/trustgraph/cores/knowledge.py @@ -0,0 +1,248 @@ + +from .. schema import KnowledgeResponse, Error, Triple +from .. knowledge import hash +from .. exceptions import RequestError +from . table_store import TableStore +from . blob_store import BlobStore +import base64 + +import uuid + +class KnowledgeManager: + + def __init__( + self, cassandra_host, cassandra_user, cassandra_password, + keyspace, + ): + + self.table_store = TableStore( + 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): + + print("Updating doc...") + + # You can't update the document ID, user or kind. + + if not await self.table_store.document_exists( + request.document_metadata.user, + request.document_metadata.id + ): + raise RuntimeError("Document does not exist") + + await self.table_store.update_document(request.document_metadata) + + print("Update complete", flush=True) + + return LibrarianResponse( + error = None, + document_metadata = None, + content = None, + document_metadatas = None, + processing_metadatas = None, + ) + + async def get_document_metadata(self, request): + + print("Get doc...") + + 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") + + # Remove doc table row + await self.table_store.remove_processing( + request.user, + request.processing_id + ) + + print("Remove complete", flush=True) + + return LibrarianResponse( + error = None, + document_metadata = None, + content = None, + document_metadatas = None, + processing_metadatas = None, + ) + + async def list_cores(self, request): + + ids = await self.table_store.list_cores(request.user) + + return KnowledgeResponse( + error = None, + document_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 new file mode 100755 index 00000000..a2ac38bc --- /dev/null +++ b/trustgraph-flow/trustgraph/cores/service.py @@ -0,0 +1,289 @@ + +""" +Knowledge core service, manages cores and exports them +""" + +from functools import partial +import asyncio +import base64 +import json + +from .. base import AsyncProcessor, Consumer, Producer, Publisher, Subscriber +from .. base import ConsumerMetrics, ProducerMetrics + +from .. schema import KnowledgeRequest, KnowledgeResponse, Error +from .. schema import knowledge_request_queue, knowledge_response_queue + +from .. schema import Document, Metadata +from .. schema import TextDocument, Metadata + +from .. exceptions import RequestError + +from . knowledge import KnowledgeManager + +default_ident = "knowledge" + +default_knowledge_request_queue = knowledge_request_queue +default_knowledge_response_queue = knowledge_response_queue + +default_cassandra_host = "cassandra" + +# FIXME: How to ensure this doesn't conflict with other usage? +keyspace = "knowledge" + +class Processor(AsyncProcessor): + + def __init__(self, **params): + + id = params.get("id") + + knowledge_request_queue = params.get( + "knowledge_request_queue", default_knowledge_request_queue + ) + + knowledge_response_queue = params.get( + "knowledge_response_queue", default_knowledge_response_queue + ) + + 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 | { + "knowledge_request_queue": knowledge_request_queue, + "knowledge_response_queue": knowledge_response_queue, + "cassandra_host": cassandra_host, + "cassandra_user": cassandra_user, + } + ) + + knowledge_request_metrics = ConsumerMetrics( + processor = self.id, flow = None, name = "knowledge-request" + ) + + knowledge_response_metrics = ProducerMetrics( + processor = self.id, flow = None, name = "knowledge-response" + ) + + self.knowledge_request_consumer = Consumer( + taskgroup = self.taskgroup, + client = self.pulsar_client, + flow = None, + topic = knowledge_request_queue, + subscriber = id, + schema = KnowledgeRequest, + handler = self.on_knowledge_request, + metrics = knowledge_request_metrics, + ) + + self.knowledge_response_producer = Producer( + client = self.pulsar_client, + topic = knowledge_response_queue, + schema = KnowledgeResponse, + metrics = knowledge_response_metrics, + ) + + self.knowledge = KnowledgeManager( + cassandra_host = cassandra_host.split(","), + cassandra_user = cassandra_user, + cassandra_password = cassandra_password, + keyspace = keyspace, + ) + + self.register_config_handler(self.on_knowledge_config) + + self.flows = {} + + print("Initialised.", flush=True) + + async def start(self): + + await super(Processor, self).start() + await self.knowledge_request_consumer.start() + await self.knowledge_response_producer.start() + + async def on_knowledge_config(self, config, version): + + print("config version", version) + + if "flows" in config: + + self.flows = { + k: json.loads(v) + for k, v in config["flows"].items() + } + + print(self.flows) + + def __del__(self): + + pass + + async def load_document(self, document, processing, content): + + print("Ready for processing...") + + print(document, processing, len(content)) + + if processing.flow not in self.flows: + raise RuntimeError("Invalid flow ID") + + flow = self.flows[processing.flow] + + if document.kind == "text/plain": + kind = "text-load" + elif document.kind == "application/pdf": + kind = "document-load" + else: + raise RuntimeError("Document with a MIME type I don't know") + + q = flow["interfaces"][kind] + + if kind == "text-load": + doc = TextDocument( + metadata = Metadata( + id = document.id, + metadata = document.metadata, + user = processing.user, + collection = processing.collection + ), + text = content, + ) + schema = TextDocument + else: + doc = Document( + metadata = Metadata( + id = document.id, + metadata = document.metadata, + user = processing.user, + collection = processing.collection + ), + data = base64.b64encode(content).decode("utf-8") + + ) + schema = Document + + print(f"Submit on queue {q}...") + + pub = Publisher( + self.pulsar_client, q, schema=schema + ) + + await pub.start() + + # FIXME: Time wait kludge? + await asyncio.sleep(1) + + await pub.send(None, doc) + + await pub.stop() + + print("Document submitted") + + async def process_request(self, v): + + if v.operation is None: + raise RequestError("Null operation") + + print("request", v.operation) + + impls = { + "list-cores": self.knowledge.list_cores, + "fetch-core": self.knowledge.fetch_core, + "delet-core": self.knowledge.delete_core, + } + + if v.operation not in impls: + raise RequestError(f"Invalid operation: {v.operation}") + + return await impls[v.operation](v) + + async def on_knowledge_request(self, msg, consumer, flow): + + v = msg.value() + + # Sender-produced ID + + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + try: + + resp = await self.process_request(v) + + await self.knowledge_response_producer.send( + resp, properties={"id": id} + ) + + return + + except RequestError as e: + resp = KnowledgeResponse( + error = Error( + type = "request-error", + message = str(e), + ) + ) + + await self.knowledge_response_producer.send( + resp, properties={"id": id} + ) + + return + except Exception as e: + resp = KnowledgeResponse( + error = Error( + type = "unexpected-error", + message = str(e), + ) + ) + + await self.knowledge_response_producer.send( + resp, properties={"id": id} + ) + + return + + print("Done.", flush=True) + + @staticmethod + def add_args(parser): + + AsyncProcessor.add_args(parser) + + parser.add_argument( + '--knowledge-request-queue', + default=default_knowledge_request_queue, + help=f'Config request queue (default: {default_knowledge_request_queue})' + ) + + parser.add_argument( + '--knowledge-response-queue', + default=default_knowledge_response_queue, + help=f'Config response queue {default_knowledge_response_queue}', + ) + + parser.add_argument( + '--cassandra-host', + default="cassandra", + help=f'Graph host (default: cassandra)' + ) + + parser.add_argument( + '--cassandra-user', + default=None, + help=f'Cassandra user' + ) + + parser.add_argument( + '--cassandra-password', + default=None, + help=f'Cassandra password' + ) + +def run(): + + Processor.launch(default_ident, __doc__) + diff --git a/trustgraph-flow/trustgraph/cores/table_store.py b/trustgraph-flow/trustgraph/cores/table_store.py new file mode 100644 index 00000000..8a75e622 --- /dev/null +++ b/trustgraph-flow/trustgraph/cores/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) + + diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index 4c18b9a5..d1ce4805 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -210,7 +210,7 @@ class Processor(AsyncProcessor): if v.operation is None: raise RequestError("Null operation") - print("requets", v.operation) + print("request", v.operation) impls = { "add-document": self.librarian.add_document,