Store service works, building management service

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 *
 

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',
+)
+

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(<key,op,value>[]) : id[]
-
 class DocumentMetadata(Record):
     id = String()
     time = Long()

trustgraph-flow/trustgraph/cores/__init__.py

@@ -0,0 +1,3 @@
+
+from . service import run
+

trustgraph-flow/trustgraph/cores/__main__.py

@@ -0,0 +1,5 @@
+
+from . service import run
+
+if __name__ == '__main__':
+    run()

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
+        )
+

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__)
+

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<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,
+                time timestamp,
+                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,
+                time timestamp,
+                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_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)
+
+        

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,