mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-20 18:51:03 +02:00
Store service works, building management service
This commit is contained in:
Cyber MacGeddon
parent
0393c7636a
commit
7f694b623f
9 changed files with 884 additions and 12 deletions
|
|
@ -13,4 +13,5 @@ from . lookup import *
|
|||
from . library import *
|
||||
from . config import *
|
||||
from . flows import *
|
||||
from . knowledge import *
|
||||
|
||||
|
|
|
|||
52
trustgraph-base/trustgraph/schema/knowledge.py
Normal file
52
trustgraph-base/trustgraph/schema/knowledge.py
Normal file
|
|
@ -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',
|
||||
)
|
||||
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
3
trustgraph-flow/trustgraph/cores/__init__.py
Normal file
3
trustgraph-flow/trustgraph/cores/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
|
||||
from . service import run
|
||||
|
||||
5
trustgraph-flow/trustgraph/cores/__main__.py
Normal file
5
trustgraph-flow/trustgraph/cores/__main__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
from . service import run
|
||||
|
||||
if __name__ == '__main__':
|
||||
run()
|
||||
248
trustgraph-flow/trustgraph/cores/knowledge.py
Normal file
248
trustgraph-flow/trustgraph/cores/knowledge.py
Normal file
|
|
@ -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
|
||||
)
|
||||
|
||||
289
trustgraph-flow/trustgraph/cores/service.py
Executable file
289
trustgraph-flow/trustgraph/cores/service.py
Executable file
|
|
@ -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__)
|
||||
|
||||
285
trustgraph-flow/trustgraph/cores/table_store.py
Normal file
285
trustgraph-flow/trustgraph/cores/table_store.py
Normal file
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue