mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-21 11:11:03 +02:00
Update librarian to new API
This commit is contained in:
Cyber MacGeddon
parent
6bf485788a
commit
6e31d24fb1
8 changed files with 136 additions and 221 deletions
|
|
@ -28,6 +28,7 @@ from . documents import Document, TextDocument
|
|||
|
||||
class DocumentPackage(Record):
|
||||
id = String()
|
||||
flow = String()
|
||||
document = Bytes()
|
||||
kind = String()
|
||||
user = String()
|
||||
|
|
@ -39,6 +40,7 @@ class DocumentPackage(Record):
|
|||
|
||||
class DocumentInfo(Record):
|
||||
id = String()
|
||||
flow = String()
|
||||
kind = String()
|
||||
user = String()
|
||||
collection = String()
|
||||
|
|
@ -55,6 +57,7 @@ class Criteria(Record):
|
|||
class LibrarianRequest(Record):
|
||||
operation = String()
|
||||
id = String()
|
||||
flow = String()
|
||||
document = DocumentPackage()
|
||||
user = String()
|
||||
collection = String()
|
||||
|
|
|
|||
|
|
@ -3,8 +3,6 @@
|
|||
Config service. Manages system global configuration state
|
||||
"""
|
||||
|
||||
from pulsar.schema import JsonSchema
|
||||
|
||||
from trustgraph.schema import Error
|
||||
|
||||
from trustgraph.schema import ConfigRequest, ConfigResponse, ConfigPush
|
||||
|
|
@ -14,7 +12,6 @@ from trustgraph.schema import config_push_queue
|
|||
from trustgraph.schema import FlowRequest, FlowResponse
|
||||
from trustgraph.schema import flow_request_queue, flow_response_queue
|
||||
|
||||
from trustgraph.log_level import LogLevel
|
||||
from trustgraph.base import AsyncProcessor, Consumer, Producer
|
||||
|
||||
from . config import Configuration
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ class LibrarianRequestor(ServiceRequestor):
|
|||
return LibrarianRequest(
|
||||
operation = body.get("operation", None),
|
||||
id = body.get("id", None),
|
||||
flow = body.get("flow", None),
|
||||
document = dp,
|
||||
user = body.get("user", None),
|
||||
collection = body.get("collection", None),
|
||||
|
|
|
|||
|
|
@ -99,6 +99,9 @@ def serialize_document_package(message):
|
|||
if message.kind:
|
||||
ret["kind"] = message.kind
|
||||
|
||||
if message.flow:
|
||||
ret["flow"] = message.flow
|
||||
|
||||
if message.user:
|
||||
ret["user"] = message.user
|
||||
|
||||
|
|
@ -117,6 +120,9 @@ def serialize_document_info(message):
|
|||
if message.kind:
|
||||
ret["kind"] = message.kind
|
||||
|
||||
if message.flow:
|
||||
ret["flow"] = message.flow
|
||||
|
||||
if message.user:
|
||||
ret["user"] = message.user
|
||||
|
||||
|
|
@ -141,6 +147,7 @@ def to_document_package(x):
|
|||
|
||||
return DocumentPackage(
|
||||
id = x.get("id", None),
|
||||
flow = x.get("flow", None),
|
||||
kind = x.get("kind", None),
|
||||
user = x.get("user", None),
|
||||
collection = x.get("collection", None),
|
||||
|
|
@ -155,6 +162,7 @@ def to_document_info(x):
|
|||
|
||||
return DocumentInfo(
|
||||
id = x.get("id", None),
|
||||
flow = x.get("flow", None),
|
||||
kind = x.get("kind", None),
|
||||
user = x.get("user", None),
|
||||
collection = x.get("collection", None),
|
||||
|
|
|
|||
|
|
@ -95,7 +95,6 @@ class Api:
|
|||
|
||||
await self.config_receiver.start()
|
||||
|
||||
|
||||
for ep in self.endpoints:
|
||||
ep.add_routes(self.app)
|
||||
|
||||
|
|
|
|||
|
|
@ -67,22 +67,3 @@ class Librarian:
|
|||
info = info,
|
||||
)
|
||||
|
||||
def handle_triples(self, m):
|
||||
self.table_store.add_triples(m)
|
||||
|
||||
def handle_graph_embeddings(self, m):
|
||||
self.table_store.add_graph_embeddings(m)
|
||||
|
||||
def handle_document_embeddings(self, m):
|
||||
self.table_store.add_document_embeddings(m)
|
||||
|
||||
|
||||
def handle_triples(self, m):
|
||||
self.table_store.add_triples(m)
|
||||
|
||||
def handle_graph_embeddings(self, m):
|
||||
self.table_store.add_graph_embeddings(m)
|
||||
|
||||
def handle_document_embeddings(self, m):
|
||||
self.table_store.add_document_embeddings(m)
|
||||
|
||||
|
|
|
|||
|
|
@ -5,41 +5,27 @@ Librarian service, manages documents in collections
|
|||
|
||||
from functools import partial
|
||||
import asyncio
|
||||
import threading
|
||||
import queue
|
||||
import base64
|
||||
import json
|
||||
|
||||
from pulsar.schema import JsonSchema
|
||||
from .. base import AsyncProcessor, Consumer, Producer, Publisher, Subscriber
|
||||
from .. base import ConsumerMetrics, ProducerMetrics
|
||||
|
||||
from .. schema import LibrarianRequest, LibrarianResponse, Error
|
||||
from .. schema import librarian_request_queue, librarian_response_queue
|
||||
|
||||
from .. schema import GraphEmbeddings
|
||||
from .. schema import graph_embeddings_store_queue
|
||||
from .. schema import Triples
|
||||
from .. schema import triples_store_queue
|
||||
from .. schema import DocumentEmbeddings
|
||||
from .. schema import document_embeddings_store_queue
|
||||
|
||||
from .. schema import Document, Metadata
|
||||
from .. schema import document_ingest_queue
|
||||
from .. schema import TextDocument, Metadata
|
||||
from .. schema import text_ingest_queue
|
||||
|
||||
from .. base import Publisher
|
||||
from .. base import Subscriber
|
||||
|
||||
from .. log_level import LogLevel
|
||||
from .. base import ConsumerProducer
|
||||
from .. exceptions import RequestError
|
||||
|
||||
from . librarian import Librarian
|
||||
|
||||
module = "librarian"
|
||||
default_ident = "librarian"
|
||||
|
||||
default_librarian_request_queue = librarian_request_queue
|
||||
default_librarian_response_queue = librarian_response_queue
|
||||
|
||||
default_input_queue = librarian_request_queue
|
||||
default_output_queue = librarian_response_queue
|
||||
default_subscriber = module
|
||||
default_minio_host = "minio:9000"
|
||||
default_minio_access_key = "minioadmin"
|
||||
default_minio_secret_key = "minioadmin"
|
||||
|
|
@ -50,15 +36,21 @@ bucket_name = "library"
|
|||
# FIXME: How to ensure this doesn't conflict with other usage?
|
||||
keyspace = "librarian"
|
||||
|
||||
class Processor(ConsumerProducer):
|
||||
class Processor(AsyncProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
self.running = True
|
||||
id = params.get("id")
|
||||
|
||||
input_queue = params.get("input_queue", default_input_queue)
|
||||
output_queue = params.get("output_queue", default_output_queue)
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
# self.running = True
|
||||
|
||||
librarian_request_queue = params.get(
|
||||
"librarian_request_queue", default_librarian_request_queue
|
||||
)
|
||||
|
||||
librarian_response_queue = params.get(
|
||||
"librarian_response_queue", default_librarian_response_queue
|
||||
)
|
||||
|
||||
minio_host = params.get("minio_host", default_minio_host)
|
||||
minio_access_key = params.get(
|
||||
|
|
@ -74,19 +66,10 @@ class Processor(ConsumerProducer):
|
|||
cassandra_user = params.get("cassandra_user")
|
||||
cassandra_password = params.get("cassandra_password")
|
||||
|
||||
triples_queue = params.get("triples_queue")
|
||||
graph_embeddings_queue = params.get("graph_embeddings_queue")
|
||||
document_embeddings_queue = params.get("document_embeddings_queue")
|
||||
document_load_queue = params.get("document_load_queue")
|
||||
text_load_queue = params.get("text_load_queue")
|
||||
|
||||
super(Processor, self).__init__(
|
||||
**params | {
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": LibrarianRequest,
|
||||
"output_schema": LibrarianResponse,
|
||||
"librarian_request_queue": librarian_request_queue,
|
||||
"librarian_response_queue": librarian_response_queue,
|
||||
"minio_host": minio_host,
|
||||
"minio_access_key": minio_access_key,
|
||||
"cassandra_host": cassandra_host,
|
||||
|
|
@ -94,38 +77,30 @@ class Processor(ConsumerProducer):
|
|||
}
|
||||
)
|
||||
|
||||
self.document_load = Publisher(
|
||||
self.client, document_load_queue, JsonSchema(Document),
|
||||
librarian_request_metrics = ConsumerMetrics(
|
||||
processor = self.id, flow = None, name = "librarian-request"
|
||||
)
|
||||
|
||||
self.text_load = Publisher(
|
||||
self.client, text_load_queue, JsonSchema(TextDocument),
|
||||
librarian_response_metrics = ProducerMetrics(
|
||||
processor = self.id, flow = None, name = "librarian-response"
|
||||
)
|
||||
|
||||
self.triples_brk = Subscriber(
|
||||
self.client, triples_store_queue,
|
||||
"librarian", "librarian",
|
||||
schema=JsonSchema(Triples),
|
||||
)
|
||||
self.graph_embeddings_brk = Subscriber(
|
||||
self.client, graph_embeddings_store_queue,
|
||||
"librarian", "librarian",
|
||||
schema=JsonSchema(GraphEmbeddings),
|
||||
)
|
||||
self.document_embeddings_brk = Subscriber(
|
||||
self.client, document_embeddings_store_queue,
|
||||
"librarian", "librarian",
|
||||
schema=JsonSchema(DocumentEmbeddings),
|
||||
self.librarian_request_consumer = Consumer(
|
||||
taskgroup = self.taskgroup,
|
||||
client = self.pulsar_client,
|
||||
flow = None,
|
||||
topic = librarian_request_queue,
|
||||
subscriber = id,
|
||||
schema = LibrarianRequest,
|
||||
handler = self.on_librarian_request,
|
||||
metrics = librarian_request_metrics,
|
||||
)
|
||||
|
||||
self.triples_reader = threading.Thread(
|
||||
target=self.receive_triples
|
||||
)
|
||||
self.graph_embeddings_reader = threading.Thread(
|
||||
target=self.receive_graph_embeddings
|
||||
)
|
||||
self.document_embeddings_reader = threading.Thread(
|
||||
target=self.receive_document_embeddings
|
||||
self.librarian_response_producer = Producer(
|
||||
client = self.pulsar_client,
|
||||
topic = librarian_response_queue,
|
||||
schema = LibrarianResponse,
|
||||
metrics = librarian_response_metrics,
|
||||
)
|
||||
|
||||
self.librarian = Librarian(
|
||||
|
|
@ -141,90 +116,43 @@ class Processor(ConsumerProducer):
|
|||
load_text = self.load_text,
|
||||
)
|
||||
|
||||
self.register_config_handler(self.on_librarian_config)
|
||||
|
||||
self.flows = {}
|
||||
|
||||
print("Initialised.", flush=True)
|
||||
|
||||
async def start(self):
|
||||
|
||||
self.document_load.start()
|
||||
self.text_load.start()
|
||||
|
||||
self.triples_brk.start()
|
||||
self.graph_embeddings_brk.start()
|
||||
self.document_embeddings_brk.start()
|
||||
await super(Processor, self).start()
|
||||
await self.librarian_request_consumer.start()
|
||||
await self.librarian_response_producer.start()
|
||||
|
||||
self.triples_sub = self.triples_brk.subscribe_all("x")
|
||||
self.graph_embeddings_sub = self.graph_embeddings_brk.subscribe_all("x")
|
||||
self.document_embeddings_sub = self.document_embeddings_brk.subscribe_all("x")
|
||||
async def on_librarian_config(self, config, version):
|
||||
|
||||
self.triples_reader.start()
|
||||
self.graph_embeddings_reader.start()
|
||||
self.document_embeddings_reader.start()
|
||||
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):
|
||||
|
||||
self.running = False
|
||||
|
||||
if hasattr(self, "document_load"):
|
||||
self.document_load.stop()
|
||||
self.document_load.join()
|
||||
|
||||
if hasattr(self, "text_load"):
|
||||
self.text_load.stop()
|
||||
self.text_load.join()
|
||||
|
||||
if hasattr(self, "triples_sub"):
|
||||
self.triples_sub.unsubscribe_all("x")
|
||||
|
||||
if hasattr(self, "graph_embeddings_sub"):
|
||||
self.graph_embeddings_sub.unsubscribe_all("x")
|
||||
|
||||
if hasattr(self, "document_embeddings_sub"):
|
||||
self.document_embeddings_sub.unsubscribe_all("x")
|
||||
|
||||
if hasattr(self, "triples_brk"):
|
||||
self.triples_brk.stop()
|
||||
self.triples_brk.join()
|
||||
|
||||
if hasattr(self, "graph_embeddings_brk"):
|
||||
self.graph_embeddings_brk.stop()
|
||||
self.graph_embeddings_brk.join()
|
||||
|
||||
if hasattr(self, "document_embeddings_brk"):
|
||||
self.document_embeddings_brk.stop()
|
||||
self.document_embeddings_brk.join()
|
||||
|
||||
def receive_triples(self):
|
||||
|
||||
while self.running:
|
||||
try:
|
||||
msg = self.triples_sub.get(timeout=1)
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
self.librarian.handle_triples(msg)
|
||||
|
||||
def receive_graph_embeddings(self):
|
||||
|
||||
while self.running:
|
||||
try:
|
||||
msg = self.graph_embeddings_sub.get(timeout=1)
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
self.librarian.handle_graph_embeddings(msg)
|
||||
|
||||
def receive_document_embeddings(self):
|
||||
|
||||
while self.running:
|
||||
try:
|
||||
msg = self.document_embeddings_sub.get(timeout=1)
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
self.librarian.handle_document_embeddings(msg)
|
||||
pass
|
||||
# self.running = False
|
||||
|
||||
async def load_document(self, document):
|
||||
|
||||
print(document)
|
||||
print(document.flow)
|
||||
|
||||
|
||||
|
||||
doc = Document(
|
||||
metadata = Metadata(
|
||||
id = document.id,
|
||||
|
|
@ -235,6 +163,8 @@ class Processor(ConsumerProducer):
|
|||
data = document.document
|
||||
)
|
||||
|
||||
|
||||
|
||||
self.document_load.send(None, doc)
|
||||
|
||||
async def load_text(self, document):
|
||||
|
|
@ -283,12 +213,14 @@ class Processor(ConsumerProducer):
|
|||
collection = v.collection,
|
||||
)
|
||||
else:
|
||||
print("BROK")
|
||||
print("User not specified")
|
||||
raise RequestError("Invalid call")
|
||||
|
||||
raise RequestError("Invalid operation: " + v.operation)
|
||||
|
||||
async def handle(self, msg):
|
||||
async def on_librarian_request(self, msg, consumer, flow):
|
||||
|
||||
print("REQUEST")
|
||||
|
||||
v = msg.value()
|
||||
|
||||
|
|
@ -307,7 +239,11 @@ class Processor(ConsumerProducer):
|
|||
message = str(e),
|
||||
)
|
||||
)
|
||||
await self.send(resp, properties={"id": id})
|
||||
|
||||
await self.librarian_response_producer.send(
|
||||
resp, properties={"id": id}
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
try:
|
||||
|
|
@ -320,7 +256,11 @@ class Processor(ConsumerProducer):
|
|||
message = str(e),
|
||||
)
|
||||
)
|
||||
await self.send(resp, properties={"id": id})
|
||||
|
||||
await self.librarian_response_producer.send(
|
||||
resp, properties={"id": id}
|
||||
)
|
||||
|
||||
return
|
||||
except Exception as e:
|
||||
print("Exception:", e, flush=True)
|
||||
|
|
@ -330,21 +270,36 @@ class Processor(ConsumerProducer):
|
|||
message = "Unhandled error: " + str(e),
|
||||
)
|
||||
)
|
||||
await self.send(resp, properties={"id": id})
|
||||
|
||||
await self.librarian_response_producer.send(
|
||||
resp, properties={"id": id}
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
print("Send response..!.", flush=True)
|
||||
|
||||
await self.send(resp, properties={"id": id})
|
||||
await self.librarian_response_producer.send(
|
||||
resp, properties={"id": id}
|
||||
)
|
||||
|
||||
print("Done.", flush=True)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
ConsumerProducer.add_args(
|
||||
parser, default_input_queue, default_subscriber,
|
||||
default_output_queue,
|
||||
AsyncProcessor.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
'--librarian-request-queue',
|
||||
default=default_librarian_request_queue,
|
||||
help=f'Config request queue (default: {default_librarian_request_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--librarian-response-queue',
|
||||
default=default_librarian_response_queue,
|
||||
help=f'Config response queue {default_librarian_response_queue}',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
|
|
@ -385,40 +340,7 @@ class Processor(ConsumerProducer):
|
|||
help=f'Cassandra password'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--triples-queue',
|
||||
default=triples_store_queue,
|
||||
help=f'Triples queue (default: {triples_store_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--graph-embeddings-queue',
|
||||
default=graph_embeddings_store_queue,
|
||||
help=f'Graph embeddings queue (default: {triples_store_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--document-embeddings-queue',
|
||||
default=document_embeddings_store_queue,
|
||||
help='Document embeddings queue '
|
||||
f'(default: {document_embeddings_store_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--document-load-queue',
|
||||
default=document_ingest_queue,
|
||||
help='Document load queue '
|
||||
f'(default: {document_ingest_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--text-load-queue',
|
||||
default=text_ingest_queue,
|
||||
help='Text ingest queue '
|
||||
f'(default: {text_ingest_queue})'
|
||||
)
|
||||
|
||||
def run():
|
||||
|
||||
Processor.launch(module, __doc__)
|
||||
Processor.launch(default_ident, __doc__)
|
||||
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ class TableStore:
|
|||
user text,
|
||||
collection text,
|
||||
id text,
|
||||
flow text,
|
||||
time timestamp,
|
||||
title text,
|
||||
comments text,
|
||||
|
|
@ -155,22 +156,24 @@ class TableStore:
|
|||
self.insert_document_stmt = self.cassandra.prepare("""
|
||||
INSERT INTO document
|
||||
(
|
||||
id, user, collection, kind, object_id, time, title, comments,
|
||||
metadata
|
||||
id, user, collection, flow, kind, object_id, time, title,
|
||||
comments, metadata
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""")
|
||||
|
||||
self.list_document_stmt = self.cassandra.prepare("""
|
||||
SELECT
|
||||
id, kind, user, collection, title, comments, time, metadata
|
||||
id, kind, flow, user, collection, title, comments, time,
|
||||
metadata
|
||||
FROM document
|
||||
WHERE user = ?
|
||||
""")
|
||||
|
||||
self.list_document_by_collection_stmt = self.cassandra.prepare("""
|
||||
SELECT
|
||||
id, kind, user, collection, title, comments, time, metadata
|
||||
id, kind, flow, user, collection, title, comments, time,
|
||||
metadata
|
||||
FROM document
|
||||
WHERE user = ? AND collection = ?
|
||||
""")
|
||||
|
|
@ -325,18 +328,19 @@ class TableStore:
|
|||
DocumentInfo(
|
||||
id = row[0],
|
||||
kind = row[1],
|
||||
user = row[2],
|
||||
collection = row[3],
|
||||
title = row[4],
|
||||
comments = row[5],
|
||||
time = int(1000 * row[6].timestamp()),
|
||||
flow = row[2],
|
||||
user = row[3],
|
||||
collection = row[4],
|
||||
title = row[5],
|
||||
comments = row[6],
|
||||
time = int(1000 * row[7].timestamp()),
|
||||
metadata = [
|
||||
Triple(
|
||||
s=Value(value=m[0], is_uri=m[1]),
|
||||
p=Value(value=m[2], is_uri=m[3]),
|
||||
o=Value(value=m[4], is_uri=m[5])
|
||||
)
|
||||
for m in row[7]
|
||||
for m in row[8]
|
||||
],
|
||||
)
|
||||
for row in resp
|
||||
|
|
@ -344,16 +348,16 @@ class TableStore:
|
|||
|
||||
print("OK3")
|
||||
|
||||
print(info[0])
|
||||
print(info)
|
||||
|
||||
print(info[0].user)
|
||||
print(info[0].time)
|
||||
print(info[0].kind)
|
||||
print(info[0].collection)
|
||||
print(info[0].title)
|
||||
print(info[0].comments)
|
||||
print(info[0].metadata)
|
||||
print(info[0].metadata)
|
||||
# print(info[0].user)
|
||||
# print(info[0].time)
|
||||
# print(info[0].kind)
|
||||
# print(info[0].collection)
|
||||
# print(info[0].title)
|
||||
# print(info[0].comments)
|
||||
# print(info[0].metadata)
|
||||
# print(info[0].metadata)
|
||||
|
||||
return info
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue